feat: Support for Ollama keep_alive API parameter (#237)

* feat: Support for Ollama keep_alive API parameter * test fix
brainlid · Jan 22, 2025 · 0c10c33 · 0c10c33
1 parent 45708db
commit 0c10c33
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/lib/chat_models/chat_ollama_ai.ex b/lib/chat_models/chat_ollama_ai.ex
@@ -49,6 +49,7 @@ defmodule LangChain.ChatModels.ChatOllamaAI do
 
   @create_fields [
     :endpoint,
+    :keep_alive,
     :mirostat,
     :mirostat_eta,
     :mirostat_tau,
@@ -78,6 +79,10 @@ defmodule LangChain.ChatModels.ChatOllamaAI do
   embedded_schema do
     field :endpoint, :string, default: "http://localhost:11434/api/chat"
 
+    # Change Keep Alive setting for unloading the model from memory.
+    # (Default: "5m", set to a negative interval to disable)
+    field :keep_alive, :string, default: "5m"
+
     # Enable Mirostat sampling for controlling perplexity.
     # (default: 0, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0)
     field :mirostat, :integer, default: 0
@@ -199,6 +204,7 @@ defmodule LangChain.ChatModels.ChatOllamaAI do
       num_predict: model.num_predict,
       repeat_last_n: model.repeat_last_n,
       repeat_penalty: model.repeat_penalty,
+      keep_alive: model.keep_alive,
       mirostat: model.mirostat,
       mirostat_eta: model.mirostat_eta,
       mirostat_tau: model.mirostat_tau,
@@ -411,6 +417,7 @@ defmodule LangChain.ChatModels.ChatOllamaAI do
       model,
       [
         :endpoint,
+        :keep_alive,
         :model,
         :mirostat,
         :mirostat_eta,

diff --git a/test/chat_models/chat_ollama_ai_test.exs b/test/chat_models/chat_ollama_ai_test.exs
@@ -257,6 +257,7 @@ defmodule ChatModels.ChatOllamaAITest do
 
       assert result == %{
                "endpoint" => "http://localhost:11434/api/chat",
+               "keep_alive" => "5m",
                "mirostat" => 0,
                "mirostat_eta" => 0.1,
                "mirostat_tau" => 5.0,