diff --git a/docs/my-website/docs/providers/moonshot.md b/docs/my-website/docs/providers/moonshot.md index 2e00bae355..827f2fd53c 100644 --- a/docs/my-website/docs/providers/moonshot.md +++ b/docs/my-website/docs/providers/moonshot.md @@ -219,6 +219,37 @@ curl http://localhost:4000/v1/chat/completions \ For more detailed information on using the LiteLLM Proxy, see the [LiteLLM Proxy documentation](../providers/litellm_proxy). +## Image / Vision Support + +Moonshot vision models (`kimi-k2.5`, `kimi-latest`, `moonshot-v1-*-vision-preview`, etc.) accept the standard OpenAI content array with `image_url` blocks. + +LiteLLM automatically detects when your messages contain images and preserves the content array so the image payload reaches the Moonshot API. For text-only requests the content is flattened to a plain string, as required by Moonshot text models. + +```python showLineNumbers title="Moonshot Vision Example" +import os +import litellm + +os.environ["MOONSHOT_API_KEY"] = "" + +response = litellm.completion( + model="moonshot/kimi-k2.5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + }, + ], + } + ], +) + +print(response.choices[0].message.content) +``` + ## Moonshot AI Limitations & LiteLLM Handling LiteLLM automatically handles the following [Moonshot AI limitations](https://platform.moonshot.ai/docs/guide/migrating-from-openai-to-kimi#about-api-compatibility) to provide seamless OpenAI compatibility: diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 5c051797e8..e9ac1d2ad7 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -221,7 +221,9 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - return streamwrapper + return self._apply_post_stream_processing( + streamwrapper, model, custom_llm_provider + ) async def acompletion( self, *args, **kwargs @@ -300,7 +302,30 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - return streamwrapper + return self._apply_post_stream_processing( + streamwrapper, model, custom_llm_provider + ) + + @staticmethod + def _apply_post_stream_processing( + stream: "CustomStreamWrapper", + model: str, + custom_llm_provider: str, + ) -> Any: + """Apply provider-specific post-stream processing if available.""" + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + try: + provider_config = ProviderConfigManager.get_provider_chat_config( + model=model, provider=LlmProviders(custom_llm_provider) + ) + except (ValueError, KeyError): + return stream + + if provider_config is not None: + return provider_config.post_stream_processing(stream) + return stream responses_api_bridge = ResponsesToCompletionBridgeHandler() diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 35fc93bbeb..fbd0810a75 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -949,9 +949,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if provider_specific_fields: function_chunk["provider_specific_fields"] = provider_specific_fields + tool_call_index = parsed_chunk.get("output_index", 0) tool_call_chunk = ChatCompletionToolCallChunk( id=output_item.get("call_id"), - index=0, + index=tool_call_index, type="function", function=function_chunk, ) @@ -972,6 +973,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): elif event_type == "response.function_call_arguments.delta": content_part: Optional[str] = parsed_chunk.get("delta", None) if content_part: + tool_call_index = parsed_chunk.get("output_index", 0) return ModelResponseStream( choices=[ StreamingChoices( @@ -980,7 +982,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): tool_calls=[ ChatCompletionToolCallChunk( id=None, - index=0, + index=tool_call_index, type="function", function=ChatCompletionToolCallFunctionChunk(name=None, arguments=content_part), ) @@ -1012,9 +1014,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if provider_specific_fields: function_chunk["provider_specific_fields"] = provider_specific_fields + tool_call_index = parsed_chunk.get("output_index", 0) tool_call_chunk = ChatCompletionToolCallChunk( id=output_item.get("call_id"), - index=0, + index=tool_call_index, type="function", function=function_chunk, ) diff --git a/litellm/images/main.py b/litellm/images/main.py index 236266af6a..eb6aa0c209 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -469,6 +469,8 @@ def image_generation( # noqa: PLR0915 or custom_llm_provider == LlmProviders.LITELLM_PROXY.value or custom_llm_provider in litellm.openai_compatible_providers ): + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) organization: Optional[str] = kwargs.get("organization", None) model_response = openai_chat_completions.image_generation( @@ -764,6 +766,8 @@ def image_edit( # noqa: PLR0915 } # model-specific params - pass them straight to the model/provider litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + model_info = kwargs.get("model_info", None) + metadata = kwargs.get("metadata", {}) _is_async = kwargs.pop("async_call", False) is True # add images / or return a single image @@ -872,8 +876,10 @@ def image_edit( # noqa: PLR0915 user=user, optional_params=dict(image_edit_request_params), litellm_params={ - "litellm_call_id": litellm_call_id, **image_edit_request_params, + "litellm_call_id": litellm_call_id, + "model_info": model_info, + "metadata": metadata, }, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index b996813b4e..c77a1b2564 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -16,6 +16,7 @@ class HeliconeLogger: helicone_model_list = [ "gpt", "claude", + "gemini", "command-r", "command-r-plus", "command-light", @@ -127,15 +128,20 @@ class HeliconeLogger: f"Helicone Logging - Enters logging function for model {model}" ) litellm_params = kwargs.get("litellm_params", {}) + custom_llm_provider = litellm_params.get("custom_llm_provider", "") kwargs.get("litellm_call_id", None) metadata = litellm_params.get("metadata", {}) or {} metadata = self.add_metadata_from_header(litellm_params, metadata) + + # Check if model is a vertex_ai model + is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") + model = ( model if any( accepted_model in model for accepted_model in self.helicone_model_list - ) + ) or is_vertex_ai else "gpt-3.5-turbo" ) provider_request = {"model": model, "messages": messages} @@ -144,7 +150,7 @@ class HeliconeLogger: ): response_obj = response_obj.json() - if "claude" in model: + if "claude" in model and not is_vertex_ai: response_obj = self.claude_mapping( model=model, messages=messages, response_obj=response_obj ) @@ -158,9 +164,15 @@ class HeliconeLogger: # Code to be executed provider_url = self.provider_url url = f"{self.api_base}/oai/v1/log" - if "claude" in model: + if "claude" in model and not is_vertex_ai: url = f"{self.api_base}/anthropic/v1/log" provider_url = "https://api.anthropic.com/v1/messages" + elif "gemini" in model: + url = f"{self.api_base}/custom/v1/log" + provider_url = "https://generativelanguage.googleapis.com/v1beta" + elif is_vertex_ai: + url = f"{self.api_base}/custom/v1/log" + provider_url = "https://aiplatform.googleapis.com/v1" headers = { "Authorization": f"Bearer {self.key}", "Content-Type": "application/json", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 8ab4ec15b0..864405749f 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -3,6 +3,7 @@ from typing import Optional, Tuple import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.llms.openrouter.common_utils import NATIVE_OPENROUTER_MODELS from litellm.secret_managers.main import get_secret, get_secret_str from ..types.router import LiteLLM_Params @@ -171,6 +172,12 @@ def get_llm_provider( # noqa: PLR0915 dynamic_api_key=dynamic_api_key, ) + # Check native OpenRouter models before provider_list stripping. + # These models have IDs like "openrouter/free" which would be + # incorrectly stripped to just "free" by the logic below. + if model in NATIVE_OPENROUTER_MODELS: + return model, "openrouter", dynamic_api_key, api_base + # check if llm provider part of model name if ( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3b75a56fcc..1256638b23 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -161,6 +161,7 @@ class CustomStreamWrapper: ) # keep track of the returned chunks - used for calculating the input/output tokens for stream options self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: Optional[int] = None + self._last_returned_hidden_params: Optional[dict] = None def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" @@ -1230,7 +1231,7 @@ class CustomStreamWrapper: ], ) _streaming_response = StreamingChoices(delta=_delta_obj) - _model_response = ModelResponse(stream=True) + _model_response = ModelResponseStream() _model_response.choices = [_streaming_response] response_obj = {"original_chunk": _model_response} else: @@ -1835,6 +1836,7 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) response._hidden_params["usage"] = usage + self._last_returned_hidden_params = response._hidden_params # Add MCP metadata to final chunk if present response = self._add_mcp_metadata_to_final_chunk(response) # RETURN RESULT @@ -1876,6 +1878,24 @@ class CustomStreamWrapper: None, cache_hit, ) + # Update hidden_params with final usage from + # stream_chunk_builder. Some providers (e.g. OpenRouter) + # send usage in a chunk after finish_reason, which arrives + # after _hidden_params["usage"] was initially set. The + # _hidden_params dict is the same object the user received + # (shared by reference), so mutating it here also corrects + # the user's copy. + if ( + self.stream_options is None + and complete_streaming_response is not None + and self._last_returned_hidden_params is not None + ): + final_usage = getattr( + complete_streaming_response, "usage", None + ) + if final_usage is not None: + self._last_returned_hidden_params["usage"] = final_usage + if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response @@ -2005,6 +2025,7 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) processed_chunk._hidden_params["usage"] = usage + self._last_returned_hidden_params = processed_chunk._hidden_params # Call post-call streaming deployment hook for final chunk if self.sent_last_chunk is True: @@ -2069,6 +2090,19 @@ class CustomStreamWrapper: cache_hit=cache_hit, ) ) + # Update hidden_params with final usage from + # stream_chunk_builder (see sync __next__ for full comment). + if ( + self.stream_options is None + and complete_streaming_response is not None + and self._last_returned_hidden_params is not None + ): + final_usage = getattr( + complete_streaming_response, "usage", None + ) + if final_usage is not None: + self._last_returned_hidden_params["usage"] = final_usage + if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 5b5354228f..07481917af 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -31,6 +31,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): api_key: str, api_base: Optional[str] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Dict[str, Any]: """ Handle a CountTokens request using httpx. @@ -60,6 +62,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): request_body = self.transform_request_to_count_tokens( model=model, messages=messages, + tools=tools, + system=system, ) verbose_logger.debug(f"Transformed request: {request_body}") diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 266b2794fc..93989c5854 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -30,6 +30,8 @@ class AnthropicTokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Optional[TokenCountResponse]: """ Count tokens using Anthropic's CountTokens API. @@ -66,6 +68,8 @@ class AnthropicTokenCounter(BaseTokenCounter): model=model_to_use, messages=messages, api_key=api_key, + tools=tools, + system=system, ) if result is not None: diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index c3ad72436b..ea4d60a60e 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -4,7 +4,7 @@ Anthropic CountTokens API transformation logic. This module handles the transformation of requests to Anthropic's CountTokens API format. """ -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION @@ -32,27 +32,27 @@ class AnthropicCountTokensConfig: self, model: str, messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Dict[str, Any]: """ Transform request to Anthropic CountTokens format. - Input: - { - "model": "claude-3-5-sonnet-20241022", - "messages": [{"role": "user", "content": "Hello!"}] - } - - Output (Anthropic CountTokens format): - { - "model": "claude-3-5-sonnet-20241022", - "messages": [{"role": "user", "content": "Hello!"}] - } + Includes optional system and tools fields for accurate token counting. """ - return { + request: Dict[str, Any] = { "model": model, "messages": messages, } + if system is not None: + request["system"] = system + + if tools is not None: + request["tools"] = tools + + return request + def get_required_headers(self, api_key: str) -> Dict[str, str]: """ Get the required headers for the CountTokens API. diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 52a0bb8bb0..2cba27925c 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -32,6 +32,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_base: str, litellm_params: Optional[Dict[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Dict[str, Any]: """ Handle a CountTokens request using httpx with Azure authentication. @@ -62,6 +64,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): request_body = self.transform_request_to_count_tokens( model=model, messages=messages, + tools=tools, + system=system, ) verbose_logger.debug(f"Transformed request: {request_body}") diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 14f9280007..afdfe9bdee 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -32,6 +32,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Optional[TokenCountResponse]: """ Count tokens using Azure AI Anthropic's CountTokens API. @@ -79,6 +81,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): api_key=api_key, api_base=api_base, litellm_params=litellm_params, + tools=tools, + system=system, ) if result is not None: diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index b1ccfc36d0..f6c6da2409 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -121,6 +121,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): Returns: Complete URL for Azure DI analyze endpoint """ + if api_base is None: + api_base = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") + if api_base is None: raise ValueError( "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter" diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 9172a05e38..ecff9053dc 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -24,6 +24,8 @@ class BaseTokenCounter(ABC): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Optional[TokenCountResponse]: pass diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index ac209904e6..f22c8ee0d9 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -438,6 +438,10 @@ class BaseConfig(ABC): """ return True + def post_stream_processing(self, stream: Any) -> Any: + """Hook for providers to post-process streaming responses. Default: pass-through.""" + return stream + def calculate_additional_costs( self, model: str, prompt_tokens: int, completion_tokens: int ) -> Optional[dict]: diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 9ae850ad4c..fe7d4b194a 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -26,7 +26,7 @@ from litellm.types.llms.bedrock_agentcore import ( AgentCoreUsage, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices, Usage +from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -481,7 +481,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): text = delta.get("text", "") if text: - chunk = ModelResponse( + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=model, @@ -499,7 +499,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Process metadata/usage metadata = event_payload.get("metadata") if metadata and "usage" in metadata: - chunk = ModelResponse( + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=model, @@ -522,7 +522,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Process final message if "message" in data_obj and isinstance(data_obj["message"], dict): - chunk = ModelResponse( + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=model, @@ -601,7 +601,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): self, response: httpx.Response, model: str, - ) -> AsyncGenerator[ModelResponse, None]: + ) -> AsyncGenerator[ModelResponseStream, None]: """ Internal async generator that parses SSE and yields ModelResponse chunks. """ @@ -636,7 +636,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): text = delta.get("text", "") if text: - chunk = ModelResponse( + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=model, @@ -654,7 +654,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Process metadata/usage metadata = event_payload.get("metadata") if metadata and "usage" in metadata: - chunk = ModelResponse( + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=model, @@ -677,7 +677,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Process final message if "message" in data_obj and isinstance(data_obj["message"], dict): - chunk = ModelResponse( + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=model, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 88f7341ed0..9b06e19820 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -559,7 +559,7 @@ class BedrockLLM(BaseAWSLLM): "INSIDE BEDROCK STREAMING TOOL CALLING CONDITION BLOCK" ) # return an iterator - streaming_model_response = ModelResponse(stream=True) + streaming_model_response = ModelResponseStream() streaming_model_response.choices[0].finish_reason = getattr( model_response.choices[0], "finish_reason", "stop" ) @@ -696,7 +696,7 @@ class BedrockLLM(BaseAWSLLM): ) if stream and provider == "ai21": - streaming_model_response = ModelResponse(stream=True) + streaming_model_response = ModelResponseStream() streaming_model_response.choices[0].finish_reason = model_response.choices[ # type: ignore 0 ].finish_reason diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index 0260eeafe6..fe0fd40b55 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -68,13 +68,8 @@ class AmazonQwen2Config(AmazonQwen3Config): # Set the content in the existing model_response structure if hasattr(model_response, 'choices') and len(model_response.choices) > 0: choice = model_response.choices[0] - if hasattr(choice, 'message'): - choice.message.content = generated_text - choice.finish_reason = "stop" - else: - # Handle streaming choices - choice.delta.content = generated_text - choice.finish_reason = "stop" + choice.message.content = generated_text + choice.finish_reason = "stop" # Set usage information if available in response if "usage" in response_data: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index 6eddcccd63..4be3e370fa 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -190,13 +190,8 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Set the content in the existing model_response structure if hasattr(model_response, 'choices') and len(model_response.choices) > 0: choice = model_response.choices[0] - if hasattr(choice, 'message'): - choice.message.content = generated_text - choice.finish_reason = "stop" - else: - # Handle streaming choices - choice.delta.content = generated_text - choice.finish_reason = "stop" + choice.message.content = generated_text + choice.finish_reason = "stop" # Set usage information if available in response if "usage" in response_data: diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 54f8a8dbd6..772eb16968 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -30,6 +30,8 @@ class BedrockTokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + tools: Optional[List[Dict[str, Any]]] = None, + system: Optional[Any] = None, ) -> Optional[TokenCountResponse]: """ Count tokens using AWS Bedrock's CountTokens API. @@ -54,11 +56,17 @@ class BedrockTokenCounter(BaseTokenCounter): litellm_params = deployment.get("litellm_params", {}) # Build request data in the format expected by BedrockCountTokensHandler - request_data = { + request_data: Dict[str, Any] = { "model": model_to_use, "messages": messages, } + if tools: + request_data["tools"] = tools + + if system: + request_data["system"] = system + # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) resolved_model = get_bedrock_base_model(model_to_use) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index b313cc9df3..64f1098e64 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -5,7 +5,8 @@ This module handles the transformation of requests from Anthropic Messages API f to AWS Bedrock's CountTokens API format and vice versa. """ -from typing import Any, Dict, List +import re +from typing import Any, Dict, List, Optional from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model @@ -75,46 +76,81 @@ class BedrockCountTokensConfig(BaseAWSLLM): input_type = self._detect_input_type(request_data) if input_type == "converse": - return self._transform_to_converse_format(request_data.get("messages", [])) + return self._transform_to_converse_format(request_data) else: return self._transform_to_invoke_model_format(request_data) def _transform_to_converse_format( - self, messages: List[Dict[str, Any]] + self, request_data: Dict[str, Any] ) -> Dict[str, Any]: - """Transform to Converse input format.""" - # Extract system messages if present - system_messages = [] + """Transform to Converse input format, including system and tools.""" + messages = request_data.get("messages", []) + system = request_data.get("system") + tools = request_data.get("tools") + + # Transform messages user_messages = [] - for message in messages: - if message.get("role") == "system": - system_messages.append({"text": message.get("content", "")}) - else: - # Transform message content to Bedrock format - transformed_message: Dict[str, Any] = {"role": message.get("role"), "content": []} + transformed_message: Dict[str, Any] = {"role": message.get("role"), "content": []} + content = message.get("content", "") + if isinstance(content, str): + transformed_message["content"].append({"text": content}) + elif isinstance(content, list): + transformed_message["content"] = content + user_messages.append(transformed_message) - # Handle content - ensure it's in the correct array format - content = message.get("content", "") - if isinstance(content, str): - # String content -> convert to text block - transformed_message["content"].append({"text": content}) - elif isinstance(content, list): - # Already in blocks format - use as is - transformed_message["content"] = content + converse_input: Dict[str, Any] = {"messages": user_messages} - user_messages.append(transformed_message) + # Transform system prompt (string or list of blocks → Bedrock format) + system_blocks = self._transform_system(system) + if system_blocks: + converse_input["system"] = system_blocks - # Build the converse input format - converse_input = {"messages": user_messages} + # Transform tools (Anthropic format → Bedrock toolConfig) + tool_config = self._transform_tools(tools) + if tool_config: + converse_input["toolConfig"] = tool_config - # Add system messages if present - if system_messages: - converse_input["system"] = system_messages - - # Build the complete request return {"input": {"converse": converse_input}} + def _transform_system(self, system: Optional[Any]) -> List[Dict[str, Any]]: + """Transform Anthropic system prompt to Bedrock system blocks.""" + if system is None: + return [] + if isinstance(system, str): + return [{"text": system}] + if isinstance(system, list): + # Already in blocks format (e.g. [{"type": "text", "text": "..."}]) + return [{"text": block.get("text", "")} for block in system if isinstance(block, dict)] + return [] + + def _transform_tools(self, tools: Optional[List[Dict[str, Any]]]) -> Optional[Dict[str, Any]]: + """Transform Anthropic tools to Bedrock toolConfig format.""" + if not tools: + return None + + bedrock_tools = [] + for tool in tools: + name = tool.get("name", "") + # Bedrock tool names must match [a-zA-Z][a-zA-Z0-9_]* and max 64 chars + name = re.sub(r"[^a-zA-Z0-9_]", "_", name) + if name and not name[0].isalpha(): + name = "t_" + name + name = name[:64] + + description = tool.get("description") or name + input_schema = tool.get("input_schema", {"type": "object", "properties": {}}) + + bedrock_tools.append({ + "toolSpec": { + "name": name, + "description": description, + "inputSchema": {"json": input_schema}, + } + }) + + return {"tools": bedrock_tools} + def _transform_to_invoke_model_format( self, request_data: Dict[str, Any] ) -> Dict[str, Any]: diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py new file mode 100644 index 0000000000..730becb06e --- /dev/null +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -0,0 +1,83 @@ +""" +Streaming utilities for ChatGPT provider. + +Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. +""" + +from typing import Any + + +class ChatGPTToolCallNormalizer: + """ + Wraps a streaming response and fixes tool_call index/dedup issues. + + The ChatGPT backend API (chatgpt.com/backend-api) sends non-spec-compliant + streaming tool call chunks: + 1. `index` is always 0, even for multiple parallel tool calls + 2. `id` and `name` get repeated in "closing" chunks that shouldn't exist + + This wrapper normalizes the stream to match the OpenAI spec before yielding + chunks to the consumer. + """ + + def __init__(self, stream: Any): + self._stream = stream + self._seen_ids: Dict[str, int] = {} # tool_call_id -> assigned_index + self._next_index: int = 0 + self._last_id: Optional[str] = None # tracks which tool call the next delta belongs to + + def __getattr__(self, name: str) -> Any: + return getattr(self._stream, name) + + def __iter__(self): + return self + + def __aiter__(self): + return self + + def __next__(self): + while True: + chunk = next(self._stream) + result = self._normalize(chunk) + if result is not None: + return result + + async def __anext__(self): + while True: + chunk = await self._stream.__anext__() + result = self._normalize(chunk) + if result is not None: + return result + + def _normalize(self, chunk: Any) -> Any: + """Fix tool_calls in the chunk. Returns None to skip duplicate chunks.""" + if not chunk.choices: + return chunk + + delta = chunk.choices[0].delta + if delta is None or not delta.tool_calls: + return chunk + + normalized = [] + for tc in delta.tool_calls: + if tc.id and tc.id not in self._seen_ids: + # New tool call — assign correct index + self._seen_ids[tc.id] = self._next_index + tc.index = self._next_index + self._last_id = tc.id + self._next_index += 1 + normalized.append(tc) + elif tc.id and tc.id in self._seen_ids: + # Duplicate "closing" chunk — skip it + continue + else: + # Continuation delta (id=None) — fix index + if self._last_id: + tc.index = self._seen_ids[self._last_id] + normalized.append(tc) + + if not normalized: + return None # all tool_calls were duplicates, skip chunk + + delta.tool_calls = normalized + return chunk diff --git a/litellm/llms/chatgpt/chat/transformation.py b/litellm/llms/chatgpt/chat/transformation.py index 2db5eb3c58..e6480398c7 100644 --- a/litellm/llms/chatgpt/chat/transformation.py +++ b/litellm/llms/chatgpt/chat/transformation.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple from litellm.exceptions import AuthenticationError from litellm.llms.openai.openai import OpenAIConfig @@ -10,6 +10,7 @@ from ..common_utils import ( ensure_chatgpt_session_id, get_chatgpt_default_headers, ) +from .streaming_utils import ChatGPTToolCallNormalizer class ChatGPTConfig(OpenAIConfig): @@ -61,6 +62,9 @@ class ChatGPTConfig(OpenAIConfig): ) return {**default_headers, **validated_headers} + def post_stream_processing(self, stream: Any) -> Any: + return ChatGPTToolCallNormalizer(stream) + def map_openai_params( self, non_default_params: dict, diff --git a/litellm/llms/codestral/completion/transformation.py b/litellm/llms/codestral/completion/transformation.py index 646c0e8e56..31d6652f48 100644 --- a/litellm/llms/codestral/completion/transformation.py +++ b/litellm/llms/codestral/completion/transformation.py @@ -102,7 +102,7 @@ class CodestralTextCompletionConfig(OpenAITextCompletionConfig): "finish_reason": finish_reason, } - original_chunk = litellm.ModelResponse(**chunk_data_dict, stream=True) + original_chunk = litellm.ModelResponseStream(**chunk_data_dict) _choices = chunk_data_dict.get("choices", []) or [] if len(_choices) == 0: return { diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index e53829d332..f99548c2c4 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -166,6 +166,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + **kwargs, ) -> Optional[TokenCountResponse]: import copy diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index bdb32cc0fe..cf81998055 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Optional import httpx from litellm._logging import verbose_logger -from litellm.types.utils import Delta, ModelResponse, StreamingChoices +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices if TYPE_CHECKING: pass @@ -44,7 +44,7 @@ class LangGraphSSEStreamIterator: self.async_line_iterator = self.response.aiter_lines() return self - def _parse_sse_line(self, line: str) -> Optional[ModelResponse]: + def _parse_sse_line(self, line: str) -> Optional[ModelResponseStream]: """ Parse a single SSE line and return a ModelResponse chunk if applicable. @@ -71,7 +71,7 @@ class LangGraphSSEStreamIterator: return None - def _process_data(self, data) -> Optional[ModelResponse]: + def _process_data(self, data) -> Optional[ModelResponseStream]: """ Process parsed data from SSE stream. @@ -101,7 +101,7 @@ class LangGraphSSEStreamIterator: return None - def _process_messages_event(self, payload) -> Optional[ModelResponse]: + def _process_messages_event(self, payload) -> Optional[ModelResponseStream]: """ Process a messages event from the stream. @@ -128,7 +128,7 @@ class LangGraphSSEStreamIterator: return None - def _process_metadata_event(self, payload) -> Optional[ModelResponse]: + def _process_metadata_event(self, payload) -> Optional[ModelResponseStream]: """ Process a metadata event, which may signal the end of the stream. """ @@ -139,9 +139,9 @@ class LangGraphSSEStreamIterator: return self._create_final_chunk() return None - def _create_content_chunk(self, text: str) -> ModelResponse: - """Create a ModelResponse chunk with content.""" - chunk = ModelResponse( + def _create_content_chunk(self, text: str) -> ModelResponseStream: + """Create a ModelResponseStream chunk with content.""" + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=self.model, @@ -158,9 +158,9 @@ class LangGraphSSEStreamIterator: return chunk - def _create_final_chunk(self) -> ModelResponse: - """Create a final ModelResponse chunk with finish_reason.""" - chunk = ModelResponse( + def _create_final_chunk(self) -> ModelResponseStream: + """Create a final ModelResponseStream chunk with finish_reason.""" + chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", created=0, model=self.model, @@ -177,7 +177,7 @@ class LangGraphSSEStreamIterator: return chunk - def __next__(self) -> ModelResponse: + def __next__(self) -> ModelResponseStream: """Sync iteration - parse SSE events and yield ModelResponse chunks.""" try: if self.line_iterator is None: @@ -205,7 +205,7 @@ class LangGraphSSEStreamIterator: verbose_logger.error(f"Error in LangGraph SSE stream: {str(e)}") raise StopIteration - async def __anext__(self) -> ModelResponse: + async def __anext__(self) -> ModelResponseStream: """Async iteration - parse SSE events and yield ModelResponse chunks.""" try: if self.async_line_iterator is None: diff --git a/litellm/llms/moonshot/chat/transformation.py b/litellm/llms/moonshot/chat/transformation.py index 0e78e58c7f..72c51bf74f 100644 --- a/litellm/llms/moonshot/chat/transformation.py +++ b/litellm/llms/moonshot/chat/transformation.py @@ -33,9 +33,25 @@ class MoonshotChatConfig(OpenAIGPTConfig): self, messages: List[AllMessageValues], model: str, is_async: bool = False ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: """ - Moonshot AI does not support content in list format. + Moonshot text-only models don't support content in list format. + Multimodal models (kimi-k2.5, kimi-latest, etc.) accept the + standard OpenAI content array with non-text blocks (image_url, + input_audio, video_url, file, etc.). + + If any message contains a non-text content part, skip flattening + so the multimodal payload is preserved. """ - messages = handle_messages_with_content_list_to_str_conversion(messages) + has_non_text = False + for m in messages: + _content = m.get("content") + if _content and isinstance(_content, list): + if any(c.get("type") != "text" for c in _content): + has_non_text = True + break + + if not has_non_text: + messages = handle_messages_with_content_list_to_str_conversion(messages) + if is_async: return super()._transform_messages( messages=messages, model=model, is_async=True diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 05c003c8b7..6903f84a4e 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -23,6 +23,11 @@ class OpenAIGPT5Config(OpenAIGPTConfig): # Don't route it through GPT-5 reasoning-specific parameter restrictions. return "gpt-5" in model and "gpt-5-chat" not in model + @classmethod + def is_model_gpt_5_search_model(cls, model: str) -> bool: + """Check if the model is a GPT-5 search variant (e.g. gpt-5-search-api).""" + return "gpt-5" in model and "search" in model + @classmethod def is_model_gpt_5_codex_model(cls, model: str) -> bool: """Check if the model is specifically a GPT-5 Codex variant.""" @@ -60,6 +65,23 @@ class OpenAIGPT5Config(OpenAIGPTConfig): return model_name.startswith("gpt-5.2") def get_supported_openai_params(self, model: str) -> list: + if self.is_model_gpt_5_search_model(model): + return [ + "max_tokens", + "max_completion_tokens", + "stream", + "stream_options", + "web_search_options", + "service_tier", + "safety_identifier", + "response_format", + "user", + "store", + "verbosity", + "max_retries", + "extra_headers", + ] + from litellm.utils import supports_tool_choice base_gpt_series_params = super().get_supported_openai_params(model=model) @@ -69,14 +91,20 @@ class OpenAIGPT5Config(OpenAIGPTConfig): base_gpt_series_params.remove("tool_choice") non_supported_params = [ - "logprobs", - "top_p", "presence_penalty", "frequency_penalty", - "top_logprobs", "stop", + "logit_bias", + "modalities", + "prediction", + "audio", + "web_search_options", ] + # gpt-5.1/5.2 support logprobs, top_p, top_logprobs when reasoning_effort="none" + if not self.is_model_gpt_5_1_model(model): + non_supported_params.extend(["logprobs", "top_p", "top_logprobs"]) + return [ param for param in base_gpt_series_params @@ -90,6 +118,18 @@ class OpenAIGPT5Config(OpenAIGPTConfig): model: str, drop_params: bool, ) -> dict: + if self.is_model_gpt_5_search_model(model): + if "max_tokens" in non_default_params: + optional_params["max_completion_tokens"] = non_default_params.pop( + "max_tokens" + ) + return super()._map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + reasoning_effort = ( non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort") @@ -118,6 +158,24 @@ class OpenAIGPT5Config(OpenAIGPTConfig): "max_tokens" ) + # gpt-5.1/5.2 support logprobs, top_p, top_logprobs only when reasoning_effort="none" + if self.is_model_gpt_5_1_model(model): + sampling_params = ["logprobs", "top_logprobs", "top_p"] + has_sampling = any(p in non_default_params for p in sampling_params) + if has_sampling and reasoning_effort not in (None, "none"): + if litellm.drop_params or drop_params: + for p in sampling_params: + non_default_params.pop(p, None) + else: + raise litellm.utils.UnsupportedParamsError( + message=( + "gpt-5.1/5.2 only support logprobs, top_p, top_logprobs when " + "reasoning_effort='none'. Current reasoning_effort='{}'. " + "To drop unsupported params set `litellm.drop_params = True`" + ).format(reasoning_effort), + status_code=400, + ) + if "temperature" in non_default_params: temperature_value: Optional[float] = non_default_params.pop("temperature") if temperature_value is not None: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 683e165c31..67e9e42bc3 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -542,16 +542,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if len(choice.message.tool_calls) > 0: return True elif isinstance(response, ModelResponseStream): - for choice in response.choices: - if isinstance(choice, litellm.StreamingChoices): + for streaming_choice in response.choices: + if isinstance(streaming_choice, litellm.StreamingChoices): # Check for text content - if choice.delta.content and isinstance(choice.delta.content, str): + if streaming_choice.delta.content and isinstance(streaming_choice.delta.content, str): return True # Check for tool calls - if choice.delta.tool_calls and isinstance( - choice.delta.tool_calls, list + if streaming_choice.delta.tool_calls and isinstance( + streaming_choice.delta.tool_calls, list ): - if len(choice.delta.tool_calls) > 0: + if len(streaming_choice.delta.tool_calls) > 0: return True return False diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index e67bfbe0c6..b89204230a 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -16,20 +16,17 @@ from litellm.types.containers.main import ( ) from litellm.types.router import GenericLiteLLMParams +from ...base_llm.containers.transformation import BaseContainerConfig + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from ...base_llm.chat.transformation import BaseLLMException as _BaseLLMException - from ...base_llm.containers.transformation import ( - BaseContainerConfig as _BaseContainerConfig, - ) LiteLLMLoggingObj = _LiteLLMLoggingObj - BaseContainerConfig = _BaseContainerConfig BaseLLMException = _BaseLLMException else: LiteLLMLoggingObj = Any - BaseContainerConfig = Any BaseLLMException = Any diff --git a/litellm/llms/openrouter/common_utils.py b/litellm/llms/openrouter/common_utils.py index 96e53a5aae..d4054278cf 100644 --- a/litellm/llms/openrouter/common_utils.py +++ b/litellm/llms/openrouter/common_utils.py @@ -1,5 +1,15 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException +# Native OpenRouter models whose IDs start with "openrouter/". +# When used via LiteLLM (openrouter/openrouter/free), get_llm_provider() +# must not strip the inner "openrouter/" prefix on its second invocation. +# See: https://github.com/BerriAI/litellm/issues/16353 +NATIVE_OPENROUTER_MODELS = { + "openrouter/auto", + "openrouter/free", + "openrouter/bodybuilder", +} + class OpenRouterException(BaseLLMException): pass diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 02b69b94d9..244ea098cc 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1042,6 +1042,7 @@ class VertexAITokenCounter(BaseTokenCounter): contents: Optional[List[Dict[str, Any]]], deployment: Optional[Dict[str, Any]] = None, request_model: str = "", + **kwargs, ) -> Optional[TokenCountResponse]: import copy diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 7bcefc1dd8..f35fc1aa4a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2100,7 +2100,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): chat_completion_logprobs=chat_completion_logprobs, image_response=image_response, ) - model_response.choices.append(choice) + model_response.choices.append(choice) # type: ignore[arg-type] elif isinstance(model_response, ModelResponse): choice = litellm.Choices( finish_reason=VertexGeminiConfig._check_finish_reason( @@ -2111,7 +2111,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): logprobs=chat_completion_logprobs, enhancements=None, ) - model_response.choices.append(choice) + model_response.choices.append(choice) # type: ignore[arg-type] return ( grounding_metadata, diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index ba3df88be1..db5693fb3a 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -43,13 +43,20 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): def get_supported_openai_params( self, model: str - ) -> List[OpenAIImageGenerationOptionalParams]: + ) -> list: """ Gemini image generation supported parameters + + Includes native Gemini imageConfig params (aspectRatio, imageSize) + in both camelCase and snake_case variants. """ return [ "n", "size", + "aspectRatio", + "aspect_ratio", + "imageSize", + "image_size", ] def map_openai_params( @@ -71,6 +78,10 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): elif k == "size": # Map OpenAI size format to Gemini aspectRatio mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + elif k in ("aspectRatio", "aspect_ratio"): + mapped_params["aspectRatio"] = v + elif k in ("imageSize", "image_size"): + mapped_params["imageSize"] = v else: mapped_params[k] = v diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index afedb6c8e7..a1241b87f4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2861,6 +2861,9 @@ class TokenCountRequest(LiteLLMPydanticObjectBase): Google /countTokens endpoint expects contents to be a list of dicts with the following structure: """ + tools: Optional[List[dict]] = None + system: Optional[Any] = None + class CallInfo(LiteLLMPydanticObjectBase): """Used for slack budget alerting""" diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 77bb1f53e6..5b23b47923 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -204,7 +204,12 @@ async def count_tokens( # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest - token_request = TokenCountRequest(model=model_name, messages=messages) + token_request = TokenCountRequest( + model=model_name, + messages=messages, + tools=data.get("tools"), + system=data.get("system"), + ) # Call the internal token counter function with direct request flag set to False token_response = await internal_token_counter( diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 34fbf47253..6786815583 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -61,7 +61,6 @@ from litellm.utils import ( ImageResponse, ModelResponse, ModelResponseStream, - StreamingChoices, ) @@ -863,9 +862,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if self.output_parse_pii is False and litellm.output_parse_pii is False: return response - if isinstance(response, ModelResponse) and not isinstance( - response.choices[0], StreamingChoices - ): # /chat/completions requests + if isinstance(response, ModelResponse): # /chat/completions requests if isinstance(response.choices[0].message.content, str): verbose_proxy_logger.debug( f"self.pii_tokens: {self.pii_tokens}; initial response: {response.choices[0].message.content}" @@ -888,7 +885,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return response # skip streaming here; handled in async_post_call_streaming_iterator_hook - if response.choices and isinstance(response.choices[0], StreamingChoices): + if isinstance(response, ModelResponseStream): return response presidio_config = self.get_presidio_settings_from_request_data( @@ -896,10 +893,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) for choice in response.choices: - # Type narrowing: StreamingChoices doesn't have .message attribute - if not hasattr(choice, "message"): - continue - content = getattr(choice.message, "content", None) # type: ignore + content = getattr(choice.message, "content", None) if content is None: continue if isinstance(content, str): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bd5b5309e0..3a3ce06f72 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8361,6 +8361,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) prompt = request.prompt messages = request.messages contents = request.contents + tools = request.tools + system = request.system ######################################################### # Validate request @@ -8421,6 +8423,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) contents=contents, deployment=deployment, request_model=request.model, + tools=tools, + system=system, ) ######################################################### # Transfrom the Response to the well known format diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dda32d9838..7290057c95 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1651,6 +1651,7 @@ class StreamingChatCompletionChunk(OpenAIChatCompletionChunk): super().__init__(**kwargs) + class ModelResponseBase(OpenAIObject): id: str """A unique identifier for the completion.""" @@ -1759,7 +1760,7 @@ class ModelResponseStream(ModelResponseBase): class ModelResponse(ModelResponseBase): - choices: List[Union[Choices, StreamingChoices]] + choices: List[Choices] """The list of completion choices the model generated for the input prompt.""" def __init__( # noqa: PLR0915 @@ -1778,44 +1779,27 @@ class ModelResponse(ModelResponseBase): _response_headers=None, **params, ) -> None: - if stream is not None and stream is True: - object = "chat.completion.chunk" - if choices is not None and isinstance(choices, list): - new_choices = [] - for choice in choices: - _new_choice = None - if isinstance(choice, StreamingChoices): - _new_choice = choice - elif isinstance(choice, dict): - _new_choice = StreamingChoices(**choice) - elif isinstance(choice, BaseModel): - _new_choice = StreamingChoices(**choice.model_dump()) - new_choices.append(_new_choice) - choices = new_choices - else: - choices = [StreamingChoices()] + object = "chat.completion" + if choices is not None and isinstance(choices, list): + new_choices = [] + for choice in choices: + if isinstance(choice, Choices): + _new_choice = choice # type: ignore + elif isinstance(choice, dict): + _new_choice = Choices(**choice) # type: ignore + elif isinstance(choice, BaseModel): + dump = ( + choice.model_dump() + if hasattr(choice, "model_dump") + else choice.dict() + ) + _new_choice = Choices(**dump) # type: ignore + else: + _new_choice = choice + new_choices.append(_new_choice) + choices = new_choices else: - object = "chat.completion" - if choices is not None and isinstance(choices, list): - new_choices = [] - for choice in choices: - if isinstance(choice, Choices): - _new_choice = choice # type: ignore - elif isinstance(choice, dict): - _new_choice = Choices(**choice) # type: ignore - elif isinstance(choice, BaseModel): - dump = ( - choice.model_dump() - if hasattr(choice, "model_dump") - else choice.dict() - ) - _new_choice = Choices(**dump) # type: ignore - else: - _new_choice = choice - new_choices.append(_new_choice) - choices = new_choices - else: - choices = [Choices()] + choices = [Choices()] if id is None: id = _generate_id() else: diff --git a/litellm/utils.py b/litellm/utils.py index 81e772b176..24da930b96 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2802,8 +2802,8 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 litellm.anthropic_models.add(key) elif value.get("litellm_provider") == "openrouter": split_string = key.split("/", 1) - if key not in litellm.openrouter_models: - litellm.openrouter_models.add(split_string[1]) + if split_string[-1] not in litellm.openrouter_models: + litellm.openrouter_models.add(split_string[-1]) elif value.get("litellm_provider") == "vercel_ai_gateway": if key not in litellm.vercel_ai_gateway_models: litellm.vercel_ai_gateway_models.add(key) @@ -4963,9 +4963,7 @@ 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[Choices], List[StreamingChoices]] = response_obj.choices # Use list accumulation to avoid O(n^2) string concatenation across choices response_parts: List[str] = [] @@ -7384,9 +7382,9 @@ def _get_base_model_from_metadata(model_call_details=None): class ModelResponseIterator: def __init__(self, model_response: ModelResponse, convert_to_delta: bool = False): if convert_to_delta is True: - self.model_response = ModelResponse(stream=True) - _delta = self.model_response.choices[0].delta # type: ignore - _delta.content = model_response.choices[0].message.content # type: ignore + _stream_response = ModelResponseStream() + _stream_response.choices[0].delta.content = model_response.choices[0].message.content # type: ignore + self.model_response: Union[ModelResponse, ModelResponseStream] = _stream_response else: self.model_response = model_response self.is_done = False diff --git a/tests/litellm/integrations/helicone/test_helicone_gemini.py b/tests/litellm/integrations/helicone/test_helicone_gemini.py new file mode 100644 index 0000000000..f42a701613 --- /dev/null +++ b/tests/litellm/integrations/helicone/test_helicone_gemini.py @@ -0,0 +1,64 @@ +""" +Test HeliconeLogger Gemini/Vertex AI support. +Fixes: https://github.com/BerriAI/litellm/issues/19093 +""" + +import pytest + + +def test_helicone_gemini_model_in_list(): + """ + Test that Gemini models are in the helicone_model_list. + """ + from litellm.integrations.helicone import HeliconeLogger + + logger = HeliconeLogger() + + # Test that "gemini" is in the model list + assert "gemini" in logger.helicone_model_list, "gemini should be in helicone_model_list" + + +def test_helicone_gemini_models_recognized(): + """ + Test that Gemini models are recognized and not replaced with gpt-3.5-turbo. + """ + from litellm.integrations.helicone import HeliconeLogger + + logger = HeliconeLogger() + + test_models = ["gemini-1.5-pro", "gemini-2.0-flash", "vertex_ai/gemini-1.5-flash"] + for model in test_models: + is_recognized = any( + accepted_model in model + for accepted_model in logger.helicone_model_list + ) + assert is_recognized, f"{model} should be recognized by helicone_model_list" + + +def test_helicone_vertex_ai_models_recognized(): + """ + Test that Vertex AI models (GLM, DeepSeek, etc.) are recognized via custom_llm_provider. + """ + # Test models that don't contain "gemini" but are vertex_ai + test_models = [ + "vertex_ai/zai-org/glm-4.7-maas", + "vertex_ai/deepseek-ai/deepseek-v3", + "vertex_ai/meta/llama-3.1-405b", + ] + for model in test_models: + is_vertex_ai = model.startswith("vertex_ai/") + assert is_vertex_ai, f"{model} should be recognized as vertex_ai model" + + +def test_helicone_vertex_ai_via_custom_llm_provider(): + """ + Test that vertex_ai models are recognized when custom_llm_provider is set. + """ + # Models without vertex_ai/ prefix but with custom_llm_provider="vertex_ai" + test_cases = [ + ("zai-org/glm-4.7-maas", "vertex_ai"), + ("deepseek-ai/deepseek-v3", "vertex_ai"), + ] + for model, custom_llm_provider in test_cases: + is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") + assert is_vertex_ai, f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai" diff --git a/tests/litellm/llms/openrouter/__init__.py b/tests/litellm/llms/openrouter/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/litellm/llms/openrouter/test_openrouter_native_models.py b/tests/litellm/llms/openrouter/test_openrouter_native_models.py new file mode 100644 index 0000000000..a6d2608a78 --- /dev/null +++ b/tests/litellm/llms/openrouter/test_openrouter_native_models.py @@ -0,0 +1,75 @@ +""" +Tests for native OpenRouter model name handling in get_llm_provider. + +OpenRouter's native models (openrouter/auto, openrouter/free, +openrouter/bodybuilder) should not have their "openrouter/" prefix +stripped when passed to get_llm_provider(), since that prefix is part +of the actual model ID on OpenRouter's API. + +""" + +import pytest + +import litellm + + +class TestNativeOpenRouterModelsNotStripped: + """get_llm_provider must preserve native OpenRouter model names.""" + + @pytest.mark.parametrize( + "model", + [ + "openrouter/auto", + "openrouter/free", + "openrouter/bodybuilder", + ], + ) + def test_native_model_not_stripped(self, model): + """Native OpenRouter model IDs are returned as-is.""" + result_model, provider, _, _ = litellm.get_llm_provider(model=model) + assert result_model == model + assert provider == "openrouter" + + @pytest.mark.parametrize( + "model,expected_model", + [ + ("openrouter/openrouter/free", "openrouter/free"), + ("openrouter/openrouter/auto", "openrouter/auto"), + ("openrouter/openrouter/bodybuilder", "openrouter/bodybuilder"), + ], + ) + def test_double_prefixed_model_strips_once_to_native(self, model, expected_model): + """openrouter/openrouter/free strips to openrouter/free (not further).""" + result_model, provider, _, _ = litellm.get_llm_provider(model=model) + assert result_model == expected_model + assert provider == "openrouter" + + @pytest.mark.parametrize( + "model,expected_model", + [ + ("openrouter/openrouter/free", "openrouter/free"), + ("openrouter/openrouter/auto", "openrouter/auto"), + ], + ) + def test_full_round_trip_no_double_strip(self, model, expected_model): + """Simulates the bridge flow: two consecutive get_llm_provider calls.""" + # First call (in adapter/handler) + model_after_first, provider, _, _ = litellm.get_llm_provider(model=model) + assert model_after_first == expected_model + + # Second call (inside litellm.completion) + model_after_second, provider2, _, _ = litellm.get_llm_provider( + model=model_after_first + ) + # Should stay as native model, not stripped further + assert model_after_second == expected_model + assert provider2 == "openrouter" + + def test_regular_openrouter_model_still_strips_normally(self): + """Non-native models like openrouter/anthropic/claude-3-haiku still strip normally.""" + model, provider, _, _ = litellm.get_llm_provider( + model="openrouter/anthropic/claude-3-haiku" + ) + assert provider == "openrouter" + # Should strip the openrouter/ prefix + assert model == "anthropic/claude-3-haiku" diff --git a/tests/litellm/test_stream_chunk_builder_images.py b/tests/litellm/test_stream_chunk_builder_images.py index c51a14ede6..92fb0f93aa 100644 --- a/tests/litellm/test_stream_chunk_builder_images.py +++ b/tests/litellm/test_stream_chunk_builder_images.py @@ -72,7 +72,7 @@ def test_stream_chunk_builder_preserves_images(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) @@ -163,7 +163,7 @@ def test_stream_chunk_builder_preserves_multiple_images(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) @@ -230,7 +230,7 @@ def test_stream_chunk_builder_no_images(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py index 8224773aa4..ddb1546097 100644 --- a/tests/local_testing/test_stream_chunk_builder.py +++ b/tests/local_testing/test_stream_chunk_builder.py @@ -542,7 +542,7 @@ def test_stream_chunk_builder_multiple_tool_calls(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) print(f"Returned response: {response}") @@ -616,7 +616,7 @@ def test_stream_chunk_builder_openai_prompt_caching(): chunks: List[litellm.ModelResponse] = [] usage_obj = None for chunk in chat_completion: - chunks.append(litellm.ModelResponse(**chunk.model_dump(), stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk.model_dump())) print(f"chunks: {chunks}") @@ -661,7 +661,7 @@ def test_stream_chunk_builder_openai_audio_output_usage(): chunks = [] for chunk in completion: - chunks.append(litellm.ModelResponse(**chunk.model_dump(), stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk.model_dump())) usage_obj: Optional[litellm.Usage] = None diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index b3f13e8a4b..50915e9993 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -393,7 +393,7 @@ def test_completion_azure_stream_content_filter_no_delta(): chunk_list = [] for chunk in chunks: - new_chunk = litellm.ModelResponse(stream=True, id=chunk["id"]) + new_chunk = litellm.ModelResponseStream(id=chunk["id"]) if "choices" in chunk and isinstance(chunk["choices"], list): new_choices = [] for choice in chunk["choices"]: @@ -3026,7 +3026,7 @@ def test_unit_test_custom_stream_wrapper(): {"index": 0, "delta": {"content": "How are you?"}, "finish_reason": "stop"} ], } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) @@ -3223,7 +3223,7 @@ def test_unit_test_custom_stream_wrapper_openai(): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) @@ -3457,7 +3457,7 @@ def test_aamazing_unit_test_custom_stream_wrapper_n(): chunk_list = [] for chunk in chunks: - new_chunk = litellm.ModelResponse(stream=True, id=chunk["id"]) + new_chunk = litellm.ModelResponseStream(id=chunk["id"]) if "choices" in chunk and isinstance(chunk["choices"], list): print("INSIDE CHUNK CHOICES!") new_choices = [] @@ -3541,7 +3541,7 @@ def test_unit_test_custom_stream_wrapper_function_call(): "system_fingerprint": "fp_44709d6fcb", "choices": [{"index": 0, "delta": delta, "finish_reason": "stop"}], } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) @@ -3651,7 +3651,7 @@ def test_unit_test_perplexity_citations_chunk(): } ], } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 3021fff9a2..f696859b08 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1318,3 +1318,94 @@ def test_transform_response_preserves_annotations(): assert result.usage.total_tokens == 30 print("✓ Annotations from Responses API are correctly preserved in Chat Completions format") + + +# ============================================================================= +# Tests for issue #21331: Parallel tool call indices in streaming +# ============================================================================= + + +def test_streaming_parallel_tool_calls_have_distinct_indices(): + """ + Test that parallel tool calls get distinct indices matching output_index + from the Responses API streaming chunks. + + Regression test for issue #21331 where all tool calls were emitted with + index=0, making it impossible to distinguish parallel calls. + """ + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + # Simulate two parallel tool calls with output_index 0 and 1 + chunks = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "type": "function_call", + "id": "fc_001", + "call_id": "call_abc", + "name": "get_weather", + "arguments": "", + }, + }, + { + "type": "response.function_call_arguments.delta", + "output_index": 0, + "item_id": "fc_001", + "delta": '{"city": "SF"}', + }, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "function_call", + "id": "fc_001", + "call_id": "call_abc", + "name": "get_weather", + "arguments": '{"city": "SF"}', + }, + }, + { + "type": "response.output_item.added", + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_002", + "call_id": "call_def", + "name": "get_weather", + "arguments": "", + }, + }, + { + "type": "response.function_call_arguments.delta", + "output_index": 1, + "item_id": "fc_002", + "delta": '{"city": "NY"}', + }, + { + "type": "response.output_item.done", + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_002", + "call_id": "call_def", + "name": "get_weather", + "arguments": '{"city": "NY"}', + }, + }, + ] + + for chunk in chunks: + result = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( + chunk + ) + expected_index = chunk["output_index"] + for choice in result.choices: + if choice.delta.tool_calls: + for tc in choice.delta.tool_calls: + assert tc.index == expected_index, ( + f"Event {chunk['type']}: expected tool_call.index={expected_index}, " + f"got {tc.index}" + ) diff --git a/tests/test_litellm/images/test_image_edit_utils.py b/tests/test_litellm/images/test_image_edit_utils.py index 56d8e48405..7a950375d3 100644 --- a/tests/test_litellm/images/test_image_edit_utils.py +++ b/tests/test_litellm/images/test_image_edit_utils.py @@ -5,6 +5,7 @@ import pytest import litellm from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.types.images.main import ImageEditOptionalRequestParams @@ -168,3 +169,92 @@ class TestImageEditRequestUtilsDropParams: assert "size" in result assert "quality" not in result assert "unsupported_param" not in result + + +class TestImageEditCustomPricing: + """ + Regression tests for https://github.com/BerriAI/litellm/issues/22244 + + image_edit must forward model_info and metadata into litellm_params + when calling update_environment_variables, so that custom pricing + detection works after PR #20679 stripped custom pricing fields from + the shared backend model key. + """ + + def test_image_edit_passes_model_info_to_logging(self): + """ + When the router provides model_info with custom pricing fields, + image_edit should include model_info and metadata in litellm_params. + """ + from litellm.images.main import image_edit + + custom_model_info = { + "id": "test-deployment-id", + "input_cost_per_image": 0.00676128, + "mode": "image_generation", + } + custom_metadata = { + "model_info": custom_model_info, + } + + captured_litellm_params = {} + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + + original_update = mock_logging_obj.update_environment_variables + + def capturing_update(**kwargs): + captured_litellm_params.update(kwargs.get("litellm_params", {})) + return original_update(**kwargs) + + mock_logging_obj.update_environment_variables = capturing_update + + with patch( + "litellm.images.main.get_llm_provider", + return_value=("test-model", "openai", None, None), + ), patch( + "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config", + return_value=MagicMock(), + ), patch( + "litellm.images.main._get_ImageEditRequestUtils", + return_value=MagicMock( + get_requested_image_edit_optional_param=MagicMock(return_value={}), + get_optional_params_image_edit=MagicMock(return_value={}), + ), + ), patch( + "litellm.images.main.base_llm_http_handler" + ) as mock_handler: + mock_handler.image_edit_handler.return_value = MagicMock() + + try: + image_edit( + image=b"fake-image-data", + prompt="test prompt", + model="openai/test-model", + litellm_logging_obj=mock_logging_obj, + model_info=custom_model_info, + metadata=custom_metadata, + ) + except Exception: + pass + + assert "model_info" in captured_litellm_params + assert captured_litellm_params["model_info"] == custom_model_info + assert "metadata" in captured_litellm_params + assert captured_litellm_params["metadata"] == custom_metadata + + def test_custom_pricing_detected_from_model_info_in_metadata(self): + litellm_params = { + "metadata": { + "model_info": { + "id": "deployment-id", + "input_cost_per_image": 0.00676128, + }, + }, + } + assert use_custom_pricing_for_model(litellm_params) is True + + def test_custom_pricing_not_detected_without_model_info(self): + litellm_params = {"litellm_call_id": "test-call-id"} + assert use_custom_pricing_for_model(litellm_params) is False diff --git a/tests/test_litellm/images/test_image_generation_extra_headers.py b/tests/test_litellm/images/test_image_generation_extra_headers.py new file mode 100644 index 0000000000..d1cbe5fc69 --- /dev/null +++ b/tests/test_litellm/images/test_image_generation_extra_headers.py @@ -0,0 +1,84 @@ +""" +Unit test for https://github.com/BerriAI/litellm/issues/22285 + +Verifies that extra_headers passed to image_generation() are forwarded +to the OpenAI SDK on the openai/litellm_proxy/openai_compatible_providers +code paths. +""" + +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.images.main import image_generation + + +class TestImageGenerationExtraHeaders: + """Test that extra_headers are forwarded on the OpenAI code path.""" + + @patch("litellm.images.main.openai_chat_completions") + def test_extra_headers_forwarded_to_openai_image_generation( + self, mock_openai_chat_completions + ): + """ + extra_headers passed to image_generation() should appear in + optional_params["extra_headers"] when the provider is openai. + """ + mock_image_response = litellm.utils.ImageResponse( + created=1234567890, + data=[{"url": "https://example.com/image.png"}], + ) + mock_openai_chat_completions.image_generation.return_value = ( + mock_image_response + ) + + extra_headers = {"traceparent": "00-abc123-def456-01", "X-Custom": "value"} + + image_generation( + model="openai/dall-e-3", + prompt="A red circle", + extra_headers=extra_headers, + ) + + mock_openai_chat_completions.image_generation.assert_called_once() + call_kwargs = mock_openai_chat_completions.image_generation.call_args + optional_params = call_kwargs.kwargs.get( + "optional_params", call_kwargs[1].get("optional_params", {}) + ) + + assert "extra_headers" in optional_params + assert optional_params["extra_headers"] == extra_headers + + @patch("litellm.images.main.openai_chat_completions") + def test_no_extra_headers_when_not_provided( + self, mock_openai_chat_completions + ): + """ + When extra_headers is not passed, optional_params should not + contain extra_headers. + """ + mock_image_response = litellm.utils.ImageResponse( + created=1234567890, + data=[{"url": "https://example.com/image.png"}], + ) + mock_openai_chat_completions.image_generation.return_value = ( + mock_image_response + ) + + image_generation( + model="openai/dall-e-3", + prompt="A red circle", + ) + + mock_openai_chat_completions.image_generation.assert_called_once() + call_kwargs = mock_openai_chat_completions.image_generation.call_args + optional_params = call_kwargs.kwargs.get( + "optional_params", call_kwargs[1].get("optional_params", {}) + ) + + assert "extra_headers" not in optional_params diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 73ecaa20a2..153ca5d5aa 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1187,6 +1187,93 @@ def test_is_chunk_non_empty_with_valid_tool_calls( ) +def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): + """ + Test that provider-reported usage from a post-finish_reason chunk + is surfaced in _hidden_params even when stream_options is NOT set. + + Reproduces issue #20760: OpenRouter sends a final chunk with usage data + after the finish_reason chunk. The hidden_params["usage"] on the last + user-visible chunk was being calculated before this usage chunk arrived, + resulting in zeros. The fix recalculates it in the StopIteration handler + after stream_chunk_builder processes all chunks. + """ + # Simulate OpenRouter's actual streaming pattern: + # 1) content chunk + # 2) finish_reason chunk (content="") + # 3) usage chunk (content="", finish_reason=None, usage={...}) + chunks = [ + ModelResponseStream( + id="gen-abc", + object="chat.completion.chunk", + created=1000000, + model="openrouter/openai/gpt-4o-mini", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content="Hello"), + finish_reason=None, + ) + ], + ), + ModelResponseStream( + id="gen-abc", + object="chat.completion.chunk", + created=1000000, + model="openrouter/openai/gpt-4o-mini", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=""), + finish_reason="stop", + ) + ], + ), + ModelResponseStream( + id="gen-abc", + object="chat.completion.chunk", + created=1000000, + model="openrouter/openai/gpt-4o-mini", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=""), + finish_reason=None, + ) + ], + usage=Usage( + prompt_tokens=20, + completion_tokens=135, + total_tokens=155, + ), + ), + ] + + # Create a CustomStreamWrapper with NO stream_options + wrapper = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="openrouter/openai/gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openrouter", + stream_options=None, + ) + + # Consume the stream + collected = [] + for chunk in wrapper: + collected.append(chunk) + + # The last user-visible chunk's _hidden_params["usage"] should + # contain the provider-reported values, not zeros. + last_chunk = collected[-1] + hidden_usage = last_chunk._hidden_params.get("usage") + assert hidden_usage is not None, "Expected usage in _hidden_params" + assert hidden_usage.prompt_tokens == 20, ( + f"Expected prompt_tokens=20 from provider, got {hidden_usage.prompt_tokens}" + ) + assert hidden_usage.completion_tokens == 135, ( + f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" + ) @pytest.mark.asyncio async def test_custom_stream_wrapper_aclose(): """Test that aclose() delegates to the underlying completion_stream's aclose()""" diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py new file mode 100644 index 0000000000..e982f735fd --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -0,0 +1,92 @@ +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) + + +def test_transform_basic_request(): + """Test basic request with only model and messages.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + ) + + assert result == { + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "Hello"}], + } + + +def test_transform_includes_system(): + """Test that system prompt is included when provided.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + system="You are a helpful assistant.", + ) + + assert result["system"] == "You are a helpful assistant." + assert result["model"] == "claude-3-5-sonnet" + assert result["messages"] == [{"role": "user", "content": "Hello"}] + + +def test_transform_includes_tools(): + """Test that tools are included when provided.""" + config = AnthropicCountTokensConfig() + + tools = [ + { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ] + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + tools=tools, + ) + + assert result["tools"] == tools + + +def test_transform_includes_system_and_tools(): + """Test that both system and tools are included together.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + system="Be helpful", + tools=[{"name": "my_tool", "input_schema": {"type": "object"}}], + ) + + assert "system" in result + assert "tools" in result + assert "messages" in result + assert "model" in result + + +def test_transform_no_system_no_tools(): + """Test that None system/tools are not included.""" + config = AnthropicCountTokensConfig() + + result = config.transform_request_to_count_tokens( + model="claude-3-5-sonnet", + messages=[{"role": "user", "content": "Hello"}], + system=None, + tools=None, + ) + + assert "system" not in result + assert "tools" not in result diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index ed8d6e1b35..699b67911d 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -34,3 +34,123 @@ def test_transform_anthropic_to_bedrock_request(): assert "input" in result assert "converse" in result["input"] assert "messages" in result["input"]["converse"] + + +def test_transform_includes_system_prompt(): + """Test that system prompt is included in Bedrock converse format.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "system": "You are a helpful assistant.", + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "system" in converse + assert converse["system"] == [{"text": "You are a helpful assistant."}] + + +def test_transform_includes_system_prompt_as_list(): + """Test that system prompt as list of blocks is handled.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "system": [{"type": "text", "text": "Block 1"}, {"type": "text", "text": "Block 2"}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert converse["system"] == [{"text": "Block 1"}, {"text": "Block 2"}] + + +def test_transform_includes_tools(): + """Test that tools are transformed to Bedrock toolConfig format.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "tools": [ + { + "name": "read_file", + "description": "Read a file", + "input_schema": { + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + }, + } + ], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "toolConfig" in converse + tools = converse["toolConfig"]["tools"] + assert len(tools) == 1 + assert tools[0]["toolSpec"]["name"] == "read_file" + assert tools[0]["toolSpec"]["description"] == "Read a file" + assert tools[0]["toolSpec"]["inputSchema"]["json"]["type"] == "object" + + +def test_transform_includes_system_and_tools_together(): + """Test that both system and tools are included together.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "system": "Be helpful", + "tools": [ + {"name": "my_tool", "description": "A tool", "input_schema": {"type": "object", "properties": {}}}, + ], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "system" in converse + assert "toolConfig" in converse + assert "messages" in converse + + +def test_transform_no_system_no_tools(): + """Test that missing system and tools don't add extra keys.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + converse = result["input"]["converse"] + assert "system" not in converse + assert "toolConfig" not in converse + + +def test_tool_name_sanitization(): + """Test that tool names are sanitized for Bedrock requirements.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "tools": [ + {"name": "my-tool!", "description": "A tool", "input_schema": {"type": "object", "properties": {}}}, + ], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + tool_name = result["input"]["converse"]["toolConfig"]["tools"][0]["toolSpec"]["name"] + # Should be sanitized: only [a-zA-Z0-9_] + assert tool_name == "my_tool_" diff --git a/tests/test_litellm/llms/chatgpt/__init__.py b/tests/test_litellm/llms/chatgpt/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/chatgpt/chat/__init__.py b/tests/test_litellm/llms/chatgpt/chat/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py b/tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py new file mode 100644 index 0000000000..0e6e4580e4 --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py @@ -0,0 +1,195 @@ +""" +Tests for ChatGPTToolCallNormalizer. + +Verifies that non-spec-compliant tool_call chunks from the ChatGPT backend API +are normalized to match the OpenAI streaming spec: +- Correct index assignment for parallel tool calls +- Deduplication of "closing" chunks with repeated id/name +""" + +import pytest + +from litellm.llms.chatgpt.chat.streaming_utils import ChatGPTToolCallNormalizer +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, +) + + +def _make_chunk(tool_calls=None, content=None): + """Helper to build a ModelResponseStream chunk with tool_calls on the delta.""" + delta = Delta( + content=content, + role="assistant", + tool_calls=tool_calls, + ) + choice = StreamingChoices(delta=delta, index=0) + return ModelResponseStream(choices=[choice]) + + +def _make_tc(index=0, id=None, name=None, arguments=None): + """Helper to build a ChatCompletionDeltaToolCall.""" + func = Function(name=name, arguments=arguments) + return ChatCompletionDeltaToolCall( + index=index, + id=id, + function=func, + type="function" if id else None, + ) + + +class TestChatGPTToolCallNormalizer: + """Test that the normalizer fixes ChatGPT-style tool_call streaming issues.""" + + def test_single_tool_call_index_preserved(self): + """A single tool call should get index=0.""" + chunks = [ + _make_chunk(tool_calls=[_make_tc(index=0, id="call_1", name="get_weather")]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"loc')]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='ation": "NYC"}')]), + ] + normalizer = ChatGPTToolCallNormalizer(iter(chunks)) + results = list(normalizer) + + assert len(results) == 3 + assert results[0].choices[0].delta.tool_calls[0].index == 0 + assert results[0].choices[0].delta.tool_calls[0].id == "call_1" + assert results[1].choices[0].delta.tool_calls[0].index == 0 + assert results[2].choices[0].delta.tool_calls[0].index == 0 + + def test_parallel_tool_calls_get_correct_indices(self): + """ + ChatGPT sends all tool_calls with index=0. The normalizer should assign + sequential indices: 0 for the first, 1 for the second. + """ + chunks = [ + # First tool call: intro chunk with id + name + _make_chunk(tool_calls=[_make_tc(index=0, id="call_aaa", name="get_weather")]), + # First tool call: arguments streaming + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"location": "NYC"}')]), + # First tool call: duplicate closing chunk (id repeated) — should be skipped + _make_chunk(tool_calls=[_make_tc(index=0, id="call_aaa", name="get_weather")]), + # Second tool call: intro chunk with id + name (index=0 from ChatGPT) + _make_chunk(tool_calls=[_make_tc(index=0, id="call_bbb", name="get_time")]), + # Second tool call: arguments streaming + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"tz": "EST"}')]), + # Second tool call: duplicate closing chunk — should be skipped + _make_chunk(tool_calls=[_make_tc(index=0, id="call_bbb", name="get_time")]), + ] + + normalizer = ChatGPTToolCallNormalizer(iter(chunks)) + results = list(normalizer) + + # 2 duplicate chunks should be skipped → 4 results + assert len(results) == 4 + + # First tool call chunks should have index=0 + assert results[0].choices[0].delta.tool_calls[0].index == 0 + assert results[0].choices[0].delta.tool_calls[0].id == "call_aaa" + assert results[1].choices[0].delta.tool_calls[0].index == 0 + + # Second tool call chunks should have index=1 + assert results[2].choices[0].delta.tool_calls[0].index == 1 + assert results[2].choices[0].delta.tool_calls[0].id == "call_bbb" + assert results[3].choices[0].delta.tool_calls[0].index == 1 + + def test_non_tool_call_chunks_pass_through(self): + """Chunks without tool_calls should pass through unchanged.""" + chunks = [ + _make_chunk(content="Hello"), + _make_chunk(content=" world"), + ] + normalizer = ChatGPTToolCallNormalizer(iter(chunks)) + results = list(normalizer) + + assert len(results) == 2 + assert results[0].choices[0].delta.content == "Hello" + assert results[1].choices[0].delta.content == " world" + + def test_empty_choices_pass_through(self): + """Chunks with empty choices should pass through.""" + chunk = ModelResponseStream(choices=[]) + normalizer = ChatGPTToolCallNormalizer(iter([chunk])) + results = list(normalizer) + + assert len(results) == 1 + + def test_three_parallel_tool_calls(self): + """Three parallel tool calls should get indices 0, 1, 2.""" + chunks = [ + _make_chunk(tool_calls=[_make_tc(index=0, id="call_1", name="fn_a")]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"a":1}')]), + _make_chunk(tool_calls=[_make_tc(index=0, id="call_2", name="fn_b")]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"b":2}')]), + _make_chunk(tool_calls=[_make_tc(index=0, id="call_3", name="fn_c")]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"c":3}')]), + ] + + normalizer = ChatGPTToolCallNormalizer(iter(chunks)) + results = list(normalizer) + + assert len(results) == 6 + # First tool call + assert results[0].choices[0].delta.tool_calls[0].index == 0 + assert results[1].choices[0].delta.tool_calls[0].index == 0 + # Second tool call + assert results[2].choices[0].delta.tool_calls[0].index == 1 + assert results[3].choices[0].delta.tool_calls[0].index == 1 + # Third tool call + assert results[4].choices[0].delta.tool_calls[0].index == 2 + assert results[5].choices[0].delta.tool_calls[0].index == 2 + + def test_all_duplicates_skipped(self): + """If a chunk contains only duplicate tool_calls, the entire chunk is skipped.""" + chunks = [ + _make_chunk(tool_calls=[_make_tc(index=0, id="call_x", name="fn")]), + # Duplicate — same id seen before + _make_chunk(tool_calls=[_make_tc(index=0, id="call_x", name="fn")]), + ] + + normalizer = ChatGPTToolCallNormalizer(iter(chunks)) + results = list(normalizer) + + assert len(results) == 1 + assert results[0].choices[0].delta.tool_calls[0].id == "call_x" + + @pytest.mark.asyncio + async def test_async_iteration(self): + """The normalizer should work with async iteration.""" + + async def async_gen(): + chunks = [ + _make_chunk(tool_calls=[_make_tc(index=0, id="call_a", name="fn_a")]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"x":1}')]), + _make_chunk(tool_calls=[_make_tc(index=0, id="call_b", name="fn_b")]), + _make_chunk(tool_calls=[_make_tc(index=0, arguments='{"y":2}')]), + ] + for c in chunks: + yield c + + normalizer = ChatGPTToolCallNormalizer(async_gen()) + results = [] + async for chunk in normalizer: + results.append(chunk) + + assert len(results) == 4 + assert results[0].choices[0].delta.tool_calls[0].index == 0 + assert results[2].choices[0].delta.tool_calls[0].index == 1 + + def test_getattr_proxies_to_stream(self): + """Attribute access should be proxied to the underlying stream.""" + + class FakeStream: + custom_attr = "test_value" + + def __iter__(self): + return iter([]) + + def __next__(self): + raise StopIteration + + normalizer = ChatGPTToolCallNormalizer(FakeStream()) + assert normalizer.custom_attr == "test_value" diff --git a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py index 62fcec04c1..345186e8a6 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py @@ -309,4 +309,99 @@ class TestMoonshotConfig: # Check that no extra message was added assert len(result["messages"]) == 1 - assert result["messages"][0]["content"] == "What's the weather?" \ No newline at end of file + assert result["messages"][0]["content"] == "What's the weather?" + + def test_transform_messages_preserves_image_url_content(self): + """Test that messages with image_url blocks are NOT flattened to strings. + + Multimodal models like kimi-k2.5 accept the standard OpenAI content + array with non-text blocks. When any message contains a non-text part, + the content array must be preserved so the payload reaches the API. + """ + config = MoonshotChatConfig() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + }, + ], + } + ] + + result = config.transform_request( + model="kimi-k2.5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + # Content must remain a list (not flattened to a string) + assert isinstance(result["messages"][0]["content"], list) + assert len(result["messages"][0]["content"]) == 2 + assert result["messages"][0]["content"][0]["type"] == "text" + assert result["messages"][0]["content"][1]["type"] == "image_url" + + def test_transform_messages_preserves_non_text_content(self): + """Test that any non-text content type (input_audio, video_url, file, + etc.) also prevents flattening, matching the OpenAI content spec.""" + config = MoonshotChatConfig() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Transcribe this audio"}, + { + "type": "input_audio", + "input_audio": {"data": "base64data", "format": "wav"}, + }, + ], + } + ] + + result = config.transform_request( + model="kimi-k2.5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert isinstance(result["messages"][0]["content"], list) + assert len(result["messages"][0]["content"]) == 2 + assert result["messages"][0]["content"][1]["type"] == "input_audio" + + def test_transform_messages_flattens_text_only_content(self): + """Test that text-only content arrays ARE flattened to strings. + + For text-only requests, Moonshot expects plain string content. + The content list should be converted to a single string. + """ + config = MoonshotChatConfig() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello, how are you?"}, + ], + } + ] + + result = config.transform_request( + model="moonshot-v1-8k", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + # Content should be flattened to a plain string + assert isinstance(result["messages"][0]["content"], str) + assert result["messages"][0]["content"] == "Hello, how are you?" \ No newline at end of file diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 386f264a4d..776e519ffc 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -414,3 +414,174 @@ def test_gpt5_2_allows_reasoning_effort_xhigh(config: OpenAIConfig): drop_params=False, ) assert params["reasoning_effort"] == "xhigh" + + +# GPT-5-Search specific tests +def test_gpt5_search_model_detection(gpt5_config: OpenAIGPT5Config): + """Test that GPT-5 search models are correctly detected.""" + assert gpt5_config.is_model_gpt_5_search_model("gpt-5-search-api") + assert gpt5_config.is_model_gpt_5_search_model("gpt-5-search-mini-api") + + assert not gpt5_config.is_model_gpt_5_search_model("gpt-5") + assert not gpt5_config.is_model_gpt_5_search_model("gpt-5-codex") + assert not gpt5_config.is_model_gpt_5_search_model("gpt-5-mini") + + +def test_gpt5_search_supported_params(gpt5_config: OpenAIGPT5Config): + """Test that search models do NOT list reasoning/tool params as supported.""" + supported = gpt5_config.get_supported_openai_params(model="gpt-5-search-api") + rejected = [ + "logit_bias", + "modalities", + "prediction", + "n", + "seed", + "temperature", + "tools", + "tool_choice", + "function_call", + "functions", + "parallel_tool_calls", + "audio", + "reasoning_effort", + ] + for param in rejected: + assert param not in supported, f"{param} should not be supported for search models" + + +def test_gpt5_search_has_expected_params(gpt5_config: OpenAIGPT5Config): + """Test that search models DO list the correct supported params.""" + supported = gpt5_config.get_supported_openai_params(model="gpt-5-search-api") + expected = [ + "max_tokens", + "max_completion_tokens", + "stream", + "stream_options", + "web_search_options", + "service_tier", + "response_format", + "user", + "store", + "verbosity", + "extra_headers", + ] + for param in expected: + assert param in supported, f"{param} should be supported for search models" + + +def test_gpt5_search_maps_max_tokens(config: OpenAIConfig): + """Test that search models map max_tokens -> max_completion_tokens.""" + params = config.map_openai_params( + non_default_params={"max_tokens": 200}, + optional_params={}, + model="gpt-5-search-api", + drop_params=False, + ) + assert params["max_completion_tokens"] == 200 + assert "max_tokens" not in params + + +def test_gpt5_search_drops_unsupported_params(config: OpenAIConfig): + """Test that search models drop unsupported params via map_openai_params.""" + params = config.map_openai_params( + non_default_params={"n": 2, "temperature": 0.7, "tools": [{"type": "function"}]}, + optional_params={}, + model="gpt-5-search-api", + drop_params=True, + ) + assert "n" not in params + assert "temperature" not in params + assert "tools" not in params +# GPT-5 unsupported params audit (validated via direct API calls) +def test_gpt5_rejects_params_unsupported_by_openai(config: OpenAIConfig): + """Params that OpenAI rejects for all GPT-5 reasoning models.""" + rejected_params = [ + "logit_bias", + "modalities", + "prediction", + "audio", + "web_search_options", + ] + for model in ["gpt-5", "gpt-5-mini", "gpt-5-codex", "gpt-5.1", "gpt-5.2"]: + supported = config.get_supported_openai_params(model=model) + for param in rejected_params: + assert param not in supported, ( + f"{param} should not be supported for {model}" + ) + + +def test_gpt5_1_supports_logprobs_top_p(config: OpenAIConfig): + """gpt-5.1/5.2 support logprobs, top_p, top_logprobs when reasoning_effort='none'.""" + for model in ["gpt-5.1", "gpt-5.2"]: + supported = config.get_supported_openai_params(model=model) + assert "logprobs" in supported, f"logprobs should be supported for {model}" + assert "top_p" in supported, f"top_p should be supported for {model}" + assert "top_logprobs" in supported, f"top_logprobs should be supported for {model}" + + +def test_gpt5_base_does_not_support_logprobs_top_p(config: OpenAIConfig): + """Base gpt-5/gpt-5-mini do NOT support logprobs, top_p, top_logprobs.""" + for model in ["gpt-5", "gpt-5-mini", "gpt-5-codex"]: + supported = config.get_supported_openai_params(model=model) + assert "logprobs" not in supported, f"logprobs should not be supported for {model}" + assert "top_p" not in supported, f"top_p should not be supported for {model}" + assert "top_logprobs" not in supported, f"top_logprobs should not be supported for {model}" + + +def test_gpt5_1_logprobs_passthrough(config: OpenAIConfig): + """Test that logprobs passes through for gpt-5.1.""" + params = config.map_openai_params( + non_default_params={"logprobs": True, "top_logprobs": 3}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + assert params["logprobs"] is True + assert params["top_logprobs"] == 3 + + +def test_gpt5_1_top_p_passthrough(config: OpenAIConfig): + """Test that top_p passes through for gpt-5.1.""" + params = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + assert params["top_p"] == 0.9 + + +def test_gpt5_1_logprobs_rejected_with_reasoning_effort(config: OpenAIConfig): + """logprobs/top_p/top_logprobs are rejected when reasoning_effort != 'none'.""" + for effort in ["low", "medium", "high"]: + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"logprobs": True, "reasoning_effort": effort}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +def test_gpt5_1_top_p_rejected_with_reasoning_effort(config: OpenAIConfig): + """top_p is rejected when reasoning_effort != 'none'.""" + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"top_p": 0.9, "reasoning_effort": "high"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +def test_gpt5_1_logprobs_dropped_with_reasoning_effort(config: OpenAIConfig): + """logprobs/top_p are dropped when reasoning_effort != 'none' and drop_params=True.""" + params = config.map_openai_params( + non_default_params={"logprobs": True, "top_p": 0.9, "reasoning_effort": "high"}, + optional_params={}, + model="gpt-5.1", + drop_params=True, + ) + assert "logprobs" not in params + assert "top_p" not in params + assert params["reasoning_effort"] == "high" diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 6736eaffeb..350fd75d3d 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -65,6 +65,42 @@ class TestVertexAIGeminiImageGenerationConfig: assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + def test_get_supported_openai_params_includes_native_gemini_params(self): + """Test that native Gemini imageConfig params are supported""" + supported = self.config.get_supported_openai_params("gemini-3-pro-image-preview") + assert "aspectRatio" in supported + assert "aspect_ratio" in supported + assert "imageSize" in supported + assert "image_size" in supported + + def test_map_openai_params_aspect_ratio_camel_case(self): + """Test mapping native aspectRatio parameter""" + result = self.config.map_openai_params( + {"aspectRatio": "9:16"}, {}, "gemini-3-pro-image-preview", False + ) + assert result["aspectRatio"] == "9:16" + + def test_map_openai_params_aspect_ratio_snake_case(self): + """Test mapping native aspect_ratio parameter""" + result = self.config.map_openai_params( + {"aspect_ratio": "16:9"}, {}, "gemini-3-pro-image-preview", False + ) + assert result["aspectRatio"] == "16:9" + + def test_map_openai_params_image_size_camel_case(self): + """Test mapping native imageSize parameter""" + result = self.config.map_openai_params( + {"imageSize": "4K"}, {}, "gemini-3-pro-image-preview", False + ) + assert result["imageSize"] == "4K" + + def test_map_openai_params_image_size_snake_case(self): + """Test mapping native image_size parameter""" + result = self.config.map_openai_params( + {"image_size": "2K"}, {}, "gemini-3-pro-image-preview", False + ) + assert result["imageSize"] == "2K" + def test_transform_image_generation_request_basic(self): """Test basic request transformation""" request = self.config.transform_image_generation_request( diff --git a/tests/test_litellm/test_model_response_normalization.py b/tests/test_litellm/test_model_response_normalization.py index 57281d3c1f..85b9fc1450 100644 --- a/tests/test_litellm/test_model_response_normalization.py +++ b/tests/test_litellm/test_model_response_normalization.py @@ -2,7 +2,14 @@ import warnings import pytest -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) def test_modelresponse_normalizes_openai_base_models() -> None: @@ -59,3 +66,63 @@ def test_modelresponse_serialization_avoids_pydantic_warnings() -> None: or "Pydantic serializer warnings" in str(w.message) for w in captured ) + + +def test_modelresponse_model_dump_json_no_pydantic_warnings() -> None: + """model_dump_json() and model_dump() should not trigger any Pydantic + serialization warnings now that choices is List[Choices] (no Union).""" + response = ModelResponse( + model="test-model", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="hello", role="assistant"), + ) + ], + ) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + _ = response.model_dump_json() + _ = response.model_dump() + _ = response.model_dump(exclude_none=True) + + pydantic_warnings = [ + w + for w in captured + if "PydanticSerializationUnexpectedValue" in str(w.message) + or "Pydantic serializer warnings" in str(w.message) + ] + assert pydantic_warnings == [], ( + f"Unexpected Pydantic serialization warnings: {pydantic_warnings}" + ) + + +def test_streaming_modelresponsestream_no_pydantic_warnings() -> None: + """Streaming responses use ModelResponseStream with List[StreamingChoices] + and should serialize without warnings.""" + response = ModelResponseStream( + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + _ = response.model_dump_json() + _ = response.model_dump() + + pydantic_warnings = [ + w + for w in captured + if "PydanticSerializationUnexpectedValue" in str(w.message) + or "Pydantic serializer warnings" in str(w.message) + ] + assert pydantic_warnings == [], ( + f"Unexpected Pydantic serialization warnings: {pydantic_warnings}" + ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 35cb290fcc..7f0b3b5b50 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2377,6 +2377,64 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() +def test_register_model_openrouter_without_slash(): + """ + Test that register_model handles openrouter models without '/' in the name. + + Fixes https://github.com/BerriAI/litellm/issues/18936 + + Previously, the code did `split_string[1]` which would fail with IndexError + when the model name didn't contain '/'. Now it uses `split_string[-1]` which + always works. + """ + # Clear any existing entries + litellm.openrouter_models.discard("my-custom-alias") + litellm.openrouter_models.discard("gpt-4") + litellm.openrouter_models.discard("openai/gpt-4") + + # Test 1: Model name without '/' (this was the bug - would raise IndexError) + litellm.register_model( + { + "my-custom-alias": { + "max_tokens": 8192, + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "litellm_provider": "openrouter", + "mode": "chat", + }, + } + ) + assert "my-custom-alias" in litellm.openrouter_models + + # Test 2: Model name with single '/' (openrouter/model format) + litellm.register_model( + { + "openrouter/gpt-4": { + "max_tokens": 8192, + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "litellm_provider": "openrouter", + "mode": "chat", + }, + } + ) + assert "gpt-4" in litellm.openrouter_models + + # Test 3: Model name with double '/' (openrouter/provider/model format) + litellm.register_model( + { + "openrouter/openai/gpt-4-turbo": { + "max_tokens": 8192, + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00002, + "litellm_provider": "openrouter", + "mode": "chat", + }, + } + ) + assert "openai/gpt-4-turbo" in litellm.openrouter_models + + def test_reasoning_content_preserved_in_text_completion_wrapper(): """Ensure reasoning_content is copied from delta to text_choices.""" chunk = ModelResponseStream(