scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(param_utils): allow metadata passing and fix test cache pollution

Browse source
This commit is contained in:
mutnale_sushant 2026-07-11 07:00:43 +05:30
parent 59c8e36add
commit 8df2302427
5 changed files with 102 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/litellm_core_utils/param_utils.py
View file

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

3
litellm/llms/azure/azure.py
View file

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

3
litellm/llms/openai/openai.py
View file

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

3
litellm/llms/openai_like/chat/handler.py
View file

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

93
tests/test_litellm/test_openai_params_strip.py
View file

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