scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

Merge 92ef50bc35 into 8b9cd12a89

Browse source
This commit is contained in:
Mrunmayee Daware 2026-10-20 01:15:42 +00:00 • committed by GitHub
commit 906b99cc76
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 103 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/llms/ollama/completion/handler.py
View file

@ -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(

3
litellm/main.py
View file

@ -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(

6
tests/unit/llms/ollama/test_ollama_chat_transformation.py
View file

@ -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}")

80
tests/unit/llms/ollama/test_ollama_embedding.py
View file

@ -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