Add support for Ollama Generate API for using vision models

Signed-off-by: Demis Bellot <demis.bellot@gmail.com>

Author: mythz

AiServer.ServiceInterface/AiProviderFactory.cs

@@ -11,14 +11,22 @@ public interface IOpenAiProvider
     Task<OpenAiChatResult> ChatAsync(AiProvider provider, OpenAiChat request, CancellationToken token = default);
 }
 
-public class AiProviderFactory(OpenAiProvider openAiProvider, GoogleAiProvider googleProvider, AnthropicAiProvider anthropicAiProvider)
+public record OllamaGenerationResult(OllamaGenerateResponse Response, int DurationMs);
+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 IOpenAiProvider GetOpenAiProvider(AiProviderType aiProviderType=AiProviderType.OpenAiProvider)
     {
-        return aiProviderType == AiProviderType.GoogleAiProvider
-            ? googleProvider
-            : aiProviderType == AiProviderType.AnthropicAiProvider
-                ? anthropicAiProvider 
-                : openAiProvider;
+        return aiProviderType switch
+        {
+            AiProviderType.OllamaAiProvider => ollamaAiProvider,
+            AiProviderType.GoogleAiProvider => googleProvider,
+            AiProviderType.AnthropicAiProvider => anthropicAiProvider,
+            _ => openAiProvider
+        };
     }
 }

AiServer.ServiceInterface/AppDb/CompleteOllamaGenerateCommand.cs

@@ -0,0 +1,39 @@
+using System.Data;
+using AiServer.ServiceModel;
+using ServiceStack;
+using ServiceStack.Jobs;
+using ServiceStack.OrmLite;
+
+namespace AiServer.ServiceInterface.AppDb;
+
+public record class CompleteOllamaGenerate(QueueOllamaGeneration Request, OllamaGenerateResponse Response, BackgroundJob Job);
+
+[Worker(Workers.AppDb)]
+public class CompleteOllamaGenerateCommand(IDbConnection db) : SyncCommand<CompleteOllamaGenerate>
+{
+    protected override void Run(CompleteOllamaGenerate ctx)
+    {
+        var summary = new ChatSummary
+        {
+            Id = ctx.Job.Id,
+            RefId = ctx.Job.RefId!,
+            CreatedDate = ctx.Job.CreatedDate,
+            DurationMs = ctx.Job.DurationMs,
+            Tag = ctx.Job.Tag,
+            Model = ctx.Request.Request.Model,
+            Provider = ctx.Job.Worker!,
+            PromptTokens = ctx.Response?.PromptTokens ?? 0,
+            CompletionTokens = ctx.Response?.EvalCount ?? 0,
+        };
+        try
+        {
+            db.Insert(summary);
+        }
+        catch (Exception e)
+        {
+            // completing failed jobs could fail with unique constraint
+            db.DeleteById<ChatSummary>(summary.Id);
+            db.Insert(summary);
+        }
+    }
+}

AiServer.ServiceInterface/Jobs/CreateOllamaGenerationCommand.cs

@@ -0,0 +1,60 @@
+using AiServer.ServiceInterface.AppDb;
+using AiServer.ServiceModel;
+using Microsoft.Extensions.Logging;
+using ServiceStack;
+using ServiceStack.Jobs;
+
+namespace AiServer.ServiceInterface.Jobs;
+
+public class CreateOllamaGenerationCommand(ILogger<CreateOllamaGenerationCommand> logger, IBackgroundJobs jobs, AppData appData, AiProviderFactory aiFactory, IHttpClientFactory clientFactory) 
+    : AsyncCommandWithResult<QueueOllamaGeneration, OllamaGenerateResponse>
+{
+    protected override async Task<OllamaGenerateResponse> RunAsync(QueueOllamaGeneration request, CancellationToken token)
+    {
+        var job = Request.GetBackgroundJob();
+        var log = Request.CreateJobLogger(jobs,logger);
+        var apiProvider = appData.AssertAiProvider(job.Worker!);
+        var chatProvider = aiFactory.GetOpenAiProvider(apiProvider.AiType.Provider);
+        if (chatProvider is not IOllamaAiProvider generateProvider)
+            throw new NotSupportedException($"{chatProvider.GetType()} is not an IOllamaAiProvider");
+
+        try
+        {
+            var origModel = request.Request.Model;
+            request.Request.Model = appData.GetQualifiedModel(origModel) ?? origModel;
+            log.LogInformation("GENERATE Ollama #{JobId} request for {OriginalModel}, using {Model}", job.Id, origModel, request.Request.Model);
+            var (response, durationMs) = await generateProvider.GenerateAsync(apiProvider, request.Request, token);
+            request.Request.Model = origModel;
+
+            job.DurationMs = durationMs;
+            jobs.RunCommand<CompleteOllamaGenerateCommand>(
+                new CompleteOllamaGenerate(Request: request, Response: response, Job: job));
+
+            log.LogInformation("GENERATE Ollama #{JobId} request finished in {Ms} ms{ReplyMessage}", 
+                job.Id, job.DurationMs, job.ReplyTo == null ? "" : $", sending response to {job.ReplyTo}");
+            if (job.ReplyTo != null)
+            {
+                await clientFactory.SendJsonCallbackAsync(Request.GetBackgroundJob().ReplyTo!, request, token:token);
+                // jobs.EnqueueCommand<NotifyOpenAiChatResponseCommand>(response, new() {
+                //     ParentId = job.Id,
+                //     ReplyTo = job.ReplyTo,
+                // });
+            }
+            return response;
+        }
+        catch (Exception e)
+        {
+            var offline = !await chatProvider.IsOnlineAsync(apiProvider, token);
+            log.LogError("CHAT OpenAi #{JobId} request failed after {Ms} with: {Message} (offline:{Offline})", 
+                job.Id, job.DurationMs, e.Message, offline);
+            if (offline)
+            {
+                jobs.RunCommand<ChangeProviderStatusCommand>(new ChangeProviderStatus {
+                    Name = apiProvider.Name,
+                    OfflineDate = DateTime.UtcNow,
+                });
+            }
+            throw;
+        }
+    }
+}

AiServer.ServiceInterface/Jobs/CreateOpenAiChatCommand.cs

@@ -39,7 +39,7 @@ protected override async Task<OpenAiChatResponse> RunAsync(QueueOpenAiChatComple
             }
             return response;
         }
-        catch(Exception e)
+        catch (Exception e)
         {
             var offline = !await chatProvider.IsOnlineAsync(apiProvider, token);
             log.LogError("CHAT OpenAi #{JobId} request failed after {Ms} with: {Message} (offline:{Offline})", 

AiServer.ServiceInterface/OpenAiChatServices.cs

@@ -108,6 +108,19 @@ public object Any(GetModelImage request)
         return GetModelImage(request.Model);
     }
 
+    public async Task<object> Post(OllamaGeneration request)
+    {
+        var generateRequest = new QueueOllamaGeneration
+        {
+            Request = request,
+            RefId = request.RefId,
+            Tag = request.Tag,
+            Provider = request.Provider
+        };
+        
+        return await generateRequest.ProcessSync(jobs, this);
+    }
+
     public async Task<object> Post(OpenAiChatCompletion request)
     {
         var chatRequest = new QueueOpenAiChatCompletion
@@ -121,6 +134,70 @@ public async Task<object> Post(OpenAiChatCompletion request)
         return await chatRequest.ProcessSync(jobs, this);
     }
 
+    public QueueOllamaGenerationResponse Any(QueueOllamaGeneration request)
+    {
+        if (request.Request == null)
+            throw new ArgumentNullException(nameof(request.Request));
+        
+        if (request.Request.Prompt.IsNullOrEmpty())
+            throw new ArgumentNullException(nameof(request.Request.Prompt));
+    
+        var qualifiedModel = appData.GetQualifiedModel(request.Request.Model);
+        if (qualifiedModel == null)
+            throw HttpError.NotFound($"Model {request.Request.Model} not found");
+
+        var queueCounts = jobs.GetWorkerQueueCounts();
+        var providerQueueCount = int.MaxValue;
+        AiProvider? useProvider = null;
+        var candidates = appData.AiProviders
+            .Where(x => x is { Enabled: true, AiType.Provider: AiProviderType.OllamaAiProvider }
+                && x.Models.Any(m => m.Model == qualifiedModel)).ToList();
+        foreach (var candidate in candidates)
+        {
+            if (candidate.OfflineDate != null)
+                continue;
+            var pendingJobs = queueCounts.GetValueOrDefault(candidate.Name, 0);
+            if (useProvider == null)
+            {
+                useProvider = candidate;
+                providerQueueCount = pendingJobs;
+                continue;
+            }
+            if (pendingJobs < providerQueueCount || (pendingJobs == providerQueueCount && candidate.Priority > useProvider.Priority))
+            {
+                useProvider = candidate;
+                providerQueueCount = pendingJobs;
+            }
+        }
+
+        useProvider ??= candidates.FirstOrDefault(x => x.Name == qualifiedModel); // Allow selecting offline models
+        if (useProvider == null)
+            throw new NotSupportedException("No active AI Providers support this model");
+
+        var jobRef = jobs.EnqueueCommand<CreateOllamaGenerationCommand>(request, new()
+        {
+            RefId = request.RefId,
+            ReplyTo = request.ReplyTo,
+            Tag = request.Tag,
+            Args = request.Provider == null ? null : new() {
+                [nameof(request.Provider)] = request.Provider
+            },
+            Worker = useProvider.Name,
+        });
+        
+        var jobStatusUrl = AppConfig.Instance.ApplicationBaseUrl
+            .CombineWith($"/api/{nameof(GetOllamaGenerationStatus)}?RefId=" + jobRef.RefId);
+    
+        var response = new QueueOllamaGenerationResponse
+        {
+            Id = jobRef.Id,
+            RefId = jobRef.RefId,
+            StatusUrl = jobStatusUrl
+        };
+
+        return response;
+    }
+    
     public QueueOpenAiChatResponse Any(QueueOpenAiChatCompletion request)
     {
         if (request.Request == null)
@@ -292,6 +369,40 @@ private async Task<JobResult> WaitForJobCompletion(long jobId)
         return HttpError.NotFound("Job not found");
     }
 
+    public async Task<object> Get(GetOllamaGenerationStatus request)
+    {
+        var summary = GetJobSummary((int)request.JobId, request.RefId);
+        if (summary == null)
+            return HttpError.NotFound("JobSummary not found");
+
+        var response = GetOpenAiChat(summary);
+        if (response == null)
+            return HttpError.NotFound("Job not found");
+
+        var job = response.Result;
+
+        var generateResponse = response.Result?.ResponseBody.FromJson<OllamaGenerateResponse>();
+        if (generateResponse == null)
+        {
+            return new GetOllamaGenerationStatusResponse
+            {
+                JobId = request.JobId,
+                RefId = request.RefId,
+                JobState = job.State,
+                Status = job.State.ToString(),
+            };
+        }
+        
+        return new GetOllamaGenerationStatusResponse
+        {
+            JobId = request.JobId,
+            RefId = request.RefId,
+            JobState = job.State,
+            Status = job.State.ToString(),
+            Result = generateResponse,
+        };
+    }
+
     public async Task<object> Get(GetOpenAiChatStatus request)
     {
         var summary = GetJobSummary((int)request.JobId, request.RefId);
@@ -428,6 +539,81 @@ public object Any(DeleteAiProvider request)
 
 public static class OpenAiChatServiceExtensions
 {
+    public static async Task<OllamaGenerateResponse> ProcessSync(this QueueOllamaGeneration generateRequest,
+        IBackgroundJobs jobs, OpenAiChatServices chatService)
+    {
+        QueueOllamaGenerationResponse? generateResponse = null;
+        try
+        {
+            var response = chatService.Any(generateRequest);
+            generateResponse = response;
+        }
+        catch (Exception e)
+        {
+            Console.WriteLine(e);
+            throw;
+        }
+        
+        if (generateResponse == null)
+            throw new Exception("Failed to start chat request");
+        
+        var job = jobs.GetJob(generateResponse.Id);
+        // For all requests, wait for the job to be created
+        while (job == null)
+        {
+            await Task.Delay(1000);
+            job = jobs.GetJob(generateResponse.Id);
+        }
+        
+        // We know at this point, we definitely have a job
+        JobResult queuedJob = job;
+        
+        var completedResponse = new OllamaGenerateResponse();
+
+        // Handle failed jobs
+        if (job.Failed != null)
+        {
+            throw new Exception(job.Failed.Error!.Message);
+        }
+        
+        // Wait for the job to complete max 2 minutes
+        var timeout = DateTime.UtcNow.AddMinutes(2);
+        while (queuedJob?.Job?.State is not (BackgroundJobState.Completed or BackgroundJobState.Cancelled
+               or BackgroundJobState.Failed) && DateTime.UtcNow < timeout)
+        {
+            await Task.Delay(1000);
+            queuedJob = jobs.GetJob(generateResponse.Id);
+        }
+        
+        // Check if the job is still not completed
+        if (queuedJob?.Job?.State is not (BackgroundJobState.Completed or BackgroundJobState.Cancelled
+               or BackgroundJobState.Failed))
+        {
+            throw new Exception("Job did not complete within the specified timeout.");
+        }
+        
+        // Process successful job results
+        var jobResponseBody = queuedJob.Completed?.ResponseBody;
+        var jobRes = jobResponseBody.FromJson<OllamaGenerateResponse>();
+        if (jobRes != null)
+        {
+            completedResponse.Model = jobRes.Model;
+            completedResponse.CreatedAt = jobRes.CreatedAt;
+            completedResponse.Response = jobRes.Response;
+            completedResponse.Done = jobRes.Done;
+            completedResponse.Context = jobRes.Context;
+            completedResponse.DoneReason = jobRes.DoneReason;
+            completedResponse.TotalDuration = jobRes.TotalDuration;
+            completedResponse.LoadDuration = jobRes.LoadDuration;
+            completedResponse.PromptEvalCount = jobRes.PromptEvalCount;
+            completedResponse.EvalCount = jobRes.EvalCount;
+            completedResponse.PromptTokens = jobRes.PromptTokens;
+            completedResponse.ResponseStatus = jobRes.ResponseStatus;
+        }
+
+        return completedResponse;
+    }
+
     public static async Task<OpenAiChatResponse> ProcessSync(this QueueOpenAiChatCompletion chatRequest,
         IBackgroundJobs jobs, OpenAiChatServices chatService)
     {
@@ -443,7 +629,7 @@ public static async Task<OpenAiChatResponse> ProcessSync(this QueueOpenAiChatCom
             throw;
         }
 
-        if(chatResponse == null)
+        if (chatResponse == null)
             throw new Exception("Failed to start chat request");
 
         var job = jobs.GetJob(chatResponse.Id);