From cfb76b9c69999937fd7c53e53a1c075ea9c130cc Mon Sep 17 00:00:00 2001
From: Alen Joses R <50145957+ALENJOSESR@users.noreply.github.com>
Date: Wed, 5 Mar 2025 22:00:36 +0530
Subject: [PATCH] feat: add client_kwargs Parameter to OllamaEmbedding Class
(#18012)
---
.../llama_index/embeddings/ollama/base.py | 6 ++++--
.../embeddings/llama-index-embeddings-ollama/pyproject.toml | 2 +-
.../tests/test_embeddings_ollama.py | 4 +++-
3 files changed, 8 insertions(+), 4 deletions(-)
diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-ollama/llama_index/embeddings/ollama/base.py b/llama-index-integrations/embeddings/llama-index-embeddings-ollama/llama_index/embeddings/ollama/base.py
index c4112f2c7..83f936449 100644
--- a/llama-index-integrations/embeddings/llama-index-embeddings-ollama/llama_index/embeddings/ollama/base.py
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-ollama/llama_index/embeddings/ollama/base.py
@@ -34,6 +34,7 @@ class OllamaEmbedding(BaseEmbedding):
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 @@ class OllamaEmbedding(BaseEmbedding):
**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
index f0ea73fb7..7cad77517 100644
--- 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/llama-index-integrations/embeddings/llama-index-embeddings-ollama/tests/test_embeddings_ollama.py b/llama-index-integrations/embeddings/llama-index-embeddings-ollama/tests/test_embeddings_ollama.py
index 74ec40ece..8c9ac3953 100644
--- a/llama-index-integrations/embeddings/llama-index-embeddings-ollama/tests/test_embeddings_ollama.py
+++ b/llama-index-integrations/embeddings/llama-index-embeddings-ollama/tests/test_embeddings_ollama.py
@@ -3,5 +3,7 @@ from llama_index.embeddings.ollama import OllamaEmbedding
def test_embedding_class():
- emb = OllamaEmbedding(model_name="")
+ emb = OllamaEmbedding(
+ model_name="", client_kwargs={"headers": {"Authorization": "Bearer token"}}
+ )
assert isinstance(emb, BaseEmbedding)
--
GitLab