feat: Added OllamaProvider.CanPullModelsAutomatically property.

tryAGI · Jan 14, 2025 · e8bb29a · e8bb29a
1 parent 50644fd
commit e8bb29a
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 29 deletions.
diff --git a/src/Ollama/src/OllamaChatModel.cs b/src/Ollama/src/OllamaChatModel.cs
@@ -26,20 +26,7 @@ public override async IAsyncEnumerable<ChatResponse> GenerateAsync(
     {
         request = request ?? throw new ArgumentNullException(nameof(request));
 
-        try
-        {
-            var runningModels = await Provider.Api.Models.ListRunningModelsAsync(cancellationToken).ConfigureAwait(false);
-            if (runningModels.Models != null &&
-                runningModels.Models.All(x => x.Model?.Contains(Id) != true))
-            {
-                await Provider.Api.Models.PullModelAsync(Id, cancellationToken: cancellationToken)
-                    .EnsureSuccessAsync().ConfigureAwait(false);
-            }
-        }
-        catch (HttpRequestException)
-        {
-            // Ignore
-        }
+        await Provider.PullModelIfRequiredAndAllowedAsync(Id, cancellationToken).ConfigureAwait(false);
 
         var usedSettings = OllamaChatSettings.Calculate(
             requestSettings: settings,

diff --git a/src/Ollama/src/OllamaEmbeddingModel.cs b/src/Ollama/src/OllamaEmbeddingModel.cs
@@ -25,21 +25,7 @@ public async Task<EmbeddingResponse> CreateEmbeddingsAsync(
     {
         request = request ?? throw new ArgumentNullException(nameof(request));
 
-        try
-        {
-            var runningModels =
-                await Provider.Api.Models.ListRunningModelsAsync(cancellationToken).ConfigureAwait(false);
-            if (runningModels.Models != null &&
-                runningModels.Models.All(x => x.Model?.Contains(Id) != true))
-            {
-                await Provider.Api.Models.PullModelAsync(Id, cancellationToken: cancellationToken)
-                    .EnsureSuccessAsync().ConfigureAwait(false);
-            }
-        }
-        catch (HttpRequestException)
-        {
-            // Ignore
-        }
+        await Provider.PullModelIfRequiredAndAllowedAsync(Id, cancellationToken).ConfigureAwait(false);
 
         var results = new List<IList<double>>(capacity: request.Strings.Count);
         foreach (var prompt in request.Strings)

diff --git a/src/Ollama/src/OllamaProvider.cs b/src/Ollama/src/OllamaProvider.cs
@@ -17,9 +17,40 @@ public sealed class OllamaProvider(
     {
         Timeout = TimeSpan.FromHours(1),
     }, baseUri: new Uri(url));
+
+    /// <summary>
+    /// OllamaChatModel and OllamaEmbeddingModel will pull models automatically if this is true.
+    /// </summary>
+    public bool CanPullModelsAutomatically { get; set; } = true;
 
     public void Dispose()
     {
         Api.Dispose();
     }
+
+    public async Task PullModelIfRequiredAndAllowedAsync(
+        string id,
+        CancellationToken cancellationToken = default)
+    {
+        if (!CanPullModelsAutomatically)
+        {
+            return;
+        }
+
+        try
+        {
+            // Pull the model if it is not running
+            var runningModels = await Api.Models.ListRunningModelsAsync(cancellationToken).ConfigureAwait(false);
+            if (runningModels.Models != null &&
+                runningModels.Models.All(x => x.Model?.Contains(id) != true))
+            {
+                await Api.Models.PullModelAsync(id, cancellationToken: cancellationToken)
+                    .EnsureSuccessAsync().ConfigureAwait(false);
+            }
+        }
+        catch (HttpRequestException)
+        {
+            // Ignore
+        }
+    }
 }