Fix: Custom Provider Support Image Generation

Author: ShenKSPZ

src/components/settings/ApiManagement.tsx

@@ -166,24 +166,24 @@ export const ApiManagement: React.FC<ApiManagementProps> = ({
   };
 
   // Handle capability change
-  // const handleCapabilityChange = (capability: AIServiceCapability, isChecked: boolean) => {
-  //   if (!currentEditModel) return;
+  const handleCapabilityChange = (capability: AIServiceCapability, isChecked: boolean) => {
+    if (!currentEditModel) return;
 
-  //   let updatedCapabilities = [...currentEditModel.modelCapabilities];
+    let updatedCapabilities = [...currentEditModel.modelCapabilities];
 
-  //   if (isChecked) {
-  //     if (!updatedCapabilities.includes(capability)) {
-  //       updatedCapabilities.push(capability);
-  //     }
-  //   } else {
-  //     updatedCapabilities = updatedCapabilities.filter(cap => cap !== capability);
-  //   }
+    if (isChecked) {
+      if (!updatedCapabilities.includes(capability)) {
+        updatedCapabilities.push(capability);
+      }
+    } else {
+      updatedCapabilities = updatedCapabilities.filter(cap => cap !== capability);
+    }
 
-  //   setCurrentEditModel({
-  //     ...currentEditModel,
-  //     modelCapabilities: updatedCapabilities
-  //   });
-  // };
+    setCurrentEditModel({
+      ...currentEditModel,
+      modelCapabilities: updatedCapabilities
+    });
+  };
 
   // Get capability icon
   const getCapabilityIcon = (capability: AIServiceCapability) => {
@@ -288,17 +288,14 @@ export const ApiManagement: React.FC<ApiManagementProps> = ({
               ></textarea>
             </div> */}
 
-            {/* <div>
+            <div>
               <label className="block mb-2 text-sm font-medium text-gray-700">
                 Capabilities
               </label>
               <div className="space-y-2">
                 {[
                   {capability: AIServiceCapability.TextCompletion, label: 'Text Completion'},
-                  {capability: AIServiceCapability.Reasoning, label: 'Reasoning'},
-                  {capability: AIServiceCapability.VisionAnalysis, label: 'Vision'},
-                  {capability: AIServiceCapability.ToolUsage, label: 'Tool Usage'},
-                  {capability: AIServiceCapability.Embedding, label: 'Embedding'}
+                  {capability: AIServiceCapability.ImageGeneration, label: 'Image Generation'},
                 ].map(({capability, label}) => (
                   <div key={capability} className="flex items-center">
                     <input
@@ -317,7 +314,7 @@ export const ApiManagement: React.FC<ApiManagementProps> = ({
                   </div>
                 ))}
               </div>
-            </div> */}
+            </div>
           </div>
 
           <div className="flex justify-end mt-6 space-x-2">

src/services/providers/custom-service.ts

@@ -143,13 +143,29 @@ export class CustomService implements AiServiceProvider {
    */
   public async getImageGeneration(
     prompt: string,
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
     options: {
       size?: `${number}x${number}`;
+      aspectRatio?: `${number}:${number}`;
       style?: string;
       quality?: string;
-    } = {}
-  ): Promise<string[]> {
-    throw new Error('Not implemented');
+    }
+  ): Promise<string[] | Uint8Array<ArrayBufferLike>[]> {
+
+    const imageModel = this.openAIProvider.imageModel('dall-e-3');
+
+    const result = await imageModel.doGenerate({
+      prompt: prompt,
+      n: 1,
+      size: options.size || '1024x1024',
+      aspectRatio: options.aspectRatio || '1:1',
+      seed: 42,
+      providerOptions: {
+        "openai": {
+          "style": options.style || 'vivid'
+        }
+      }
+    });
+
+    return result.images;
   }
 } 

src/services/settings-service.ts

@@ -9,31 +9,31 @@ import { v4 as uuidv4 } from 'uuid';
  */
 const DEFAULT_SETTINGS: UserSettings = {
   providers: {
-    ['TensorBlock']: {
-      providerId: 'TensorBlock',
-      providerName: 'TensorBlock',
-      apiKey: '',
-      baseUrl: 'http://54.177.123.202:8000/v1',
-      customProvider: false,
-      models:[
-        {
-          modelName: 'GPT-4o',
-          modelId: 'gpt-4o',
-          modelCategory: 'GPT 4',
-          modelDescription: 'GPT-4o is the latest and most powerful model from OpenAI.',
-          modelCapabilities: [AIServiceCapability.TextCompletion, AIServiceCapability.WebSearch],
-          modelRefUUID: uuidv4(),
-        },
-        {
-          modelName: 'DALL-E 3',
-          modelId: 'dall-e-3',
-          modelCategory: 'Image Generation',
-          modelDescription: 'DALL-E 3 is OpenAI\'s advanced image generation model.',
-          modelCapabilities: [AIServiceCapability.ImageGeneration],
-          modelRefUUID: uuidv4(),
-        },
-      ]
-    },
+    // ['TensorBlock']: {
+    //   providerId: 'TensorBlock',
+    //   providerName: 'TensorBlock',
+    //   apiKey: '',
+    //   baseUrl: 'http://54.177.123.202:8000/v1',
+    //   customProvider: false,
+    //   models:[
+    //     {
+    //       modelName: 'GPT-4o',
+    //       modelId: 'gpt-4o',
+    //       modelCategory: 'GPT 4',
+    //       modelDescription: 'GPT-4o is the latest and most powerful model from OpenAI.',
+    //       modelCapabilities: [AIServiceCapability.TextCompletion, AIServiceCapability.WebSearch],
+    //       modelRefUUID: uuidv4(),
+    //     },
+    //     {
+    //       modelName: 'DALL-E 3',
+    //       modelId: 'dall-e-3',
+    //       modelCategory: 'Image Generation',
+    //       modelDescription: 'DALL-E 3 is OpenAI\'s advanced image generation model.',
+    //       modelCapabilities: [AIServiceCapability.ImageGeneration],
+    //       modelRefUUID: uuidv4(),
+    //     },
+    //   ]
+    // },
     ['OpenAI']: {
       providerId: 'OpenAI',
       providerName: 'OpenAI',
@@ -348,6 +348,10 @@ export class SettingsService {
           });
         }
       }
+      else if(provider === 'TensorBlock'){
+        // Remove Default TensorBlock provider
+        delete this.settings.providers[provider];
+      }
     }
   }