From 1dc2d2ddacc4bf54335a8023ecd4645858f9b3e6 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:30:23 +0530 Subject: [PATCH] fix(utils.py): correctly extract messages from google genai contents (#19156) * fix(utils.py): correctly extract messages from google genai contents * refactor use shared utilities --- litellm/utils.py | 392 ++++++++++++------ .../test_google_api_endpoints.py | 257 ++++++------ 2 files changed, 392 insertions(+), 257 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index f5b5c50468f..9198b8e2ecc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -87,6 +87,7 @@ def _get_cached_custom_logger(): global _CustomLogger if _CustomLogger is None: from litellm.integrations.custom_logger import CustomLogger + _CustomLogger = CustomLogger return _CustomLogger @@ -100,6 +101,7 @@ def _get_cached_custom_guardrail(): global _CustomGuardrail if _CustomGuardrail is None: from litellm.integrations.custom_guardrail import CustomGuardrail + _CustomGuardrail = CustomGuardrail return _CustomGuardrail @@ -113,6 +115,7 @@ def _get_cached_caching_handler_response(): global _CachingHandlerResponse if _CachingHandlerResponse is None: from litellm.caching.caching_handler import CachingHandlerResponse + _CachingHandlerResponse = CachingHandlerResponse return _CachingHandlerResponse @@ -126,6 +129,7 @@ def _get_cached_llm_caching_handler(): global _LLMCachingHandler if _LLMCachingHandler is None: from litellm.caching.caching_handler import LLMCachingHandler + _LLMCachingHandler = LLMCachingHandler return _LLMCachingHandler @@ -144,9 +148,11 @@ def _get_cached_audio_utils(): global _audio_utils_module if _audio_utils_module is None: import litellm.litellm_core_utils.audio_utils.utils + _audio_utils_module = litellm.litellm_core_utils.audio_utils.utils return _audio_utils_module + from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -203,10 +209,6 @@ from litellm.types.utils import ( # Thank you users! We ❤️ you! - Krrish & Ishaan - - - - try: # Python 3.9+ with resources.files("litellm.litellm_core_utils.tokenizers").joinpath( @@ -250,10 +252,14 @@ from litellm.llms.base_llm.base_utils import ( if TYPE_CHECKING: # Heavy types that are only needed for type checking; avoid importing # their modules at runtime during `litellm` import. - from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler + from litellm.caching.caching_handler import ( + CachingHandlerResponse, + LLMCachingHandler, + ) from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.proxy._types import AllowedModelRegion + # Type stubs for lazy-loaded functions to help mypy understand their types # These imports allow mypy to understand the types when these are accessed via __getattr__ from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -288,10 +294,13 @@ if TYPE_CHECKING: ) from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig from litellm.llms.base_llm.search.transformation import BaseSearchConfig - from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig + from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig, + ) from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + # Type stubs for lazy-loaded functions and classes from litellm.litellm_core_utils.cached_imports import ( get_coroutine_checker, @@ -335,6 +344,7 @@ if TYPE_CHECKING: reset_retry_policy, ) from litellm.secret_managers.main import get_secret + # Type stubs for lazy-loaded config classes and types from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.containers.transformation import BaseContainerConfig @@ -618,8 +628,8 @@ def load_credentials_from_list(kwargs: dict): Updates kwargs with the credentials if credential_name in kwarg """ # Access CredentialAccessor via module to trigger lazy loading if needed - CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor') - + CredentialAccessor = getattr(sys.modules[__name__], "CredentialAccessor") + credential_name = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: credential_accessor = CredentialAccessor.get_credential_values(credential_name) @@ -646,7 +656,7 @@ def _is_gemini_model(model: Optional[str], custom_llm_provider: Optional[str]) - if custom_llm_provider in ["vertex_ai", "vertex_ai_beta"]: return model is not None and "gemini" in model.lower() return True - + # Check if model name contains gemini return model is not None and "gemini" in model.lower() @@ -668,7 +678,7 @@ def _process_assistant_message_tool_calls( """ role = msg_copy.get("role") tool_calls = msg_copy.get("tool_calls") - + if role == "assistant" and isinstance(tool_calls, list): new_tool_calls = [] for tc in tool_calls: @@ -681,17 +691,17 @@ def _process_assistant_message_tool_calls( else: new_tool_calls.append(tc) continue - + # Remove thought signature from ID if present if isinstance(tc_dict.get("id"), str): if thought_signature_separator in tc_dict["id"]: tc_dict["id"] = _remove_thought_signature_from_id( tc_dict["id"], thought_signature_separator ) - + new_tool_calls.append(tc_dict) msg_copy["tool_calls"] = new_tool_calls - + return msg_copy @@ -699,14 +709,12 @@ def _process_tool_message_id(msg_copy: dict, thought_signature_separator: str) - """ Process tool message to remove thought signature from tool_call_id. """ - if msg_copy.get("role") == "tool" and isinstance( - msg_copy.get("tool_call_id"), str - ): + if msg_copy.get("role") == "tool" and isinstance(msg_copy.get("tool_call_id"), str): if thought_signature_separator in msg_copy["tool_call_id"]: msg_copy["tool_call_id"] = _remove_thought_signature_from_id( msg_copy["tool_call_id"], thought_signature_separator ) - + return msg_copy @@ -717,7 +725,7 @@ def _remove_thought_signatures_from_messages( Remove thought signatures from tool call IDs in all messages. """ processed_messages = [] - + for msg in messages: # Handle Pydantic models (convert to dict) if hasattr(msg, "model_dump"): @@ -728,17 +736,17 @@ def _remove_thought_signatures_from_messages( # Unknown type, keep as is processed_messages.append(msg) continue - + # Process assistant messages with tool_calls msg_dict = _process_assistant_message_tool_calls( msg_dict, thought_signature_separator ) - + # Process tool messages with tool_call_id msg_dict = _process_tool_message_id(msg_dict, thought_signature_separator) - + processed_messages.append(msg_dict) - + return processed_messages @@ -763,12 +771,12 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## LAZY LOAD COROUTINE CHECKER ## - get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') + get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = ( - kwargs.pop("callbacks", None) - ) + dynamic_callbacks: Optional[ + List[Union[str, Callable, "CustomLogger"]] + ] = kwargs.pop("callbacks", None) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -811,7 +819,7 @@ def function_setup( # noqa: PLR0915 + litellm.failure_callback ) ) - get_set_callbacks = getattr(sys.modules[__name__], 'get_set_callbacks') + get_set_callbacks = getattr(sys.modules[__name__], "get_set_callbacks") get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS if len(litellm.input_callback) > 0: @@ -939,7 +947,7 @@ def function_setup( # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] ### PRE-CALL RULES ### - Rules = getattr(sys.modules[__name__], 'Rules') + Rules = getattr(sys.modules[__name__], "Rules") if ( Rules.has_pre_call_rules() and isinstance(messages, list) @@ -947,7 +955,6 @@ def function_setup( # noqa: PLR0915 and isinstance(messages[0], dict) and "content" in messages[0] ): - buffer = StringIO() for m in messages: content = m.get("content", "") @@ -958,7 +965,7 @@ def function_setup( # noqa: PLR0915 input=buffer.getvalue(), model=model, ) - + ### REMOVE THOUGHT SIGNATURES FROM TOOL CALL IDS FOR NON-GEMINI MODELS ### # Gemini models embed thought signatures in tool call IDs. When sending # messages with tool calls to non-Gemini providers, we need to remove these @@ -974,7 +981,7 @@ def function_setup( # noqa: PLR0915 # Get custom_llm_provider to determine target provider custom_llm_provider = kwargs.get("custom_llm_provider") - + # If custom_llm_provider not in kwargs, try to determine it from the model if not custom_llm_provider and model: try: @@ -985,18 +992,18 @@ def function_setup( # noqa: PLR0915 except Exception: # If we can't determine the provider, skip this processing pass - + # Only process if target is NOT a Gemini model if not _is_gemini_model(model, custom_llm_provider): verbose_logger.debug( "Removing thought signatures from tool call IDs for non-Gemini model" ) - + # Process messages to remove thought signatures processed_messages = _remove_thought_signatures_from_messages( messages, THOUGHT_SIGNATURE_SEPARATOR ) - + # Update messages in kwargs or args if "messages" in kwargs: kwargs["messages"] = processed_messages @@ -1041,9 +1048,7 @@ def function_setup( # noqa: PLR0915 _file_obj: FileTypes = args[1] if len(args) > 1 else kwargs["file"] # Lazy import audio_utils.utils only when needed for transcription calls audio_utils = _get_cached_audio_utils() - file_checksum = audio_utils.get_audio_file_content_hash( - file_obj=_file_obj - ) + file_checksum = audio_utils.get_audio_file_content_hash(file_obj=_file_obj) if "metadata" in kwargs: kwargs["metadata"]["file_checksum"] = file_checksum else: @@ -1064,6 +1069,42 @@ def function_setup( # noqa: PLR0915 else kwargs.get("input") or kwargs.get("messages", "default-message-value") ) + elif ( + call_type == CallTypes.generate_content.value + or call_type == CallTypes.agenerate_content.value + or call_type == CallTypes.generate_content_stream.value + or call_type == CallTypes.agenerate_content_stream.value + ): + try: + from litellm.google_genai.adapters.transformation import ( + GoogleGenAIAdapter, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_last_user_message, + ) + + contents_param = args[1] if len(args) > 1 else kwargs.get("contents") + model_param = args[0] if len(args) > 0 else kwargs.get("model", "") + + if contents_param: + adapter = GoogleGenAIAdapter() + transformed = adapter.translate_generate_content_to_completion( + model=model_param, + contents=contents_param, + config=kwargs.get("config"), + ) + transformed_messages = transformed.get("messages", []) + messages = ( + get_last_user_message(transformed_messages) + or "default-message-value" + ) + else: + messages = "default-message-value" + except Exception as e: + verbose_logger.debug( + f"Error extracting messages from Google contents: {str(e)}" + ) + messages = "default-message-value" else: messages = "default-message-value" stream = False @@ -1072,7 +1113,9 @@ def function_setup( # noqa: PLR0915 call_type=call_type, ): stream = True - get_litellm_logging_class = getattr(sys.modules[__name__], 'get_litellm_logging_class') + get_litellm_logging_class = getattr( + sys.modules[__name__], "get_litellm_logging_class" + ) logging_obj = get_litellm_logging_class()( # Victim for object pool model=model, # type: ignore messages=messages, @@ -1156,8 +1199,10 @@ def _get_wrapper_num_retries( if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') - reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') + get_num_retries_from_retry_policy = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") retry_policy_num_retries = get_num_retries_from_retry_policy( exception=exception, retry_policy=kwargs.get("retry_policy"), @@ -1185,7 +1230,7 @@ def _get_wrapper_timeout( def check_coroutine(value) -> bool: - get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') + get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") return get_coroutine_checker().is_async_callable(value) @@ -1344,7 +1389,7 @@ def post_call_processing( def client(original_function): # noqa: PLR0915 - Rules = getattr(sys.modules[__name__], 'Rules') + Rules = getattr(sys.modules[__name__], "Rules") rules_obj = Rules() @wraps(original_function) @@ -1530,7 +1575,9 @@ def client(original_function): # noqa: PLR0915 ) else: # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1574,7 +1621,7 @@ def client(original_function): # noqa: PLR0915 # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx = contextvars.copy_context() - executor = getattr(sys.modules[__name__], 'executor') + executor = getattr(sys.modules[__name__], "executor") executor.submit( ctx.run, logging_obj.success_handler, @@ -1583,7 +1630,9 @@ def client(original_function): # noqa: PLR0915 end_time, ) # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1600,15 +1649,19 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') - reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') + get_num_retries_from_retry_policy = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy = getattr( + sys.modules[__name__], "reset_retry_policy" + ) num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -1645,15 +1698,19 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') - reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') + get_num_retries_from_retry_policy = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy = getattr( + sys.modules[__name__], "reset_retry_policy" + ) num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -1804,7 +1861,9 @@ def client(original_function): # noqa: PLR0915 chunks, messages=kwargs.get("messages", None) ) else: - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1869,7 +1928,9 @@ def client(original_function): # noqa: PLR0915 end_time=end_time, ) - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1969,7 +2030,7 @@ def client(original_function): # noqa: PLR0915 setattr(e, "timeout", timeout) raise e - get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') + get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") is_coroutine = get_coroutine_checker().is_async_callable(original_function) # Return the appropriate wrapper based on the original function type @@ -2330,7 +2391,7 @@ def supports_response_schema( """ ## GET LLM PROVIDER ## try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider ) @@ -2694,10 +2755,10 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 ## override / add new keys to the existing model cost dictionary updated_dictionary = _update_dictionary(existing_model, value) litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) - + # Invalidate case-insensitive lookup map since model_cost was modified _invalidate_model_cost_lowercase_map() - + verbose_logger.debug( f"added/updated model={model_cost_key} in litellm.model_cost: {model_cost_key}" ) @@ -3034,8 +3095,10 @@ def get_optional_params_embeddings( # noqa: PLR0915 **kwargs, ): # Lazy load get_supported_openai_params - get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') - + get_supported_openai_params = getattr( + sys.modules[__name__], "get_supported_openai_params" + ) + # retrieve all parameters passed to the function passed_params = locals() custom_llm_provider = passed_params.pop("custom_llm_provider", None) @@ -3308,8 +3371,8 @@ def get_optional_params_embeddings( # noqa: PLR0915 ) elif custom_llm_provider == "ollama": - if 'dimensions' in non_default_params: - optional_params['dimensions']=non_default_params.pop('dimensions') + if "dimensions" in non_default_params: + optional_params["dimensions"] = non_default_params.pop("dimensions") if len(non_default_params.keys()) > 0: if ( litellm.drop_params is True or drop_params is True @@ -3575,10 +3638,10 @@ def pre_process_non_default_params( if "response_format" in non_default_params: if provider_config is not None: - non_default_params["response_format"] = ( - provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] - ) + non_default_params[ + "response_format" + ] = provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -3707,16 +3770,16 @@ def pre_process_optional_params( True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("tools") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("tools") non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("functions") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("functions") elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -3840,7 +3903,9 @@ def get_optional_params( # noqa: PLR0915 message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) - get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') + get_supported_openai_params = getattr( + sys.modules[__name__], "get_supported_openai_params" + ) supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) @@ -4095,7 +4160,7 @@ def get_optional_params( # noqa: PLR0915 ), ) elif custom_llm_provider == "bedrock": - BedrockModelInfo = getattr(sys.modules[__name__], 'BedrockModelInfo') + BedrockModelInfo = getattr(sys.modules[__name__], "BedrockModelInfo") bedrock_route = BedrockModelInfo.get_bedrock_route(model) bedrock_base_model = BedrockModelInfo.get_base_model(model) if bedrock_route == "converse" or bedrock_route == "converse_like": @@ -4520,8 +4585,8 @@ def get_optional_params( # noqa: PLR0915 # Apply nested drops from additional_drop_params if additional_drop_params: - is_nested_path = getattr(sys.modules[__name__], 'is_nested_path') - delete_nested_value = getattr(sys.modules[__name__], 'delete_nested_value') + is_nested_path = getattr(sys.modules[__name__], "is_nested_path") + delete_nested_value = getattr(sys.modules[__name__], "delete_nested_value") nested_paths = [p for p in additional_drop_params if is_nested_path(p)] for path in nested_paths: optional_params = delete_nested_value(optional_params, path) @@ -4571,7 +4636,9 @@ def add_provider_specific_params_to_optional_params( else: processed_extra_body = initial_extra_body - _ensure_extra_body_is_safe = getattr(sys.modules[__name__], '_ensure_extra_body_is_safe') + _ensure_extra_body_is_safe = getattr( + sys.modules[__name__], "_ensure_extra_body_is_safe" + ) optional_params["extra_body"] = _ensure_extra_body_is_safe( extra_body=processed_extra_body ) @@ -4862,9 +4929,9 @@ def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) return delta if isinstance(delta, str) else "" # Handle standard ModelResponse and ModelResponseStream - _choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = ( - response_obj.choices - ) + _choices: Union[ + List[Union[Choices, StreamingChoices]], List[StreamingChoices] + ] = response_obj.choices # Use list accumulation to avoid O(n^2) string concatenation across choices response_parts: List[str] = [] @@ -4982,7 +5049,7 @@ def get_max_tokens(model: str) -> Optional[int]: return litellm.model_cost[model]["max_output_tokens"] elif "max_tokens" in litellm.model_cost[model]: return litellm.model_cost[model]["max_tokens"] - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) @@ -5058,7 +5125,7 @@ _model_cost_lowercase_map: Optional[Dict[str, str]] = None def _invalidate_model_cost_lowercase_map() -> None: """Invalidate the case-insensitive lookup map for model_cost. - + Call this whenever litellm.model_cost is modified to ensure the map is rebuilt. """ global _model_cost_lowercase_map @@ -5067,7 +5134,7 @@ def _invalidate_model_cost_lowercase_map() -> None: def _rebuild_model_cost_lowercase_map() -> Dict[str, str]: """Rebuild the case-insensitive lookup map from the current model_cost. - + Returns: The rebuilt map (guaranteed to be not None). """ @@ -5081,9 +5148,9 @@ def _handle_stale_map_entry_rebuild( ) -> Optional[str]: """ Handle stale _model_cost_lowercase_map entry (key was popped). - + Rebuilds the map and retries the lookup. - + Returns: The matched key if found after rebuild, None otherwise. """ @@ -5100,9 +5167,9 @@ def _handle_new_key_with_scan( ) -> Optional[str]: """ Handle new key added to model_cost without invalidating _model_cost_lowercase_map. - + Scans model_cost for case-insensitive match and rebuilds the map if found. - + Returns: The matched key if found, None otherwise. """ @@ -5117,20 +5184,20 @@ def _handle_new_key_with_scan( def _get_model_cost_key(potential_key: str) -> Optional[str]: """ Get the actual key from model_cost, with case-insensitive fallback. - + WARNING: Only O(1) lookup operations are acceptable. O(n) lookups will cause severe CPU overhead. This function is called frequently during router operations. - + ALLOWED HELPER FUNCTIONS (conditionally called, O(n) operations are acceptable): - _rebuild_model_cost_lowercase_map: Rebuilds the lookup map (only when map is None) - _handle_stale_map_entry_rebuild: Rebuilds map when stale entry detected (rare case) - + If you need to add a new helper function with O(n) operations that is conditionally called and confirmed not to cause performance issues, add it to the allowed_helpers list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py """ global _model_cost_lowercase_map - + # Exact match (O(1)) if potential_key in litellm.model_cost: return potential_key @@ -5138,20 +5205,20 @@ def _get_model_cost_key(potential_key: str) -> Optional[str]: # Case-insensitive lookup via map (O(1)) if _model_cost_lowercase_map is None: _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() - + potential_key_lower = potential_key.lower() matched_key = _model_cost_lowercase_map.get(potential_key_lower) - + # Verify key exists (O(1) - handles model_cost.pop() case) if matched_key is not None and matched_key in litellm.model_cost: return matched_key - + # Rebuild map if stale entry detected (O(n) rebuild, but only when stale entry found) if matched_key is not None: matched_key = _handle_stale_map_entry_rebuild(potential_key_lower) if matched_key is not None: return matched_key - + return None @@ -5183,9 +5250,12 @@ def _check_provider_match(model_info: dict, custom_llm_provider: Optional[str]) custom_llm_provider == "litellm_proxy" ): # litellm_proxy is a special case, it's not a provider, it's a proxy for the provider return True - elif custom_llm_provider == "azure_ai" and model_info["litellm_provider"] in ("azure", "openai"): - # Azure AI also works with azure models - # as a last attempt if the model is not on Azure AI, Azure then fallback to OpenAI cost + elif custom_llm_provider == "azure_ai" and model_info["litellm_provider"] in ( + "azure", + "openai", + ): + # Azure AI also works with azure models + # as a last attempt if the model is not on Azure AI, Azure then fallback to OpenAI cost # tracking the cost is better than attributing 0 cost to it. return True else: @@ -5211,7 +5281,7 @@ def _get_potential_model_names( if custom_llm_provider is None: # Get custom_llm_provider try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") split_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: split_model = model @@ -5941,7 +6011,7 @@ def validate_environment( # noqa: PLR0915 } ## EXTRACT LLM PROVIDER - if model name provided try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") _, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: custom_llm_provider = None @@ -6504,7 +6574,7 @@ def register_prompt_template( complete_model = model potential_models = [complete_model] try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") model = get_llm_provider(model=model)[0] potential_models.append(model) except Exception: @@ -6590,7 +6660,7 @@ class TextCompletionStreamWrapper: except StopIteration: raise StopIteration except Exception as e: - exception_type = getattr(sys.modules[__name__], 'exception_type') + exception_type = getattr(sys.modules[__name__], "exception_type") raise exception_type( model=self.model, custom_llm_provider=self.custom_llm_provider or "", @@ -7084,7 +7154,7 @@ def get_valid_models( # init litellm_params ################################# from litellm.types.router import LiteLLM_Params - + if litellm_params is None: litellm_params = LiteLLM_Params(model="") if api_key is not None: @@ -7214,14 +7284,20 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata = litellm_params.get("metadata", {}) - _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') - base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata) + _get_base_model_from_litellm_call_metadata = getattr( + sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" + ) + base_model_from_metadata = _get_base_model_from_litellm_call_metadata( + metadata=metadata + ) if base_model_from_metadata is not None: return base_model_from_metadata # Also check litellm_metadata (used by Responses API and other generic API calls) litellm_metadata = litellm_params.get("litellm_metadata", {}) - _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') + _get_base_model_from_litellm_call_metadata = getattr( + sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" + ) return _get_base_model_from_litellm_call_metadata(metadata=litellm_metadata) return None @@ -7618,7 +7694,7 @@ class ProviderConfigManager: @staticmethod def _build_provider_config_map() -> dict[LlmProviders, tuple[Callable, bool]]: """Build the provider-to-config mapping dictionary. - + Returns a dict mapping provider to (factory_function, needs_model_parameter). This avoids expensive inspect.signature() calls at runtime. """ @@ -7627,12 +7703,30 @@ class ProviderConfigManager: # Format: (factory_function, needs_model_parameter: bool) LlmProviders.OPENAI: (lambda: litellm.OpenAIGPTConfig(), False), LlmProviders.ANTHROPIC: (lambda: litellm.AnthropicConfig(), False), - LlmProviders.AZURE: (lambda model: ProviderConfigManager._get_azure_config(model), True), - LlmProviders.AZURE_AI: (lambda model: ProviderConfigManager._get_azure_ai_config(model), True), - LlmProviders.VERTEX_AI: (lambda model: ProviderConfigManager._get_vertex_ai_config(model), True), - LlmProviders.BEDROCK: (lambda model: ProviderConfigManager._get_bedrock_config(model), True), - LlmProviders.COHERE: (lambda model: ProviderConfigManager._get_cohere_config(model), True), - LlmProviders.COHERE_CHAT: (lambda model: ProviderConfigManager._get_cohere_config(model), True), + LlmProviders.AZURE: ( + lambda model: ProviderConfigManager._get_azure_config(model), + True, + ), + LlmProviders.AZURE_AI: ( + lambda model: ProviderConfigManager._get_azure_ai_config(model), + True, + ), + LlmProviders.VERTEX_AI: ( + lambda model: ProviderConfigManager._get_vertex_ai_config(model), + True, + ), + LlmProviders.BEDROCK: ( + lambda model: ProviderConfigManager._get_bedrock_config(model), + True, + ), + LlmProviders.COHERE: ( + lambda model: ProviderConfigManager._get_cohere_config(model), + True, + ), + LlmProviders.COHERE_CHAT: ( + lambda model: ProviderConfigManager._get_cohere_config(model), + True, + ), # Simple provider mappings (no model parameter needed) LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False), LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False), @@ -7642,7 +7736,10 @@ class ProviderConfigManager: LlmProviders.ZAI: (lambda: litellm.ZAIChatConfig(), False), LlmProviders.LAMBDA_AI: (lambda: litellm.LambdaAIChatConfig(), False), LlmProviders.LLAMA: (lambda: litellm.LlamaAPIConfig(), False), - LlmProviders.TEXT_COMPLETION_OPENAI: (lambda: litellm.OpenAITextCompletionConfig(), False), + LlmProviders.TEXT_COMPLETION_OPENAI: ( + lambda: litellm.OpenAITextCompletionConfig(), + False, + ), LlmProviders.SNOWFLAKE: (lambda: litellm.SnowflakeConfig(), False), LlmProviders.CLARIFAI: (lambda: litellm.ClarifaiConfig(), False), LlmProviders.ANTHROPIC_TEXT: (lambda: litellm.AnthropicTextConfig(), False), @@ -7665,7 +7762,10 @@ class ProviderConfigManager: LlmProviders.CUSTOM: (lambda: litellm.OpenAILikeChatConfig(), False), LlmProviders.CUSTOM_OPENAI: (lambda: litellm.OpenAILikeChatConfig(), False), LlmProviders.OPENAI_LIKE: (lambda: litellm.OpenAILikeChatConfig(), False), - LlmProviders.AIOHTTP_OPENAI: (lambda: litellm.AiohttpOpenAIChatConfig(), False), + LlmProviders.AIOHTTP_OPENAI: ( + lambda: litellm.AiohttpOpenAIChatConfig(), + False, + ), LlmProviders.HOSTED_VLLM: (lambda: litellm.HostedVLLMChatConfig(), False), LlmProviders.LLAMAFILE: (lambda: litellm.LlamafileChatConfig(), False), LlmProviders.LM_STUDIO: (lambda: litellm.LMStudioChatConfig(), False), @@ -7674,7 +7774,10 @@ class ProviderConfigManager: LlmProviders.HUGGINGFACE: (lambda: litellm.HuggingFaceChatConfig(), False), LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIConfig(), False), LlmProviders.OPENROUTER: (lambda: litellm.OpenrouterConfig(), False), - LlmProviders.VERCEL_AI_GATEWAY: (lambda: litellm.VercelAIGatewayConfig(), False), + LlmProviders.VERCEL_AI_GATEWAY: ( + lambda: litellm.VercelAIGatewayConfig(), + False, + ), LlmProviders.COMETAPI: (lambda: litellm.CometAPIConfig(), False), LlmProviders.DATAROBOT: (lambda: litellm.DataRobotConfig(), False), LlmProviders.GEMINI: (lambda: litellm.GoogleAIStudioGeminiConfig(), False), @@ -7692,7 +7795,10 @@ class ProviderConfigManager: LlmProviders.CEREBRAS: (lambda: litellm.CerebrasConfig(), False), LlmProviders.BASETEN: (lambda: litellm.BasetenConfig(), False), LlmProviders.VOLCENGINE: (lambda: litellm.VolcEngineConfig(), False), - LlmProviders.TEXT_COMPLETION_CODESTRAL: (lambda: litellm.CodestralTextCompletionConfig(), False), + LlmProviders.TEXT_COMPLETION_CODESTRAL: ( + lambda: litellm.CodestralTextCompletionConfig(), + False, + ), LlmProviders.SAMBANOVA: (lambda: litellm.SambanovaConfig(), False), LlmProviders.MARITALK: (lambda: litellm.MaritalkConfig(), False), LlmProviders.VLLM: (lambda: litellm.VLLMConfig(), False), @@ -7700,17 +7806,26 @@ class ProviderConfigManager: LlmProviders.PREDIBASE: (lambda: litellm.PredibaseConfig(), False), LlmProviders.TRITON: (lambda: litellm.TritonConfig(), False), LlmProviders.PETALS: (lambda: litellm.PetalsConfig(), False), - LlmProviders.SAP_GENERATIVE_AI_HUB: (lambda: litellm.GenAIHubOrchestrationConfig(), False), + LlmProviders.SAP_GENERATIVE_AI_HUB: ( + lambda: litellm.GenAIHubOrchestrationConfig(), + False, + ), LlmProviders.FEATHERLESS_AI: (lambda: litellm.FeatherlessAIConfig(), False), LlmProviders.NOVITA: (lambda: litellm.NovitaConfig(), False), LlmProviders.NEBIUS: (lambda: litellm.NebiusConfig(), False), LlmProviders.WANDB: (lambda: litellm.WandbConfig(), False), LlmProviders.DASHSCOPE: (lambda: litellm.DashScopeChatConfig(), False), LlmProviders.MOONSHOT: (lambda: litellm.MoonshotChatConfig(), False), - LlmProviders.DOCKER_MODEL_RUNNER: (lambda: litellm.DockerModelRunnerChatConfig(), False), + LlmProviders.DOCKER_MODEL_RUNNER: ( + lambda: litellm.DockerModelRunnerChatConfig(), + False, + ), LlmProviders.V0: (lambda: litellm.V0ChatConfig(), False), LlmProviders.MORPH: (lambda: litellm.MorphChatConfig(), False), - LlmProviders.LITELLM_PROXY: (lambda: litellm.LiteLLMProxyChatConfig(), False), + LlmProviders.LITELLM_PROXY: ( + lambda: litellm.LiteLLMProxyChatConfig(), + False, + ), LlmProviders.GRADIENT_AI: (lambda: litellm.GradientAIConfig(), False), LlmProviders.NSCALE: (lambda: litellm.NscaleConfig(), False), LlmProviders.HEROKU: (lambda: litellm.HerokuChatConfig(), False), @@ -7718,7 +7833,10 @@ class ProviderConfigManager: LlmProviders.HYPERBOLIC: (lambda: litellm.HyperbolicChatConfig(), False), LlmProviders.OVHCLOUD: (lambda: litellm.OVHCloudChatConfig(), False), LlmProviders.AMAZON_NOVA: (lambda: litellm.AmazonNovaChatConfig(), False), - LlmProviders.LANGGRAPH: (lambda: ProviderConfigManager._get_langgraph_config(), False), + LlmProviders.LANGGRAPH: ( + lambda: ProviderConfigManager._get_langgraph_config(), + False, + ), } @staticmethod @@ -7748,6 +7866,7 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.vertex_ai_partner_models.gpt_oss.transformation import ( VertexAIGPTOSSTransformation, ) + return VertexAIGPTOSSTransformation() elif model in litellm.vertex_mistral_models: if "codestral" in model: @@ -7762,12 +7881,13 @@ class ProviderConfigManager: def _get_bedrock_config(model: str) -> BaseConfig: """Get Bedrock config based on model.""" from litellm.llms.bedrock.common_utils import get_bedrock_chat_config + return get_bedrock_chat_config(model=model) @staticmethod def _get_cohere_config(model: str) -> BaseConfig: """Get Cohere config based on route.""" - CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') + CohereModelInfo = getattr(sys.modules[__name__], "CohereModelInfo") route = CohereModelInfo.get_cohere_route(model) if route == "v2": return litellm.CohereV2ChatConfig() @@ -7777,6 +7897,7 @@ class ProviderConfigManager: def _get_langgraph_config() -> BaseConfig: """Get LangGraph config.""" from litellm.llms.langgraph.chat.transformation import LangGraphConfig + return LangGraphConfig() @staticmethod @@ -7785,7 +7906,7 @@ class ProviderConfigManager: ) -> Optional[BaseConfig]: """ Returns the provider config for a given provider. - + Uses O(1) dictionary lookup for fast provider resolution. """ # Check JSON providers FIRST (these override standard mappings) @@ -7807,7 +7928,9 @@ class ProviderConfigManager: # Initialize provider config map lazily (avoids circular imports) if ProviderConfigManager._PROVIDER_CONFIG_MAP is None: - ProviderConfigManager._PROVIDER_CONFIG_MAP = ProviderConfigManager._build_provider_config_map() + ProviderConfigManager._PROVIDER_CONFIG_MAP = ( + ProviderConfigManager._build_provider_config_map() + ) # O(1) dictionary lookup config_entry = ProviderConfigManager._PROVIDER_CONFIG_MAP.get(provider) @@ -7877,6 +8000,7 @@ class ProviderConfigManager: from litellm.llms.openrouter.embedding.transformation import ( OpenrouterEmbeddingConfig, ) + return OpenrouterEmbeddingConfig() elif litellm.LlmProviders.GIGACHAT == provider: return litellm.GigaChatEmbeddingConfig() @@ -8490,8 +8614,8 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.ocr.common_utils import get_vertex_ai_ocr_config return get_vertex_ai_ocr_config(model=model) - - MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') + + MistralOCRConfig = getattr(sys.modules[__name__], "MistralOCRConfig") PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -8632,7 +8756,9 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ - get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], 'get_litellm_metadata_from_kwargs') + get_litellm_metadata_from_kwargs = getattr( + sys.modules[__name__], "get_litellm_metadata_from_kwargs" + ) _metadata = cast( dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)) ) @@ -8928,12 +9054,12 @@ def __getattr__(name: str) -> Any: """Lazy import handler for utils module with cached registry for improved performance.""" # Use cached registry from _lazy_imports instead of importing tuples every time from litellm._lazy_imports import _get_lazy_import_registry - + registry = _get_lazy_import_registry() - + # Check if name is in registry and call the cached handler function if name in registry: handler_func = registry[name] return handler_func(name) - + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 26c60b2c303..2f2eaa905be 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -26,30 +26,28 @@ def test_google_generate_content_endpoint(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock the router's agenerate_content method with patch("litellm.proxy.proxy_server.llm_router") as mock_router: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:generateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - } + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - + # Verify the response assert response.status_code == 200 assert response.json() == {"test": "response"} - + # Verify that agenerate_content was called mock_router.agenerate_content.assert_called_once() @@ -64,40 +62,42 @@ def test_google_stream_generate_content_endpoint(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock the router's agenerate_content_stream method to return a stream async def mock_stream_generator(): yield 'data: {"test": "stream_chunk_1"}\n\n' yield 'data: {"test": "stream_chunk_2"}\n\n' yield "data: [DONE]\n\n" - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router: - mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream_generator()) - + mock_router.agenerate_content_stream = AsyncMock( + return_value=mock_stream_generator() + ) + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:streamGenerateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - } + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that agenerate_content_stream was called with correct parameters mock_router.agenerate_content_stream.assert_called_once() call_args = mock_router.agenerate_content_stream.call_args assert call_args[1]["stream"] is True assert call_args[1]["model"] == "test-model" - assert call_args[1]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + assert call_args[1]["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] def test_google_generate_content_with_cost_tracking_metadata(): @@ -110,25 +110,28 @@ def test_google_generate_content_with_cost_tracking_metadata(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): # Simulate adding user metadata data["litellm_metadata"] = { "user_api_key_user_id": "test-user-id", @@ -136,29 +139,27 @@ def test_google_generate_content_with_cost_tracking_metadata(): "user_api_key": "hashed-key", } return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:generateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - }, - headers={"Authorization": "Bearer sk-test-key"} + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that add_litellm_data_to_request was called mock_add_data.assert_called_once() - + # Verify that agenerate_content was called with metadata mock_router.agenerate_content.assert_called_once() call_args = mock_router.agenerate_content.call_args called_data = call_args[1] - + # Verify that litellm_metadata exists and contains user information assert "litellm_metadata" in called_data assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" @@ -174,30 +175,33 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock the router's agenerate_content_stream method to return a stream mock_stream = AsyncMock() mock_stream.__aiter__ = lambda self: mock_stream mock_stream.__anext__.side_effect = StopAsyncIteration - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) - + # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): # Simulate adding user metadata data["litellm_metadata"] = { "user_api_key_user_id": "test-user-id", @@ -205,29 +209,27 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): "user_api_key": "hashed-key", } return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:streamGenerateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - }, - headers={"Authorization": "Bearer sk-test-key"} + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that add_litellm_data_to_request was called mock_add_data.assert_called_once() - + # Verify that agenerate_content_stream was called with metadata mock_router.agenerate_content_stream.assert_called_once() call_args = mock_router.agenerate_content_stream.call_args called_data = call_args[1] - + # Verify that litellm_metadata exists and contains user information assert "litellm_metadata" in called_data assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" @@ -239,7 +241,7 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): def test_google_generate_content_with_system_instruction(): """ Test that systemInstruction is correctly passed through from the endpoint to the router. - + This test verifies the fix for systemInstruction being dropped when forwarding requests to Vertex AI through the Google GenAI endpoint. """ @@ -250,62 +252,63 @@ def test_google_generate_content_with_system_instruction(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Define the systemInstruction to test - system_instruction = { - "parts": [{"text": "Your name is Doodle."}] - } - + system_instruction = {"parts": [{"text": "Your name is Doodle."}]} + # Send a request with systemInstruction response = client.post( "/v1beta/models/gemini-2.5-pro:generateContent", json={ "systemInstruction": system_instruction, "contents": [ - { - "parts": [{"text": "What is your name?"}], - "role": "user" - } - ] + {"parts": [{"text": "What is your name?"}], "role": "user"} + ], }, - headers={"Authorization": "Bearer sk-test-key"} + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that agenerate_content was called mock_router.agenerate_content.assert_called_once() call_args = mock_router.agenerate_content.call_args called_data = call_args[1] - + # Verify that systemInstruction is present in the call arguments assert "systemInstruction" in called_data assert called_data["systemInstruction"] == system_instruction - assert called_data["systemInstruction"]["parts"][0]["text"] == "Your name is Doodle." - + assert ( + called_data["systemInstruction"]["parts"][0]["text"] + == "Your name is Doodle." + ) + # Verify contents are also present assert "contents" in called_data assert len(called_data["contents"]) == 1 @@ -315,7 +318,7 @@ def test_google_generate_content_with_system_instruction(): def test_google_generate_content_with_image_config(): """ Test that imageConfig is correctly passed through from generationConfig to the router. - + This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved when forwarding requests to Google GenAI through the endpoint. """ @@ -326,69 +329,75 @@ def test_google_generate_content_with_image_config(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Send a request with generationConfig containing imageConfig response = client.post( "/v1beta/models/gemini-3-pro-image-preview:generateContent", json={ - "contents": [{ - "role": "user", - "parts": [{"text": "Create a vibrant infographic about photosynthesis"}] - }], + "contents": [ + { + "role": "user", + "parts": [ + { + "text": "Create a vibrant infographic about photosynthesis" + } + ], + } + ], "generationConfig": { "responseModalities": ["TEXT", "IMAGE"], - "imageConfig": { - "aspectRatio": "9:16", - "imageSize": "4K" - } - } + "imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"}, + }, }, - headers={"Authorization": "Bearer sk-test-key"} + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that agenerate_content was called mock_router.agenerate_content.assert_called_once() call_args = mock_router.agenerate_content.call_args called_data = call_args[1] - + # Verify that config is present in the call arguments assert "config" in called_data - + # Verify that imageConfig is preserved in the config assert "imageConfig" in called_data["config"] assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16" assert called_data["config"]["imageConfig"]["imageSize"] == "4K" - + # Verify that responseModalities is also preserved assert "responseModalities" in called_data["config"] assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"] - + # Verify contents are also present assert "contents" in called_data assert len(called_data["contents"]) == 1 - assert called_data["contents"][0]["role"] == "user" \ No newline at end of file + assert called_data["contents"][0]["role"] == "user"