From 6078738d34645388fa2dca0e0215fa60639e4c32 Mon Sep 17 00:00:00 2001
From: Felipe Amaral <97637336+felipeavilis@users.noreply.github.com>
Date: Mon, 4 Mar 2024 23:09:45 -0300
Subject: [PATCH] Feature: Custom Ollama endpoint base_url (#1301)

---
 configs/ollama.yaml           | 1 +
 embedchain/config/llm/base.py | 2 ++
 embedchain/llm/ollama.py      | 1 +
 embedchain/utils/misc.py      | 1 +
 4 files changed, 5 insertions(+)

diff --git a/configs/ollama.yaml b/configs/ollama.yaml
index 5b01cc18..da4be9c2 100644
--- a/configs/ollama.yaml
+++ b/configs/ollama.yaml
@@ -5,6 +5,7 @@ llm:
     temperature: 0.5
     top_p: 1
     stream: true
+    base_url: http://localhost:11434
 
 embedder:
   provider: huggingface
diff --git a/embedchain/config/llm/base.py b/embedchain/config/llm/base.py
index fb6491d1..c6d3d640 100644
--- a/embedchain/config/llm/base.py
+++ b/embedchain/config/llm/base.py
@@ -96,6 +96,7 @@ class BaseLlmConfig(BaseConfig):
         endpoint: Optional[str] = None,
         model_kwargs: Optional[dict[str, Any]] = None,
         local: Optional[bool] = False,
+        base_url: Optional[str] = None,
     ):
         """
         Initializes a configuration class instance for the LLM.
@@ -169,6 +170,7 @@ class BaseLlmConfig(BaseConfig):
         self.endpoint = endpoint
         self.model_kwargs = model_kwargs
         self.local = local
+        self.base_url = base_url
 
         if isinstance(prompt, str):
             prompt = Template(prompt)
diff --git a/embedchain/llm/ollama.py b/embedchain/llm/ollama.py
index 6d7f802b..221753d3 100644
--- a/embedchain/llm/ollama.py
+++ b/embedchain/llm/ollama.py
@@ -31,6 +31,7 @@ class OllamaLlm(BaseLlm):
             temperature=config.temperature,
             top_p=config.top_p,
             callback_manager=CallbackManager(callback_manager),
+            base_url=config.base_url,
         )
 
         return llm.invoke(prompt)
diff --git a/embedchain/utils/misc.py b/embedchain/utils/misc.py
index 6ad7ad2b..53802d81 100644
--- a/embedchain/utils/misc.py
+++ b/embedchain/utils/misc.py
@@ -427,6 +427,7 @@ def validate_config(config_data):
                     Optional("endpoint"): str,
                     Optional("model_kwargs"): dict,
                     Optional("local"): bool,
+                    Optional("base_url"): str,
                 },
             },
             Optional("vectordb"): {