From 0c10c33f2cbeed7742b7ff1481a66908b57aa518 Mon Sep 17 00:00:00 2001
From: Matthew Pope <mpope9@users.noreply.github.com>
Date: Wed, 22 Jan 2025 11:04:55 -0800
Subject: [PATCH] feat: Support for Ollama keep_alive API parameter (#237)

* feat: Support for Ollama keep_alive API parameter

* test fix
---
 lib/chat_models/chat_ollama_ai.ex        | 7 +++++++
 test/chat_models/chat_ollama_ai_test.exs | 1 +
 2 files changed, 8 insertions(+)

diff --git a/lib/chat_models/chat_ollama_ai.ex b/lib/chat_models/chat_ollama_ai.ex
index 8b948eca..63725eff 100644
--- 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
index b1fbd3f1..68c702d4 100644
--- 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,