diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 816028499e8..6c324ccaa60 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -88,6 +88,15 @@ def _process_ollama_embedding_response( return model_response +def _ollama_embedding_request_headers( + api_key: str | None, extra_headers: Mapping[str, str] | None +) -> dict[str, str] | None: # mutable-ok: headers dict is built fresh per request and handed to httpx + headers: Final = dict(extra_headers) if extra_headers else {} + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + return headers or None + + async def ollama_aembeddings( api_base: str, model: str, @@ -96,10 +105,16 @@ async def ollama_aembeddings( optional_params: dict, logging_obj: object, encoding: TokenEncoder | None, + api_key: str | None = None, + extra_headers: Mapping[str, str] | None = None, + timeout: float | None = None, ): data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = await litellm.module_level_aclient.post(url=_ollama_embed_url(api_base), json=data) + headers: Final = _ollama_embedding_request_headers(api_key=api_key, extra_headers=extra_headers) + response: Final = await litellm.module_level_aclient.post( + url=_ollama_embed_url(api_base), json=data, headers=headers, timeout=timeout + ) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( @@ -120,10 +135,16 @@ def ollama_embeddings( model_response: EmbeddingResponse, logging_obj: object, encoding: TokenEncoder | None = None, + api_key: str | None = None, + extra_headers: Mapping[str, str] | None = None, + timeout: float | None = None, ): data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = litellm.module_level_client.post(url=_ollama_embed_url(api_base), json=data) + headers: Final = _ollama_embedding_request_headers(api_key=api_key, extra_headers=extra_headers) + response: Final = litellm.module_level_client.post( + url=_ollama_embed_url(api_base), json=data, headers=headers, timeout=timeout + ) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( diff --git a/litellm/main.py b/litellm/main.py index 9430af34401..9a572017aee 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7222,6 +7222,9 @@ def embedding( logging_obj=logging, optional_params=optional_params, model_response=EmbeddingResponse(), + api_key=api_key, + extra_headers=cast("dict[str, str] | None", headers), # cast-ok: untyped kwargs headers + timeout=timeout, ) elif custom_llm_provider == "sagemaker": response = sagemaker_llm.embedding( diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 1a0bbe171a9..6664b10ecb6 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1064,6 +1064,9 @@ def test_ollama_embeddings(mock_embeddings): logging_obj=mock.ANY, model_response=mock.ANY, encoding=mock.ANY, + api_key=None, + extra_headers={}, + timeout=600, ) except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -1088,6 +1091,9 @@ def test_ollama_aembeddings(mock_aembeddings): logging_obj=mock.ANY, model_response=mock.ANY, encoding=mock.ANY, + api_key=None, + extra_headers={}, + timeout=600, ) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index 6a1dfebc236..e0a89c28a26 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -35,7 +35,6 @@ def test_ollama_embeddings(mock_response_data, mock_embedding_response, mock_enc patch("litellm.module_level_client.post") as mock_post, patch("litellm.OllamaConfig.get_config", return_value={"truncate": 512}), ): - mock_response = MagicMock() mock_response.json.return_value = mock_response_data mock_post.return_value = mock_response @@ -57,19 +56,14 @@ def test_ollama_embeddings(mock_response_data, mock_embedding_response, mock_enc @pytest.mark.asyncio -async def test_ollama_aembeddings( - mock_response_data, mock_embedding_response, mock_encoding -): +async def test_ollama_aembeddings(mock_response_data, mock_embedding_response, mock_encoding): mock_response = AsyncMock() # Make json() a regular synchronous method, not async mock_response.json = MagicMock(return_value=mock_response_data) with ( - patch( - "litellm.module_level_aclient.post", return_value=mock_response - ) as mock_post, + patch("litellm.module_level_aclient.post", return_value=mock_response) as mock_post, patch("litellm.OllamaConfig.get_config", return_value={"truncate": 512}), ): - response = await ollama_aembeddings( api_base="http://localhost:11434", model="test-model", @@ -96,7 +90,6 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin patch("litellm.module_level_client.post") as mock_post, patch("litellm.OllamaConfig.get_config", return_value={}), ): - mock_response = MagicMock() mock_response.json.return_value = response_data mock_post.return_value = mock_response @@ -136,3 +129,72 @@ def test_ollama_embeddings_ignore_a_trailing_slash_on_api_base(api_base: str) -> assert route.called assert [item["embedding"] for item in response.data] == [[0.1, 0.2], [0.3, 0.4]] + + +def test_ollama_embeddings_sends_api_key_and_timeout(mock_response_data, mock_embedding_response, mock_encoding): + with patch("litellm.module_level_client.post") as mock_post: + mock_response = MagicMock() + mock_response.json.return_value = mock_response_data + mock_post.return_value = mock_response + + ollama_embeddings( + api_base="http://localhost:11434", + model="test-model", + prompts=["hello"], + optional_params={}, + model_response=mock_embedding_response, + logging_obj=None, + encoding=mock_encoding, + api_key="secret", + timeout=5, + ) + + _, kwargs = mock_post.call_args + assert kwargs["headers"] == {"Authorization": "Bearer secret"} + assert kwargs["timeout"] == 5 + + +def test_ollama_embeddings_extra_headers_and_no_clobber(mock_response_data, mock_embedding_response, mock_encoding): + with patch("litellm.module_level_client.post") as mock_post: + mock_response = MagicMock() + mock_response.json.return_value = mock_response_data + mock_post.return_value = mock_response + + ollama_embeddings( + api_base="http://localhost:11434", + model="test-model", + prompts=["hello"], + optional_params={}, + model_response=mock_embedding_response, + logging_obj=None, + encoding=mock_encoding, + api_key="secret", + extra_headers={"Authorization": "Token custom", "X-Tenant": "acme"}, + ) + + _, kwargs = mock_post.call_args + assert kwargs["headers"] == {"Authorization": "Token custom", "X-Tenant": "acme"} + + +@pytest.mark.asyncio +async def test_ollama_aembeddings_sends_api_key_and_timeout(mock_response_data, mock_embedding_response, mock_encoding): + with patch("litellm.module_level_aclient.post", new=AsyncMock()) as mock_post: + mock_response = MagicMock() + mock_response.json.return_value = mock_response_data + mock_post.return_value = mock_response + + await ollama_aembeddings( + api_base="http://localhost:11434", + model="test-model", + prompts=["hello"], + model_response=mock_embedding_response, + optional_params={}, + logging_obj=None, + encoding=mock_encoding, + api_key="secret", + timeout=5, + ) + + _, kwargs = mock_post.call_args + assert kwargs["headers"] == {"Authorization": "Bearer secret"} + assert kwargs["timeout"] == 5