diff --git a/litellm/main.py b/litellm/main.py index 80a2f74c571..316e5a811fd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -249,6 +249,7 @@ from .types.utils import ( FileTypes, HiddenParams, LlmProviders, + LlmProvidersSet, PromptTokensDetails, ProviderSpecificHeader, all_litellm_params, @@ -1104,8 +1105,6 @@ def completion( # type: ignore # noqa: PLR0915 stop = validate_openai_optional_params(stop=stop) ######### unpacking kwargs ##################### - args = locals() - skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False) if not skip_mcp_handler and tools: from litellm.responses.mcp.chat_completions_handler import ( @@ -1284,12 +1283,12 @@ def completion( # type: ignore # noqa: PLR0915 logging: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, litellm_logging_obj) fallbacks = fallbacks or litellm.model_fallbacks if fallbacks is not None: - return completion_with_fallbacks(**args) + return completion_with_fallbacks(**locals()) if model_list is not None: deployments = [ m["litellm_params"] for m in model_list if m["model_name"] == model ] - return litellm.batch_completion_models(deployments=deployments, **args) + return litellm.batch_completion_models(deployments=deployments, **locals()) if litellm.model_alias_map and model in litellm.model_alias_map: model = litellm.model_alias_map[ model @@ -1394,9 +1393,7 @@ def completion( # type: ignore # noqa: PLR0915 ) provider_config: Optional[BaseConfig] = None - if custom_llm_provider is not None and custom_llm_provider in [ - provider.value for provider in LlmProviders - ]: + if custom_llm_provider is not None and custom_llm_provider in LlmProvidersSet: provider_config = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider) ) @@ -4306,11 +4303,13 @@ def completion( # type: ignore # noqa: PLR0915 model=model, custom_llm_provider=custom_llm_provider, original_exception=e, - completion_kwargs=args, + completion_kwargs={"messages": messages}, extra_kwargs=kwargs, ) + + def completion_with_retries(*args, **kwargs): """ Executes a litellm.completion() with 3 retries diff --git a/litellm/utils.py b/litellm/utils.py index 2a29c08904c..3d12ad5ea36 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3874,9 +3874,7 @@ def get_optional_params( # noqa: PLR0915 custom_llm_provider=custom_llm_provider, ) provider_config: Optional[BaseConfig] = None - if custom_llm_provider is not None and custom_llm_provider in [ - provider.value for provider in LlmProviders - ]: + if custom_llm_provider is not None and custom_llm_provider in LlmProvidersSet: provider_config = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider) ) diff --git a/tests/litellm/test_completion_perf_optimizations.py b/tests/litellm/test_completion_perf_optimizations.py new file mode 100644 index 00000000000..72a8e964675 --- /dev/null +++ b/tests/litellm/test_completion_perf_optimizations.py @@ -0,0 +1,47 @@ +""" +Tests for completion() performance optimizations. + +Covers: +1. LlmProvidersSet used instead of list comprehension for provider lookup +2. Deferred locals() - get_first_chars_messages works with minimal dict +""" + +import litellm +from litellm.types.utils import LlmProviders, LlmProvidersSet + + +class TestLlmProvidersSetConsistency: + """Verify LlmProvidersSet matches the enum values exactly.""" + + def test_set_matches_enum(self): + expected = {p.value for p in LlmProviders} + assert LlmProvidersSet == expected + + def test_known_providers_in_set(self): + for provider in ["openai", "anthropic", "azure", "bedrock", "vertex_ai"]: + assert provider in LlmProvidersSet, f"{provider} not in LlmProvidersSet" + + def test_unknown_provider_not_in_set(self): + assert "not_a_real_provider_xyz" not in LlmProvidersSet + + +class TestGetFirstCharsMessagesWithMinimalKwargs: + """Verify get_first_chars_messages works when completion_kwargs only has 'messages'.""" + + def test_with_messages_key(self): + """get_first_chars_messages should work with just {"messages": ...}.""" + messages = [{"role": "user", "content": "hello"}] + result = litellm.get_first_chars_messages(kwargs={"messages": messages}) + assert "hello" in result + + def test_with_empty_dict(self): + """get_first_chars_messages should handle empty dict gracefully.""" + result = litellm.get_first_chars_messages(kwargs={}) + # Should return empty string or "None" — not crash + assert isinstance(result, str) + + def test_truncates_long_messages(self): + """get_first_chars_messages truncates to 100 chars.""" + messages = [{"role": "user", "content": "x" * 200}] + result = litellm.get_first_chars_messages(kwargs={"messages": messages}) + assert len(result) <= 100