Add support for custom models

Author: mythz

AiServer.ServiceInterface/AiProviderFactory.cs

@@ -17,7 +17,12 @@ public interface IOllamaAiProvider
     Task<OllamaGenerationResult> GenerateAsync(AiProvider provider, OllamaGenerate request, CancellationToken token = default);
 }
 
-public class AiProviderFactory(OpenAiProvider openAiProvider, OllamaAiProvider ollamaAiProvider, GoogleAiProvider googleProvider, AnthropicAiProvider anthropicAiProvider)
+public class AiProviderFactory(
+    OpenAiProvider openAiProvider, 
+    OllamaAiProvider ollamaAiProvider, 
+    GoogleAiProvider googleProvider, 
+    AnthropicAiProvider anthropicAiProvider,
+    CustomAiProvider customAiProvider)
 {
     public IOpenAiProvider GetOpenAiProvider(AiProviderType aiProviderType=AiProviderType.OpenAiProvider)
     {
@@ -26,6 +31,7 @@ public IOpenAiProvider GetOpenAiProvider(AiProviderType aiProviderType=AiProvide
             AiProviderType.OllamaAiProvider => ollamaAiProvider,
             AiProviderType.GoogleAiProvider => googleProvider,
             AiProviderType.AnthropicAiProvider => anthropicAiProvider,
+            AiProviderType.CustomOpenAiProvider => customAiProvider,
             _ => openAiProvider
         };
     }

AiServer.ServiceInterface/OpenAiChatServices.cs

@@ -42,9 +42,7 @@ public object Any(ActiveAiModels request)
             .SelectMany(x => x.Models.Select(m => appData.GetQualifiedModel(m.Model)))
             .Where(x => x != null)
             .Select(x => x!)  // Non-null assertion after filtering out null values
-            .Distinct()
-            .OrderBy(x => x)
-            .ToList();
+            .ToSet();
 
         if (request.Vision == true)
         {
@@ -53,12 +51,32 @@ public object Any(ActiveAiModels request)
                 .Where(x => x.Vision == true)
                 .Select(x => x.Id)
                 .ToSet();
-            activeModels = activeModels.Where(x => allVisionModels.Contains(x.LeftPart(':'))).ToList();
+            activeModels = activeModels.Where(x => allVisionModels.Contains(x.LeftPart(':'))).ToSet();
+        }
+
+        var customModels = appData.AiProviders
+            .Where(x => x.AiTypeId == "Custom")
+            .SelectMany(x => x.SelectedModels);
+        foreach (var customModel in customModels)
+        {
+            activeModels.Add(customModel);
         }
 
         return new StringsResponse
         {
-            Results = activeModels 
+            Results = activeModels.OrderBy(x => x).ToList() 
+        };
+    }
+
+    public object Any(ActiveCustomAiModels request)
+    {
+        return new StringsResponse
+        {
+            Results = appData.AiProviders
+                .Where(x => x.AiTypeId == "Custom")
+                .SelectMany(x => x.SelectedModels)
+                .OrderBy(x => x)
+                .ToList() 
         };
     }
 
@@ -104,6 +122,23 @@ public object GetModelImage(string model)
             }
         }
 
+        var customModels = appData.AiProviders
+            .Where(x => x.AiTypeId == "Custom")
+            .SelectMany(x => x.SelectedModels)
+            .ToSet(StringComparer.OrdinalIgnoreCase);
+
+        if (customModels.Contains(model))
+        {
+            return new HttpResult(
+                """
+                <svg xmlns="http://www.w3.org/2000/svg" width="32" height="32" viewBox="0 0 32 32">
+                    <path fill="currentColor" d="M19 22v-2h1v-7h-1v-2h4v2h-1v7h1v2zm-3.5 0h2L14 11h-3L7.503 22h2l.601-2h4.778zm-4.794-4l1.628-5.411l.256-.003L14.264 18zM32 4h-4V0h-2v4h-4v2h4v4h2V6h4zm-2 8h2v2h-2zM18 0h2v2h-2z"/>
+                    <path fill="currentColor" d="M32 32H0V0h14v2H2v28h28V18h2z"/>
+                </svg>
+                """, 
+                MimeTypes.ImageSvg);
+        }
+        
         return new HttpResult(
             """
             <svg xmlns="http://www.w3.org/2000/svg" width="1em" height="1em" viewBox="0 0 48 48">
@@ -217,11 +252,19 @@ public QueueOpenAiChatResponse Any(QueueOpenAiChatCompletion request)
 
         if (request.Request.Messages.IsNullOrEmpty())
             throw new ArgumentNullException(nameof(request.Request.Messages));
-    
+
         var qualifiedModel = appData.GetQualifiedModel(request.Request.Model);
-        if (qualifiedModel == null)
+
+        var customModel = appData.AiProviders
+            .Where(x => x.AiTypeId == "Custom")
+            .SelectMany(x => x.SelectedModels)
+            .FirstOrDefault(x => x == request.Request.Model);
+        
+        if (qualifiedModel == null && customModel == null)
             throw HttpError.NotFound($"Model {request.Request.Model} not found");
 
+        qualifiedModel ??= customModel;
+        
         var queueCounts = jobs.GetWorkerQueueCounts();
         var providerQueueCount = int.MaxValue;
         AiProvider? useProvider = null;
@@ -497,13 +540,23 @@ public object Any(CreateAiProvider request)
             request.Models ??= [];
             foreach (var selectedModel in request.SelectedModels)
             {
-                var qualifiedModel = appData.GetQualifiedModel(selectedModel);
-                if (qualifiedModel == null)
-                    continue;
-                request.Models.Add(new()
+                if (request.AiTypeId == "Custom")
+                {
+                    request.Models.Add(new()
+                    {
+                        Model = selectedModel
+                    });
+                }
+                else
                 {
-                    Model = qualifiedModel
-                });
+                    var qualifiedModel = appData.GetQualifiedModel(selectedModel);
+                    if (qualifiedModel == null)
+                        continue;
+                    request.Models.Add(new()
+                    {
+                        Model = qualifiedModel
+                    });
+                }
             }
         }
 
@@ -515,20 +568,34 @@ public object Any(CreateAiProvider request)
     public object Any(UpdateAiProvider request)
     {
         var ignore = new[] { nameof(request.Id), nameof(request.SelectedModels) };
+        var provider = Db.SingleById<AiProvider>(request.Id);
+        if (provider == null)
+            throw HttpError.NotFound("Provider not found");
+        
         // Only call AutoQuery Update if there's something to update
         IdResponse? response = null;
         if (request.SelectedModels is { Count: > 0 })
         {
             request.Models ??= [];
             foreach (var selectedModel in request.SelectedModels)
             {
-                var qualifiedModel = appData.GetQualifiedModel(selectedModel);
-                if (qualifiedModel == null)
-                    continue;
-                request.Models.Add(new()
+                if (provider.AiTypeId == "Custom")
                 {
-                    Model = qualifiedModel
-                });
+                    request.Models.Add(new()
+                    {
+                        Model = selectedModel
+                    });
+                }
+                else
+                {
+                    var qualifiedModel = appData.GetQualifiedModel(selectedModel);
+                    if (qualifiedModel == null)
+                        continue;
+                    request.Models.Add(new()
+                    {
+                        Model = qualifiedModel
+                    });
+                }
             }
         }
         if (request.ToObjectDictionary().HasNonDefaultValues(ignoreKeys:ignore) || Request!.QueryString[Keywords.Reset] != null)

AiServer.ServiceInterface/OpenAiProvider.cs

@@ -11,6 +11,14 @@ public class OpenAiProvider(ILogger<OpenAiProvider> log) : OpenAiProviderBase(lo
 {
 }
 
+public class CustomAiProvider(ILogger<OpenAiProvider> log) : OpenAiProviderBase(log)
+{
+    public override Task<OpenAiChatResult> ChatAsync(AiProvider provider, OpenAiChat request, CancellationToken token = default)
+    {
+        return base.ChatAsync(provider, request, token);
+    }
+}
+
 public class OllamaAiProvider(ILogger<OllamaAiProvider> log) : OpenAiProviderBase(log), IOllamaAiProvider
 {
     protected virtual async Task<OllamaGenerateResponse> SendOllamaGenerateRequestAsync(AiProvider provider, OllamaGenerate request, 
@@ -141,6 +149,8 @@ public string GetApiEndpointUrlFor(AiProvider aiProvider, TaskType taskType)
     {
         var apiBaseUrl = aiProvider.ApiBaseUrl ?? aiProvider.AiType?.ApiBaseUrl
             ?? throw new NotSupportedException($"[{aiProvider.Name}] No ApiBaseUrl found in AiProvider or AiType");
+        if (aiProvider.AiTypeId == "Custom")
+            return apiBaseUrl;
         if (taskType == TaskType.OllamaGenerate)
             return apiBaseUrl.CombineWith("/api/generate");
         if (taskType == TaskType.OpenAiChat)

AiServer.ServiceModel/AiProvider.cs

@@ -107,6 +107,7 @@ public enum AiProviderType
     OpenAiProvider,
     GoogleAiProvider,
     AnthropicAiProvider,
+    CustomOpenAiProvider,
 }
 
 /// <summary>

AiServer.ServiceModel/ApiAdmin.cs

@@ -19,6 +19,11 @@ public class ActiveAiModels : IGet, IReturn<StringsResponse>
     public bool? Vision { get; set; }
 }
 
+[Tag(Tags.AiInfo)]
+[Api("Active Custom AI Worker Models available in AI Server")]
+public class ActiveCustomAiModels : IGet, IReturn<StringsResponse>
+{
+}
 
 [Tag(Tags.AiInfo)]
 [ValidateApiKey]

AiServer/Configure.AppHost.cs

@@ -54,6 +54,7 @@ public void Configure(IWebHostBuilder builder) => builder
             services.AddSingleton<OllamaAiProvider>();
             services.AddSingleton<GoogleAiProvider>();
             services.AddSingleton<AnthropicAiProvider>();
+            services.AddSingleton<CustomAiProvider>();
             services.AddSingleton<AiProviderFactory>();
 
             services.AddSingleton(new ComfyMediaProviderOptions

AiServer/wwwroot/img/providers/custom.svg

@@ -222,5 +222,12 @@
       "claude-3-5-sonnet": "claude-3-5-sonnet-latest",
       "claude-3-5-haiku": "claude-3-5-haiku-latest"
     }
+  },
+  {
+    "id": "Custom",
+    "provider": "CustomOpenAiProvider",
+    "website": "https://platform.openai.com/docs/api-reference/chat",
+    "icon": "/img/providers/custom.svg",
+    "apiBaseUrl": "http://localhost:8080/v1/chat/completions"
   }
 ]

AiServer/wwwroot/lib/data/ai-types.json

@@ -148,7 +148,11 @@ export default {
             </template>
             <template #formfooter="{ form, formInstance, apis, type, model, id, updateModel }">
                 <div class="pl-6">
-                    <SelectModels v-if="model" :aiTypeId="model?.aiTypeId" :edit="model" @update:modelValue="updateModel(model)" />
+                    <div v-if="providerType=='Custom'" class="pr-6">
+                        <TagInput v-if="model" v-model="model.selectedModels" label="Custom Models" @update:modelValue="updateModel(model)"  />                         
+                        <TagInput v-else v-model="formModel.selectedModels" label="Custom Models" @update:modelValue="updateModel(model)"  />                         
+                    </div>
+                    <SelectModels v-else-if="model" :aiTypeId="model?.aiTypeId" :edit="model" @update:modelValue="updateModel(model)" />
                     <SelectModels v-else-if="formModel" :aiTypeId="providerType" :edit="formInstance.model" @update:modelValue="formInstance.setModel" />
                 </div>
             </template>