feat: add client_kwargs Parameter to OllamaEmbedding Class (#18012)

run-llama · Mar 5, 2025 · cfb76b9 · cfb76b9
1 parent e486157
commit cfb76b9
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/...tegrations/embeddings/llama-index-embeddings-ollama/llama_index/embeddings/ollama/base.py b/...tegrations/embeddings/llama-index-embeddings-ollama/llama_index/embeddings/ollama/base.py
@@ -34,6 +34,7 @@ def __init__(
         embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
         ollama_additional_kwargs: Optional[Dict[str, Any]] = None,
         callback_manager: Optional[CallbackManager] = None,
+        client_kwargs: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -45,8 +46,9 @@ def __init__(
             **kwargs,
         )
 
-        self._client = Client(host=self.base_url)
-        self._async_client = AsyncClient(host=self.base_url)
+        client_kwargs = client_kwargs or {}
+        self._client = Client(host=self.base_url, **client_kwargs)
+        self._async_client = AsyncClient(host=self.base_url, **client_kwargs)
 
     @classmethod
     def class_name(cls) -> str:

diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-ollama/pyproject.toml b/llama-index-integrations/embeddings/llama-index-embeddings-ollama/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-embeddings-ollama"
 readme = "README.md"
-version = "0.5.0"
+version = "0.6.0"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"

diff --git a/...dex-integrations/embeddings/llama-index-embeddings-ollama/tests/test_embeddings_ollama.py b/...dex-integrations/embeddings/llama-index-embeddings-ollama/tests/test_embeddings_ollama.py
@@ -3,5 +3,7 @@
 
 
 def test_embedding_class():
-    emb = OllamaEmbedding(model_name="")
+    emb = OllamaEmbedding(
+        model_name="", client_kwargs={"headers": {"Authorization": "Bearer token"}}
+    )
     assert isinstance(emb, BaseEmbedding)