diff --git a/litellm/main.py b/litellm/main.py index 7d457d9cdd1..4f263398d72 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2850,6 +2850,7 @@ def _complete_cohere_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatc def _complete_maritalk(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key custom_prompt_dict = ctx.custom_prompt_dict @@ -2879,10 +2880,12 @@ def _complete_maritalk(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe logging_obj=logging, custom_llm_provider="maritalk", custom_prompt_dict=custom_prompt_dict, + acompletion=acompletion, ) def _complete_amazon_nova(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key custom_llm_provider = ctx.custom_llm_provider @@ -2915,6 +2918,7 @@ def _complete_amazon_nova(ctx: _CompletionDispatchContext) -> _CompletionDispatc timeout=timeout, custom_llm_provider=custom_llm_provider, custom_prompt_dict=custom_prompt_dict, + acompletion=acompletion, ) diff --git a/tests/test_litellm/test_openai_params_strip.py b/tests/test_litellm/test_openai_params_strip.py index a5f9131e95d..b2c2a2614a1 100644 --- a/tests/test_litellm/test_openai_params_strip.py +++ b/tests/test_litellm/test_openai_params_strip.py @@ -229,3 +229,285 @@ async def test_openai_metadata_preview_feature(): mock_create_no_preview.assert_called_once() call_kwargs = mock_create_no_preview.call_args[1] assert "metadata" not in call_kwargs + + +@pytest.mark.asyncio +async def test_strip_litellm_internal_params_non_dict(): + """ + Test that strip_litellm_internal_params returns the input directly + if it is not a dictionary. + """ + from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params + + assert strip_litellm_internal_params("not a dict") == "not a dict" + assert strip_litellm_internal_params(None) is None + assert strip_litellm_internal_params(123) == 123 + + +@pytest.mark.asyncio +async def test_azure_chat_completion_params_strip(): + """ + Test that litellm_params and _litellm_* prefixed params are stripped + from Azure OpenAI completion calls. + """ + mock_choice = MagicMock() + mock_choice.finish_reason = "stop" + mock_choice.index = 0 + mock_choice.message = MagicMock(content="Mock response", role="assistant") + mock_choice.message.tool_calls = None + mock_choice.message.function_call = None + mock_choice.message.provider_specific_fields = {} + + mock_response_data = MagicMock() + mock_response_data.choices = [mock_choice] + mock_response_data.id = "chatcmpl-123" + mock_response_data.created = 1677858242 + mock_response_data.model = "gpt-4" + mock_response_data.object = "chat.completion" + mock_response_data.usage = MagicMock(completion_tokens=10, prompt_tokens=5, total_tokens=15) + + mock_create = MagicMock() + mock_raw_resp = MagicMock() + mock_raw_resp.headers = {"x-test-header": "test"} + mock_raw_resp.parse.return_value = mock_response_data + mock_create.return_value = mock_raw_resp + + with patch("openai.resources.chat.completions.Completions.create", mock_create): + completion( + model="azure/gpt-4", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + api_base="https://mock.openai.azure.com", + api_version="2023-05-15", + # internal params that should be stripped + litellm_params={"metadata": {"some_internal_key": "some_value"}}, + _litellm_test_param="test_value", + ) + + mock_create.assert_called_once() + call_kwargs = mock_create.call_args[1] + + assert "litellm_params" not in call_kwargs + assert "_litellm_test_param" not in call_kwargs + + +@pytest.mark.asyncio +async def test_azure_chat_acompletion_params_strip(): + """ + Test that litellm_params and _litellm_* prefixed params are stripped + from Azure OpenAI async completion calls. + """ + mock_choice = MagicMock() + mock_choice.finish_reason = "stop" + mock_choice.index = 0 + mock_choice.message = MagicMock(content="Mock response", role="assistant") + mock_choice.message.tool_calls = None + mock_choice.message.function_call = None + mock_choice.message.provider_specific_fields = {} + + mock_response_data = MagicMock() + mock_response_data.choices = [mock_choice] + mock_response_data.id = "chatcmpl-123" + mock_response_data.created = 1677858242 + mock_response_data.model = "gpt-4" + mock_response_data.object = "chat.completion" + mock_response_data.usage = MagicMock(completion_tokens=10, prompt_tokens=5, total_tokens=15) + + mock_raw_resp = MagicMock() + mock_raw_resp.headers = {"x-test-header": "test"} + mock_raw_resp.parse.return_value = mock_response_data + + mock_acreate = AsyncMock(return_value=mock_raw_resp) + + with patch("openai.resources.chat.completions.AsyncCompletions.create", mock_acreate): + await acompletion( + model="azure/gpt-4", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + api_base="https://mock.openai.azure.com", + api_version="2023-05-15", + litellm_params={"metadata": {"some_internal_key": "some_value"}}, + _litellm_test_param="test_value", + ) + + mock_acreate.assert_called_once() + call_kwargs = mock_acreate.call_args[1] + + assert "litellm_params" not in call_kwargs + assert "_litellm_test_param" not in call_kwargs + + +@pytest.mark.asyncio +async def test_azure_embedding_params_strip(): + """ + Test that litellm_params and _litellm_* prefixed params are stripped + from Azure OpenAI embedding calls. + """ + mock_response_data = MagicMock() + mock_response_data.model = "text-embedding-ada-002" + mock_response_data.object = "list" + mock_response_data.data = [MagicMock(embedding=[0.1, 0.2])] + mock_response_data.usage = MagicMock(prompt_tokens=5, total_tokens=5) + + mock_create = MagicMock() + mock_raw_resp = MagicMock() + mock_raw_resp.headers = {"x-test-header": "test"} + mock_raw_resp.parse.return_value = mock_response_data + mock_create.return_value = mock_raw_resp + + with patch("openai.resources.embeddings.Embeddings.create", mock_create): + embedding( + model="azure/text-embedding-ada-002", + input=["hello"], + api_key="mock-key", + api_base="https://mock.openai.azure.com", + api_version="2023-05-15", + litellm_params={"metadata": {"some_internal_key": "some_value"}}, + _litellm_test_param="test_value", + ) + + mock_create.assert_called_once() + call_kwargs = mock_create.call_args[1] + + assert "litellm_params" not in call_kwargs + assert "_litellm_test_param" not in call_kwargs + + +@pytest.mark.asyncio +async def test_azure_aembedding_params_strip(): + """ + Test that litellm_params and _litellm_* prefixed params are stripped + from Azure OpenAI async embedding calls. + """ + from litellm import aembedding + + mock_response_data = MagicMock() + mock_response_data.model = "text-embedding-ada-002" + mock_response_data.object = "list" + mock_response_data.data = [MagicMock(embedding=[0.1, 0.2])] + mock_response_data.usage = MagicMock(prompt_tokens=5, total_tokens=5) + + mock_raw_resp = MagicMock() + mock_raw_resp.headers = {"x-test-header": "test"} + mock_raw_resp.parse.return_value = mock_response_data + + mock_acreate = AsyncMock(return_value=mock_raw_resp) + + with patch("openai.resources.embeddings.AsyncEmbeddings.create", mock_acreate): + await aembedding( + model="azure/text-embedding-ada-002", + input=["hello"], + api_key="mock-key", + api_base="https://mock.openai.azure.com", + api_version="2023-05-15", + litellm_params={"metadata": {"some_internal_key": "some_value"}}, + _litellm_test_param="test_value", + ) + + mock_acreate.assert_called_once() + call_kwargs = mock_acreate.call_args[1] + + assert "litellm_params" not in call_kwargs + assert "_litellm_test_param" not in call_kwargs + + +@pytest.mark.asyncio +async def test_openai_like_chat_completion_params_strip(): + """ + Test that litellm_params and _litellm_* prefixed params are stripped + from OpenAI-like completion calls. + """ + import json + import httpx + + # Mock Response object + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers({"content-type": "application/json"}) + + mock_body = { + "choices": [ + {"finish_reason": "stop", "index": 0, "message": {"content": "Mock response", "role": "assistant"}} + ], + "id": "chatcmpl-123", + "created": 1677858242, + "model": "maritalk", + "object": "chat.completion", + "usage": {"completion_tokens": 10, "prompt_tokens": 5, "total_tokens": 15}, + } + + mock_response.text = json.dumps(mock_body) + mock_response.json.return_value = mock_body + mock_response.content = json.dumps(mock_body).encode("utf-8") + + mock_post = MagicMock(return_value=mock_response) + + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post): + completion( + model="maritalk/maritalk", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + api_base="https://example.com", + # internal params that should be stripped + litellm_params={"metadata": {"some_internal_key": "some_value"}}, + _litellm_test_param="test_value", + ) + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args[1] + + # Verify that data parameter passed to post has been stripped + sent_data = json.loads(call_kwargs.get("data", "{}")) + assert "litellm_params" not in sent_data + assert "_litellm_test_param" not in sent_data + + +@pytest.mark.asyncio +async def test_openai_like_chat_acompletion_params_strip(): + """ + Test that litellm_params and _litellm_* prefixed params are stripped + from OpenAI-like async completion calls. + """ + import json + import httpx + + # Mock Response object + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = httpx.Headers({"content-type": "application/json"}) + + mock_body = { + "choices": [ + {"finish_reason": "stop", "index": 0, "message": {"content": "Mock response", "role": "assistant"}} + ], + "id": "chatcmpl-123", + "created": 1677858242, + "model": "maritalk", + "object": "chat.completion", + "usage": {"completion_tokens": 10, "prompt_tokens": 5, "total_tokens": 15}, + } + + mock_response.text = json.dumps(mock_body) + mock_response.json.return_value = mock_body + mock_response.content = json.dumps(mock_body).encode("utf-8") + + mock_apost = AsyncMock(return_value=mock_response) + + with patch.object(litellm.module_level_aclient, "post", mock_apost): + await acompletion( + model="maritalk/maritalk", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + api_base="https://example.com", + # internal params that should be stripped + litellm_params={"metadata": {"some_internal_key": "some_value"}}, + _litellm_test_param="test_value", + ) + + mock_apost.assert_called_once() + call_kwargs = mock_apost.call_args[1] + + # Verify that data parameter passed to post has been stripped + sent_data = json.loads(call_kwargs.get("data", "{}")) + assert "litellm_params" not in sent_data + assert "_litellm_test_param" not in sent_data