From 542abc0429ab482131d0a1665c70b07dfed2979b Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 17 Jan 2026 18:00:00 -0300 Subject: [PATCH 01/32] fix(helicone): add Gemini/Vertex AI support to HeliconeLogger - Add "gemini" to helicone_model_list so Gemini models are recognized - Use /custom/v1/log endpoint for Gemini models instead of /oai/v1/log - Set correct provider_url for Google's generativelanguage API - Add unit test for Gemini model recognition Previously, Gemini models were logged as "gpt-3.5-turbo" with OpenAI as the provider, corrupting analytics. Now they log correctly with their actual model name and CUSTOM provider. --- litellm/integrations/helicone.py | 4 ++++ .../test_helicone_integration.py | 24 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index 198cbaf405..d64a589951 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -11,6 +11,7 @@ class HeliconeLogger: helicone_model_list = [ "gpt", "claude", + "gemini", "command-r", "command-r-plus", "command-light", @@ -151,6 +152,9 @@ class HeliconeLogger: if "claude" in model: 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" headers = { "Authorization": f"Bearer {self.key}", "Content-Type": "application/json", diff --git a/tests/local_testing/test_helicone_integration.py b/tests/local_testing/test_helicone_integration.py index ad8fe92d1e..6773a70ac6 100644 --- a/tests/local_testing/test_helicone_integration.py +++ b/tests/local_testing/test_helicone_integration.py @@ -162,3 +162,27 @@ def test_helicone_removes_otel_span_from_metadata(): assert result_metadata["other_metadata"] == "some_value" print("✅ Test passed: litellm_parent_otel_span was successfully removed from metadata") + + +def test_helicone_gemini_model_in_list(): + """ + Test that Gemini models are in the helicone_model_list and use the correct endpoint. + Fixes: https://github.com/BerriAI/litellm/issues/19093 + """ + 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" + + # Test that gemini models are recognized (not replaced with gpt-3.5-turbo) + 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" + + print("✅ Test passed: Gemini models are properly supported in HeliconeLogger") From cc39c71ac69daf586b3789acb9bd5ce06c6b77e5 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sat, 17 Jan 2026 18:07:40 -0300 Subject: [PATCH 02/32] test: move helicone gemini test to tests/litellm/ Move test to correct directory per PR template requirements. --- .../helicone/test_helicone_gemini.py | 35 +++++++++++++++++++ .../test_helicone_integration.py | 24 ------------- 2 files changed, 35 insertions(+), 24 deletions(-) create mode 100644 tests/litellm/integrations/helicone/test_helicone_gemini.py 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..37c50241a1 --- /dev/null +++ b/tests/litellm/integrations/helicone/test_helicone_gemini.py @@ -0,0 +1,35 @@ +""" +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" diff --git a/tests/local_testing/test_helicone_integration.py b/tests/local_testing/test_helicone_integration.py index 6773a70ac6..ad8fe92d1e 100644 --- a/tests/local_testing/test_helicone_integration.py +++ b/tests/local_testing/test_helicone_integration.py @@ -162,27 +162,3 @@ def test_helicone_removes_otel_span_from_metadata(): assert result_metadata["other_metadata"] == "some_value" print("✅ Test passed: litellm_parent_otel_span was successfully removed from metadata") - - -def test_helicone_gemini_model_in_list(): - """ - Test that Gemini models are in the helicone_model_list and use the correct endpoint. - Fixes: https://github.com/BerriAI/litellm/issues/19093 - """ - 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" - - # Test that gemini models are recognized (not replaced with gpt-3.5-turbo) - 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" - - print("✅ Test passed: Gemini models are properly supported in HeliconeLogger") From f76719750ba7bea62927b89a5a3b6756397fe4ee Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 19 Jan 2026 11:35:34 -0300 Subject: [PATCH 03/32] feat(helicone): add Vertex AI support for non-Gemini models Extends HeliconeLogger to properly log Vertex AI partner models (GLM, DeepSeek, etc.) that don't contain "gemini" in their name. Uses custom_llm_provider to detect vertex_ai. --- litellm/integrations/helicone.py | 14 +++++++-- .../helicone/test_helicone_gemini.py | 29 +++++++++++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index d64a589951..e5dc05e479 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -118,15 +118,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} @@ -135,7 +140,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 ) @@ -149,12 +154,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/tests/litellm/integrations/helicone/test_helicone_gemini.py b/tests/litellm/integrations/helicone/test_helicone_gemini.py index 37c50241a1..f42a701613 100644 --- a/tests/litellm/integrations/helicone/test_helicone_gemini.py +++ b/tests/litellm/integrations/helicone/test_helicone_gemini.py @@ -33,3 +33,32 @@ def test_helicone_gemini_models_recognized(): 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" From 25ff25d1efec311219f3fa29239f577fc2b2b394 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 26 Jan 2026 13:43:17 -0300 Subject: [PATCH 04/32] fix(containers): Fix Python 3.10 compatibility for OpenAIContainerConfig LiteLLM's pyproject.toml specifies `python = ">=3.9,<4.0"`, supporting Python 3.9 and above. However, the OpenAIContainerConfig class was failing to load on Python 3.10 with: TypeError: Cannot subclass typing.Any This occurred because `BaseContainerConfig = Any` was used as a runtime placeholder, and the class then inherited from it. In Python 3.11+, subclassing `typing.Any` is allowed (added to support dynamic classes like unittest.Mock where type checkers should skip verification). However, in Python 3.10 and earlier, this raises a TypeError. The fix imports the actual BaseContainerConfig class at runtime instead of using Any as a placeholder. Fixes #19727 --- litellm/llms/openai/containers/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index e67bfbe0c6..e8ab6190c7 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: BaseLLMException = _BaseLLMException else: LiteLLMLoggingObj = Any - BaseContainerConfig = Any + from ...base_llm.containers.transformation import BaseContainerConfig BaseLLMException = Any From bf3e6bb9f420d5ec070810fd66688b022b5d1c7c Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 26 Jan 2026 14:35:57 -0300 Subject: [PATCH 05/32] fix(register_model): handle openrouter models without '/' in name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The issue was in `register_model()` where `split_string[1]` assumed the model name always contained '/'. For custom model names like "glm" or UUIDs, the split would only produce one element, causing an IndexError. Changed `split_string[1]` to `split_string[-1]` which always returns the last element, working correctly for both cases: - "openrouter/gpt-4" → ["openrouter", "gpt-4"] → [-1] = "gpt-4" - "my-custom-alias" → ["my-custom-alias"] → [-1] = "my-custom-alias" --- litellm/utils.py | 2 +- tests/test_litellm/test_utils.py | 58 ++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 584ab8805a..7576e3bb83 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2780,7 +2780,7 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 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]) + 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) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f6c24d19df..6d99d64f7b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2300,6 +2300,64 @@ def test_register_model_with_scientific_notation(): assert registered_model["mode"] == "chat" +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( From 9c3de581823424e7defda4479b41ee868139343b Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 30 Jan 2026 14:34:10 -0300 Subject: [PATCH 06/32] refactor(containers): simplify BaseContainerConfig import per review Move BaseContainerConfig import outside TYPE_CHECKING block since it's needed at runtime for class inheritance, not just for type hints. --- litellm/llms/openai/containers/transformation.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index e8ab6190c7..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 - from ...base_llm.containers.transformation import BaseContainerConfig BaseLLMException = Any From 0664ec51d83a00cf74c75f7988ccecfccbd76d1c Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 16 Feb 2026 18:00:42 -0300 Subject: [PATCH 07/32] fix(responses): use output_index for parallel tool call streaming indices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #21331 — the Responses API streaming bridge hardcoded index=0 for all tool call chunks, making parallel tool calls indistinguishable. Now reads output_index from the Responses API chunk instead. --- .../transformation.py | 9 +- ...responses_transformation_transformation.py | 91 +++++++++++++++++++ 2 files changed, 97 insertions(+), 3 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 753a94295b..b307b2e90d 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -960,9 +960,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): 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, ) @@ -983,6 +984,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( @@ -991,7 +993,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): tool_calls=[ ChatCompletionToolCallChunk( id=None, - index=0, + index=tool_call_index, type="function", function=ChatCompletionToolCallFunctionChunk( name=None, arguments=content_part @@ -1033,9 +1035,10 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): 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/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 f8a082ee30..243c1f79ca 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 @@ -1278,3 +1278,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}" + ) From 0ebbec4d0ecf8f0f84470adf9e705d24cf76671f Mon Sep 17 00:00:00 2001 From: Chesars Date: Wed, 18 Feb 2026 18:02:36 -0300 Subject: [PATCH 08/32] fix(chatgpt): normalize streaming tool_call indices and deduplicate closing chunks The ChatGPT backend API sends non-spec-compliant streaming tool call chunks where index is always 0 for parallel tool calls and id/name get repeated in duplicate closing chunks. Add ChatGPTToolCallNormalizer to fix indices and filter duplicates before they reach the consumer. Fixes #21482 --- .../handler.py | 29 ++- litellm/llms/base_llm/chat/transformation.py | 4 + litellm/llms/chatgpt/chat/streaming_utils.py | 83 ++++++++ litellm/llms/chatgpt/chat/transformation.py | 6 +- tests/test_litellm/llms/chatgpt/__init__.py | 0 .../llms/chatgpt/chat/__init__.py | 0 .../llms/chatgpt/chat/test_streaming_utils.py | 195 ++++++++++++++++++ 7 files changed, 314 insertions(+), 3 deletions(-) create mode 100644 litellm/llms/chatgpt/chat/streaming_utils.py create mode 100644 tests/test_litellm/llms/chatgpt/__init__.py create mode 100644 tests/test_litellm/llms/chatgpt/chat/__init__.py create mode 100644 tests/test_litellm/llms/chatgpt/chat/test_streaming_utils.py 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/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/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py new file mode 100644 index 0000000000..953309266e --- /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: str | None = 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/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" From f3f731a678355f905eb6790cb29373378bed8aca Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 14:37:43 -0300 Subject: [PATCH 09/32] fix(openai): restrict supported params for gpt-5-search models gpt-5-search-api models were routed through OpenAIGPT5Config which listed params like n, temperature, tools, reasoning_effort as supported, but OpenAI rejects all of these for search models. Fixes #21572 --- .../llms/openai/chat/gpt_5_transformation.py | 34 ++++++++ .../llms/openai/test_gpt5_transformation.py | 78 +++++++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 05c003c8b7..9c5c7db791 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) @@ -90,6 +112,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") diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 386f264a4d..0a676fe87c 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -414,3 +414,81 @@ 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 From 5c4c085353f9cef5be57ab87a9ecafc5edebee80 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 14:56:49 -0300 Subject: [PATCH 10/32] fix(openai): correct supported_openai_params for GPT-5 model family Remove logit_bias, modalities, prediction, audio, web_search_options from supported params for all GPT-5 reasoning models (OpenAI rejects them). Add logprobs, top_p, top_logprobs for gpt-5.1/5.2 which support them when reasoning_effort="none". Related to #21572 --- .../llms/openai/chat/gpt_5_transformation.py | 12 +++- .../llms/openai/test_gpt5_transformation.py | 59 +++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 05c003c8b7..ffb39bfaaa 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -69,14 +69,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 diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 386f264a4d..b569c39c56 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -414,3 +414,62 @@ def test_gpt5_2_allows_reasoning_effort_xhigh(config: OpenAIConfig): drop_params=False, ) assert params["reasoning_effort"] == "xhigh" + + +# 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 From c6240ff621f1a53e837fe2bb57dbabf97da4a812 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 15:14:09 -0300 Subject: [PATCH 11/32] fix(azure_ai): resolve api_base from env var in get_complete_url for Document Intelligence OCR validate_environment() resolved AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT from the env var but only returned headers. get_complete_url() still received None and raised ValueError. Now get_complete_url() also resolves the env var as a fallback. Fixes #21034 --- .../llms/azure_ai/ocr/document_intelligence/transformation.py | 3 +++ 1 file changed, 3 insertions(+) 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" From a18518208632096940fac562c5b4d8a9fc39530d Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 15:20:14 -0300 Subject: [PATCH 12/32] fix(openai): validate logprobs/top_p against reasoning_effort for gpt-5.1/5.2 logprobs, top_p, top_logprobs are only accepted by OpenAI when reasoning_effort="none". Add validation matching the existing temperature logic: raise UnsupportedParamsError or drop when reasoning_effort is set to other values. --- .../llms/openai/chat/gpt_5_transformation.py | 18 ++++++++++ .../llms/openai/test_gpt5_transformation.py | 36 +++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index ffb39bfaaa..5e85c547aa 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -124,6 +124,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/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index b569c39c56..decc19c866 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -473,3 +473,39 @@ def test_gpt5_1_top_p_passthrough(config: OpenAIConfig): 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" From 155bed57e8bae17762a62da98956d9900af5659e Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 15:39:27 -0300 Subject: [PATCH 13/32] fix(vertex_ai): pass through native Gemini imageConfig params for image generation aspectRatio and imageSize were silently dropped because they weren't listed in get_supported_openai_params(), so the validation layer filtered them out before they could reach transform_image_generation_request(). Fixes #21070 --- .../vertex_gemini_transformation.py | 13 ++++++- ...rtex_ai_image_generation_transformation.py | 36 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) 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/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( From 3aea9c81c9f7bdebdc544b6f14d688059f75cc72 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 15:50:39 -0300 Subject: [PATCH 14/32] fix(openrouter): prevent double-stripping of native model names in get_llm_provider Move the fix to the OpenRouter level: define native OpenRouter models (openrouter/auto, openrouter/free, openrouter/bodybuilder) and check them in get_llm_provider() before the provider_list stripping logic. This prevents the second strip across all bridges without modifying each adapter/handler individually. Fixes #16353 --- .../get_llm_provider_logic.py | 7 ++ litellm/llms/openrouter/common_utils.py | 10 +++ tests/litellm/llms/openrouter/__init__.py | 0 .../test_openrouter_native_models.py | 75 +++++++++++++++++++ 4 files changed, 92 insertions(+) create mode 100644 tests/litellm/llms/openrouter/__init__.py create mode 100644 tests/litellm/llms/openrouter/test_openrouter_native_models.py diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 718773a1b1..eeae959aeb 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 @@ -165,6 +166,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/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/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" From 27413790e6155c97ea188e3063f787e7c8bb4a34 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 16:07:30 -0300 Subject: [PATCH 15/32] fix(openrouter): use provider-reported usage in streaming without stream_options When providers like OpenRouter send a usage chunk after the finish_reason chunk, _hidden_params["usage"] was already calculated (with zeros) before the usage data arrived. The StopIteration handler now recalculates usage from stream_chunk_builder and updates the shared _hidden_params dict so the user's copy reflects the real provider-reported token counts. Fixes #20760 --- .../litellm_core_utils/streaming_handler.py | 34 +++++++ .../test_streaming_handler.py | 89 +++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7a6752fbff..df8095e64d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -149,6 +149,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 __iter__(self): return self @@ -1787,6 +1788,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 @@ -1828,6 +1830,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 @@ -1951,6 +1971,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: @@ -2017,6 +2038,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/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index ec2f528a35..b592237864 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1185,3 +1185,92 @@ def test_is_chunk_non_empty_with_valid_tool_calls( ) is True ) + + +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}" + ) From 3b7236d6181f4f200a48262541fcbfa7800db7d0 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Thu, 19 Feb 2026 16:24:19 -0300 Subject: [PATCH 16/32] Update litellm/llms/chatgpt/chat/streaming_utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/llms/chatgpt/chat/streaming_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index 953309266e..730becb06e 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -22,9 +22,9 @@ class ChatGPTToolCallNormalizer: def __init__(self, stream: Any): self._stream = stream - self._seen_ids: dict[str, int] = {} # tool_call_id -> assigned_index + self._seen_ids: Dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 - self._last_id: str | None = None # tracks which tool call the next delta belongs to + 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) From c5cec60fd053b16077bd5a6f91d2dba2cd1ce7d4 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 16:40:58 -0300 Subject: [PATCH 17/32] fix(moonshot): preserve image_url blocks in multimodal messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Moonshot's _transform_messages unconditionally flattened content arrays to plain text, dropping image_url blocks. Vision models like kimi-k2.5 accept the standard OpenAI content array format. Now checks for image_url blocks before flattening — if any message contains images the content array is preserved intact. Fixes #20862 --- docs/my-website/docs/providers/moonshot.md | 31 ++++++ litellm/llms/moonshot/chat/transformation.py | 20 +++- .../test_moonshot_chat_transformation.py | 97 ++++++++++++++++++- 3 files changed, 145 insertions(+), 3 deletions(-) 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/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/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 From 3564b8d83b095d3ee4d6549bb3bd7924c650822b Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 19 Feb 2026 21:43:19 -0300 Subject: [PATCH 18/32] fix(types): suppress Pydantic serialization warnings on ModelResponse choices Pydantic v2's Union serializer for `List[Union[Choices, StreamingChoices]]` tries both branches when serializing, emitting spurious `PydanticSerializationUnexpectedValue` warnings (field count mismatch on `Message` and type mismatch `Expected StreamingChoices but got Choices`). Add a `WrapSerializer` on the `choices` field that serializes each item individually via its own `model_dump()`, bypassing the Union dispatch entirely while correctly propagating `exclude_none`, `exclude_unset`, and `exclude_defaults` from the parent serialization context. --- litellm/types/utils.py | 38 ++++++++++- .../test_model_response_normalization.py | 66 ++++++++++++++++++- 2 files changed, 100 insertions(+), 4 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9228b25b03..2b64231ca7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -22,8 +22,9 @@ from openai.types.moderation_create_response import Moderation as Moderation from openai.types.moderation_create_response import ( ModerationCreateResponse as ModerationCreateResponse, ) -from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator -from typing_extensions import Required, TypedDict +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, SerializationInfo, model_validator +from pydantic.functional_serializers import WrapSerializer +from typing_extensions import Annotated, Required, TypedDict from litellm._uuid import uuid from litellm.types.llms.base import ( @@ -1640,6 +1641,34 @@ class StreamingChatCompletionChunk(OpenAIChatCompletionChunk): super().__init__(**kwargs) +def _serialize_choices_list( + choices: list, handler, info: SerializationInfo +) -> list: + """Serialize each choice individually to avoid Union serializer warnings. + + Pydantic's Union serializer for ``List[Union[Choices, StreamingChoices]]`` + may try the wrong branch first, emitting spurious + ``PydanticSerializationUnexpectedValue`` warnings. By serializing each + item via its own ``model_dump()`` we bypass the Union dispatch entirely. + """ + kwargs: Dict[str, Any] = {} + if info.exclude_none: + kwargs["exclude_none"] = True + if info.exclude_unset: + kwargs["exclude_unset"] = True + if info.exclude_defaults: + kwargs["exclude_defaults"] = True + result = [] + for choice in choices: + if hasattr(choice, "model_dump"): + result.append(choice.model_dump(**kwargs)) + elif isinstance(choice, dict): + result.append(choice) + else: + result.append(choice) + return result + + class ModelResponseBase(OpenAIObject): id: str """A unique identifier for the completion.""" @@ -1748,7 +1777,10 @@ class ModelResponseStream(ModelResponseBase): class ModelResponse(ModelResponseBase): - choices: List[Union[Choices, StreamingChoices]] + choices: Annotated[ + List[Union[Choices, StreamingChoices]], + WrapSerializer(_serialize_choices_list, return_type=list), + ] """The list of completion choices the model generated for the input prompt.""" def __init__( # noqa: PLR0915 diff --git a/tests/test_litellm/test_model_response_normalization.py b/tests/test_litellm/test_model_response_normalization.py index 57281d3c1f..52e6b8bc9e 100644 --- a/tests/test_litellm/test_model_response_normalization.py +++ b/tests/test_litellm/test_model_response_normalization.py @@ -2,7 +2,7 @@ import warnings import pytest -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices def test_modelresponse_normalizes_openai_base_models() -> None: @@ -59,3 +59,67 @@ 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() bypasses the Python model_dump() override and uses + Pydantic's C-level serializer directly. The Union[Choices, StreamingChoices] + field previously triggered PydanticSerializationUnexpectedValue warnings via + this path.""" + 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_modelresponse_no_pydantic_warnings() -> None: + """Streaming responses use StreamingChoices in the Union field and should + also serialize without warnings.""" + response = ModelResponse( + model="test-model", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + stream=True, + ) + + 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}" + ) From 0f20976efa8e6fc00b973a4c9801e2e901b0afbb Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 20 Feb 2026 17:47:42 -0300 Subject: [PATCH 19/32] fix(types): remove StreamingChoices from ModelResponse, use ModelResponseStream ModelResponse.choices was typed as List[Union[Choices, StreamingChoices]] which caused Pydantic serialization warnings and false linting errors. Now that ModelResponseStream exists for streaming, narrow ModelResponse.choices to List[Choices] and migrate all ModelResponse(stream=True) call sites to use ModelResponseStream() instead. --- .../litellm_core_utils/streaming_handler.py | 2 +- litellm/llms/bedrock/chat/invoke_handler.py | 4 +- .../codestral/completion/transformation.py | 2 +- .../guardrails/guardrail_hooks/presidio.py | 9 +- litellm/types/utils.py | 94 +++++-------------- litellm/utils.py | 6 +- .../test_stream_chunk_builder_images.py | 6 +- .../test_stream_chunk_builder.py | 6 +- tests/local_testing/test_streaming.py | 12 +-- .../test_model_response_normalization.py | 25 ++--- 10 files changed, 59 insertions(+), 107 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7a6752fbff..772a7a2894 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1183,7 +1183,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: diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 1c58a11eeb..6d8b9c5a16 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -558,7 +558,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" ) @@ -695,7 +695,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/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/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 29d57ee473..06e0b38f53 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -808,7 +808,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return response if isinstance(response, ModelResponse) and not isinstance( - response.choices[0], StreamingChoices + response, ModelResponseStream ): # /chat/completions requests if isinstance(response.choices[0].message.content, str): verbose_proxy_logger.debug( @@ -832,7 +832,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( @@ -840,10 +840,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/types/utils.py b/litellm/types/utils.py index 2b64231ca7..7ffd20e827 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -22,9 +22,8 @@ from openai.types.moderation_create_response import Moderation as Moderation from openai.types.moderation_create_response import ( ModerationCreateResponse as ModerationCreateResponse, ) -from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, SerializationInfo, model_validator -from pydantic.functional_serializers import WrapSerializer -from typing_extensions import Annotated, Required, TypedDict +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator +from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.llms.base import ( @@ -1641,33 +1640,6 @@ class StreamingChatCompletionChunk(OpenAIChatCompletionChunk): super().__init__(**kwargs) -def _serialize_choices_list( - choices: list, handler, info: SerializationInfo -) -> list: - """Serialize each choice individually to avoid Union serializer warnings. - - Pydantic's Union serializer for ``List[Union[Choices, StreamingChoices]]`` - may try the wrong branch first, emitting spurious - ``PydanticSerializationUnexpectedValue`` warnings. By serializing each - item via its own ``model_dump()`` we bypass the Union dispatch entirely. - """ - kwargs: Dict[str, Any] = {} - if info.exclude_none: - kwargs["exclude_none"] = True - if info.exclude_unset: - kwargs["exclude_unset"] = True - if info.exclude_defaults: - kwargs["exclude_defaults"] = True - result = [] - for choice in choices: - if hasattr(choice, "model_dump"): - result.append(choice.model_dump(**kwargs)) - elif isinstance(choice, dict): - result.append(choice) - else: - result.append(choice) - return result - class ModelResponseBase(OpenAIObject): id: str @@ -1777,10 +1749,7 @@ class ModelResponseStream(ModelResponseBase): class ModelResponse(ModelResponseBase): - choices: Annotated[ - List[Union[Choices, StreamingChoices]], - WrapSerializer(_serialize_choices_list, return_type=list), - ] + choices: List[Choices] """The list of completion choices the model generated for the input prompt.""" def __init__( # noqa: PLR0915 @@ -1799,44 +1768,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 241b9d217b..de8fa9c1e1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7369,9 +7369,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/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 ee208b5e0e..7f7a43095c 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/test_model_response_normalization.py b/tests/test_litellm/test_model_response_normalization.py index 52e6b8bc9e..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, Delta, Message, ModelResponse, StreamingChoices +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) def test_modelresponse_normalizes_openai_base_models() -> None: @@ -62,10 +69,8 @@ def test_modelresponse_serialization_avoids_pydantic_warnings() -> None: def test_modelresponse_model_dump_json_no_pydantic_warnings() -> None: - """model_dump_json() bypasses the Python model_dump() override and uses - Pydantic's C-level serializer directly. The Union[Choices, StreamingChoices] - field previously triggered PydanticSerializationUnexpectedValue warnings via - this path.""" + """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=[ @@ -94,11 +99,10 @@ def test_modelresponse_model_dump_json_no_pydantic_warnings() -> None: ) -def test_streaming_modelresponse_no_pydantic_warnings() -> None: - """Streaming responses use StreamingChoices in the Union field and should - also serialize without warnings.""" - response = ModelResponse( - model="test-model", +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", @@ -106,7 +110,6 @@ def test_streaming_modelresponse_no_pydantic_warnings() -> None: delta=Delta(content="hello", role="assistant"), ) ], - stream=True, ) with warnings.catch_warnings(record=True) as captured: From a2cae0070e352fa8b895bb1ee472992134b59412 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 20 Feb 2026 17:58:06 -0300 Subject: [PATCH 20/32] fix(lint): remove unused StreamingChoices import in presidio guardrail --- litellm/proxy/guardrails/guardrail_hooks/presidio.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 06e0b38f53..d3f1fd1781 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, ) From b370fcd8de0993084c219ad6dd0146a80e6095cf Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 20 Feb 2026 17:59:19 -0300 Subject: [PATCH 21/32] fix(presidio): remove redundant isinstance check for ModelResponseStream ModelResponseStream and ModelResponse are sibling classes (both inherit from ModelResponseBase), so the guard was always True. Simplify to just isinstance(response, ModelResponse). --- litellm/proxy/guardrails/guardrail_hooks/presidio.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index d3f1fd1781..b4b25e2d90 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -806,9 +806,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, ModelResponseStream - ): # /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}" From 5eeee88ff8e33c506f1d3c36c5e42f92d64d2b2e Mon Sep 17 00:00:00 2001 From: Chesars Date: Sun, 22 Feb 2026 09:03:12 -0300 Subject: [PATCH 22/32] fix(lint): migrate remaining StreamingChoices callers to ModelResponseStream After narrowing ModelResponse.choices to List[Choices], several files still assigned StreamingChoices into ModelResponse. Migrate streaming call sites to ModelResponseStream and remove dead streaming branches in qwen2/qwen3 transform_response methods. --- .../llms/bedrock/chat/agentcore/transformation.py | 14 +++++++------- .../amazon_qwen2_transformation.py | 9 ++------- .../amazon_qwen3_transformation.py | 9 ++------- litellm/llms/langgraph/chat/sse_iterator.py | 14 +++++++------- .../openai/chat/guardrail_translation/handler.py | 12 ++++++------ litellm/utils.py | 4 +--- 6 files changed, 25 insertions(+), 37 deletions(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 9ae850ad4c..b85a0b70f7 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, @@ -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_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index c532d8ea27..90adc44a49 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 b3a957ce0f..faf8bb4ac6 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/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index bdb32cc0fe..aff1c2f2db 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 @@ -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, 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/utils.py b/litellm/utils.py index de8fa9c1e1..eff5bd58ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4960,9 +4960,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] = [] From 3e00c833531033aa4749307df45778e99dbd4f1f Mon Sep 17 00:00:00 2001 From: Chesars Date: Sun, 22 Feb 2026 09:09:49 -0300 Subject: [PATCH 23/32] fix(lint): restore ModelResponse import in langgraph sse_iterator --- litellm/llms/langgraph/chat/sse_iterator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index aff1c2f2db..8d946c4e52 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, ModelResponseStream, StreamingChoices +from litellm.types.utils import Delta, ModelResponse, ModelResponseStream, StreamingChoices if TYPE_CHECKING: pass From 43319b562a4e2da5979a2b1c35548fa18ae498e7 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sun, 22 Feb 2026 09:17:36 -0300 Subject: [PATCH 24/32] fix(lint): update return/yield types to ModelResponseStream - langgraph sse_iterator: update return types from ModelResponse to ModelResponseStream across all methods - bedrock agentcore: fix async generator yield type annotation - vertex gemini: add type: ignore for MyPy narrowing false positive --- .../llms/bedrock/chat/agentcore/transformation.py | 2 +- litellm/llms/langgraph/chat/sse_iterator.py | 14 +++++++------- .../gemini/vertex_and_google_ai_studio_gemini.py | 2 +- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index b85a0b70f7..fe7d4b194a 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -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. """ diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 8d946c4e52..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, ModelResponseStream, 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. """ @@ -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/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index cf3461a996..1783b5e630 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 @@ -2093,7 +2093,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( From eaf38ed3a20df28e7368a9da79fdf145af441418 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sun, 22 Feb 2026 09:22:44 -0300 Subject: [PATCH 25/32] fix(lint): add type: ignore for MyPy narrowing issue in vertex gemini --- .../llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 1783b5e630..377801ea59 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 @@ -2104,7 +2104,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, From 7e4f07c45dcfbdc36df56cae4ab9109f79af1bf5 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sun, 22 Feb 2026 09:37:26 -0300 Subject: [PATCH 26/32] fix(vertex): use module-level import for ModelResponseStream instead of type: ignore Move ModelResponseStream to module-level import so MyPy can properly narrow the Union type after isinstance checks, removing the need for type: ignore suppressions. --- .../vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 377801ea59..28e6919065 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 @@ -80,6 +80,7 @@ from litellm.types.utils import ( ChatCompletionTokenLogprob, ChoiceLogprobs, CompletionTokensDetailsWrapper, + ModelResponseStream, PromptTokensDetailsWrapper, TopLogprob, Usage, @@ -1931,7 +1932,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _process_candidates( # noqa: PLR0915 _candidates: List[Candidates], - model_response: Union[ModelResponse, "ModelResponseStream"], + model_response: Union[ModelResponse, ModelResponseStream], standard_optional_params: dict, cumulative_tool_call_index: int = 0, ) -> Tuple[List[dict], List[dict], List, List, int]: @@ -2093,7 +2094,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): chat_completion_logprobs=chat_completion_logprobs, image_response=image_response, ) - model_response.choices.append(choice) # type: ignore[arg-type] + model_response.choices.append(choice) elif isinstance(model_response, ModelResponse): choice = litellm.Choices( finish_reason=VertexGeminiConfig._check_finish_reason( @@ -2104,7 +2105,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): logprobs=chat_completion_logprobs, enhancements=None, ) - model_response.choices.append(choice) # type: ignore[arg-type] + model_response.choices.append(choice) return ( grounding_metadata, From ca5fd1ecec859f8717de36800b9fff68fffa8987 Mon Sep 17 00:00:00 2001 From: Chesars Date: Sun, 22 Feb 2026 09:43:08 -0300 Subject: [PATCH 27/32] Revert "fix(vertex): use module-level import for ModelResponseStream instead of type: ignore" This reverts commit 7e4f07c45dcfbdc36df56cae4ab9109f79af1bf5. --- .../vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 28e6919065..377801ea59 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 @@ -80,7 +80,6 @@ from litellm.types.utils import ( ChatCompletionTokenLogprob, ChoiceLogprobs, CompletionTokensDetailsWrapper, - ModelResponseStream, PromptTokensDetailsWrapper, TopLogprob, Usage, @@ -1932,7 +1931,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _process_candidates( # noqa: PLR0915 _candidates: List[Candidates], - model_response: Union[ModelResponse, ModelResponseStream], + model_response: Union[ModelResponse, "ModelResponseStream"], standard_optional_params: dict, cumulative_tool_call_index: int = 0, ) -> Tuple[List[dict], List[dict], List, List, int]: @@ -2094,7 +2093,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( @@ -2105,7 +2104,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, From dc4f713c6d6c48fd730514073516244c7eb6e3eb Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 27 Feb 2026 15:17:58 -0300 Subject: [PATCH 28/32] fix(images): forward extra_headers on OpenAI code path in image_generation() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #22285 — extra_headers passed to litellm.image_generation() were silently dropped on the openai/litellm_proxy/openai_compatible_providers code path. The azure and azure_ai paths already forwarded them correctly. --- litellm/images/main.py | 2 + .../test_image_generation_extra_headers.py | 84 +++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 tests/test_litellm/images/test_image_generation_extra_headers.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 6c4c502a7b..494d741f93 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( 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 From c4458c09fe4b423eec26ce9b5c3bb0b4abe6d821 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 27 Feb 2026 15:39:35 -0300 Subject: [PATCH 29/32] fix(count_tokens): include system and tools in token counting API requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /v1/messages/count_tokens proxy endpoint was only passing `messages` to provider token counting APIs, discarding `system` and `tools`. This caused clients like Claude Code to receive artificially low token counts (e.g. 10 instead of 531), preventing proper context window management and leading to context overflow errors. Pass system and tools through the full chain: - TokenCountRequest → proxy_server → provider counters → API handlers - Bedrock: transform tools to toolConfig format, system to text blocks - Anthropic/Azure AI: pass through directly (same API format) --- .../llms/anthropic/count_tokens/handler.py | 4 + .../anthropic/count_tokens/token_counter.py | 4 + .../anthropic/count_tokens/transformation.py | 26 ++-- .../anthropic/count_tokens/handler.py | 4 + .../anthropic/count_tokens/token_counter.py | 4 + litellm/llms/base_llm/base_utils.py | 2 + .../count_tokens/bedrock_token_counter.py | 10 +- .../bedrock/count_tokens/transformation.py | 92 ++++++++++---- litellm/llms/gemini/common_utils.py | 1 + litellm/llms/vertex_ai/common_utils.py | 1 + litellm/proxy/_types.py | 3 + .../proxy/anthropic_endpoints/endpoints.py | 7 +- litellm/proxy/proxy_server.py | 4 + ...t_anthropic_count_tokens_transformation.py | 92 ++++++++++++++ ...est_bedrock_count_tokens_transformation.py | 120 ++++++++++++++++++ 15 files changed, 331 insertions(+), 43 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/test_anthropic_count_tokens_transformation.py 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/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/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/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/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/proxy/_types.py b/litellm/proxy/_types.py index aeb9950b11..ea60c1e2ba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2832,6 +2832,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/proxy_server.py b/litellm/proxy/proxy_server.py index f0b1e66818..53627250fe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8321,6 +8321,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 @@ -8381,6 +8383,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/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_" From 727adb0117ec88d78a825aacca5bc18a1af9b2d6 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 27 Feb 2026 16:23:23 -0300 Subject: [PATCH 30/32] fix(images): pass model_info and metadata in image_edit for custom pricing image_edit was not forwarding model_info/metadata to the logging object, so custom_pricing was never detected. After PR #20679 stripped custom pricing fields from the shared backend key, image_edit cost became 0. Fixes #22244 --- litellm/images/main.py | 4 + .../images/test_image_edit_utils.py | 90 +++++++++++++++++++ 2 files changed, 94 insertions(+) diff --git a/litellm/images/main.py b/litellm/images/main.py index 6c4c502a7b..f4e65290ca 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -763,6 +763,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,6 +874,8 @@ def image_edit( # noqa: PLR0915 optional_params=dict(image_edit_request_params), litellm_params={ "litellm_call_id": litellm_call_id, + "model_info": model_info, + "metadata": metadata, **image_edit_request_params, }, custom_llm_provider=custom_llm_provider, 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 From a08e5195e50bba0985865c3259715d250cd786de Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 27 Feb 2026 16:26:11 -0300 Subject: [PATCH 31/32] fix: put image_edit_request_params spread first to avoid overwriting model_info/metadata --- litellm/images/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index f4e65290ca..8f6a983aa0 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -873,10 +873,10 @@ def image_edit( # noqa: PLR0915 user=user, optional_params=dict(image_edit_request_params), litellm_params={ + **image_edit_request_params, "litellm_call_id": litellm_call_id, "model_info": model_info, "metadata": metadata, - **image_edit_request_params, }, custom_llm_provider=custom_llm_provider, ) From 77496776c15aa29dbca670288fd581f8cbe17b95 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 27 Feb 2026 18:48:28 -0300 Subject: [PATCH 32/32] fix(register_model): align membership check with stored value for openrouter models The guard checked `key` (full key like "openrouter/gpt-4") but the set stores `split_string[-1]` ("gpt-4"), so the duplicate check never matched. --- litellm/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 7576e3bb83..2e186bb160 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2779,7 +2779,7 @@ 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: + 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: