tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 00:22:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update mistral 'supports_response_schema' field + Fix ollama embedding (#12024)

Browse Source 
* build(model_prices_and_context_window.json): update all mistral models (besides codestral-mamba) to indicate support for response schema

Closes https://github.com/BerriAI/litellm/issues/12012

* fix(route_llm_request.py): if llm router is not initialized, go straight through to litellm sdk

Fixes https://github.com/BerriAI/litellm/issues/12008

* test: add unit test

* fix(ollama_embeddings): fix unecessary await

Fixes https://github.com/BerriAI/litellm/issues/11997

* test: update ollama embedding tests
This commit is contained in:
Krish Dholakia
2025-06-25 07:20:13 -07:00
committed by GitHub
parent 24c2cd1bd9
commit 1a4ad8bf18
7 changed files with 178 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/ollama/completion/handler.py
+17 -10
View File
@@ -5,15 +5,15 @@ Ollama /chat/completion calls handled in llm_http_handler.py
"""
from typing import Any, Dict, List
import litellm
from litellm.types.utils import EmbeddingResponse
def _prepare_ollama_embedding_payload(
model: str,
prompts: List[str],
optional_params: Dict[str, Any]
model: str, prompts: List[str], optional_params: Dict[str, Any]
) -> Dict[str, Any]:
data: Dict[str, Any] = {"model": model, "input": prompts}
special_optional_params = ["truncate", "options", "keep_alive"]
@@ -26,13 +26,14 @@ def _prepare_ollama_embedding_payload(
data["options"].update({k: v})
return data
def _process_ollama_embedding_response(
response_json: dict,
prompts: List[str],
model: str,
model_response: EmbeddingResponse,
logging_obj: Any,
encoding: Any
encoding: Any,
) -> EmbeddingResponse:
output_data = []
embeddings: List[List[float]] = response_json["embeddings"]
@@ -46,11 +47,15 @@ def _process_ollama_embedding_response(
if encoding is not None:
input_tokens = len(encoding.encode("".join(prompts)))
if logging_obj:
logging_obj.debug("Ollama response missing prompt_eval_count; estimated with encoding.")
logging_obj.debug(
"Ollama response missing prompt_eval_count; estimated with encoding."
)
else:
input_tokens = 0
if logging_obj:
logging_obj.warning("Missing prompt_eval_count and no encoding provided; defaulted to 0.")
logging_obj.warning(
"Missing prompt_eval_count and no encoding provided; defaulted to 0."
)
model_response.object = "list"
model_response.data = output_data
@@ -64,6 +69,7 @@ def _process_ollama_embedding_response(
)
return model_response
async def ollama_aembeddings(
api_base: str,
model: str,
@@ -79,7 +85,7 @@ async def ollama_aembeddings(
data = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response = await litellm.module_level_aclient.post(url=api_base, json=data)
response_json = await response.json()
response_json = response.json()
return _process_ollama_embedding_response(
response_json=response_json,
@@ -87,9 +93,10 @@ async def ollama_aembeddings(
model=model,
model_response=model_response,
logging_obj=logging_obj,
encoding=encoding
encoding=encoding,
)
def ollama_embeddings(
api_base: str,
model: str,
@@ -113,5 +120,5 @@ def ollama_embeddings(
model=model,
model_response=model_response,
logging_obj=logging_obj,
encoding=encoding
encoding=encoding,
)
litellm/model_prices_and_context_window_backup.json
+52 -26
View File
@@ -4046,7 +4046,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-small": {
"max_tokens": 8191,
@@ -4058,7 +4059,8 @@
"supports_function_calling": true,
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-small-latest": {
"max_tokens": 8191,
@@ -4070,7 +4072,8 @@
"supports_function_calling": true,
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium": {
"max_tokens": 8191,
@@ -4081,7 +4084,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium-latest": {
"max_tokens": 8191,
@@ -4093,7 +4097,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium-2505": {
"max_tokens": 8191,
@@ -4105,7 +4110,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium-2312": {
"max_tokens": 8191,
@@ -4116,7 +4122,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-latest": {
"max_tokens": 128000,
@@ -4128,7 +4135,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-2411": {
"max_tokens": 128000,
@@ -4140,7 +4148,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-2402": {
"max_tokens": 8191,
@@ -4152,7 +4161,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-2407": {
"max_tokens": 128000,
@@ -4164,7 +4174,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/pixtral-large-latest": {
"max_tokens": 128000,
@@ -4177,7 +4188,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_vision": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/pixtral-large-2411": {
"max_tokens": 128000,
@@ -4190,7 +4202,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_vision": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/pixtral-12b-2409": {
"max_tokens": 128000,
@@ -4203,7 +4216,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_vision": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mistral-7b": {
"max_tokens": 8191,
@@ -4214,7 +4228,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mixtral-8x7b": {
"max_tokens": 8191,
@@ -4226,7 +4241,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mixtral-8x22b": {
"max_tokens": 8191,
@@ -4238,7 +4254,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/codestral-latest": {
"max_tokens": 8191,
@@ -4249,7 +4266,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/codestral-2405": {
"max_tokens": 8191,
@@ -4260,7 +4278,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mistral-nemo": {
"max_tokens": 128000,
@@ -4272,7 +4291,8 @@
"mode": "chat",
"source": "https://mistral.ai/technology/",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mistral-nemo-2407": {
"max_tokens": 128000,
@@ -4284,7 +4304,8 @@
"mode": "chat",
"source": "https://mistral.ai/technology/",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-codestral-mamba": {
"max_tokens": 256000,
@@ -4321,7 +4342,8 @@
"source": "https://mistral.ai/news/devstral",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/magistral-medium-latest": {
"max_tokens": 40000,
@@ -4335,7 +4357,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/magistral-medium-2506": {
"max_tokens": 40000,
@@ -4349,7 +4372,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/magistral-small-latest": {
"max_tokens": 40000,
@@ -4363,7 +4387,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/magistral-small-2506": {
"max_tokens": 40000,
@@ -4377,7 +4402,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/mistral-embed": {
"max_tokens": 8192,
litellm/proxy/route_llm_request.py
+4 -1
View File
@@ -81,7 +81,10 @@ async def route_request(
team_id = get_team_id_from_data(data)
router_model_names = llm_router.model_names if llm_router is not None else []
if "api_key" in data or "api_base" in data:
return getattr(llm_router, f"{route_type}")(**data)
if llm_router is not None:
return getattr(llm_router, f"{route_type}")(**data)
else:
return getattr(litellm, f"{route_type}")(**data)
elif "user_config" in data:
router_config = data.pop("user_config")
model_prices_and_context_window.json
+52 -26
View File
@@ -4046,7 +4046,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-small": {
"max_tokens": 8191,
@@ -4058,7 +4059,8 @@
"supports_function_calling": true,
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-small-latest": {
"max_tokens": 8191,
@@ -4070,7 +4072,8 @@
"supports_function_calling": true,
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium": {
"max_tokens": 8191,
@@ -4081,7 +4084,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium-latest": {
"max_tokens": 8191,
@@ -4093,7 +4097,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium-2505": {
"max_tokens": 8191,
@@ -4105,7 +4110,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-medium-2312": {
"max_tokens": 8191,
@@ -4116,7 +4122,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-latest": {
"max_tokens": 128000,
@@ -4128,7 +4135,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-2411": {
"max_tokens": 128000,
@@ -4140,7 +4148,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-2402": {
"max_tokens": 8191,
@@ -4152,7 +4161,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/mistral-large-2407": {
"max_tokens": 128000,
@@ -4164,7 +4174,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/pixtral-large-latest": {
"max_tokens": 128000,
@@ -4177,7 +4188,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_vision": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/pixtral-large-2411": {
"max_tokens": 128000,
@@ -4190,7 +4202,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_vision": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/pixtral-12b-2409": {
"max_tokens": 128000,
@@ -4203,7 +4216,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_vision": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mistral-7b": {
"max_tokens": 8191,
@@ -4214,7 +4228,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mixtral-8x7b": {
"max_tokens": 8191,
@@ -4226,7 +4241,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mixtral-8x22b": {
"max_tokens": 8191,
@@ -4238,7 +4254,8 @@
"mode": "chat",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/codestral-latest": {
"max_tokens": 8191,
@@ -4249,7 +4266,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/codestral-2405": {
"max_tokens": 8191,
@@ -4260,7 +4278,8 @@
"litellm_provider": "mistral",
"mode": "chat",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mistral-nemo": {
"max_tokens": 128000,
@@ -4272,7 +4291,8 @@
"mode": "chat",
"source": "https://mistral.ai/technology/",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-mistral-nemo-2407": {
"max_tokens": 128000,
@@ -4284,7 +4304,8 @@
"mode": "chat",
"source": "https://mistral.ai/technology/",
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/open-codestral-mamba": {
"max_tokens": 256000,
@@ -4321,7 +4342,8 @@
"source": "https://mistral.ai/news/devstral",
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true
"supports_tool_choice": true,
"supports_response_schema": true
},
"mistral/magistral-medium-latest": {
"max_tokens": 40000,
@@ -4335,7 +4357,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/magistral-medium-2506": {
"max_tokens": 40000,
@@ -4349,7 +4372,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/magistral-small-latest": {
"max_tokens": 40000,
@@ -4363,7 +4387,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/magistral-small-2506": {
"max_tokens": 40000,
@@ -4377,7 +4402,8 @@
"supports_function_calling": true,
"supports_assistant_prefill": true,
"supports_tool_choice": true,
"supports_reasoning": true
"supports_reasoning": true,
"supports_response_schema": true
},
"mistral/mistral-embed": {
"max_tokens": 8192,
test_script.py
+8 -12
View File
@@ -1,24 +1,20 @@
import os
import json
gemini_model_cost_map = json.load(open("model_prices_and_context_window.json"))
mistral_model_cost_map = json.load(open("model_prices_and_context_window.json"))
for model, model_info in gemini_model_cost_map.items():
for model, model_info in mistral_model_cost_map.items():
if (
(
model_info.get("litellm_provider") == "gemini"
or model_info.get("litellm_provider") == "vertex_ai-language-models"
)
(model_info.get("litellm_provider") == "mistral")
and model_info.get("mode") == "chat"
and ("gemini-2.5" in model and "tts" not in model)
and model_info.get("supports_pdf_input") is None
and ("codestral-mamba" not in model)
):
"""
Update all gemini chat models to support pdf input
Update all mistral models to supports_response_schema
"""
model_info["supports_pdf_input"] = True
print(f"Updated {model} to support pdf input")
model_info["supports_response_schema"] = True
print(f"Updated {model} to support response schema")
json.dump(
gemini_model_cost_map, open("model_prices_and_context_window.json", "w"), indent=4
mistral_model_cost_map, open("model_prices_and_context_window.json", "w"), indent=4
)
tests/test_litellm/llms/ollama/test_ollama_embedding.py
+23 -14
View File
@@ -1,8 +1,9 @@
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.types.utils import EmbeddingResponse
from litellm.llms.ollama.completion.handler import ollama_embeddings, ollama_aembeddings
import pytest
from litellm.llms.ollama.completion.handler import ollama_aembeddings, ollama_embeddings
from litellm.types.utils import EmbeddingResponse
@pytest.fixture
@@ -26,8 +27,9 @@ def mock_encoding():
def test_ollama_embeddings(mock_response_data, mock_embedding_response, mock_encoding):
with patch("litellm.module_level_client.post") as mock_post, \
patch("litellm.OllamaConfig.get_config", return_value={"truncate": 512}):
with 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
@@ -50,11 +52,17 @@ 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):
with patch("litellm.module_level_aclient.post", new_callable=AsyncMock) as mock_post, \
patch("litellm.OllamaConfig.get_config", return_value={"truncate": 512}):
mock_post.return_value.json.return_value = mock_response_data
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.OllamaConfig.get_config", return_value={"truncate": 512}
):
response = await ollama_aembeddings(
api_base="http://localhost:11434",
@@ -78,9 +86,10 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin
# No "prompt_eval_count"
}
with patch("litellm.module_level_client.post") as mock_post, \
patch("litellm.OllamaConfig.get_config", return_value={}):
with 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
@@ -99,4 +108,4 @@ def test_prompt_eval_fallback_when_missing(mock_embedding_response, mock_encodin
assert response.usage.prompt_tokens == 5
assert response.usage.total_tokens == 5
assert response.usage.completion_tokens == 0
assert response.data[0]['embedding'] == [0.1, 0.2, 0.3]
assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
tests/test_litellm/proxy/test_route_llm_request.py
+22
View File
@@ -102,3 +102,25 @@ async def test_route_request_no_model_required_with_router_settings():
# Reset the mock for the next route
llm_router.reset_mock()
@pytest.mark.asyncio
async def test_route_request_no_model_required_with_router_settings_and_no_router():
"""Test route types that don't require model parameter with router settings and no router"""
from unittest.mock import patch
import litellm
from litellm.proxy.route_llm_request import route_request
data = {
"model": "my-model-id",
"api_key": "my-api-key",
"messages": [{"role": "user", "content": "what llm are you"}],
}
with patch.object(
litellm, "acompletion", return_value="fake_response"
) as mock_completion:
response = await route_request(data, None, "gpt-3.5-turbo", "acompletion")
mock_completion.assert_called_once_with(**data)