From 1a4ad8bf18f3fa4e07e5c5b2cc901653a7e3654c Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Wed, 25 Jun 2025 07:20:13 -0700 Subject: [PATCH] Update mistral 'supports_response_schema' field + Fix ollama embedding (#12024) * 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 --- litellm/llms/ollama/completion/handler.py | 27 ++++--- ...odel_prices_and_context_window_backup.json | 78 ++++++++++++------- litellm/proxy/route_llm_request.py | 5 +- model_prices_and_context_window.json | 78 ++++++++++++------- test_script.py | 20 ++--- .../llms/ollama/test_ollama_embedding.py | 37 +++++---- .../proxy/test_route_llm_request.py | 22 ++++++ 7 files changed, 178 insertions(+), 89 deletions(-) diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 9f507aded0..9e6497e66a 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -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, ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d441115f06..8d72806c32 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 970526afad..ce623aa7f0 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -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") diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d441115f06..8d72806c32 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/test_script.py b/test_script.py index 0dce4b6256..d0fec5035b 100644 --- a/test_script.py +++ b/test_script.py @@ -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 ) diff --git a/tests/test_litellm/llms/ollama/test_ollama_embedding.py b/tests/test_litellm/llms/ollama/test_ollama_embedding.py index 30248f26b6..a5cbd36cee 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_embedding.py +++ b/tests/test_litellm/llms/ollama/test_ollama_embedding.py @@ -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] diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index c99815bcd9..d4a65d8847 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -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)