diff --git a/litellm/litellm_core_utils/param_utils.py b/litellm/litellm_core_utils/param_utils.py index 832f885573d..ced45937a1d 100644 --- a/litellm/litellm_core_utils/param_utils.py +++ b/litellm/litellm_core_utils/param_utils.py @@ -1,13 +1,8 @@ -import logging - -logger = logging.getLogger(__name__) - LITELLM_INTERNAL_PARAM_NAMES = frozenset( ( "litellm_params", "proxy_server_request", "model_info", - "metadata", "preset_cache_key", "litellm_metadata", "acompletion", @@ -27,23 +22,19 @@ def strip_litellm_internal_params( if not isinstance(data, dict): # pyright: ignore[reportUnnecessaryIsInstance] # runtime guard for unsanitized input return data # pyright: ignore[reportUnreachable] # runtime guard - try: - # Create a shallow copy so we don't modify the input dictionary in-place - cleaned_data: dict[str, object] = {} # mutable-ok: building cleaned payload dict - for key, value in data.items(): - if key in LITELLM_INTERNAL_PARAM_NAMES or key.startswith("_litellm_"): - continue - if key == "extra_body" and isinstance(value, dict): - cleaned_extra_body: dict[str, object] = {} # mutable-ok: building cleaned extra_body dict - extra_body_dict: dict[object, object] = value # pyright: ignore[reportUnknownVariableType] # mutable-ok: reading from extra_body - for k, v in extra_body_dict.items(): - if isinstance(k, str) and (k in LITELLM_INTERNAL_PARAM_NAMES or k.startswith("_litellm_")): - continue - cleaned_extra_body[str(k)] = v - cleaned_data["extra_body"] = cleaned_extra_body - else: - cleaned_data[key] = value - return cleaned_data - except Exception as e: - logger.warning(f"Error in strip_litellm_internal_params: {str(e)}") - return data + # Create a shallow copy so we don't modify the input dictionary in-place + cleaned_data: dict[str, object] = {} # mutable-ok: building cleaned payload dict + for key, value in data.items(): + if key in LITELLM_INTERNAL_PARAM_NAMES or key.startswith("_litellm_"): + continue + if key == "extra_body" and isinstance(value, dict): + cleaned_extra_body: dict[str, object] = {} # mutable-ok: building cleaned extra_body dict + extra_body_dict: dict[object, object] = value # pyright: ignore[reportUnknownVariableType] # mutable-ok: reading from extra_body + for k, v in extra_body_dict.items(): + if isinstance(k, str) and (k in LITELLM_INTERNAL_PARAM_NAMES or k.startswith("_litellm_")): + continue + cleaned_extra_body[str(k)] = v + cleaned_data["extra_body"] = cleaned_extra_body + else: + cleaned_data[key] = value + return cleaned_data diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 67efac19499..8c5dafbd60c 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -16,6 +16,7 @@ import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -33,8 +34,6 @@ from litellm.utils import ( convert_to_model_response_object, modify_url, ) -from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params - from ...types.llms.openai import HttpxBinaryResponseContent from ..base import BaseLLM diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 8048e13068b..46295585cbb 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -35,6 +35,7 @@ from litellm.constants import DEFAULT_MAX_RETRIES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator @@ -50,8 +51,6 @@ from litellm.utils import ( ProviderConfigManager, convert_to_model_response_object, ) -from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params - from ...types.llms.openai import * from ..base import BaseLLM diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 11d798a9aa2..1c14e76959b 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -11,6 +11,7 @@ import httpx import litellm from litellm import LlmProviders +from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.databricks.streaming_utils import ModelResponseIterator @@ -18,8 +19,6 @@ from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai.openai import OpenAIConfig from litellm.types.utils import CustomStreamingDecoder, ModelResponse from litellm.utils import CustomStreamWrapper, ProviderConfigManager -from litellm.litellm_core_utils.param_utils import strip_litellm_internal_params - from ..common_utils import OpenAILikeBase, OpenAILikeError from .transformation import OpenAILikeChatConfig diff --git a/tests/test_litellm/test_openai_params_strip.py b/tests/test_litellm/test_openai_params_strip.py index f145b022f0d..a5f9131e95d 100644 --- a/tests/test_litellm/test_openai_params_strip.py +++ b/tests/test_litellm/test_openai_params_strip.py @@ -13,6 +13,20 @@ from litellm import acompletion, completion, embedding litellm.return_response_headers = False +@pytest.fixture(autouse=True) +def clear_client_cache(): + """ + Clear the HTTP client cache before each test to ensure mocks are used. + This prevents cached real clients from being reused across tests. + """ + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is not None: + cache.flush_cache() + yield + if cache is not None: + cache.flush_cache() + + @pytest.mark.asyncio async def test_openai_chat_completion_params_strip(): """ @@ -96,16 +110,13 @@ async def test_openai_chat_acompletion_params_strip(): mock_acreate = AsyncMock(return_value=mock_raw_resp) with patch("openai.resources.chat.completions.AsyncCompletions.create", mock_acreate): - try: - await acompletion( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - api_key="mock-key", - litellm_params={"metadata": {"some_internal_key": "some_value"}}, - _litellm_test_param="test_value", - ) - except Exception: - pass + await acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + 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] @@ -156,3 +167,65 @@ async def test_openai_embedding_params_strip(): if extra_body: assert "litellm_params" not in extra_body assert "_litellm_test_param" not in extra_body + + +@pytest.mark.asyncio +async def test_openai_metadata_preview_feature(): + """ + Test that when litellm.enable_preview_features = True, + metadata is successfully forwarded and not stripped. + And when it is False, metadata is not in the request. + """ + 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-4o" + 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 + + # 1. Test with enable_preview_features = True (metadata should be forwarded) + mock_create_preview = MagicMock(return_value=mock_raw_resp) + litellm.enable_preview_features = True + try: + with patch("openai.resources.chat.completions.Completions.create", mock_create_preview): + completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + metadata={"user_api_key_user_id": "test_user_id"}, + ) + mock_create_preview.assert_called_once() + call_kwargs = mock_create_preview.call_args[1] + assert "metadata" in call_kwargs + assert call_kwargs["metadata"] == {"user_api_key_user_id": "test_user_id"} + finally: + litellm.enable_preview_features = False + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is not None: + cache.flush_cache() + + # 2. Test with enable_preview_features = False (metadata should be stripped/omitted) + mock_create_no_preview = MagicMock(return_value=mock_raw_resp) + with patch("openai.resources.chat.completions.Completions.create", mock_create_no_preview): + completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + api_key="mock-key", + metadata={"user_api_key_user_id": "test_user_id"}, + ) + mock_create_no_preview.assert_called_once() + call_kwargs = mock_create_no_preview.call_args[1] + assert "metadata" not in call_kwargs