mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
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
This commit is contained in:
Harshit Jain
GitHub
parent
98e87c3e67
commit
1dc2d2ddac
2 changed files with 392 additions and 257 deletions
392
litellm/utils.py
392
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}")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
assert called_data["contents"][0]["role"] == "user"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue