From 18101345a015540c3b5643f2e0c4206c54700f72 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Tue, 7 Jul 2026 17:30:19 -0400 Subject: [PATCH 001/393] fix(responses): propagate message cache_control safely through objects and models --- .../context_caching/transformation.py | 22 +- .../transformation.py | 65 +- litellm/utils.py | 53 +- .../test_litellm_completion_responses.py | 559 ++++++++---------- 4 files changed, 348 insertions(+), 351 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index f0ce3323ef6..183d9743f13 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -62,23 +62,27 @@ def extract_ttl_from_cached_messages(messages: List[AllMessageValues]) -> Option if not is_cached_message(message): continue - content = message.get("content") + content = message.get("content") if isinstance(message, dict) else getattr(message, "content", None) if not content or isinstance(content, str): continue for content_item in content: - # Type check to ensure content_item is a dictionary before calling .get() - if not isinstance(content_item, dict): + # Check if content_item is dict or object model + if isinstance(content_item, dict): + cache_control = content_item.get("cache_control") + else: + cache_control = getattr(content_item, "cache_control", None) + + if not cache_control: continue - cache_control = content_item.get("cache_control") - if not cache_control or not isinstance(cache_control, dict): + cc_type = ( + cache_control.get("type") if isinstance(cache_control, dict) else getattr(cache_control, "type", None) + ) + if cc_type != "ephemeral": continue - if cache_control.get("type") != "ephemeral": - continue - - ttl = cache_control.get("ttl") + ttl = cache_control.get("ttl") if isinstance(cache_control, dict) else getattr(cache_control, "ttl", None) if ttl and _is_valid_ttl_format(ttl): return str(ttl) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 6b1ca3564e3..8052f69f22b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -892,26 +892,38 @@ class LiteLLMCompletionResponsesConfig: function_call=input_item ) else: - content = input_item.get("content") + content = ( + input_item.get("content") if isinstance(input_item, dict) else getattr(input_item, "content", None) + ) # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: return [] - return [ - GenericChatCompletionMessage( - role=input_item.get("role") or "user", - content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ), - ) - ] + + role = input_item.get("role") if isinstance(input_item, dict) else getattr(input_item, "role", None) + cache_control = ( + input_item.get("cache_control") + if isinstance(input_item, dict) + else getattr(input_item, "cache_control", None) + ) + + msg = GenericChatCompletionMessage( + role=role or "user", + content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( + content + ), + ) + if cache_control is not None: + msg["cache_control"] = cache_control + return [msg] @staticmethod def _is_input_item_tool_call_output(input_item: Any) -> bool: """ Check if the input item is a tool call output """ - return input_item.get("type") in [ + val = input_item.get("type") if isinstance(input_item, dict) else getattr(input_item, "type", None) + return val in [ "function_call_output", "custom_tool_call_output", "web_search_call", @@ -926,7 +938,8 @@ class LiteLLMCompletionResponsesConfig: Both need to be reconstructed as assistant tool_calls for Chat Completions providers. """ - return input_item.get("type") in ("function_call", "custom_tool_call") + val = input_item.get("type") if isinstance(input_item, dict) else getattr(input_item, "type", None) + return val in ("function_call", "custom_tool_call") @staticmethod def _transform_responses_api_tool_call_output_to_chat_completion_message( @@ -1155,6 +1168,31 @@ class LiteLLMCompletionResponsesConfig: return ChatCompletionImageObject(type="image_url", image_url=image_url_obj) + @staticmethod + def _normalize_responses_api_object_to_dict(item: Any) -> dict[str, Any]: + """ + Normalize a Responses API object (Pydantic model or custom class) to a dictionary + """ + if hasattr(item, "model_dump"): + return item.model_dump() + elif hasattr(item, "dict"): + return item.dict() + + item_dict = {} + for attr in [ + "type", + "text", + "cache_control", + "file_id", + "file_data", + "file_url", + "image_url", + "detail", + ]: + if hasattr(item, attr): + item_dict[attr] = getattr(item, attr) + return item_dict + @staticmethod def _transform_responses_api_content_to_chat_completion_content( content: Any, @@ -1176,7 +1214,10 @@ class LiteLLMCompletionResponsesConfig: for item in content: if isinstance(item, str): content_list.append(item) - elif isinstance(item, dict): + elif isinstance(item, dict) or (item is not None and not isinstance(item, (str, int, float, bool))): + if not isinstance(item, dict): + item = LiteLLMCompletionResponsesConfig._normalize_responses_api_object_to_dict(item) + if item.get("type") == "input_file": content_list.append( LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) diff --git a/litellm/utils.py b/litellm/utils.py index a11c5500503..3fdf0a5eee6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7209,36 +7209,51 @@ def is_cached_message(message: AllMessageValues) -> bool: return False # Check message-level cache_control (set by cache_control_injection_points hook for string content) - message_level_cache_control = message.get("cache_control") - if ( - message_level_cache_control is not None - and isinstance(message_level_cache_control, dict) - and message_level_cache_control.get("type") == "ephemeral" - ): - return True + message_level_cache_control = ( + message.get("cache_control") + if isinstance(message, dict) + else getattr(message, "cache_control", None) + ) + if message_level_cache_control is not None: + cc_type = ( + message_level_cache_control.get("type") + if isinstance(message_level_cache_control, dict) + else getattr(message_level_cache_control, "type", None) + ) + if cc_type == "ephemeral": + return True - if "content" not in message: - return False - - content = message["content"] + if isinstance(message, dict): + if "content" not in message: + return False + content = message["content"] + else: + content = getattr(message, "content", None) # Handle non-list content types (None, str, etc.) if not isinstance(content, list): return False for content_item in content: - # Ensure content_item is a dictionary before accessing keys - if not isinstance(content_item, dict): - continue + # Check if content_item is dict or object model + if isinstance(content_item, dict): + cache_control = content_item.get("cache_control") + item_type = content_item.get("type") + else: + cache_control = getattr(content_item, "cache_control", None) + item_type = getattr(content_item, "type", None) - cache_control = content_item.get("cache_control") if ( - content_item.get("type") == "text" + item_type == "text" and cache_control is not None - and isinstance(cache_control, dict) - and cache_control.get("type") == "ephemeral" ): - return True + cc_type = ( + cache_control.get("type") + if isinstance(cache_control, dict) + else getattr(cache_control, "type", None) + ) + if cc_type == "ephemeral": + return True return False diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index d8e3f495ced..4b07d638261 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -3,9 +3,7 @@ import sys import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, @@ -34,11 +32,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_file", "file_id": "file-abc123xyz"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {"file_id": "file-abc123xyz"}} @@ -53,11 +47,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_file", "file_data": file_data} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {"file_data": file_data}} @@ -75,11 +65,7 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = { @@ -97,11 +83,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_file"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {}} @@ -120,11 +102,7 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(input_item) # Assert expected = {"type": "file", "file": {"file_id": "file-abc123xyz"}} @@ -134,10 +112,8 @@ class TestLiteLLMCompletionResponsesConfig: def test_transform_input_file_item_to_file_item_with_file_url(self): """file_url should be mapped to file_id for downstream URL handling""" - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - {"type": "input_file", "file_url": "https://example.com/doc.pdf"} - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( + {"type": "input_file", "file_url": "https://example.com/doc.pdf"} ) assert result == { "type": "file", @@ -146,14 +122,12 @@ class TestLiteLLMCompletionResponsesConfig: def test_transform_input_file_item_file_id_takes_precedence_over_file_url(self): """explicit file_id should not be overwritten by file_url""" - result = ( - LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( - { - "type": "input_file", - "file_id": "file-abc123", - "file_url": "https://example.com/doc.pdf", - } - ) + result = LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( + { + "type": "input_file", + "file_id": "file-abc123", + "file_url": "https://example.com/doc.pdf", + } ) assert result == {"type": "file", "file": {"file_id": "file-abc123"}} @@ -164,11 +138,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image", "image_url": image_url, "detail": "high"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -187,11 +157,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image", "image_url": image_url, "detail": "high"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -210,11 +176,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image", "image_url": image_url} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -232,11 +194,7 @@ class TestLiteLLMCompletionResponsesConfig: input_item = {"type": "input_image"} # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = {"type": "image_url", "image_url": {"url": "", "detail": "auto"}} @@ -256,11 +214,7 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - result = ( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item( - input_item - ) - ) + result = LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(input_item) # Assert expected = { @@ -296,28 +250,26 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="What is the meaning of life?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="What is the meaning of life?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert assert hasattr(responses_api_response, "output") assert len(responses_api_response.output) >= 2 - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1, "Should have exactly one reasoning item" reasoning_item = reasoning_items[0] # Note: ID auto-generation was disabled, so reasoning items may not have IDs # Only assert ID format if an ID is present if hasattr(reasoning_item, "id") and reasoning_item.id: - assert reasoning_item.id.startswith( - "rs_" - ), f"Expected ID to start with 'rs_', got: {reasoning_item.id}" + assert reasoning_item.id.startswith("rs_"), f"Expected ID to start with 'rs_', got: {reasoning_item.id}" assert reasoning_item.status == "completed" assert reasoning_item.role == "assistant" assert len(reasoning_item.content) == 1 @@ -325,9 +277,7 @@ class TestLiteLLMCompletionResponsesConfig: assert "step by step" in reasoning_item.content[0].text assert "42" in reasoning_item.content[0].text - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) == 1, "Should have exactly one message item" message_item = message_items[0] @@ -354,21 +304,19 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="A simple question?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="A simple question?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 0, "Should have no reasoning items" - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) == 1, "Should have exactly one message item" assert message_items[0].content[0].text == "Just a regular answer." assert responses_api_response.object == "response" @@ -404,22 +352,20 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="A question with multiple answers?", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="A question with multiple answers?", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert - reasoning_items = [ - item for item in responses_api_response.output if item.type == "reasoning" - ] + reasoning_items = [item for item in responses_api_response.output if item.type == "reasoning"] assert len(reasoning_items) == 1, "Should have exactly one reasoning item" assert reasoning_items[0].content[0].text == "First reasoning process." - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) == 2, "Should have two message items" def test_transform_chat_completion_response_status_with_stop(self): @@ -446,10 +392,12 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="this is a test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="this is a test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) assert responses_api_response.status == "completed" @@ -485,15 +433,15 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="this is a test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="this is a test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) - message_items = [ - item for item in responses_api_response.output if item.type == "message" - ] + message_items = [item for item in responses_api_response.output if item.type == "message"] assert len(message_items) > 0 for item in message_items: @@ -528,10 +476,12 @@ class TestLiteLLMCompletionResponsesConfig: ], ) - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="this is a test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="this is a test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) assert responses_api_response.status == "incomplete" @@ -563,10 +513,12 @@ class TestLiteLLMCompletionResponsesConfig: } # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert @@ -598,10 +550,12 @@ class TestLiteLLMCompletionResponsesConfig: ) # Execute - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="Test", - responses_api_request={}, - chat_completion_response=chat_completion_response, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) ) # Assert - should default to empty dict @@ -630,26 +584,14 @@ class TestFunctionCallTransformation: regular_message = {"type": "message", "role": "user", "content": "Hello"} # Test function_call detection - assert LiteLLMCompletionResponsesConfig._is_input_item_function_call( - function_call_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call( - function_call_output_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call( - regular_message - ) + assert LiteLLMCompletionResponsesConfig._is_input_item_function_call(function_call_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call(function_call_output_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_function_call(regular_message) # Test function_call_output detection (should still work) - assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - function_call_output_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - function_call_item - ) - assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( - regular_message - ) + assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(function_call_output_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(function_call_item) + assert not LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(regular_message) def test_function_call_transformation(self): """Test that function_call items are correctly transformed to assistant messages with tool calls""" @@ -733,9 +675,7 @@ class TestFunctionCallTransformation: tool_msg = messages[2] assert tool_msg.get("role") == "tool" assert tool_msg.get("content") == "Rainy" - assert ( - tool_msg.get("tool_call_id") == "call_1fe70e2a-a596-45ef-b72c-9b8567c460e5" - ) + assert tool_msg.get("tool_call_id") == "call_1fe70e2a-a596-45ef-b72c-9b8567c460e5" def test_complete_request_transformation_with_function_calls(self): """Test the complete request transformation that would be used by the responses API""" @@ -940,9 +880,7 @@ class TestToolChoiceTransformation: Test that {"type": "tool"} is transformed to "required". This fixes the Anthropic error: "tool_choice.tool.name: Field required" """ - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "tool"} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "tool"}) assert result == "required" def test_transform_tool_choice_preserves_function_with_name(self): @@ -954,23 +892,17 @@ class TestToolChoiceTransformation: def test_transform_tool_choice_responses_flat_function_name(self): """Responses-API forced-function with a top-level name maps to the nested Chat Completions shape instead of degrading to required and dropping the name""" - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "function", "name": "get_weather"} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "function", "name": "get_weather"}) assert result == {"type": "function", "function": {"name": "get_weather"}} def test_transform_tool_choice_function_without_name_falls_back_to_required(self): """A function-type dict with no name still falls back to required""" - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "function"} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "function"}) assert result == "required" def test_transform_tool_choice_function_empty_name_falls_back_to_required(self): """An empty top-level name is falsy and must not produce an empty function name""" - result = LiteLLMCompletionResponsesConfig._transform_tool_choice( - {"type": "function", "name": ""} - ) + result = LiteLLMCompletionResponsesConfig._transform_tool_choice({"type": "function", "name": ""}) assert result == "required" @@ -982,20 +914,12 @@ class TestContentTypeTransformation: Test that 'tool_result' content type is transformed to 'text'. This fixes: Invalid user message - content type 'tool_result' not valid. """ - result = ( - LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( - "tool_result" - ) - ) + result = LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type("tool_result") assert result == "text" def test_input_text_content_type_transformed_to_text(self): """Test that 'input_text' content type is transformed to 'text'""" - result = ( - LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( - "input_text" - ) - ) + result = LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type("input_text") assert result == "text" def test_none_text_blocks_filtered_out(self): @@ -1009,9 +933,7 @@ class TestContentTypeTransformation: {"type": "text", "text": None}, # Should be filtered out {"type": "text", "text": "another valid"}, ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert len(result) == 2 assert result[0]["text"] == "valid text" assert result[1]["text"] == "another valid" @@ -1033,9 +955,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1057,9 +977,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1087,9 +1005,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert - computer_use has no Chat Completions equivalent, so it is dropped assert len(result_tools) == 0 @@ -1114,9 +1030,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert - custom tool is converted to a function tool assert len(result_tools) == 1 @@ -1141,9 +1055,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1167,9 +1079,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1190,9 +1100,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1211,9 +1119,7 @@ class TestToolTransformation: tools = [custom_tool] with pytest.raises(ValueError, match="allowed_callers must be a list of strings"): - LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) def test_transform_web_search_tools_to_web_search_options(self): """Test that web_search tools are converted to web_search_options""" @@ -1229,9 +1135,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 0 # Web search is not added to tools @@ -1262,9 +1166,7 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1294,9 +1196,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1322,9 +1222,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1350,9 +1248,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1376,9 +1272,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 2 @@ -1410,14 +1304,10 @@ class TestToolTransformation: ( result_tools, web_search_options, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert - assert ( - len(result_tools) == 3 - ) # function, mcp, vertex (web_search becomes options) + assert len(result_tools) == 3 # function, mcp, vertex (web_search becomes options) assert web_search_options is not None # Check function tool @@ -1447,9 +1337,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1472,9 +1360,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1495,9 +1381,7 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 @@ -1519,19 +1403,14 @@ class TestToolTransformation: ( result_tools, _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - tools=tools - ) + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools=tools) # Assert assert len(result_tools) == 1 result_tool = result_tools[0] assert result_tool["function"]["parameters"]["type"] == "object" assert "properties" in result_tool["function"]["parameters"] - assert ( - result_tool["function"]["parameters"]["properties"]["arg"]["type"] - == "string" - ) + assert result_tool["function"]["parameters"]["properties"]["arg"]["type"] == "string" def test_bedrock_anthropic_drops_derived_web_search_options(self): """ @@ -1657,11 +1536,7 @@ class TestToolTransformation: assert "web_search_options" not in result bedrock_tool_blocks = _bedrock_tools_pt(tools=result["tools"], model=model) - names = [ - block["toolSpec"]["name"] - for block in bedrock_tool_blocks - if "toolSpec" in block - ] + names = [block["toolSpec"]["name"] for block in bedrock_tool_blocks if "toolSpec" in block] assert names == ["noop"] assert not any(name.startswith("litellm_unnamed_tool_") for name in names) @@ -1940,9 +1815,7 @@ class TestUsageTransformation: Choices( finish_reason="stop", index=0, - message=Message( - content="Here is the generated image.", role="assistant" - ), + message=Message(content="Here is the generated image.", role="assistant"), ) ], ) @@ -2166,9 +2039,7 @@ class TestStreamingIDConsistency: # Verify the cached ID is set and matches assert iterator._cached_item_id is not None, "Iterator should cache the item_id" assert iterator._cached_item_id == item_id_1, "Cached ID should match event IDs" - assert ( - iterator._cached_item_id == "chatcmpl-first-id" - ), "Should use the first chunk's ID" + assert iterator._cached_item_id == "chatcmpl-first-id", "Should use the first chunk's ID" def test_streaming_iterator_initial_events_use_cached_id(self): """ @@ -2258,9 +2129,7 @@ class TestStreamingIDConsistency: # Create done events text_done_event = iterator.create_output_text_done_event(complete_response) - content_done_event = iterator.create_output_content_part_done_event( - complete_response - ) + content_done_event = iterator.create_output_content_part_done_event(complete_response) item_done_event = iterator.create_output_item_done_event(complete_response) # Extract IDs @@ -2320,27 +2189,19 @@ class TestStreamingIDConsistency: input=input_items ) - roles = [ - m.get("role") if isinstance(m, dict) else getattr(m, "role", None) - for m in messages - ] + roles = [m.get("role") if isinstance(m, dict) else getattr(m, "role", None) for m in messages] # Must not have two consecutive assistant messages for i in range(len(roles) - 1): - assert not ( - roles[i] == "assistant" and roles[i + 1] == "assistant" - ), f"Consecutive assistant messages at indices {i} and {i+1}: {roles}" + assert not (roles[i] == "assistant" and roles[i + 1] == "assistant"), ( + f"Consecutive assistant messages at indices {i} and {i + 1}: {roles}" + ) # The single assistant message must contain BOTH tool_calls assistant_messages = [ - m - for m in messages - if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) - == "assistant" + m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "assistant" ] - assert ( - len(assistant_messages) == 1 - ), f"Expected 1 assistant message, got {len(assistant_messages)}" + assert len(assistant_messages) == 1, f"Expected 1 assistant message, got {len(assistant_messages)}" assistant_msg = assistant_messages[0] tool_calls = ( @@ -2348,27 +2209,19 @@ class TestStreamingIDConsistency: if isinstance(assistant_msg, dict) else getattr(assistant_msg, "tool_calls", None) ) - assert ( - tool_calls is not None and len(tool_calls) == 2 - ), f"Expected 2 tool_calls in the merged assistant message, got: {tool_calls}" + assert tool_calls is not None and len(tool_calls) == 2, ( + f"Expected 2 tool_calls in the merged assistant message, got: {tool_calls}" + ) - call_ids = [ - (tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)) - for tc in tool_calls - ] + call_ids = [(tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)) for tc in tool_calls] assert "toolu_01" in call_ids, f"toolu_01 missing from tool_calls: {call_ids}" assert "toolu_02" in call_ids, f"toolu_02 missing from tool_calls: {call_ids}" # Both tool messages must be present tool_messages = [ - m - for m in messages - if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) - == "tool" + m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "tool" ] - assert ( - len(tool_messages) == 2 - ), f"Expected 2 tool messages, got {len(tool_messages)}" + assert len(tool_messages) == 2, f"Expected 2 tool messages, got {len(tool_messages)}" def test_single_tool_call_still_works_after_merge_fix(self): """ @@ -2390,20 +2243,14 @@ class TestStreamingIDConsistency: input=input_items ) - roles = [ - m.get("role") if isinstance(m, dict) else getattr(m, "role", None) - for m in messages - ] + roles = [m.get("role") if isinstance(m, dict) else getattr(m, "role", None) for m in messages] assert "user" in roles assert "assistant" in roles assert "tool" in roles assistant_messages = [ - m - for m in messages - if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) - == "assistant" + m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "assistant" ] assert len(assistant_messages) == 1 @@ -2576,9 +2423,7 @@ class TestEnsureOutputItemContentPartAdded: LiteLLMCompletionStreamingIterator, ) - iterator = LiteLLMCompletionStreamingIterator.__new__( - LiteLLMCompletionStreamingIterator - ) + iterator = LiteLLMCompletionStreamingIterator.__new__(LiteLLMCompletionStreamingIterator) iterator.sent_output_item_added_event = False iterator.sent_content_part_added_event = False iterator._sequence_number = 0 @@ -2671,9 +2516,7 @@ class TestEnsureOutputItemContentPartAdded: usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), ) - completed_event = iterator._emit_response_completed_event( - litellm_model_response - ) + completed_event = iterator._emit_response_completed_event(litellm_model_response) assert completed_event is not None assert completed_event.response.status == "incomplete" @@ -2716,9 +2559,7 @@ class TestCacheControlPreservation: "cache_control": {"type": "ephemeral"}, } ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} @@ -2726,9 +2567,7 @@ class TestCacheControlPreservation: def test_content_without_cache_control_unaffected(self): """Content blocks that don't have cache_control should be unaffected.""" content = [{"type": "text", "text": "hello"}] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert "cache_control" not in result[0] @@ -2750,9 +2589,7 @@ class TestCacheControlPreservation: ) assert len(messages) == 1 msg_content = ( - messages[0].get("content") - if isinstance(messages[0], dict) - else getattr(messages[0], "content", None) + messages[0].get("content") if isinstance(messages[0], dict) else getattr(messages[0], "content", None) ) assert isinstance(msg_content, list) assert msg_content[0]["cache_control"] == {"type": "ephemeral"} @@ -2765,9 +2602,7 @@ class TestCacheControlPreservation: "cache_control": {"type": "ephemeral"}, } ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} @@ -2780,13 +2615,121 @@ class TestCacheControlPreservation: "cache_control": {"type": "ephemeral"}, } ] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content - ) + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} + def test_cache_control_preserved_for_object_input_item(self): + """Test that cache_control is preserved when input_item is a custom object / model.""" + + class MockInputItem: + def __init__(self): + self.role = "user" + self.content = "hello" + self.cache_control = {"type": "ephemeral"} + + input_item = MockInputItem() + messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item + ) + assert len(messages) == 1 + assert messages[0].get("cache_control") == {"type": "ephemeral"} + + def test_cache_control_preserved_for_object_content_item(self): + """Test that cache_control is preserved when content items are custom objects.""" + + class MockContentBlock: + def __init__(self): + self.type = "text" + self.text = "hello" + self.cache_control = {"type": "ephemeral"} + + class MockPydanticV2Block: + def __init__(self): + self.type = "text" + self.text = "hello v2" + self.cache_control = {"type": "ephemeral"} + + def model_dump(self): + return {"type": self.type, "text": self.text, "cache_control": self.cache_control} + + class MockPydanticV1Block: + def __init__(self): + self.type = "text" + self.text = "hello v1" + self.cache_control = {"type": "ephemeral"} + + def dict(self): + return {"type": self.type, "text": self.text, "cache_control": self.cache_control} + + content = [MockContentBlock(), MockPydanticV2Block(), MockPydanticV1Block()] + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) + assert isinstance(result, list) + assert len(result) == 3 + assert result[0]["cache_control"] == {"type": "ephemeral"} + assert result[1]["cache_control"] == {"type": "ephemeral"} + assert result[2]["cache_control"] == {"type": "ephemeral"} + assert result[1]["text"] == "hello v2" + assert result[2]["text"] == "hello v1" + + def test_is_cached_message_for_object_message_and_content_item(self): + """Test is_cached_message on custom objects / models.""" + from litellm.utils import is_cached_message + + # Test message level cache_control object + class MockCacheControl: + def __init__(self): + self.type = "ephemeral" + + class MockMessageLevelObj: + def __init__(self): + self.role = "system" + self.content = "hello" + self.cache_control = MockCacheControl() + + msg = MockMessageLevelObj() + assert is_cached_message(msg) is True + + # Test content level cache_control object + class MockContentItem: + def __init__(self): + self.type = "text" + self.text = "hello" + self.cache_control = MockCacheControl() + + class MockContentLevelObj: + def __init__(self): + self.role = "system" + self.content = [MockContentItem()] + + msg = MockContentLevelObj() + assert is_cached_message(msg) is True + + def test_extract_ttl_from_cached_messages_for_object_models(self): + """Test extract_ttl_from_cached_messages with object-based messages and content items.""" + from litellm.llms.vertex_ai.context_caching.transformation import extract_ttl_from_cached_messages + + class MockCacheControl: + def __init__(self): + self.type = "ephemeral" + self.ttl = "3600s" + + class MockContentItem: + def __init__(self): + self.type = "text" + self.text = "hello" + self.cache_control = MockCacheControl() + + class MockMessageObj: + def __init__(self): + self.role = "system" + self.content = [MockContentItem()] + + messages = [MockMessageObj()] + ttl = extract_ttl_from_cached_messages(messages) + assert ttl == "3600s" + def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): """Bedrock Mantle returns a non-unique, index-based ``call_id`` (``call_0`` that @@ -2798,16 +2741,10 @@ def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): for the bedrock-mantle gpt-5.5 non-streaming path.""" from types import SimpleNamespace - convert = ( - LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call - ) + convert = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call - mantle = SimpleNamespace( - id="fc_unique_abc123", call_id="call_0", name="get_weather", arguments="{}" - ) + mantle = SimpleNamespace(id="fc_unique_abc123", call_id="call_0", name="get_weather", arguments="{}") assert convert(mantle)["id"] == "fc_unique_abc123" - openai = SimpleNamespace( - id="fc_2", call_id="call_tokyo", name="get_weather", arguments="{}" - ) + openai = SimpleNamespace(id="fc_2", call_id="call_tokyo", name="get_weather", arguments="{}") assert convert(openai)["id"] == "call_tokyo" From 42a39dd81988273c88835ebc44725660c7048c62 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Tue, 7 Jul 2026 20:05:42 -0400 Subject: [PATCH 002/393] fix(transformations): implementing greptile feedback --- .../transformation.py | 4 +++- litellm/utils.py | 13 +++---------- .../test_litellm_completion_responses.py | 11 +++++++++-- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 8052f69f22b..217ed711d37 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1190,7 +1190,9 @@ class LiteLLMCompletionResponsesConfig: "detail", ]: if hasattr(item, attr): - item_dict[attr] = getattr(item, attr) + val = getattr(item, attr) + if val is not None: + item_dict[attr] = val return item_dict @staticmethod diff --git a/litellm/utils.py b/litellm/utils.py index 3fdf0a5eee6..d1c7af7050d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7210,9 +7210,7 @@ def is_cached_message(message: AllMessageValues) -> bool: # Check message-level cache_control (set by cache_control_injection_points hook for string content) message_level_cache_control = ( - message.get("cache_control") - if isinstance(message, dict) - else getattr(message, "cache_control", None) + message.get("cache_control") if isinstance(message, dict) else getattr(message, "cache_control", None) ) if message_level_cache_control is not None: cc_type = ( @@ -7243,14 +7241,9 @@ def is_cached_message(message: AllMessageValues) -> bool: cache_control = getattr(content_item, "cache_control", None) item_type = getattr(content_item, "type", None) - if ( - item_type == "text" - and cache_control is not None - ): + if item_type == "text" and cache_control is not None: cc_type = ( - cache_control.get("type") - if isinstance(cache_control, dict) - else getattr(cache_control, "type", None) + cache_control.get("type") if isinstance(cache_control, dict) else getattr(cache_control, "type", None) ) if cc_type == "ephemeral": return True diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 4b07d638261..537753ec3e1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2663,15 +2663,22 @@ class TestCacheControlPreservation: def dict(self): return {"type": self.type, "text": self.text, "cache_control": self.cache_control} - content = [MockContentBlock(), MockPydanticV2Block(), MockPydanticV1Block()] + class MockBlockWithNoneCacheControl: + def __init__(self): + self.type = "text" + self.text = "hello none" + self.cache_control = None + + content = [MockContentBlock(), MockPydanticV2Block(), MockPydanticV1Block(), MockBlockWithNoneCacheControl()] result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) assert isinstance(result, list) - assert len(result) == 3 + assert len(result) == 4 assert result[0]["cache_control"] == {"type": "ephemeral"} assert result[1]["cache_control"] == {"type": "ephemeral"} assert result[2]["cache_control"] == {"type": "ephemeral"} assert result[1]["text"] == "hello v2" assert result[2]["text"] == "hello v1" + assert "cache_control" not in result[3] def test_is_cached_message_for_object_message_and_content_item(self): """Test is_cached_message on custom objects / models.""" From f09b8ce34af3e2521faf87a386835c6a42a32b35 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Wed, 8 Jul 2026 21:59:31 -0400 Subject: [PATCH 003/393] fix(responses): propagate message cache_control safely through objects and models --- .../context_caching/transformation.py | 50 ++++-- .../transformation.py | 45 +++-- litellm/types/llms/openai.py | 1 + .../test_context_caching_ttl.py | 157 ++++++++++++++++-- .../test_litellm_completion_responses.py | 148 +++++------------ 5 files changed, 248 insertions(+), 153 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index 183d9743f13..bc6ee47d3b8 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -59,32 +59,52 @@ def extract_ttl_from_cached_messages(messages: List[AllMessageValues]) -> Option Optional[str]: TTL string in format "3600s" or None if not found/invalid """ for message in messages: - if not is_cached_message(message): - continue + # Check message-level cache_control first + msg_cache_control = ( + message.get("cache_control") if isinstance(message, dict) else getattr(message, "cache_control", None) + ) + if msg_cache_control is not None: + cc_type = ( + msg_cache_control.get("type") + if isinstance(msg_cache_control, dict) + else getattr(msg_cache_control, "type", None) + ) + if cc_type == "ephemeral": + ttl = ( + msg_cache_control.get("ttl") + if isinstance(msg_cache_control, dict) + else getattr(msg_cache_control, "ttl", None) + ) + if ttl and _is_valid_ttl_format(ttl): + return str(ttl) content = message.get("content") if isinstance(message, dict) else getattr(message, "content", None) - if not content or isinstance(content, str): + if not isinstance(content, list): continue for content_item in content: # Check if content_item is dict or object model if isinstance(content_item, dict): cache_control = content_item.get("cache_control") + item_type = content_item.get("type") else: cache_control = getattr(content_item, "cache_control", None) + item_type = getattr(content_item, "type", None) - if not cache_control: - continue - - cc_type = ( - cache_control.get("type") if isinstance(cache_control, dict) else getattr(cache_control, "type", None) - ) - if cc_type != "ephemeral": - continue - - ttl = cache_control.get("ttl") if isinstance(cache_control, dict) else getattr(cache_control, "ttl", None) - if ttl and _is_valid_ttl_format(ttl): - return str(ttl) + if item_type == "text" and cache_control is not None: + cc_type = ( + cache_control.get("type") + if isinstance(cache_control, dict) + else getattr(cache_control, "type", None) + ) + if cc_type == "ephemeral": + ttl = ( + cache_control.get("ttl") + if isinstance(cache_control, dict) + else getattr(cache_control, "ttl", None) + ) + if ttl and _is_valid_ttl_format(ttl): + return str(ttl) return None diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 217ed711d37..4105fc0bcf5 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -912,9 +912,8 @@ class LiteLLMCompletionResponsesConfig: content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( content ), + **({"cache_control": cache_control} if cache_control is not None else {}), ) - if cache_control is not None: - msg["cache_control"] = cache_control return [msg] @staticmethod @@ -948,6 +947,11 @@ class LiteLLMCompletionResponsesConfig: """ ChatCompletionToolMessage is used to indicate the output from a tool call """ + if not isinstance(tool_call_output, dict): + tool_call_output = LiteLLMCompletionResponsesConfig._normalize_responses_api_object_to_dict( + tool_call_output + ) + call_id = tool_call_output.get("call_id") # If call_id is missing or empty, skip this message # Empty call_id means we can't create a valid tool message @@ -1097,6 +1101,9 @@ class LiteLLMCompletionResponsesConfig: } ``` """ + if not isinstance(function_call, dict): + function_call = LiteLLMCompletionResponsesConfig._normalize_responses_api_object_to_dict(function_call) + # Create a tool call for the function call. Custom tool calls # store their payload in "input" (raw string) rather than # "arguments" (JSON string), so normalize to arguments here. @@ -1178,8 +1185,7 @@ class LiteLLMCompletionResponsesConfig: elif hasattr(item, "dict"): return item.dict() - item_dict = {} - for attr in [ + target_attrs = ( "type", "text", "cache_control", @@ -1188,12 +1194,19 @@ class LiteLLMCompletionResponsesConfig: "file_url", "image_url", "detail", - ]: - if hasattr(item, attr): - val = getattr(item, attr) - if val is not None: - item_dict[attr] = val - return item_dict + "call_id", + "arguments", + "input", + "name", + "id", + "output", + "status", + ) + return { + attr: getattr(item, attr) + for attr in target_attrs + if hasattr(item, attr) and getattr(item, attr) is not None + } @staticmethod def _transform_responses_api_content_to_chat_completion_content( @@ -1225,11 +1238,10 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) - ) - if "cache_control" in item: - image_block["cache_control"] = item["cache_control"] + image_block = { + **LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item), + **({"cache_control": item["cache_control"]} if "cache_control" in item else {}), + } content_list.append(image_block) else: # Skip text blocks with None text to avoid downstream errors @@ -1241,9 +1253,8 @@ class LiteLLMCompletionResponsesConfig: item.get("type") or "text" ), "text": text_value, + **({"cache_control": item["cache_control"]} if "cache_control" in item else {}), } - if "cache_control" in item: - content_block["cache_control"] = item["cache_control"] content_list.append(content_block) return content_list else: diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 9f689a2dd31..77d1680eefd 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -791,6 +791,7 @@ class ChatCompletionDeveloperMessage(OpenAIChatCompletionDeveloperMessage, total class GenericChatCompletionMessage(TypedDict, total=False): role: Required[str] content: Required[Union[str, List]] + cache_control: ChatCompletionCachedContent ValidUserMessageContentTypes = [ diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py index 8a13baa0006..2b786820f8f 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py @@ -91,9 +91,7 @@ class TestTTLExtraction: messages = [ { "role": "user", - "content": [ - {"type": "text", "text": "Regular message without cache control"} - ], + "content": [{"type": "text", "text": "Regular message without cache control"}], } ] @@ -175,9 +173,7 @@ class TestTTLExtraction: class TestTransformationWithTTL: """Test the complete transformation with TTL support""" - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) def test_transform_with_valid_ttl(self, custom_llm_provider): """Test transformation includes TTL when provided""" messages = [ @@ -218,9 +214,7 @@ class TestTransformationWithTTL: assert result["displayName"] == "test-cache-key" - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) def test_transform_without_ttl(self, custom_llm_provider): """Test transformation without TTL""" messages = [ @@ -260,9 +254,7 @@ class TestTransformationWithTTL: assert result["displayName"] == "test-cache-key" - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) def test_transform_with_invalid_ttl(self, custom_llm_provider): """Test transformation with invalid TTL (should be ignored)""" messages = [ @@ -301,9 +293,7 @@ class TestTransformationWithTTL: assert result["displayName"] == "test-cache-key" - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) def test_transform_with_system_message_and_ttl(self, custom_llm_provider): """Test transformation with system message and TTL""" messages = [ @@ -388,6 +378,143 @@ class TestEdgeCases: assert isinstance(ttl, str) assert ttl == "3600s" + def test_cache_control_preserved_for_object_content_items(self): + """Test that cache_control is preserved when content items are real Pydantic models.""" + from pydantic import BaseModel, Field + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + class MockContentBlock: + def __init__(self): + self.type = "text" + self.text = "hello" + self.cache_control = {"type": "ephemeral"} + + class RealPydanticV2Block(BaseModel): + type: str = "text" + text: str = "hello v2" + cache_control: dict = Field(default_factory=lambda: {"type": "ephemeral"}) + + class MockBlockWithNoneCacheControl: + def __init__(self): + self.type = "text" + self.text = "hello none" + self.cache_control = None + + content = [ + MockContentBlock(), + RealPydanticV2Block(), + MockBlockWithNoneCacheControl(), + ] + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) + assert result == [ + {"type": "text", "text": "hello", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "hello v2", "cache_control": {"type": "ephemeral"}}, + {"type": "text", "text": "hello none"}, + ] + + def test_is_cached_message_for_object_message_and_content_item(self): + """Test is_cached_message on custom objects / models.""" + from litellm.utils import is_cached_message + + # Test message level cache_control object + class MockCacheControl: + def __init__(self): + self.type = "ephemeral" + + class MockMessageLevelObj: + def __init__(self): + self.role = "system" + self.content = "hello" + self.cache_control = MockCacheControl() + + msg = MockMessageLevelObj() + assert is_cached_message(msg) is True + + # Test content level cache_control object + class MockContentItem: + def __init__(self): + self.type = "text" + self.text = "hello" + self.cache_control = MockCacheControl() + + class MockContentLevelObj: + def __init__(self): + self.role = "system" + self.content = [MockContentItem()] + + msg = MockContentLevelObj() + assert is_cached_message(msg) is True + + def test_extract_ttl_from_cached_messages_for_object_models(self): + """Test extract_ttl_from_cached_messages with object-based messages and content items.""" + + class MockCacheControl: + def __init__(self): + self.type = "ephemeral" + self.ttl = "3600s" + + class MockContentItem: + def __init__(self): + self.type = "text" + self.text = "hello" + self.cache_control = MockCacheControl() + + class MockMessageObj: + def __init__(self): + self.role = "system" + self.content = [MockContentItem()] + + messages = [MockMessageObj()] + ttl = extract_ttl_from_cached_messages(messages) + assert ttl == "3600s" + + def test_extract_ttl_from_cached_messages_with_message_level_object_cache_control(self): + """Test extract_ttl_from_cached_messages with message-level object cache_control.""" + + class MockCacheControl: + def __init__(self): + self.type = "ephemeral" + self.ttl = "7200s" + + class MockMessageObj: + def __init__(self): + self.role = "system" + self.content = "hello" + self.cache_control = MockCacheControl() + + messages = [MockMessageObj()] + ttl = extract_ttl_from_cached_messages(messages) + assert ttl == "7200s" + + def test_is_cached_message_for_dict_message_with_dict_content_items(self): + """Test is_cached_message with dict message and dict content list items.""" + from litellm.utils import is_cached_message + + # Dictionary message without content should return False + assert is_cached_message({"role": "user"}) is False + + msg = { + "role": "user", + "content": [ + {"type": "text", "text": "hello", "cache_control": {"type": "ephemeral"}} + ], + } + assert is_cached_message(msg) is True + + def test_normalize_responses_api_object_to_dict_pydantic_v1(self): + """Test _normalize_responses_api_object_to_dict with Pydantic v1 dict fallback.""" + from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig + + class MockPydanticV1Model: + def dict(self): + return {"type": "text", "text": "hello", "cache_control": {"type": "ephemeral"}} + + item = MockPydanticV1Model() + res = LiteLLMCompletionResponsesConfig._normalize_responses_api_object_to_dict(item) + assert res == {"type": "text", "text": "hello", "cache_control": {"type": "ephemeral"}} + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 537753ec3e1..b34bdd812d6 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2621,121 +2621,57 @@ class TestCacheControlPreservation: assert result[0]["cache_control"] == {"type": "ephemeral"} def test_cache_control_preserved_for_object_input_item(self): - """Test that cache_control is preserved when input_item is a custom object / model.""" + """Test that cache_control is preserved when input_item is a real Pydantic model.""" + from pydantic import BaseModel, Field - class MockInputItem: - def __init__(self): - self.role = "user" - self.content = "hello" - self.cache_control = {"type": "ephemeral"} + class RealInputItem(BaseModel): + role: str = "user" + content: str = "hello" + cache_control: dict = Field(default_factory=lambda: {"type": "ephemeral"}) - input_item = MockInputItem() + input_item = RealInputItem() messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item ) + assert messages == [{"role": "user", "content": "hello", "cache_control": {"type": "ephemeral"}}] + + def test_tool_call_output_as_custom_object(self): + """Test _transform_responses_api_tool_call_output_to_chat_completion_message with a custom object.""" + class MockToolCallOutput: + def __init__(self): + self.call_id = "call_abc123" + self.output = "tool output content" + self.status = "completed" + + item = MockToolCallOutput() + messages = LiteLLMCompletionResponsesConfig._transform_responses_api_tool_call_output_to_chat_completion_message( + item + ) assert len(messages) == 1 - assert messages[0].get("cache_control") == {"type": "ephemeral"} + assert messages[0]["role"] == "tool" + assert messages[0]["tool_call_id"] == "call_abc123" + assert messages[0]["content"] == "tool output content" - def test_cache_control_preserved_for_object_content_item(self): - """Test that cache_control is preserved when content items are custom objects.""" - - class MockContentBlock: + def test_function_call_as_custom_object(self): + """Test _transform_responses_api_function_call_to_chat_completion_message with a custom object.""" + class MockFunctionCall: def __init__(self): - self.type = "text" - self.text = "hello" - self.cache_control = {"type": "ephemeral"} + self.type = "function_call" + self.arguments = '{"location": "Boston"}' + self.call_id = "call_xyz789" + self.name = "get_weather" + self.id = "fc_12345" + self.status = "completed" - class MockPydanticV2Block: - def __init__(self): - self.type = "text" - self.text = "hello v2" - self.cache_control = {"type": "ephemeral"} - - def model_dump(self): - return {"type": self.type, "text": self.text, "cache_control": self.cache_control} - - class MockPydanticV1Block: - def __init__(self): - self.type = "text" - self.text = "hello v1" - self.cache_control = {"type": "ephemeral"} - - def dict(self): - return {"type": self.type, "text": self.text, "cache_control": self.cache_control} - - class MockBlockWithNoneCacheControl: - def __init__(self): - self.type = "text" - self.text = "hello none" - self.cache_control = None - - content = [MockContentBlock(), MockPydanticV2Block(), MockPydanticV1Block(), MockBlockWithNoneCacheControl()] - result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content) - assert isinstance(result, list) - assert len(result) == 4 - assert result[0]["cache_control"] == {"type": "ephemeral"} - assert result[1]["cache_control"] == {"type": "ephemeral"} - assert result[2]["cache_control"] == {"type": "ephemeral"} - assert result[1]["text"] == "hello v2" - assert result[2]["text"] == "hello v1" - assert "cache_control" not in result[3] - - def test_is_cached_message_for_object_message_and_content_item(self): - """Test is_cached_message on custom objects / models.""" - from litellm.utils import is_cached_message - - # Test message level cache_control object - class MockCacheControl: - def __init__(self): - self.type = "ephemeral" - - class MockMessageLevelObj: - def __init__(self): - self.role = "system" - self.content = "hello" - self.cache_control = MockCacheControl() - - msg = MockMessageLevelObj() - assert is_cached_message(msg) is True - - # Test content level cache_control object - class MockContentItem: - def __init__(self): - self.type = "text" - self.text = "hello" - self.cache_control = MockCacheControl() - - class MockContentLevelObj: - def __init__(self): - self.role = "system" - self.content = [MockContentItem()] - - msg = MockContentLevelObj() - assert is_cached_message(msg) is True - - def test_extract_ttl_from_cached_messages_for_object_models(self): - """Test extract_ttl_from_cached_messages with object-based messages and content items.""" - from litellm.llms.vertex_ai.context_caching.transformation import extract_ttl_from_cached_messages - - class MockCacheControl: - def __init__(self): - self.type = "ephemeral" - self.ttl = "3600s" - - class MockContentItem: - def __init__(self): - self.type = "text" - self.text = "hello" - self.cache_control = MockCacheControl() - - class MockMessageObj: - def __init__(self): - self.role = "system" - self.content = [MockContentItem()] - - messages = [MockMessageObj()] - ttl = extract_ttl_from_cached_messages(messages) - assert ttl == "3600s" + item = MockFunctionCall() + messages = LiteLLMCompletionResponsesConfig._transform_responses_api_function_call_to_chat_completion_message( + item + ) + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert len(messages[0]["tool_calls"]) == 1 + assert messages[0]["tool_calls"][0]["id"] == "call_xyz789" + assert messages[0]["tool_calls"][0]["function"]["name"] == "get_weather" def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): From 5822aa87eedbab9e017683c598145498d2001600 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Thu, 9 Jul 2026 22:52:14 -0400 Subject: [PATCH 004/393] feat(caching): enhance Gemini context caching propagation and TTL normalization --- .../adapters/transformation.py | 8 +- .../context_caching/transformation.py | 54 ++++++++++++- ...al_pass_through_adapters_transformation.py | 55 ++++++++++++- .../test_context_caching_ttl.py | 81 +++++++++++++++++++ .../test_litellm_completion_responses.py | 16 ++++ 5 files changed, 207 insertions(+), 7 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 4b6617fbeac..97f98cbea7e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -305,11 +305,17 @@ class LiteLLMAnthropicMessagesAdapter: target: Dict or TypedDict to add cache_control to model: Model name to check if cache_control should be preserved """ + from litellm.utils import _is_gemini_model + # TypedDict objects are dicts at runtime, so .get() works cache_control = ( source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) ) - if cache_control and model and (self.is_anthropic_claude_model(model) or self.is_bedrock_arn_model(model)): + if cache_control and model and ( + self.is_anthropic_claude_model(model) + or self.is_bedrock_arn_model(model) + or _is_gemini_model(model, None) + ): # TypedDict objects support dict operations at runtime # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) if isinstance(target, dict): diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index bc6ee47d3b8..b325c61aeca 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -75,8 +75,9 @@ def extract_ttl_from_cached_messages(messages: List[AllMessageValues]) -> Option if isinstance(msg_cache_control, dict) else getattr(msg_cache_control, "ttl", None) ) - if ttl and _is_valid_ttl_format(ttl): - return str(ttl) + normalized = _normalize_ttl_to_seconds(ttl) + if normalized is not None: + return normalized content = message.get("content") if isinstance(message, dict) else getattr(message, "content", None) if not isinstance(content, list): @@ -103,8 +104,9 @@ def extract_ttl_from_cached_messages(messages: List[AllMessageValues]) -> Option if isinstance(cache_control, dict) else getattr(cache_control, "ttl", None) ) - if ttl and _is_valid_ttl_format(ttl): - return str(ttl) + normalized = _normalize_ttl_to_seconds(ttl) + if normalized is not None: + return normalized return None @@ -138,6 +140,50 @@ def _is_valid_ttl_format(ttl: str) -> bool: return False +def _normalize_ttl_to_seconds(ttl: object) -> Optional[str]: + """ + Normalize a cache_control TTL into Gemini's "s" format. + + Accepts Gemini-native seconds (e.g. "3600s", "1.5s") and Anthropic-style + minute/hour units (e.g. "5m", "1h") that Claude Code and the Anthropic + /v1/messages spec use. Returns None for missing or unparseable values so + Gemini falls back to its own default TTL. + """ + if not isinstance(ttl, str): + return None + + if _is_valid_ttl_format(ttl): + return ttl + + match = re.match(r"^([0-9]*\.?[0-9]+)(m|h)$", ttl) + if not match: + return None + + value = float(match.group(1)) + + if value <= 0: + return None + + seconds = value * (60 if match.group(2) == "m" else 3600) + return f"{int(seconds)}s" if seconds.is_integer() else f"{seconds}s" + + +def get_gemini_context_caching_min_tokens(model: str) -> int: + """ + Minimum input token count required to create an explicit Gemini context cache. + + Gemini rejects a cachedContents create below a per-model floor with a 400, so + the caller skips caching below this value. Figures from + https://ai.google.dev/gemini-api/docs/caching (Gemini 2.5 -> 2048, Gemini 3.x + -> 4096). Unknown Gemini models default to the highest known floor so a create + is never attempted below the real minimum. + """ + model_lower = model.lower() + if "gemini-2.5" in model_lower or "gemini-2-5" in model_lower: + return 2048 + return 4096 + + def separate_cached_messages( messages: List[AllMessageValues], ) -> Tuple[List[AllMessageValues], List[AllMessageValues]]: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index dfe7e0c3a51..e86010a00f0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1388,14 +1388,17 @@ def test_should_add_cache_control_for_anthropic_model(): def test_should_not_add_cache_control_for_non_anthropic_model(): - """Should not add cache_control for non-Anthropic models.""" + """Should not add cache_control for providers that reject an explicit cache_control field. + + OpenAI/Azure do prompt caching implicitly and 400 on an unexpected + cache_control field, so it must not be forwarded to them. + """ adapter = LiteLLMAnthropicMessagesAdapter() cache_control = {"type": "ephemeral"} for model in [ CACHE_CONTROL_NON_ANTHROPIC_MODEL, "openai/gpt-4-turbo", - "gemini-pro", ]: target = {} adapter._add_cache_control_if_applicable( @@ -1404,6 +1407,54 @@ def test_should_not_add_cache_control_for_non_anthropic_model(): assert "cache_control" not in target +def test_should_add_cache_control_for_gemini_model(): + """Should add cache_control for Gemini / Vertex Gemini targets. + + These consume anthropic-style cache_control blocks via the Gemini context + caching path, so /v1/messages requests (e.g. Claude Code) routed to a + Gemini model must keep it. Regression for the adapter dropping the field + before it reaches the Gemini transformation. + """ + adapter = LiteLLMAnthropicMessagesAdapter() + cache_control = {"type": "ephemeral", "ttl": "1h"} + + for model in [ + "gemini-3.5-flash", + "gemini/gemini-3.5-flash", + "gemini-3.1-pro-preview", + "vertex_ai/gemini-2.5-pro", + ]: + target = {} + adapter._add_cache_control_if_applicable( + {"cache_control": cache_control}, target, model + ) + assert target.get("cache_control") == cache_control + + +def test_cache_control_preserved_in_text_content_for_gemini(): + """cache_control must survive message translation for a Gemini target.""" + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[ + { + "type": "text", + "text": "This is cached content", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_anthropic_messages_to_openai( + messages=anthropic_messages, model="gemini/gemini-3.5-flash" + ) + + assert len(result) == 1 + assert result[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + def test_should_not_add_cache_control_when_none(): """Should not add cache_control when source has None or empty cache_control.""" adapter = LiteLLMAnthropicMessagesAdapter() diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py index 2b786820f8f..ec193f8b9d8 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py @@ -1,11 +1,33 @@ import pytest from litellm.llms.vertex_ai.context_caching.transformation import ( extract_ttl_from_cached_messages, + get_gemini_context_caching_min_tokens, _is_valid_ttl_format, + _normalize_ttl_to_seconds, transform_openai_messages_to_gemini_context_caching, ) +class TestGeminiContextCachingMinTokens: + """Per-model floor for explicit Gemini context cache creation.""" + + @pytest.mark.parametrize( + "model, expected", + [ + ("gemini-2.5-flash", 2048), + ("gemini-2.5-pro", 2048), + ("gemini/gemini-2.5-pro", 2048), + ("vertex_ai/gemini-2.5-flash", 2048), + ("gemini-3.5-flash", 4096), + ("gemini-3.1-pro-preview", 4096), + ("gemini/gemini-3.5-flash", 4096), + ("gemini-1.5-pro", 4096), + ], + ) + def test_min_tokens_by_model(self, model, expected): + assert get_gemini_context_caching_min_tokens(model) == expected + + class TestTTLValidation: """Test TTL format validation""" @@ -37,6 +59,65 @@ class TestTTLValidation: assert not _is_valid_ttl_format(ttl), f"TTL {ttl} should be invalid" +class TestTTLNormalization: + """Normalization of anthropic-style TTL units into Gemini's seconds format.""" + + @pytest.mark.parametrize( + "ttl, expected", + [ + ("3600s", "3600s"), + ("1.5s", "1.5s"), + ("5m", "300s"), + ("90m", "5400s"), + ("1h", "3600s"), + ("2h", "7200s"), + ("0.5h", "1800s"), + ], + ) + def test_normalizes_units_to_seconds(self, ttl, expected): + assert _normalize_ttl_to_seconds(ttl) == expected + + @pytest.mark.parametrize( + "ttl", + ["invalid", "", "0m", "0h", "-1h", "5d", "1 h", "m", None, 123, 3600], + ) + def test_rejects_unparseable_ttl(self, ttl): + assert _normalize_ttl_to_seconds(ttl) is None + + def test_extract_ttl_normalizes_anthropic_hour_unit(self): + """Claude Code / Anthropic send "1h"; Gemini must receive "3600s".""" + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "cached", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + + assert extract_ttl_from_cached_messages(messages) == "3600s" + + def test_extract_ttl_normalizes_anthropic_minute_unit(self): + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "cached", + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + ], + } + ] + + assert extract_ttl_from_cached_messages(messages) == "300s" + + class TestTTLExtraction: """Test TTL extraction from cached messages""" diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b34bdd812d6..62bc61c1a63 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2673,6 +2673,22 @@ class TestCacheControlPreservation: assert messages[0]["tool_calls"][0]["id"] == "call_xyz789" assert messages[0]["tool_calls"][0]["function"]["name"] == "get_weather" + def test_is_input_item_object_type_checks(self): + """Test _is_input_item_tool_call_output and _is_input_item_function_call with custom objects.""" + class MockObj: + def __init__(self, t): + self.type = t + + # Test tool call output + assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(MockObj("function_call_output")) is True + assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(MockObj("custom_tool_call_output")) is True + assert LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output(MockObj("text")) is False + + # Test function call + assert LiteLLMCompletionResponsesConfig._is_input_item_function_call(MockObj("function_call")) is True + assert LiteLLMCompletionResponsesConfig._is_input_item_function_call(MockObj("custom_tool_call")) is True + assert LiteLLMCompletionResponsesConfig._is_input_item_function_call(MockObj("text")) is False + def test_function_call_tool_id_falls_back_to_unique_id_for_degenerate_call_id(): """Bedrock Mantle returns a non-unique, index-based ``call_id`` (``call_0`` that From 0695f0702db4d27ce0471b8de00d5a5fb7bb130f Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Thu, 9 Jul 2026 22:53:04 -0400 Subject: [PATCH 005/393] feat(caching): enhance Gemini context caching by enforcing minimum token requirements to prevent 400s --- .../adapters/transformation.py | 12 ++-- .../vertex_ai_context_caching.py | 20 ++++--- litellm/utils.py | 2 + .../test_vertex_ai_context_caching.py | 58 +++++++++++++++++++ 4 files changed, 81 insertions(+), 11 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 97f98cbea7e..6e69d011ff1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -311,10 +311,14 @@ class LiteLLMAnthropicMessagesAdapter: cache_control = ( source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) ) - if cache_control and model and ( - self.is_anthropic_claude_model(model) - or self.is_bedrock_arn_model(model) - or _is_gemini_model(model, None) + if ( + cache_control + and model + and ( + self.is_anthropic_claude_model(model) + or self.is_bedrock_arn_model(model) + or _is_gemini_model(model, None) + ) ): # TypedDict objects support dict operations at runtime # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 0bf3715f798..bfd862966a2 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -4,7 +4,6 @@ import httpx import litellm from litellm.caching.caching import Cache, LiteLLMCacheType -from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -22,6 +21,7 @@ from litellm.types.llms.vertex_ai import ( from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase from .transformation import ( + get_gemini_context_caching_min_tokens, separate_cached_messages, transform_openai_messages_to_gemini_context_caching, ) @@ -308,17 +308,20 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None - # Gemini requires a minimum of 1024 tokens for context caching. - # Skip caching if the cached content is too small to avoid API errors. + # Gemini's explicit context caching minimum varies by model; creating a + # cache below it returns a 400. Skip caching when the cached content is + # too small to avoid the error. + min_token_count = get_gemini_context_caching_min_tokens(model) if not is_prompt_caching_valid_prompt( model=model, messages=cached_messages, custom_llm_provider=custom_llm_provider, + min_token_count=min_token_count, ): verbose_logger.debug( "Vertex AI context caching: cached content is below minimum token " "count (%d). Skipping context caching.", - MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + min_token_count, ) return messages, optional_params, None @@ -459,17 +462,20 @@ class ContextCachingEndpoints(VertexBase): if len(cached_messages) == 0: return messages, optional_params, None - # Gemini requires a minimum of 1024 tokens for context caching. - # Skip caching if the cached content is too small to avoid API errors. + # Gemini's explicit context caching minimum varies by model; creating a + # cache below it returns a 400. Skip caching when the cached content is + # too small to avoid the error. + min_token_count = get_gemini_context_caching_min_tokens(model) if not is_prompt_caching_valid_prompt( model=model, messages=cached_messages, custom_llm_provider=custom_llm_provider, + min_token_count=min_token_count, ): verbose_logger.debug( "Vertex AI context caching: cached content is below minimum token " "count (%d). Skipping context caching.", - MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + min_token_count, ) return messages, optional_params, None diff --git a/litellm/utils.py b/litellm/utils.py index d1c7af7050d..17c2b388847 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9097,6 +9097,8 @@ def is_prompt_caching_valid_prompt( nothing here and would silently fall back to the default. OpenAI's minimum is a flat 1024 across models, which the default already covers. + Pass min_token_count to override this for providers with a different floor + (e.g. Gemini, whose explicit context caching minimum varies by model). """ try: if messages is None and tools is None: diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index cf75964ddb7..8a3ef84a607 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1401,6 +1401,64 @@ class TestContextCachingEndpoints: # Restart the patcher so teardown_method can stop it cleanly self._token_check_patcher.start() + @pytest.mark.parametrize( + "model, expected_min", + [ + ("gemini-3.5-flash", 4096), + ("gemini/gemini-3.5-flash", 4096), + ("gemini-3.1-pro-preview", 4096), + ("gemini-2.5-flash", 2048), + ("gemini-2.5-pro", 2048), + ], + ) + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) + def test_check_and_create_cache_uses_model_specific_min_tokens( + self, mock_separate, model, expected_min + ): + """The Gemini per-model floor must be forwarded to the token-count guard. + + A flat 1024 floor let content between 1024 and the real minimum (2048 for + 2.5, 4096 for 3.x) reach Gemini and 400. Assert the model-derived floor is + passed so the guard skips instead of erroring. + """ + self._token_check_patcher.stop() + + cached_messages = [ + { + "role": "system", + "content": "cached", + "cache_control": {"type": "ephemeral"}, + } + ] + non_cached_messages = [{"role": "user", "content": "Hello"}] + mock_separate.return_value = (cached_messages, non_cached_messages) + + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt", + return_value=False, + ) as mock_valid: + self.context_caching.check_and_create_cache( + messages=cached_messages + non_cached_messages, + optional_params=self.sample_optional_params.copy(), + api_key="test_key", + api_base=None, + model=model, + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + cached_content=None, + custom_llm_provider="gemini", + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="test_token", + ) + + assert mock_valid.call_args.kwargs["min_token_count"] == expected_min + + self._token_check_patcher.start() + @pytest.mark.parametrize( "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] ) From a244ad63af094cc838613e7756a92eba4943b942 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Thu, 9 Jul 2026 23:25:01 -0400 Subject: [PATCH 006/393] fix(caching): updating gemini-1.5 minimum token requirements --- .../adapters/transformation.py | 2 +- .../vertex_ai/context_caching/transformation.py | 12 ++++++++---- ...mental_pass_through_adapters_transformation.py | 15 +++++++++++++++ .../context_caching/test_context_caching_ttl.py | 5 ++++- .../test_vertex_ai_context_caching.py | 1 + 5 files changed, 29 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 6e69d011ff1..789a434b518 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -326,7 +326,7 @@ class LiteLLMAnthropicMessagesAdapter: target["cache_control"] = cache_control # type: ignore[typeddict-item] else: # Fallback for non-dict objects (shouldn't happen in practice) - cast(Dict[str, Any], target)["cache_control"] = cache_control + setattr(target, "cache_control", cache_control) def translatable_anthropic_params(self) -> List: """ diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index b325c61aeca..30ad0805474 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -174,14 +174,18 @@ def get_gemini_context_caching_min_tokens(model: str) -> int: Gemini rejects a cachedContents create below a per-model floor with a 400, so the caller skips caching below this value. Figures from - https://ai.google.dev/gemini-api/docs/caching (Gemini 2.5 -> 2048, Gemini 3.x - -> 4096). Unknown Gemini models default to the highest known floor so a create - is never attempted below the real minimum. + https://ai.google.dev/gemini-api/docs/caching (Gemini 1.5 -> 32768, Gemini 2.5 + -> 2048, Gemini 3.x -> 4096). Unknown Gemini models default to the highest + known floor so a create is never attempted below the real minimum. """ model_lower = model.lower() + if "gemini-1.5" in model_lower or "gemini-1-5" in model_lower: + return 32768 if "gemini-2.5" in model_lower or "gemini-2-5" in model_lower: return 2048 - return 4096 + if "gemini-3" in model_lower: + return 4096 + return 32768 def separate_cached_messages( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e86010a00f0..03e6f7c56b3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1431,6 +1431,21 @@ def test_should_add_cache_control_for_gemini_model(): assert target.get("cache_control") == cache_control +def test_cache_control_fallback_setattr(): + """Verify cache_control is safely assigned to non-dict target objects using setattr.""" + adapter = LiteLLMAnthropicMessagesAdapter() + cache_control = {"type": "ephemeral"} + + class MockTarget: + pass + + target = MockTarget() + adapter._add_cache_control_if_applicable( + {"cache_control": cache_control}, target, "claude-3-opus-20240229" + ) + assert getattr(target, "cache_control", None) == cache_control + + def test_cache_control_preserved_in_text_content_for_gemini(): """cache_control must survive message translation for a Gemini target.""" anthropic_messages = [ diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py index ec193f8b9d8..9e31df9c27f 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py @@ -14,6 +14,9 @@ class TestGeminiContextCachingMinTokens: @pytest.mark.parametrize( "model, expected", [ + ("gemini-1.5-pro", 32768), + ("gemini-1.5-flash", 32768), + ("vertex_ai/gemini-1.5-pro-001", 32768), ("gemini-2.5-flash", 2048), ("gemini-2.5-pro", 2048), ("gemini/gemini-2.5-pro", 2048), @@ -21,7 +24,7 @@ class TestGeminiContextCachingMinTokens: ("gemini-3.5-flash", 4096), ("gemini-3.1-pro-preview", 4096), ("gemini/gemini-3.5-flash", 4096), - ("gemini-1.5-pro", 4096), + ("gemini-unknown-future-model", 32768), ], ) def test_min_tokens_by_model(self, model, expected): diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 8a3ef84a607..aa873e984f7 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1407,6 +1407,7 @@ class TestContextCachingEndpoints: ("gemini-3.5-flash", 4096), ("gemini/gemini-3.5-flash", 4096), ("gemini-3.1-pro-preview", 4096), + ("gemini-1.5-pro", 32768), ("gemini-2.5-flash", 2048), ("gemini-2.5-pro", 2048), ], From d18409ef608f849cb794fc7804f1b84f014d5367 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Fri, 10 Jul 2026 22:01:56 -0400 Subject: [PATCH 007/393] feat(transformation): capping Gemini ttl for caching --- .../vertex_ai/context_caching/transformation.py | 17 ++++++++++------- .../context_caching/test_context_caching_ttl.py | 3 +++ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index 30ad0805474..6f28427925b 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -146,16 +146,14 @@ def _normalize_ttl_to_seconds(ttl: object) -> Optional[str]: Accepts Gemini-native seconds (e.g. "3600s", "1.5s") and Anthropic-style minute/hour units (e.g. "5m", "1h") that Claude Code and the Anthropic - /v1/messages spec use. Returns None for missing or unparseable values so - Gemini falls back to its own default TTL. + /v1/messages spec use. Caps the requested TTL at 24 hours (86400s) to + prevent unbounded persistent storage costs. Returns None for missing or + unparseable values so Gemini falls back to its own default TTL. """ if not isinstance(ttl, str): return None - if _is_valid_ttl_format(ttl): - return ttl - - match = re.match(r"^([0-9]*\.?[0-9]+)(m|h)$", ttl) + match = re.match(r"^([0-9]*\.?[0-9]+)(s|m|h)$", ttl) if not match: return None @@ -164,7 +162,12 @@ def _normalize_ttl_to_seconds(ttl: object) -> Optional[str]: if value <= 0: return None - seconds = value * (60 if match.group(2) == "m" else 3600) + multiplier = {"s": 1, "m": 60, "h": 3600}[match.group(2)] + seconds = value * multiplier + + # Cap explicit caches to 24 hours to prevent unbounded billing costs + seconds = min(seconds, 86400.0) + return f"{int(seconds)}s" if seconds.is_integer() else f"{seconds}s" diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py index 9e31df9c27f..af2f0684e69 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py @@ -75,6 +75,9 @@ class TestTTLNormalization: ("1h", "3600s"), ("2h", "7200s"), ("0.5h", "1800s"), + ("48h", "86400s"), + ("1500m", "86400s"), + ("1000000s", "86400s"), ], ) def test_normalizes_units_to_seconds(self, ttl, expected): From 4b2a4363178449b4817e0906c2cb2f6bb459ac35 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Wed, 22 Jul 2026 12:23:54 -0400 Subject: [PATCH 008/393] feat(vertex_ai): look up Gemini minimum cache creation tokens from model info --- .../context_caching/transformation.py | 22 +- .../transformation.py | 11 +- model_prices_and_context_window.json | 287 ++++++++++++------ .../test_context_caching_ttl.py | 11 + tests/test_litellm/test_utils.py | 1 + 5 files changed, 226 insertions(+), 106 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index 6f28427925b..7be01cd39e5 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -140,7 +140,7 @@ def _is_valid_ttl_format(ttl: str) -> bool: return False -def _normalize_ttl_to_seconds(ttl: object) -> Optional[str]: +def _normalize_ttl_to_seconds(ttl: object) -> str | None: """ Normalize a cache_control TTL into Gemini's "s" format. @@ -168,6 +168,8 @@ def _normalize_ttl_to_seconds(ttl: object) -> Optional[str]: # Cap explicit caches to 24 hours to prevent unbounded billing costs seconds = min(seconds, 86400.0) + # Google Protobuf Duration requires up to 9 fractional digits + seconds = round(seconds, 9) return f"{int(seconds)}s" if seconds.is_integer() else f"{seconds}s" @@ -175,15 +177,19 @@ def get_gemini_context_caching_min_tokens(model: str) -> int: """ Minimum input token count required to create an explicit Gemini context cache. - Gemini rejects a cachedContents create below a per-model floor with a 400, so - the caller skips caching below this value. Figures from - https://ai.google.dev/gemini-api/docs/caching (Gemini 1.5 -> 32768, Gemini 2.5 - -> 2048, Gemini 3.x -> 4096). Unknown Gemini models default to the highest - known floor so a create is never attempted below the real minimum. + Looks up the `cache_creation_min_tokens` property from model_prices_and_context_window.json. + Defaults to string-matching fallbacks for unknown models. """ + import litellm + + try: + model_info = litellm.get_model_info(model=model) + if model_info and "cache_creation_min_tokens" in model_info: + return int(model_info["cache_creation_min_tokens"]) + except Exception: # noqa: BLE001 # fallback to string-matching heuristic if model lookup fails + pass + model_lower = model.lower() - if "gemini-1.5" in model_lower or "gemini-1-5" in model_lower: - return 32768 if "gemini-2.5" in model_lower or "gemini-2-5" in model_lower: return 2048 if "gemini-3" in model_lower: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4105fc0bcf5..4e27b756b68 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1181,9 +1181,16 @@ class LiteLLMCompletionResponsesConfig: Normalize a Responses API object (Pydantic model or custom class) to a dictionary """ if hasattr(item, "model_dump"): - return item.model_dump() + if hasattr(item, "model_dump"): + try: + return item.model_dump(exclude_none=True) + except Exception: # noqa: BLE001 # fallback if custom model_dump does not accept exclude_none + return item.model_dump() elif hasattr(item, "dict"): - return item.dict() + try: + return item.dict(exclude_none=True) + except Exception: # noqa: BLE001 # fallback if custom dict does not accept exclude_none + return item.dict() target_attrs = ( "type", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c9d871fc41d..4964a04844d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13493,7 +13493,8 @@ "output_dbu_cost_per_token": 3.5714e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_creation_min_tokens": 2048 }, "databricks/databricks-gemini-2-5-pro": { "input_cost_per_token": 1.24999e-06, @@ -13510,7 +13511,8 @@ "output_dbu_cost_per_token": 0.000142857, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_creation_min_tokens": 2048 }, "databricks/databricks-gemma-3-12b": { "input_cost_per_token": 1.5000999999999998e-07, @@ -14682,7 +14684,8 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -14693,7 +14696,8 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "cache_creation_min_tokens": 2048 }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -17338,7 +17342,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini-2.5-flash-image": { "cache_read_input_token_cost": 3e-08, @@ -17382,7 +17387,8 @@ "supports_vision": true, "supports_web_search": false, "tpm": 8000000, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini-3-pro-image": { "input_cost_per_image": 0.0011, @@ -17422,7 +17428,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -17462,7 +17469,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.1-flash-image": { "input_cost_per_image": 0.00056, @@ -17500,7 +17508,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, @@ -17538,7 +17547,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -17586,7 +17596,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.1-flash-lite": { "cache_read_input_token_cost": 2.5e-08, @@ -17642,7 +17653,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, @@ -17697,7 +17709,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, @@ -17776,7 +17789,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini-2.5-flash-lite-preview-09-2025": { "cache_read_input_token_cost": 1e-08, @@ -17821,7 +17835,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini-2.5-flash-preview-09-2025": { "cache_read_input_token_cost": 7.5e-08, @@ -17866,7 +17881,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, @@ -18002,7 +18018,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, @@ -18046,7 +18063,8 @@ "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "cache_creation_min_tokens": 2048 }, "gemini-3-pro-preview": { "deprecation_date": "2026-03-26", @@ -18102,7 +18120,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -18159,7 +18178,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.1-pro-preview-customtools": { "cache_read_input_token_cost": 2e-07, @@ -18210,7 +18230,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -18265,7 +18286,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, @@ -18313,7 +18335,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.5-flash": { "cache_read_input_token_cost": 1.5e-07, @@ -18364,7 +18387,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.6-flash": { "cache_read_input_token_cost": 1.5e-07, @@ -18418,7 +18442,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -18475,7 +18500,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "cache_read_input_token_cost": 2e-07, @@ -18532,7 +18558,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, @@ -18567,7 +18594,8 @@ "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "cache_creation_min_tokens": 2048 }, "gemini-robotics-er-1.5-preview": { "cache_read_input_token_cost": 0, @@ -18671,7 +18699,8 @@ "supports_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 2048 }, "gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -18812,7 +18841,8 @@ "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", "supports_multimodal": true, - "tpm": 10000000 + "tpm": 10000000, + "cache_creation_min_tokens": 32768 }, "gemini/gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, @@ -18973,7 +19003,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-flash-image": { "cache_read_input_token_cost": 3e-08, @@ -19023,7 +19054,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-3-pro-image": { "input_cost_per_image": 0.0011, @@ -19066,7 +19098,8 @@ "search_context_size_high": 0.014 }, "web_search_billing_unit": "per_query", - "supports_reasoning": false + "supports_reasoning": false, + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -19109,7 +19142,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.1-flash-image": { "input_cost_per_token": 2.5e-07, @@ -19151,7 +19185,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.1-flash-image-preview": { "input_cost_per_token": 2.5e-07, @@ -19193,7 +19228,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, @@ -19281,7 +19317,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { "cache_read_input_token_cost": 1e-08, @@ -19328,7 +19365,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-flash-preview-09-2025": { "cache_read_input_token_cost": 7.5e-08, @@ -19375,7 +19413,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-flash-latest": { "cache_read_input_token_cost": 7.5e-08, @@ -19515,7 +19554,8 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, @@ -19527,7 +19567,8 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, @@ -19576,7 +19617,8 @@ "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "input_cost_per_token": 1.25e-06, @@ -19606,7 +19648,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "tpm": 800000 + "tpm": 800000, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-3-pro-preview": { "deprecation_date": "2026-03-09", @@ -19662,7 +19705,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -19712,7 +19756,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_token_cost": 2.5e-08, @@ -19770,7 +19815,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, @@ -19827,7 +19873,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, @@ -19879,7 +19926,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.5-flash": { "cache_read_input_token_cost": 1.5e-07, @@ -19933,7 +19981,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.6-flash": { "cache_read_input_token_cost": 1.5e-07, @@ -19990,7 +20039,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -20080,7 +20130,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-3.1-pro-preview-customtools": { "cache_read_input_token_cost": 2e-07, @@ -20137,7 +20188,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, @@ -20187,7 +20239,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -20270,7 +20323,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini-3.6-flash": { "cache_read_input_token_cost": 1.5e-07, @@ -20325,7 +20379,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, @@ -20361,7 +20416,8 @@ "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-exp-1114": { "input_cost_per_token": 0, @@ -20750,7 +20806,8 @@ "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 2048 }, "github_copilot/gemini-3-pro-preview": { "litellm_provider": "github_copilot", @@ -20760,7 +20817,8 @@ "mode": "chat", "supports_function_calling": true, "supports_parallel_function_calling": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 4096 }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", @@ -21352,7 +21410,8 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 4096 }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, @@ -21363,7 +21422,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 4096 }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -29453,7 +29513,8 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -29467,7 +29528,8 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "cache_creation_min_tokens": 2048 }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -29482,7 +29544,8 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -30510,7 +30573,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "openrouter/google/gemini-2.5-pro": { "input_cost_per_audio_token": 7e-07, @@ -30526,7 +30590,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 2048 }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -30567,7 +30632,8 @@ "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "cache_creation_min_tokens": 4096 }, "openrouter/google/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, @@ -30608,7 +30674,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "tpm": 800000 + "tpm": 800000, + "cache_creation_min_tokens": 4096 }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -30651,7 +30718,8 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "tpm": 800000 + "tpm": 800000, + "cache_creation_min_tokens": 4096 }, "openrouter/google/gemini-3.1-flash-lite": { "cache_read_input_token_cost": 2.5e-08, @@ -30694,7 +30762,8 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "tpm": 800000 + "tpm": 800000, + "cache_creation_min_tokens": 4096 }, "openrouter/google/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -30727,7 +30796,8 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_creation_min_tokens": 4096 }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 1.875e-06, @@ -32337,21 +32407,24 @@ "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "cache_creation_min_tokens": 4096 }, "perplexity/google/gemini-3-flash-preview": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "cache_creation_min_tokens": 4096 }, "perplexity/google/gemini-2.5-pro": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "cache_creation_min_tokens": 2048 }, "perplexity/google/gemini-2.5-flash": { "litellm_provider": "perplexity", @@ -32359,7 +32432,8 @@ "supports_web_search": true, "supports_reasoning": false, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "perplexity/xai/grok-4-1-fast-non-reasoning": { "litellm_provider": "perplexity", @@ -32864,7 +32938,8 @@ "supports_vision": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "cache_creation_min_tokens": 4096 }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -32942,7 +33017,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -35730,7 +35806,8 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "vercel_ai_gateway/google/gemini-2.5-pro": { "input_cost_per_token": 2.5e-06, @@ -35743,7 +35820,8 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "cache_creation_min_tokens": 2048 }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -37435,7 +37513,8 @@ "supports_vision": true, "supports_web_search": false, "tpm": 8000000, - "supports_image_size": false + "supports_image_size": false, + "cache_creation_min_tokens": 2048 }, "vertex_ai/gemini-3-pro-image": { "input_cost_per_image": 0.0011, @@ -37451,7 +37530,8 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" + "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -37467,7 +37547,8 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" + "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.1-flash-image": { "input_cost_per_image": 0.00056, @@ -37481,7 +37562,8 @@ "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, @@ -37495,7 +37577,8 @@ "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -37543,7 +37626,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.1-flash-lite": { "cache_read_input_token_cost": 2.5e-08, @@ -37599,7 +37683,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, @@ -37654,7 +37739,8 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_creation_min_tokens": 4096 }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, @@ -44259,7 +44345,8 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "cache_creation_min_tokens": 2048 }, "gemini-2.5-flash-native-audio-preview-09-2025": { "input_cost_per_audio_token": 1e-06, @@ -44284,7 +44371,8 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "cache_creation_min_tokens": 2048 }, "gemini-2.5-flash-native-audio-preview-12-2025": { "input_cost_per_audio_token": 1e-06, @@ -44309,7 +44397,8 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "cache_creation_min_tokens": 2048 }, "gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -44342,7 +44431,8 @@ "supports_function_calling": true, "supports_vision": true, "supports_web_search": true, - "gemini_audio_only_live": true + "gemini_audio_only_live": true, + "cache_creation_min_tokens": 4096 }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 1e-06, @@ -44369,7 +44459,8 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-flash-native-audio-preview-09-2025": { "input_cost_per_audio_token": 1e-06, @@ -44396,7 +44487,8 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-2.5-flash-native-audio-preview-12-2025": { "input_cost_per_audio_token": 1e-06, @@ -44423,7 +44515,8 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "cache_creation_min_tokens": 2048 }, "gemini/gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -44458,7 +44551,8 @@ "supports_web_search": true, "tpm": 250000, "rpm": 10, - "gemini_audio_only_live": true + "gemini_audio_only_live": true, + "cache_creation_min_tokens": 4096 }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, @@ -44468,7 +44562,8 @@ "source": "https://ai.google.dev/pricing", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "cache_creation_min_tokens": 2048 }, "gemini-flash-latest": { "cache_read_input_token_cost": 3e-08, @@ -45963,4 +46058,4 @@ } ] } -} +} \ No newline at end of file diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py index af2f0684e69..4896a75ade7 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py @@ -30,6 +30,16 @@ class TestGeminiContextCachingMinTokens: def test_min_tokens_by_model(self, model, expected): assert get_gemini_context_caching_min_tokens(model) == expected + def test_min_tokens_from_model_info(self, monkeypatch): + """Should prefer cache_creation_min_tokens from model_info if present.""" + import litellm + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, **kwargs: {"cache_creation_min_tokens": 12345} + ) + assert get_gemini_context_caching_min_tokens("gemini-1.5-pro") == 12345 + class TestTTLValidation: """Test TTL format validation""" @@ -70,6 +80,7 @@ class TestTTLNormalization: [ ("3600s", "3600s"), ("1.5s", "1.5s"), + ("1.3333333333333333s", "1.333333333s"), ("5m", "300s"), ("90m", "5400s"), ("1h", "3600s"), diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index a1a9448cc58..1840f87360b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -712,6 +712,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_computer_use": {"type": "boolean"}, "cache_creation_input_audio_token_cost": {"type": "number"}, "cache_creation_input_token_cost": {"type": "number"}, + "cache_creation_min_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_1hr": {"type": "number"}, "cache_creation_input_token_cost_above_200k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, From 4a307e6759e539f7b010f6f4a94e6db964922a1a Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Wed, 22 Jul 2026 16:47:03 -0400 Subject: [PATCH 009/393] fix(responses): guard against None cache_control on content blocks --- .../transformation.py | 19 ++++++++----- .../test_litellm_completion_responses.py | 27 +++++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4e27b756b68..7e126099391 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1158,7 +1158,7 @@ class LiteLLMCompletionResponsesConfig: file_dict["file_data"] = item["file_data"] new_item: dict[str, Any] = {"type": "file", "file": file_dict} - if "cache_control" in item: + if item.get("cache_control") is not None: new_item["cache_control"] = item["cache_control"] return new_item @@ -1180,16 +1180,15 @@ class LiteLLMCompletionResponsesConfig: """ Normalize a Responses API object (Pydantic model or custom class) to a dictionary """ - if hasattr(item, "model_dump"): if hasattr(item, "model_dump"): try: return item.model_dump(exclude_none=True) - except Exception: # noqa: BLE001 # fallback if custom model_dump does not accept exclude_none + except TypeError: return item.model_dump() elif hasattr(item, "dict"): try: return item.dict(exclude_none=True) - except Exception: # noqa: BLE001 # fallback if custom dict does not accept exclude_none + except TypeError: return item.dict() target_attrs = ( @@ -1247,7 +1246,11 @@ class LiteLLMCompletionResponsesConfig: elif item.get("type") == "input_image": image_block = { **LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item), - **({"cache_control": item["cache_control"]} if "cache_control" in item else {}), + **( + {"cache_control": item["cache_control"]} + if item.get("cache_control") is not None + else {} + ), } content_list.append(image_block) else: @@ -1260,7 +1263,11 @@ class LiteLLMCompletionResponsesConfig: item.get("type") or "text" ), "text": text_value, - **({"cache_control": item["cache_control"]} if "cache_control" in item else {}), + **( + {"cache_control": item["cache_control"]} + if item.get("cache_control") is not None + else {} + ), } content_list.append(content_block) return content_list diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 62bc61c1a63..ae2e624ef70 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2635,6 +2635,33 @@ class TestCacheControlPreservation: ) assert messages == [{"role": "user", "content": "hello", "cache_control": {"type": "ephemeral"}}] + def test_none_cache_control_omitted_from_content_blocks(self): + from pydantic import BaseModel + from typing import Optional + + class MockContentItem(BaseModel): + type: str = "text" + text: str = "hello world" + cache_control: Optional[dict] = None + + item = MockContentItem() + result = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([item]) + assert isinstance(result, list) + assert len(result) == 1 + assert "cache_control" not in result[0] + + dict_item = {"type": "text", "text": "hello", "cache_control": None} + result_dict = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([dict_item]) + assert isinstance(result_dict, list) + assert len(result_dict) == 1 + assert "cache_control" not in result_dict[0] + + image_item = {"type": "input_image", "image_url": "https://example.com/a.png", "cache_control": None} + result_img = LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([image_item]) + assert isinstance(result_img, list) + assert len(result_img) == 1 + assert "cache_control" not in result_img[0] + def test_tool_call_output_as_custom_object(self): """Test _transform_responses_api_tool_call_output_to_chat_completion_message with a custom object.""" class MockToolCallOutput: From 0fd8a42972de6413c1c4467e80c582fa689885e8 Mon Sep 17 00:00:00 2001 From: Gaurav Pandey <112387553+gaurav-pandey-zocdoc@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:22:03 +0530 Subject: [PATCH 010/393] fix(alerting): clarify budget threshold messages Generated with AI Co-Authored-By: Claude Code --- litellm/integrations/SlackAlerting/slack_alerting.py | 4 ++-- .../SlackAlerting/test_slack_alerting.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 94d734546be..432cbd0917b 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -632,10 +632,10 @@ class SlackAlerting(CustomBatchLogger): event_message += f"Budget Crossed\n Total Budget:`{user_info.max_budget}`" elif percent_left <= SLACK_ALERTING_THRESHOLD_5_PERCENT: event = "threshold_crossed" - event_message += "5% Threshold Crossed " + event_message += "5% or less of budget remaining" elif percent_left <= SLACK_ALERTING_THRESHOLD_15_PERCENT: event = "threshold_crossed" - event_message += "15% Threshold Crossed" + event_message += "15% or less of budget remaining" return event, event_message diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 55e2dcdc270..6eaa3147dc3 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -10,6 +10,7 @@ import pytest import litellm from litellm.caching.caching import DualCache +from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.proxy._types import CallInfo, Litellm_EntityType from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys @@ -46,9 +47,8 @@ class TestSlackAlerting(unittest.TestCase): self.assertEqual(result, -0.2) def test_get_event_and_event_message_max_budget(self): - # Initial setup with no event event = None - event_message = "Test Message: " + event_message = get_budget_alert_type("user_budget").get_event_message() # Test case 1: When spend exceeds max_budget user_info = CallInfo( @@ -63,7 +63,7 @@ class TestSlackAlerting(unittest.TestCase): self.assertEqual(event, "budget_crossed") self.assertTrue("Budget Crossed" in event_message) - # Test case 2: When 5% of max_budget is left + event_message = get_budget_alert_type("user_budget").get_event_message() user_info = CallInfo( max_budget=100.0, spend=95.0, @@ -74,9 +74,9 @@ class TestSlackAlerting(unittest.TestCase): user_info=user_info, event=event, event_message=event_message ) self.assertEqual(event, "threshold_crossed") - self.assertTrue("5% Threshold Crossed" in event_message) + self.assertEqual(event_message, "User Budget: 5% or less of budget remaining") - # Test case 3: When 15% of max_budget is left + event_message = get_budget_alert_type("user_budget").get_event_message() user_info = CallInfo( max_budget=100.0, spend=85.0, @@ -87,7 +87,7 @@ class TestSlackAlerting(unittest.TestCase): user_info=user_info, event=event, event_message=event_message ) self.assertEqual(event, "threshold_crossed") - self.assertTrue("15% Threshold Crossed" in event_message) + self.assertEqual(event_message, "User Budget: 15% or less of budget remaining") def test_get_event_and_event_message_soft_budget(self): # Initial setup with no event From 2286bf3eca414cc24e0a03b008a7a4e6b9647c44 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 18:30:16 +0900 Subject: [PATCH 011/393] fix(router): stamp model_group when retrieving a batch Batch token usage is accounted on the retrieve call, not on create: a provider only reports token counts once the job finishes, so the usage arrives on aretrieve_batch and that is the spend log row the tokens land on. Router.acreate_batch stamps the requested model group into its metadata, but Router.aretrieve_batch never did. A batch is retrieved by id, so the request carries no model, and the router fans the lookup out over its deployments - leaving model_group unset on the one record that carries the tokens. /global/activity/model groups the spend logs by model_group, so every batch's tokens were bucketed under an empty group. Stamp the model group inside the per-deployment retrieve attempt, preferring an explicitly requested group and otherwise using the model_name of the deployment that answered, which is unambiguous even when the request named no model. An existing model_group in the metadata is left untouched, so nothing that already resolves a group changes. Scope is limited to aretrieve_batch: acompletion, aresponses and acreate_batch logging are untouched, and cost/spend attribution by model is unchanged. Signed-off-by: mynkyu --- litellm/router.py | 9 ++ .../test_router_batch_retrieve_model_group.py | 118 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py diff --git a/litellm/router.py b/litellm/router.py index 3f450661946..c6c25b6be17 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6137,6 +6137,8 @@ class Router: """ try: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) + requested_model_group: Final = model + metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch") if model is not None: filtered_model_list: ( list[DeploymentTypedDict] | list[dict] | dict | None @@ -6173,6 +6175,13 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) + ## STAMP THE MODEL GROUP FOR SPEND TRACKING ## + # A batch is retrieved by id, so the request carries no model group of its + # own - only the deployment that answered knows it. Batch token usage lands + # on this retrieve call (the provider reports counts once the job finishes), + # so without this the tokens are logged under an empty model_group. + model_group: Final = requested_model_group or model_name["model_name"] + new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) data.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py new file mode 100644 index 00000000000..ef8a23e4917 --- /dev/null +++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py @@ -0,0 +1,118 @@ +""" +model_group attribution on router batch retrieval. + +Batch token usage is accounted on the *retrieve* call, not on create: the +provider only knows the token counts once the job finishes, so +`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the +spend log tokens land on. + +`aretrieve_batch` is addressed by batch_id, so the request carries no model, +and the router fans the lookup out across its deployments. These tests lock +that the winning deployment's model group is stamped on the emitted +StandardLoggingPayload, so `/global/activity/model` - which groups the spend +logs by `model_group` - can attribute those tokens instead of bucketing every +batch under "". +""" + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +import litellm.batches.main as bm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import LiteLLMBatch, Usage + +MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev" +DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite" + + +class _PayloadCollector(CustomLogger): + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs.get("standard_logging_object")) + + +@pytest.fixture +def router(): + return Router( + model_list=[ + { + "model_name": MODEL_GROUP, + "litellm_params": { + "model": DEPLOYMENT_MODEL, + "vertex_project": "fake-project", + "vertex_location": "us-central1", + "vertex_credentials": "fake-creds", + }, + } + ] + ) + + +@pytest.fixture +def collector(): + logger = _PayloadCollector() + previous = litellm.callbacks + litellm.callbacks = [logger] + try: + yield logger + finally: + litellm.callbacks = previous + + +@pytest.fixture +def vertex_retrieve(): + """Mock the vertex provider seam - the only real network boundary.""" + batch = LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-1", + object="batch", + status="completed", + usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), + ) + seam = MagicMock(name="vertex_ai_batches_instance") + seam.retrieve_batch.return_value = batch + with patch.object(bm, "vertex_ai_batches_instance", seam): + yield seam + + +async def _collected_payload(collector) -> dict: + for _ in range(50): # the success handler runs as a background task + payloads = [p for p in collector.payloads if p is not None] + if payloads: + return payloads[-1] + await asyncio.sleep(0.05) + raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}") + + +@pytest.mark.asyncio +async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve): + """ + The proxy retrieves a managed batch by id only - no `model` in the request. + The router fans out over its deployments, so the model group is only known + from the deployment that answered. + """ + response = await router.aretrieve_batch(batch_id="batch-1") + + assert response.usage.total_tokens == 1200 + payload = await _collected_payload(collector) + assert payload["model"] == DEPLOYMENT_MODEL + assert payload["model_group"] == MODEL_GROUP + + +@pytest.mark.asyncio +async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve): + """An explicitly requested model group is what gets logged.""" + await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1") + + payload = await _collected_payload(collector) + assert payload["model_group"] == MODEL_GROUP From e630f21d16b10b78e22c28a974dee73009749167 Mon Sep 17 00:00:00 2001 From: mynkyu Date: Thu, 27 Aug 2026 19:01:18 +0900 Subject: [PATCH 012/393] test: fake the provider at the HTTP boundary in the batch model_group test The test-quality gate flagged the first version for patching an SDK internal (litellm.batches.main.vertex_ai_batches_instance) and for writing litellm.callbacks directly. Drive an openai-compatible deployment through respx instead, so the retrieve call and the usage accounting that reads the completed batch's output file both run for real, and install the collector with monkeypatch so nothing leaks into the next test. Signed-off-by: mynkyu --- .../test_router_batch_retrieve_model_group.py | 146 +++++++++++------- 1 file changed, 89 insertions(+), 57 deletions(-) diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py index ef8a23e4917..b99ec50e041 100644 --- a/tests/test_litellm/test_router_batch_retrieve_model_group.py +++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py @@ -1,35 +1,81 @@ """ model_group attribution on router batch retrieval. -Batch token usage is accounted on the *retrieve* call, not on create: the -provider only knows the token counts once the job finishes, so -`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the -spend log tokens land on. +Batch token usage is accounted on the *retrieve* call, not on create: a provider +only reports token counts once the job finishes, so the usage is read off the +completed batch's output file during retrieve logging and that is the spend log +row the tokens land on. -`aretrieve_batch` is addressed by batch_id, so the request carries no model, -and the router fans the lookup out across its deployments. These tests lock -that the winning deployment's model group is stamped on the emitted -StandardLoggingPayload, so `/global/activity/model` - which groups the spend -logs by `model_group` - can attribute those tokens instead of bucketing every -batch under "". +A batch is retrieved by id, so the request carries no model and the router fans +the lookup out across its deployments. These tests lock that the answering +deployment's model group is stamped on the emitted StandardLoggingPayload, so +`/global/activity/model` - which groups the spend logs by `model_group` - can +attribute those tokens instead of bucketing every batch under "". + +The provider is faked at the HTTP boundary, so the whole retrieve + usage +accounting path runs for real. """ import asyncio -from unittest.mock import MagicMock, patch +import json +import httpx import pytest +import respx import litellm -import litellm.batches.main as bm from litellm import Router from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import LiteLLMBatch, Usage -MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev" -DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite" +MODEL_GROUP = "gemini-batch-group" +DEPLOYMENT_MODEL = "openai/gpt-4o-mini" +API_BASE = "http://localhost:4001/v1" +BATCH_ID = "batch-1" +ROWS = 2 +TOKENS_PER_ROW = 600 + +COMPLETED_BATCH = { + "id": BATCH_ID, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-1", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-1", + "error_file_id": None, + "created_at": 0, + "completed_at": 1, + "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0}, + "metadata": None, +} + +OUTPUT_JSONL = "\n".join( + json.dumps( + { + "id": f"req-{row}", + "custom_id": f"row-{row}", + "response": { + "status_code": 200, + "body": { + "id": f"chatcmpl-{row}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW}, + }, + }, + } + ) + for row in range(ROWS) +) class _PayloadCollector(CustomLogger): + """Captures the StandardLoggingPayload the spend log is built from.""" + def __init__(self): super().__init__() self.payloads = [] @@ -37,6 +83,14 @@ class _PayloadCollector(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.payloads.append(kwargs.get("standard_logging_object")) + async def retrieve_batch_payload(self) -> dict: + for _ in range(100): # the success handler runs as a background task + for payload in self.payloads: + if payload and payload.get("call_type") == "aretrieve_batch": + return payload + await asyncio.sleep(0.05) + raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") + @pytest.fixture def router(): @@ -46,9 +100,8 @@ def router(): "model_name": MODEL_GROUP, "litellm_params": { "model": DEPLOYMENT_MODEL, - "vertex_project": "fake-project", - "vertex_location": "us-central1", - "vertex_credentials": "fake-creds", + "api_base": API_BASE, + "api_key": "sk-fake", }, } ] @@ -56,63 +109,42 @@ def router(): @pytest.fixture -def collector(): +def collector(monkeypatch): logger = _PayloadCollector() - previous = litellm.callbacks - litellm.callbacks = [logger] - try: - yield logger - finally: - litellm.callbacks = previous + monkeypatch.setattr(litellm, "callbacks", [logger]) + return logger @pytest.fixture -def vertex_retrieve(): - """Mock the vertex provider seam - the only real network boundary.""" - batch = LiteLLMBatch( - id="batch-1", - completion_window="24h", - created_at=0, - endpoint="/v1/chat/completions", - input_file_id="file-1", - object="batch", - status="completed", - usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), - ) - seam = MagicMock(name="vertex_ai_batches_instance") - seam.retrieve_batch.return_value = batch - with patch.object(bm, "vertex_ai_batches_instance", seam): - yield seam - - -async def _collected_payload(collector) -> dict: - for _ in range(50): # the success handler runs as a background task - payloads = [p for p in collector.payloads if p is not None] - if payloads: - return payloads[-1] - await asyncio.sleep(0.05) - raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}") +def provider(): + """Fake the provider at the HTTP boundary: the completed batch plus the + output file the usage accounting reads.""" + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH)) + respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL)) + yield respx_mock @pytest.mark.asyncio -async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve): +async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider): """ The proxy retrieves a managed batch by id only - no `model` in the request. The router fans out over its deployments, so the model group is only known from the deployment that answered. """ - response = await router.aretrieve_batch(batch_id="batch-1") + response = await router.aretrieve_batch(batch_id=BATCH_ID) - assert response.usage.total_tokens == 1200 - payload = await _collected_payload(collector) + assert response.id == BATCH_ID + payload = await collector.retrieve_batch_payload() + assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW assert payload["model"] == DEPLOYMENT_MODEL assert payload["model_group"] == MODEL_GROUP @pytest.mark.asyncio -async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve): +async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider): """An explicitly requested model group is what gets logged.""" - await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1") + await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID) - payload = await _collected_payload(collector) + payload = await collector.retrieve_batch_payload() assert payload["model_group"] == MODEL_GROUP From df6990c7127a0c30c77b96323e8311d9200a23be Mon Sep 17 00:00:00 2001 From: mynkyu Date: Sun, 6 Sep 2026 10:10:07 +0900 Subject: [PATCH 013/393] test: move the batch model_group regression into test_router.py CLAUDE.md asks bug fixes to extend the existing mapped test file rather than add a new one, and tests/test_litellm/test_router.py already covers Router.aretrieve_batch. Fold the two cases in next to that coverage and drop the standalone file. The helpers are prefixed so they read unambiguously in a shared file, and the respx context stays open while the payload is awaited, since the usage accounting reads the batch's output file from the success handler. Signed-off-by: mynkyu --- tests/test_litellm/test_router.py | 153 ++++++++++++++++++ .../test_router_batch_retrieve_model_group.py | 150 ----------------- 2 files changed, 153 insertions(+), 150 deletions(-) delete mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7c044310e14..9b146092927 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -958,6 +958,159 @@ async def test_arouter_aretrieve_batch(): assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base" +# --------------------------------------------------------------------------- +# Batch retrieval has to attribute its tokens to a model group. +# +# Batch token usage is accounted on the *retrieve* call, not on create: a +# provider only reports token counts once the job finishes, so the usage is read +# off the completed batch's output file during retrieve logging, and that is the +# spend log row the tokens land on. A batch is retrieved by id, so the request +# carries no model and the router fans the lookup out across its deployments - +# the group of the deployment that answered is the only one there is to stamp. +# Leaving it unset files every batch's tokens under an empty model_group, which +# is what /global/activity/model groups the spend logs by. +# +# The provider is faked at the HTTP boundary, so the retrieve call and the usage +# accounting that reads the output file both run for real. +# --------------------------------------------------------------------------- + +_BATCH_GROUP = "gemini-batch-group" +_BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini" +_BATCH_API_BASE = "http://localhost:4001/v1" +_BATCH_ID = "batch-1" +_BATCH_ROWS = 2 +_BATCH_TOKENS_PER_ROW = 600 + +_BATCH_COMPLETED = { + "id": _BATCH_ID, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-1", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-1", + "error_file_id": None, + "created_at": 0, + "completed_at": 1, + "request_counts": {"total": _BATCH_ROWS, "completed": _BATCH_ROWS, "failed": 0}, + "metadata": None, +} + +_BATCH_OUTPUT_JSONL = "\n".join( + json.dumps( + { + "id": f"req-{row}", + "custom_id": f"row-{row}", + "response": { + "status_code": 200, + "body": { + "id": f"chatcmpl-{row}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": 500, + "completion_tokens": 100, + "total_tokens": _BATCH_TOKENS_PER_ROW, + }, + }, + }, + } + ) + for row in range(_BATCH_ROWS) +) + + +class _BatchPayloadCollector(CustomLogger): + """Captures the StandardLoggingPayload the spend log row is built from.""" + + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs.get("standard_logging_object")) + + async def retrieve_batch_payload(self): + for _ in range(100): # the success handler runs as a background task + for payload in self.payloads: + if payload and payload.get("call_type") == "aretrieve_batch": + return payload + await asyncio.sleep(0.05) + raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") + + +def _batch_model_group_router(): + return litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + } + ] + ) + + +def _mock_batch_provider(respx_mock): + """The completed batch, plus the output file the usage accounting reads.""" + respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(200, json=_BATCH_COMPLETED) + ) + respx_mock.get(f"{_BATCH_API_BASE}/files/file-out-1/content").mock( + return_value=httpx.Response(200, text=_BATCH_OUTPUT_JSONL) + ) + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypatch: pytest.MonkeyPatch): + """ + The proxy retrieves a managed batch by id only - no `model` in the request. + The router fans out over its deployments, so the model group is only known + from the deployment that answered. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = _batch_model_group_router() + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + # the usage accounting reads the output file from the success handler, + # so the provider has to stay faked until that payload lands + payload = await collector.retrieve_batch_payload() + + assert response.id == _BATCH_ID + assert payload["total_tokens"] == _BATCH_ROWS * _BATCH_TOKENS_PER_ROW + assert payload["model"] == _BATCH_DEPLOYMENT_MODEL + assert payload["model_group"] == _BATCH_GROUP + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(monkeypatch: pytest.MonkeyPatch): + """An explicitly requested model group is what gets logged.""" + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = _batch_model_group_router() + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + await router.aretrieve_batch(model=_BATCH_GROUP, batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + + assert payload["model_group"] == _BATCH_GROUP + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py deleted file mode 100644 index b99ec50e041..00000000000 --- a/tests/test_litellm/test_router_batch_retrieve_model_group.py +++ /dev/null @@ -1,150 +0,0 @@ -""" -model_group attribution on router batch retrieval. - -Batch token usage is accounted on the *retrieve* call, not on create: a provider -only reports token counts once the job finishes, so the usage is read off the -completed batch's output file during retrieve logging and that is the spend log -row the tokens land on. - -A batch is retrieved by id, so the request carries no model and the router fans -the lookup out across its deployments. These tests lock that the answering -deployment's model group is stamped on the emitted StandardLoggingPayload, so -`/global/activity/model` - which groups the spend logs by `model_group` - can -attribute those tokens instead of bucketing every batch under "". - -The provider is faked at the HTTP boundary, so the whole retrieve + usage -accounting path runs for real. -""" - -import asyncio -import json - -import httpx -import pytest -import respx - -import litellm -from litellm import Router -from litellm.integrations.custom_logger import CustomLogger - -MODEL_GROUP = "gemini-batch-group" -DEPLOYMENT_MODEL = "openai/gpt-4o-mini" -API_BASE = "http://localhost:4001/v1" -BATCH_ID = "batch-1" -ROWS = 2 -TOKENS_PER_ROW = 600 - -COMPLETED_BATCH = { - "id": BATCH_ID, - "object": "batch", - "endpoint": "/v1/chat/completions", - "errors": None, - "input_file_id": "file-in-1", - "completion_window": "24h", - "status": "completed", - "output_file_id": "file-out-1", - "error_file_id": None, - "created_at": 0, - "completed_at": 1, - "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0}, - "metadata": None, -} - -OUTPUT_JSONL = "\n".join( - json.dumps( - { - "id": f"req-{row}", - "custom_id": f"row-{row}", - "response": { - "status_code": 200, - "body": { - "id": f"chatcmpl-{row}", - "object": "chat.completion", - "model": "gpt-4o-mini", - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} - ], - "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW}, - }, - }, - } - ) - for row in range(ROWS) -) - - -class _PayloadCollector(CustomLogger): - """Captures the StandardLoggingPayload the spend log is built from.""" - - def __init__(self): - super().__init__() - self.payloads = [] - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.payloads.append(kwargs.get("standard_logging_object")) - - async def retrieve_batch_payload(self) -> dict: - for _ in range(100): # the success handler runs as a background task - for payload in self.payloads: - if payload and payload.get("call_type") == "aretrieve_batch": - return payload - await asyncio.sleep(0.05) - raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") - - -@pytest.fixture -def router(): - return Router( - model_list=[ - { - "model_name": MODEL_GROUP, - "litellm_params": { - "model": DEPLOYMENT_MODEL, - "api_base": API_BASE, - "api_key": "sk-fake", - }, - } - ] - ) - - -@pytest.fixture -def collector(monkeypatch): - logger = _PayloadCollector() - monkeypatch.setattr(litellm, "callbacks", [logger]) - return logger - - -@pytest.fixture -def provider(): - """Fake the provider at the HTTP boundary: the completed batch plus the - output file the usage accounting reads.""" - with respx.mock(assert_all_called=True) as respx_mock: - respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH)) - respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL)) - yield respx_mock - - -@pytest.mark.asyncio -async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider): - """ - The proxy retrieves a managed batch by id only - no `model` in the request. - The router fans out over its deployments, so the model group is only known - from the deployment that answered. - """ - response = await router.aretrieve_batch(batch_id=BATCH_ID) - - assert response.id == BATCH_ID - payload = await collector.retrieve_batch_payload() - assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW - assert payload["model"] == DEPLOYMENT_MODEL - assert payload["model_group"] == MODEL_GROUP - - -@pytest.mark.asyncio -async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider): - """An explicitly requested model group is what gets logged.""" - await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID) - - payload = await collector.retrieve_batch_payload() - assert payload["model_group"] == MODEL_GROUP From 128cb114bdac6b8cf41a9d689f0a573a2e27eced Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:45:47 -0700 Subject: [PATCH 014/393] style: trim comments on batch retrieve model group stamp --- litellm/router.py | 7 ++----- tests/test_litellm/test_router.py | 16 ---------------- 2 files changed, 2 insertions(+), 21 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index c6c25b6be17..20c9018abee 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6175,11 +6175,8 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) - ## STAMP THE MODEL GROUP FOR SPEND TRACKING ## - # A batch is retrieved by id, so the request carries no model group of its - # own - only the deployment that answered knows it. Batch token usage lands - # on this retrieve call (the provider reports counts once the job finishes), - # so without this the tokens are logged under an empty model_group. + # A batch is retrieved by id, so only the deployment that answered knows the + # group, and batch token usage is logged on this retrieve call. model_group: Final = requested_model_group or model_name["model_name"] new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9b146092927..258d1c973d0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -958,22 +958,6 @@ async def test_arouter_aretrieve_batch(): assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base" -# --------------------------------------------------------------------------- -# Batch retrieval has to attribute its tokens to a model group. -# -# Batch token usage is accounted on the *retrieve* call, not on create: a -# provider only reports token counts once the job finishes, so the usage is read -# off the completed batch's output file during retrieve logging, and that is the -# spend log row the tokens land on. A batch is retrieved by id, so the request -# carries no model and the router fans the lookup out across its deployments - -# the group of the deployment that answered is the only one there is to stamp. -# Leaving it unset files every batch's tokens under an empty model_group, which -# is what /global/activity/model groups the spend logs by. -# -# The provider is faked at the HTTP boundary, so the retrieve call and the usage -# accounting that reads the output file both run for real. -# --------------------------------------------------------------------------- - _BATCH_GROUP = "gemini-batch-group" _BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini" _BATCH_API_BASE = "http://localhost:4001/v1" From 33d89c9814641a34cb66d357e2cc3a403677a06d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:55:11 -0700 Subject: [PATCH 015/393] style: drop redundant comments per repo comment policy --- litellm/router.py | 3 +-- tests/test_litellm/test_router.py | 3 --- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 20c9018abee..e85ca1bd7a8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6175,8 +6175,7 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) - # A batch is retrieved by id, so only the deployment that answered knows the - # group, and batch token usage is logged on this retrieve call. + # Batch token usage is logged on this retrieve call, not on create. model_group: Final = requested_model_group or model_name["model_name"] new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 258d1c973d0..dc51339cf55 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1009,8 +1009,6 @@ _BATCH_OUTPUT_JSONL = "\n".join( class _BatchPayloadCollector(CustomLogger): - """Captures the StandardLoggingPayload the spend log row is built from.""" - def __init__(self): super().__init__() self.payloads = [] @@ -1043,7 +1041,6 @@ def _batch_model_group_router(): def _mock_batch_provider(respx_mock): - """The completed batch, plus the output file the usage accounting reads.""" respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( return_value=httpx.Response(200, json=_BATCH_COMPLETED) ) From 01bdfb34aa5ed88320dbd1c9f231876df820ca29 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:24:20 -0700 Subject: [PATCH 016/393] chore: drop redundant comments in aretrieve_batch router tests --- tests/test_litellm/test_router.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index dc51339cf55..bfa9ea7e3d1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1017,7 +1017,7 @@ class _BatchPayloadCollector(CustomLogger): self.payloads.append(kwargs.get("standard_logging_object")) async def retrieve_batch_payload(self): - for _ in range(100): # the success handler runs as a background task + for _ in range(100): for payload in self.payloads: if payload and payload.get("call_type") == "aretrieve_batch": return payload @@ -1065,8 +1065,6 @@ async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypa with respx.mock(assert_all_called=True) as respx_mock: _mock_batch_provider(respx_mock) response = await router.aretrieve_batch(batch_id=_BATCH_ID) - # the usage accounting reads the output file from the success handler, - # so the provider has to stay faked until that payload lands payload = await collector.retrieve_batch_payload() assert response.id == _BATCH_ID From 9acf09f60d4f917053e1bfb9d493dce3cdd2771a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:07:55 -0700 Subject: [PATCH 017/393] fix(router): keep batch retrieves out of the per-minute tpm/rpm counters Stamping model_group let both router deployment callbacks past their `model_group is None` early return for batch retrieves. A batch reports the whole job's token total on retrieve and reports it again on every poll of the finished batch, so those tokens are not load in the current minute: three polls of one completed 1,200 token batch pushed a tpm:1000 deployment to 3,600. The fan-out also probed unrelated deployments, adding an rpm tick to each. --- litellm/router.py | 5 ++ litellm/router_utils/batch_utils.py | 18 +++++++ tests/test_litellm/test_router.py | 76 +++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index e85ca1bd7a8..9dd267d7560 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -124,6 +124,7 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, + is_batch_retrieve_call_type, replace_model_in_jsonl, should_replace_model_in_jsonl, ) @@ -7878,6 +7879,8 @@ class Router: # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object is None") @@ -8117,6 +8120,8 @@ class Router: """ Update RPM usage for a deployment """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return deployment_name: Final = kwargs["litellm_params"]["metadata"].get( "deployment", None ) # handles wildcard routes - by giving the original name sent to `litellm.completion` diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index ccb6ad95519..6e110b586fb 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -5,6 +5,7 @@ from typing import Final from litellm._logging import verbose_logger from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose +from litellm.types.utils import CallTypes class InMemoryFile(io.BytesIO): @@ -170,3 +171,20 @@ def _get_router_metadata_variable_name(function_name: str | None) -> str: return "litellm_metadata" else: return "metadata" + + +BATCH_RETRIEVE_CALL_TYPES: Final = frozenset( + { + CallTypes.aretrieve_batch.value, + CallTypes.retrieve_batch.value, + } +) + + +def is_batch_retrieve_call_type(call_type: object) -> bool: + """ + A batch retrieve reports the whole job's token usage, which the provider spent + asynchronously over the life of the batch, and reports it again on every poll of the + finished batch. Per-minute usage counters must not be fed from it. + """ + return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bfa9ea7e3d1..d8045722998 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1090,6 +1090,82 @@ async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(m assert payload["model_group"] == _BATCH_GROUP +_UNRELATED_BATCH_GROUP = "unrelated-batch-group" +_UNRELATED_BATCH_API_BASE = "http://localhost:4002/v1" + +_BATCH_NOT_FOUND = { + "error": { + "message": f"No batch found with id '{_BATCH_ID}'.", + "type": "invalid_request_error", + "code": "batch_not_found", + } +} + + +async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:")) + if keys: + return keys + await asyncio.sleep(0.05) + return [] + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch): + """ + A batch reports the whole job's tokens on retrieve, and reports them again on every + poll of the finished batch, so they are not a measure of load in the current minute. + The fan-out also probes deployments the caller never named. Neither may reach the + per-minute tpm/rpm counters that gate live traffic. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "batch-dep"}, + "tpm": 1000, + "rpm": 10, + }, + { + "model_name": _UNRELATED_BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _UNRELATED_BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "unrelated-dep"}, + "tpm": 1000, + "rpm": 10, + }, + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + usage_keys = await _router_usage_keys(router) + + assert response.id == _BATCH_ID + assert payload["model_group"] == _BATCH_GROUP + assert usage_keys == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ From 58c3d04733f2bebfbc15e8f1f6dd702a37c6e2f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:24:47 -0700 Subject: [PATCH 018/393] test: cover is_batch_retrieve_call_type in router batch utils --- .../router_unit_tests/test_router_batch_utils.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index c9f19731372..e274ac61a01 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -317,3 +317,18 @@ def test_replace_model_in_jsonl_with_embedded_newlines(): == "This is a message\nwith multiple\nlines" ) assert result_json["custom_id"] == "test123" + + +def test_is_batch_retrieve_call_type_matches_only_batch_retrieves(): + from litellm.router_utils.batch_utils import is_batch_retrieve_call_type + from litellm.types.utils import CallTypes + + assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True + assert is_batch_retrieve_call_type(CallTypes.retrieve_batch.value) is True + + for call_type in CallTypes: + if call_type in (CallTypes.aretrieve_batch, CallTypes.retrieve_batch): + continue + assert is_batch_retrieve_call_type(call_type.value) is False + + assert is_batch_retrieve_call_type(None) is False From ad2afe5e6568b69389c2258680274098b9191b6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:41:07 -0700 Subject: [PATCH 019/393] style(router): drop the inline comment on the batch retrieve stamp --- litellm/router.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 9dd267d7560..2397c6b2fc7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6176,7 +6176,6 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) - # Batch token usage is logged on this retrieve call, not on create. model_group: Final = requested_model_group or model_name["model_name"] new_kwargs[metadata_variable_name].setdefault("model_group", model_group) new_kwargs.pop("custom_llm_provider", None) From 63ea19743373fa1dd87081a66d64e5f580212f77 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:02:53 -0700 Subject: [PATCH 020/393] fix(router): keep batch retrieves out of routing strategy state Stamping model_group on a batch retrieve routed the whole batch job's token usage into the per-model-group counters that usage-based, latency-based, cost-based and least-busy routing read, so polling a finished batch could exhaust a group's TPM or RPM window and lock live chat traffic out with RouterRateLimitError. Polling also drove the least-busy in-flight counts negative once per poll per deployment, which pinned chat to whichever deployment had been polled most. The strategy callbacks now skip batch retrieve call types, so a retrieve still lands in spend logs under its model group while the numbers that pick a deployment for the next chat request stay driven by live traffic only. --- litellm/router.py | 3 +- litellm/router_strategy/least_busy.py | 11 ++++ litellm/router_strategy/lowest_cost.py | 5 ++ litellm/router_strategy/lowest_latency.py | 7 ++ litellm/router_strategy/lowest_tpm_rpm.py | 5 ++ litellm/router_utils/batch_utils.py | 3 +- tests/test_litellm/test_router.py | 79 +++++++++++++++++++++++ 7 files changed, 111 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2397c6b2fc7..8f913d96463 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6177,7 +6177,8 @@ class Router: function_name="aretrieve_batch", ) model_group: Final = requested_model_group or model_name["model_name"] - new_kwargs[metadata_variable_name].setdefault("model_group", model_group) + if not new_kwargs[metadata_variable_name].get("model_group"): + new_kwargs[metadata_variable_name]["model_group"] = model_group new_kwargs.pop("custom_llm_provider", None) data.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..e93288fd9fe 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -11,6 +11,7 @@ from typing import Final from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type class LeastBusyLoggingHandler(CustomLogger): @@ -27,6 +28,8 @@ class LeastBusyLoggingHandler(CustomLogger): Caching based on model group. """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -48,6 +51,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -76,6 +81,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass def log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -103,6 +110,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -131,6 +140,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index b927df0c438..aaad6186484 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -8,6 +8,7 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type class LowestCostLoggingHandler(CustomLogger): @@ -19,6 +20,8 @@ class LowestCostLoggingHandler(CustomLogger): self.router_cache = router_cache def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update usage on success @@ -96,6 +99,8 @@ class LowestCostLoggingHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update cost usage on success diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index a1b67eaeaf9..598ca1227ec 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -9,6 +9,7 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -35,6 +36,8 @@ class LowestLatencyLoggingHandler(CustomLogger): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update latency usage on success @@ -167,6 +170,8 @@ class LowestLatencyLoggingHandler(CustomLogger): """ Check if Timeout Error, if timeout set deployment latency -> 100 """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: metadata_field: Final = self._select_metadata_field(kwargs) _exception: Final = kwargs.get("exception", None) @@ -221,6 +226,8 @@ class LowestLatencyLoggingHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update latency usage on success diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 31c4b1d7e3f..d4abf1f8f70 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -8,6 +8,7 @@ from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose @@ -27,6 +28,8 @@ class LowestTPMLoggingHandler(CustomLogger): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success @@ -79,6 +82,8 @@ class LowestTPMLoggingHandler(CustomLogger): verbose_router_logger.debug(traceback.format_exc()) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 6e110b586fb..be20c358202 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -185,6 +185,7 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: """ A batch retrieve reports the whole job's token usage, which the provider spent asynchronously over the life of the batch, and reports it again on every poll of the - finished batch. Per-minute usage counters must not be fed from it. + finished batch. The counters that measure live traffic, per-minute rate limits and the + routing strategies' own state, must not be fed from it. """ return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d8045722998..836085c1f5d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1166,6 +1166,85 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m assert usage_keys == [] +_ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") + + +async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + keys = sorted( + key + for key in router.cache.in_memory_cache.cache_dict + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) + ) + if keys: + return keys + await asyncio.sleep(0.05) + return [] + + +def _batch_fan_out_router(routing_strategy: str): + return litellm.Router( + routing_strategy=routing_strategy, + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "batch-dep"}, + }, + { + "model_name": _UNRELATED_BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _UNRELATED_BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "unrelated-dep"}, + }, + ], + ) + + +@pytest.mark.parametrize( + "routing_strategy", + ["usage-based-routing", "latency-based-routing", "cost-based-routing", "least-busy"], +) +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( + monkeypatch: pytest.MonkeyPatch, routing_strategy: str +): + """ + Every routing strategy picks a deployment from what recent live traffic did. + A batch retrieve reports the whole job on every poll and probes deployments the + caller never named, so polling a finished batch must not move the numbers that + decide where the next chat request goes. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + monkeypatch.setattr(litellm, "input_callback", []) + router = _batch_fan_out_router(routing_strategy) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + for _ in range(3): + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + await collector.retrieve_batch_payload() + strategy_keys = await _router_strategy_keys(router) + + assert response.id == _BATCH_ID + assert strategy_keys == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ From 828a02f78f2c974f6458237c6fde1128b01b26bb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 03:03:45 -0700 Subject: [PATCH 021/393] fix(batches): stamp the model group on the proxy's model-encoded retrieve path The model-encoded batch id path calls the SDK directly, so the router never labels it. Stamp the decoded group into the request's litellm_metadata, and guard usage-based-routing-v2 the same way the other strategies already are. --- litellm/proxy/batches_endpoints/endpoints.py | 21 +++++++--- litellm/router_strategy/lowest_tpm_rpm_v2.py | 5 +++ .../proxy/batches_endpoints/test_endpoints.py | 40 ++++++++++++++++++- tests/test_litellm/test_router.py | 26 +++++++----- 4 files changed, 76 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..b6a8b421e02 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -52,6 +52,20 @@ from litellm.types.llms.openai import LiteLLMBatchCreateRequest router: Final = APIRouter() +def _litellm_metadata_of(data: dict) -> dict: + """The request's litellm_metadata mapping, created on the request when it carries none. + + The success handler reads this mapping, so a flag or a model group set here has to live + inside it rather than beside it. + """ + existing: Final = data.get("litellm_metadata") + if isinstance(existing, dict): + return existing + created: Final = {} # mutable-ok: the logging layer copies and extends this mapping, so it cannot be a read-only view + data["litellm_metadata"] = created + return created + + def _raise_not_found_when_openai_fallback_unservable( requested_provider: "str | None", data: Mapping[str, object], @@ -531,11 +545,7 @@ async def retrieve_batch( poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active() if poller_owns_accounting: - litellm_metadata = data.get("litellm_metadata") - if not isinstance(litellm_metadata, dict): - litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend - data["litellm_metadata"] = litellm_metadata - litellm_metadata["batch_ignore_default_logging"] = True + _litellm_metadata_of(data)["batch_ignore_default_logging"] = True # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info @@ -558,6 +568,7 @@ async def retrieve_batch( # so litellm.aretrieve_batch can load BedrockBatchesConfig. Without # it the call falls into the legacy provider switch and 400s. data["model"] = model_from_id + _litellm_metadata_of(data).setdefault("model_group", model_from_id) # Retrieve batch using model credentials response = await litellm.aretrieve_batch( diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 665ff69ab47..a2acce5fcb5 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -210,6 +211,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): return deployment # don't fail calls if eg. redis fails to connect def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success @@ -250,6 +253,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM usage on success diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..c6df8f2ffcf 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1233,9 +1233,11 @@ async def call_retrieve( user: Optional[UserAPIKeyAuth] = None, headers: Optional[Dict[str, str]] = None, query: Optional[Dict[str, str]] = None, + enriched_data: Optional[Dict[str, Any]] = None, ): - # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...). - harness.data["data"] = {"batch_id": batch_id} + # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...), + # then pre-call enrichment adds key/team metadata to it. + harness.data["data"] = {"batch_id": batch_id, **(enriched_data or {})} return await endpoints.retrieve_batch( request=FakeRequest(headers=headers, query=query), fastapi_response=Response(), @@ -1271,6 +1273,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness): "api_key": "sk-azure", "api_base": "https://azure.test", "model": "azure/gpt-4o", + "litellm_metadata": {"model_group": "azure/gpt-4o"}, } # 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip. @@ -1293,6 +1296,39 @@ async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment assert retrieve_harness.aretrieve_kwargs()["model"] == "azure/gpt-4o" +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__stamps_model_group(retrieve_harness): + """This path never goes through the router, so nothing else labels the call. + Without the stamp the spend log lands under a blank model group and the batch + disappears from per-model usage.""" + await call_retrieve(retrieve_harness, AZURE_BATCH_ID) + + litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"] + + assert litellm_metadata["model_group"] == "azure/gpt-4o" + + +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__stamps_model_group_beside_existing_metadata( + retrieve_harness, +): + """The stamp joins the metadata pre-call enrichment already built. Replacing + that dict instead of adding to it drops the key and team labels the spend log + is attributed with.""" + await call_retrieve( + retrieve_harness, + AZURE_BATCH_ID, + enriched_data={"litellm_metadata": {"user_api_key_alias": "team-a-key"}}, + ) + + litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"] + + assert litellm_metadata == { + "user_api_key_alias": "team-a-key", + "model_group": "azure/gpt-4o", + } + + @pytest.mark.asyncio async def test_retrieve__model_encoded_id__encodes_output_and_error_ids( retrieve_harness, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 836085c1f5d..e4e0dafd750 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1169,17 +1169,19 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") -async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]: +async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: loop = asyncio.get_event_loop() deadline = loop.time() + timeout while loop.time() < deadline: - keys = sorted( - key - for key in router.cache.in_memory_cache.cache_dict + cache_dict = router.cache.in_memory_cache.cache_dict + moved = sorted( + f"{key}={cache_dict[key]}" + for key in cache_dict if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) + and cache_dict[key] ) - if keys: - return keys + if moved: + return moved await asyncio.sleep(0.05) return [] @@ -1212,7 +1214,13 @@ def _batch_fan_out_router(routing_strategy: str): @pytest.mark.parametrize( "routing_strategy", - ["usage-based-routing", "latency-based-routing", "cost-based-routing", "least-busy"], + [ + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", + "least-busy", + ], ) @pytest.mark.asyncio async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( @@ -1239,10 +1247,10 @@ async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( for _ in range(3): response = await router.aretrieve_batch(batch_id=_BATCH_ID) await collector.retrieve_batch_payload() - strategy_keys = await _router_strategy_keys(router) + moved_counters = await _moved_routing_counters(router) assert response.id == _BATCH_ID - assert strategy_keys == [] + assert moved_counters == [] @pytest.mark.asyncio From 91761d984d55b2e4729e38c36ef7c7510c9ffc76 Mon Sep 17 00:00:00 2001 From: abhirup7 Date: Tue, 8 Sep 2026 00:30:00 +0530 Subject: [PATCH 022/393] fix(azure): send the resolved Entra ID token on image generation requests Azure image generation calls initialize_azure_sdk_client like the chat path does, but then sends the request through httpx with the headers it was given, so a credential resolved from litellm_params or the environment (Entra ID client credentials, managed or workload identity, OIDC, a static azure_ad_token) never reached the wire and Azure answered 401. Only an explicitly passed azure_ad_token_provider was applied Add get_azure_request_auth_headers, which turns the credential in azure_client_params into an Authorization: Bearer header (or api-key, following the SDK's precedence) while keeping any auth header the caller already set, and use it for both the sync and async image requests. The pre-call logging metadata receives a redacted copy of those headers so the token never reaches logging callbacks Fixes #16422 --- litellm/llms/azure/azure.py | 31 ++- litellm/llms/azure/common_utils.py | 35 +++ .../test_azure_image_generation_init.py | 249 ++++++++++++++++++ 3 files changed, 301 insertions(+), 14 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 46a9dd1a531..bfbaffd972c 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -46,7 +46,9 @@ from .common_utils import ( AzureOpenAIError, BaseAzureLLM, get_azure_ad_token_from_oidc, + get_azure_request_auth_headers, process_azure_headers, + redact_azure_auth_headers, select_azure_base_url_or_endpoint, ) from .image_generation import ( @@ -1142,7 +1144,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key: str, input: list, logging_obj: LiteLLMLoggingObj, - headers: dict, + headers: dict[str, str], client=None, timeout=None, model: str | None = None, @@ -1167,7 +1169,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={ "complete_input_dict": data, "api_base": img_gen_api_base, - "headers": headers, + "headers": redact_azure_auth_headers(headers), }, ) httpx_response: Final[httpx.Response] = await self.make_async_azure_httpx_request( @@ -1226,7 +1228,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): timeout: float, optional_params: dict, logging_obj: LiteLLMLoggingObj, - headers: dict, + headers: dict[str, str], model: str | None = None, api_key: str | None = None, api_base: str | None = None, @@ -1261,21 +1263,22 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if not isinstance(max_retries, int): raise AzureOpenAIError(status_code=422, message="max retries must be an int") - if api_key is None and azure_ad_token_provider is not None: - azure_ad_token = azure_ad_token_provider() - if azure_ad_token: - headers.pop("api-key", None) - headers["Authorization"] = f"Bearer {azure_ad_token}" - - # init AzureOpenAI Client + auth_params: Final[dict[str, object]] = {**(litellm_params or {})} # mutable-ok: SDK init takes a dict + if azure_ad_token is not None: + auth_params["azure_ad_token"] = azure_ad_token + if azure_ad_token_provider is not None: + auth_params["azure_ad_token_provider"] = azure_ad_token_provider azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client( - litellm_params=litellm_params or {}, + litellm_params=auth_params, api_key=api_key, model_name=model or "", api_version=api_version, api_base=api_base, is_async=False, ) + request_headers: Final = dict( # mutable-ok: the httpx request helpers take a dict + get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) + ) if aimg_generation is True: return self.aimage_generation( data=data, @@ -1286,7 +1289,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): client=client, azure_client_params=azure_client_params, timeout=timeout, - headers=headers, + headers=request_headers, model=model, ) @@ -1303,7 +1306,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={ "complete_input_dict": data, "api_base": img_gen_api_base, - "headers": headers, + "headers": redact_azure_auth_headers(request_headers), }, ) httpx_response: Final[httpx.Response] = self.make_sync_azure_httpx_request( @@ -1313,7 +1316,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_version=api_version or "", api_key=api_key or "", data=data, - headers=headers, + headers=request_headers, deployment_name=model, ) provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6cb7d09cec4..f276d8b18d1 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -405,6 +405,41 @@ def get_azure_ad_token( return azure_ad_token +_AZURE_AUTH_HEADER_NAMES: Final = frozenset(("api-key", "authorization")) +_REDACTED_AZURE_HEADER_VALUE: Final = "***REDACTED***" + + +def _resolve_azure_ad_token(azure_client_params: Mapping[str, object]) -> str | None: + azure_ad_token: Final = azure_client_params.get("azure_ad_token") + if isinstance(azure_ad_token, str) and azure_ad_token: + return azure_ad_token + token_provider: Final = azure_client_params.get("azure_ad_token_provider") + provided_token: Final = token_provider() if callable(token_provider) else None + return provided_token if isinstance(provided_token, str) and provided_token else None + + +def get_azure_request_auth_headers( + headers: Mapping[str, str], + azure_client_params: Mapping[str, object], +) -> Mapping[str, str]: + if any(name.lower() in _AZURE_AUTH_HEADER_NAMES for name in headers): + return headers + azure_ad_token: Final = _resolve_azure_ad_token(azure_client_params) + if azure_ad_token is not None: + return MappingProxyType({**headers, "Authorization": f"Bearer {azure_ad_token}"}) + api_key: Final = azure_client_params.get("api_key") + if isinstance(api_key, str) and api_key: + return MappingProxyType({**headers, "api-key": api_key}) + return headers + + +def redact_azure_auth_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return { # mutable-ok: logging callbacks JSON-serialize this copy + name: (_REDACTED_AZURE_HEADER_VALUE if name.lower() in _AZURE_AUTH_HEADER_NAMES else value) + for name, value in headers.items() + } + + class BaseAzureLLM(BaseOpenAILLM): @staticmethod def _try_get_default_azure_credential_provider( diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 70b5eab5c37..a30aa277f3d 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -10,6 +10,11 @@ import respx import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.azure.common_utils import ( + _cached_entra_id_token_provider, + get_azure_request_auth_headers, + redact_azure_auth_headers, +) from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, ) @@ -587,3 +592,247 @@ def test_azure_image_generation_v1_route_base_model_vs_deployment_name(respx_moc sent_body = json.loads(request.content) assert sent_body["model"] == model assert sent_body["prompt"] == prompt + + +@pytest.fixture +def fake_entra_id(monkeypatch: pytest.MonkeyPatch): + built_credentials = [] + + class FakeClientSecretCredential: + def __init__(self, tenant_id: str, client_id: str, client_secret: str) -> None: + built_credentials.append((tenant_id, client_id, client_secret)) + + monkeypatch.setattr("azure.identity.ClientSecretCredential", FakeClientSecretCredential) + monkeypatch.setattr("azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "entra-id-token") + _cached_entra_id_token_provider.cache_clear() + yield built_credentials + _cached_entra_id_token_provider.cache_clear() + + +def _mock_image_generation_route(respx_mock: respx.MockRouter, api_base: str, model: str) -> respx.Route: + return respx_mock.post(f"{api_base}/openai/deployments/{model}/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]}) + ) + + +@pytest.mark.parametrize("credentials_in_litellm_params", [False, True]) +def test_azure_image_generation_keyless_entra_id_sends_bearer_token( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fake_entra_id: list, + credentials_in_litellm_params: bool, +): + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + litellm_params = {"api_base": api_base, "api_version": api_version} + if credentials_in_litellm_params: + litellm_params.update( + tenant_id="tenant-from-params", client_id="client-from-params", client_secret="secret-from-params" + ) + expected_credential = ("tenant-from-params", "client-from-params", "secret-from-params") + else: + monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env") + expected_credential = ("tenant-from-env", "client-from-env", "secret-from-env") + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params=litellm_params, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-id-token" + assert "api-key" not in request.headers + assert fake_entra_id == [expected_credential] + assert response.data[0].b64_json == "aaaa" + logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"] + assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"} + assert "entra-id-token" not in str(logging_obj.pre_call.call_args) + + +@pytest.mark.asyncio +async def test_azure_aimage_generation_keyless_entra_id_sends_bearer_token( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = await AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + aimg_generation=True, + litellm_params={ + "api_base": api_base, + "api_version": api_version, + "tenant_id": "tenant-from-params", + "client_id": "client-from-params", + "client_secret": "secret-from-params", + }, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-id-token" + assert "api-key" not in request.headers + assert fake_entra_id == [("tenant-from-params", "client-from-params", "secret-from-params")] + assert response.data[0].b64_json == "aaaa" + logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"] + assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"} + assert "entra-id-token" not in str(logging_obj.pre_call.call_args) + + +@pytest.mark.parametrize( + "credential_kwargs, expected_authorization", + [ + ({"azure_ad_token": "static-ad-token"}, "Bearer static-ad-token"), + ({"azure_ad_token_provider": lambda: "provider-token"}, "Bearer provider-token"), + ], +) +def test_azure_image_generation_explicit_azure_ad_credential_sends_bearer_token( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + credential_kwargs: dict, + expected_authorization: str, +): + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + **credential_kwargs, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == expected_authorization + assert "api-key" not in request.headers + assert response.data[0].b64_json == "aaaa" + + +def test_azure_image_generation_with_api_key_keeps_api_key_header( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list +): + monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env") + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json", "api-key": "sk-test"}, + model="gpt-image-1", + api_key="sk-test", + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + ) + + request = route.calls.last.request + assert request.headers["api-key"] == "sk-test" + assert "Authorization" not in request.headers + assert fake_entra_id == [] + assert response.data[0].b64_json == "aaaa" + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***" + + +@pytest.mark.parametrize( + "caller_auth_header", + [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}], +) +def test_get_azure_request_auth_headers_keeps_caller_auth_header(caller_auth_header: dict): + headers = {"Content-Type": "application/json", **caller_auth_header} + azure_client_params = { + "api_key": "sk-resolved", + "azure_ad_token": "resolved-token", + "azure_ad_token_provider": lambda: "provider-token", + } + assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers + + +def test_get_azure_request_auth_headers_prefers_azure_ad_token_over_provider_and_api_key(): + headers = {"Content-Type": "application/json"} + azure_client_params = { + "api_key": "sk-resolved", + "azure_ad_token": "static-token", + "azure_ad_token_provider": lambda: "provider-token", + } + out = get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) + assert dict(out) == {"Content-Type": "application/json", "Authorization": "Bearer static-token"} + assert headers == {"Content-Type": "application/json"} + + +def test_get_azure_request_auth_headers_uses_token_provider_over_api_key(): + azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": lambda: "pt"} + out = get_azure_request_auth_headers(headers={}, azure_client_params=azure_client_params) + assert dict(out) == {"Authorization": "Bearer pt"} + + +def test_get_azure_request_auth_headers_falls_back_to_api_key(): + azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": None} + out = get_azure_request_auth_headers(headers={"Content-Type": "application/json"}, azure_client_params=azure_client_params) + assert dict(out) == {"Content-Type": "application/json", "api-key": "sk-resolved"} + + +@pytest.mark.parametrize( + "azure_client_params", + [ + {}, + {"api_key": "", "azure_ad_token": "", "azure_ad_token_provider": None}, + {"azure_ad_token_provider": lambda: None}, + {"azure_ad_token_provider": lambda: ""}, + ], +) +def test_get_azure_request_auth_headers_without_credential_leaves_headers_unchanged(azure_client_params: dict): + headers = {"Content-Type": "application/json"} + assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers + + +def test_redact_azure_auth_headers_masks_only_credential_values(): + headers = {"Content-Type": "application/json", "api-key": "sk-secret", "authorization": "Bearer secret"} + assert redact_azure_auth_headers(headers) == { + "Content-Type": "application/json", + "api-key": "***REDACTED***", + "authorization": "***REDACTED***", + } + assert headers["api-key"] == "sk-secret" + assert headers["authorization"] == "Bearer secret" From 5ef05b97c567f6163d6a20d960fc3d66e70e0c98 Mon Sep 17 00:00:00 2001 From: Aidan Sinclair Date: Tue, 8 Sep 2026 18:25:22 -0400 Subject: [PATCH 023/393] feat(websearch): let the model emit objective + multi-query search shape for providers that support it The intercepted web search tool only carries a single query string, so search providers whose APIs take a natural-language objective plus multiple keyword queries (documented best practice for Parallel AI's v1 search) always receive a degraded single-query request. Widen the tool's input schema with optional objective and search_queries fields (query stays required), and forward the richer shape from the interception handler only to providers whose search config reports supports_rich_search_input(). Every other provider, and every model that keeps emitting just query, is byte-for-byte unchanged. - BaseSearchConfig.supports_rich_search_input() defaults False; ParallelAISearchConfig overrides True - handler trims search_queries to five (the provider cap) and never overrides an objective configured on the search tool's litellm_params - mocked tests cover schema exposure, extraction validation, provider gating, and the unchanged single-string path Co-Authored-By: Claude Fable 5 --- .../websearch_interception/handler.py | 573 ++++++++++++++---- .../websearch_interception/tools.py | 92 +-- .../llms/base_llm/search/transformation.py | 36 +- .../llms/parallel_ai/search/transformation.py | 23 +- .../integrations/websearch_interception.py | 16 + .../test_websearch_rich_query_shape.py | 188 ++++++ 6 files changed, 750 insertions(+), 178 deletions(-) create mode 100644 tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 587da997f94..4fca0a36797 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -44,6 +44,7 @@ from litellm.types.integrations.custom_logger import ( from litellm.types.integrations.websearch_interception import ( AnthropicSearchQuery, AnthropicServerToolUseBlock, + RichWebSearchInput, WebSearchInterceptionConfig, ) from litellm.types.llms.anthropic import AnthropicThinkingParam @@ -173,7 +174,9 @@ class _AcompletionNamedParams(TypedDict, total=False): logprobs: ReadOnly[bool | None] top_logprobs: ReadOnly[int | None] deployment_id: ReadOnly[str | None] - reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None] + reasoning_effort: ReadOnly[ + Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None + ] verbosity: ReadOnly[Literal["low", "medium", "high"] | None] safety_identifier: ReadOnly[str | None] service_tier: ReadOnly[str | None] @@ -231,7 +234,9 @@ class WebSearchInterceptionLogger(CustomLogger): if enabled_providers is None: self.enabled_providers = [LlmProviders.BEDROCK.value] else: - self.enabled_providers = [p.value if isinstance(p, LlmProviders) else p for p in enabled_providers] + self.enabled_providers = [ + p.value if isinstance(p, LlmProviders) else p for p in enabled_providers + ] self.search_tool_name = search_tool_name self.max_agentic_loops = self._validated_max_agentic_loops(max_agentic_loops) self._request_has_websearch = False # Track if current request has web search @@ -241,7 +246,9 @@ class WebSearchInterceptionLogger(CustomLogger): """ Reject loop ceilings the agentic loop cannot honor, at config load time. """ - return validated_max_agentic_loops(max_agentic_loops, field="websearch_interception_params.max_agentic_loops") + return validated_max_agentic_loops( + max_agentic_loops, field="websearch_interception_params.max_agentic_loops" + ) async def try_short_circuit_search( self, @@ -276,7 +283,10 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if provider is in enabled list provider_str: Final = custom_llm_provider or "" - if self.enabled_providers is not None and provider_str not in self.enabled_providers: + if ( + self.enabled_providers is not None + and provider_str not in self.enabled_providers + ): return None # Only short-circuit for providers whose Anthropic Messages agentic loop @@ -292,10 +302,15 @@ class WebSearchInterceptionLogger(CustomLogger): # web-search-only requests against it. try: provider_enum: Final = LlmProviders(provider_str) - anthropic_config: Final = ProviderConfigManager.get_provider_anthropic_messages_config( - model=model, provider=provider_enum + anthropic_config: Final = ( + ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=provider_enum + ) ) - if anthropic_config is not None and anthropic_config.handles_web_search_natively(): + if ( + anthropic_config is not None + and anthropic_config.handles_web_search_natively() + ): verbose_logger.debug( "WebSearchInterception: Skipping short-circuit for %s (provider handles web search natively via the agentic loop)", provider_str, @@ -318,7 +333,9 @@ class WebSearchInterceptionLogger(CustomLogger): return None verbose_logger.debug( - "WebSearchInterception: Short-circuit search detected (provider=%s, query='%s')", provider_str, query + "WebSearchInterception: Short-circuit search detected (provider=%s, query='%s')", + provider_str, + query, ) # Native clients (Claude Desktop / Cowork / Anthropic SDK) make a @@ -338,9 +355,13 @@ class WebSearchInterceptionLogger(CustomLogger): if kwargs is None: search_result_text, structured = await self._execute_search(query) else: - search_result_text, structured = await self._execute_search(query, kwargs=kwargs) + search_result_text, structured = await self._execute_search( + query, kwargs=kwargs + ) except Exception as e: - verbose_logger.error("WebSearchInterception: Short-circuit search failed: %s", e) + verbose_logger.error( + "WebSearchInterception: Short-circuit search failed: %s", e + ) search_result_text, structured = f"Search failed: {e}", None content: Final[list[dict[str, object]]] = [] @@ -400,12 +421,14 @@ class WebSearchInterceptionLogger(CustomLogger): "litellm_params": kwargs.get("litellm_params", {}), "model": kwargs.get("model", ""), } - custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get( - "custom_llm_provider", "" - ) + custom_llm_provider = call_kwargs_view[ + "custom_llm_provider" + ] or call_kwargs_view["litellm_params"].get("custom_llm_provider", "") if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"]) + _, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=call_kwargs_view["model"] + ) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -424,7 +447,9 @@ class WebSearchInterceptionLogger(CustomLogger): if not has_websearch: return None - verbose_logger.debug("WebSearchInterception: Converting native web_search tools to LiteLLM standard") + verbose_logger.debug( + "WebSearchInterception: Converting native web_search tools to LiteLLM standard" + ) # If the client sent an Anthropic-native web_search_* tool, mark the # request so the agentic loop emits native web_search_tool_result @@ -454,7 +479,9 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs["tools"] = converted_tools if kwargs.get("stream"): - verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") + verbose_logger.debug( + "WebSearchInterception: deployment hook converting stream=True to stream=False" + ) kwargs["stream"] = False kwargs["_websearch_interception_converted_stream"] = True @@ -467,23 +494,34 @@ class WebSearchInterceptionLogger(CustomLogger): if not any(is_web_search_tool_responses(tool) for tool in tools): return None - verbose_logger.debug("WebSearchInterception: Converting Responses web_search tools to LiteLLM standard") + verbose_logger.debug( + "WebSearchInterception: Converting Responses web_search tools to LiteLLM standard" + ) converted_tools: Final = [ - get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools + ( + get_litellm_web_search_tool_responses() + if is_web_search_tool_responses(tool) + else tool + ) + for tool in tools ] converted_kwargs: Final = {**kwargs, "tools": converted_tools} if kwargs.get("stream"): - verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") + verbose_logger.debug( + "WebSearchInterception: deployment hook converting stream=True to stream=False" + ) converted_kwargs["stream"] = False converted_kwargs["_websearch_interception_converted_stream"] = True return converted_kwargs @classmethod - def from_config_yaml(cls, config: WebSearchInterceptionConfig) -> "WebSearchInterceptionLogger": + def from_config_yaml( + cls, config: WebSearchInterceptionConfig + ) -> "WebSearchInterceptionLogger": """ Initialize WebSearchInterceptionLogger from proxy config.yaml parameters. @@ -538,7 +576,9 @@ class WebSearchInterceptionLogger(CustomLogger): return tool.get("name") @classmethod - def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]]) -> object: + def _sync_forced_tool_choice( + cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]] + ) -> object: """Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it names a web-search tool that was just converted away. @@ -555,7 +595,9 @@ class WebSearchInterceptionLogger(CustomLogger): return tool_choice return {**tool_choice, "name": LITELLM_WEB_SEARCH_TOOL_NAME} - async def async_pre_request_hook(self, model: str, messages: list[dict], kwargs: dict) -> dict | None: + async def async_pre_request_hook( + self, model: str, messages: list[dict], kwargs: dict + ) -> dict | None: """ Pre-request hook to convert native web search tools to LiteLLM standard. @@ -571,7 +613,9 @@ class WebSearchInterceptionLogger(CustomLogger): Modified kwargs dict with converted tools, or None if no modifications needed """ # Check if this request is for an enabled provider - custom_llm_provider: Final = kwargs.get("litellm_params", {}).get("custom_llm_provider", "") + custom_llm_provider: Final = kwargs.get("litellm_params", {}).get( + "custom_llm_provider", "" + ) verbose_logger.debug( "WebSearchInterception: Pre-request hook called - custom_llm_provider=%s - enabled_providers=%s", @@ -579,9 +623,14 @@ class WebSearchInterceptionLogger(CustomLogger): self.enabled_providers or "ALL", ) - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( - "WebSearchInterception: Skipping - provider %s not in %s", custom_llm_provider, self.enabled_providers + "WebSearchInterception: Skipping - provider %s not in %s", + custom_llm_provider, + self.enabled_providers, ) return None @@ -595,11 +644,16 @@ class WebSearchInterceptionLogger(CustomLogger): if not has_websearch: return None - verbose_logger.debug("WebSearchInterception: Pre-request hook triggered for provider=%s", custom_llm_provider) + verbose_logger.debug( + "WebSearchInterception: Pre-request hook triggered for provider=%s", + custom_llm_provider, + ) deployment_max_agentic_loops: Final = kwargs.get("max_agentic_loops") if self.max_agentic_loops is not None and deployment_max_agentic_loops is None: - kwargs["max_agentic_loops"] = self.max_agentic_loops # rebind-ok: this hook returns the kwargs it edits + kwargs["max_agentic_loops"] = ( + self.max_agentic_loops + ) # rebind-ok: this hook returns the kwargs it edits # If the client sent an Anthropic-native web_search_* tool, mark the # request so the agentic loop emits native web_search_tool_result @@ -626,15 +680,20 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs["tools"] = converted_tools verbose_logger.debug( - "WebSearchInterception: Tools after conversion: %s", [t.get("name") for t in converted_tools] + "WebSearchInterception: Tools after conversion: %s", + [t.get("name") for t in converted_tools], ) if "tool_choice" in kwargs: - kwargs["tool_choice"] = self._sync_forced_tool_choice(kwargs.get("tool_choice"), converted_tools) + kwargs["tool_choice"] = self._sync_forced_tool_choice( + kwargs.get("tool_choice"), converted_tools + ) # Also convert here for direct callers that bypass the deployment hook. if kwargs.get("stream"): - verbose_logger.debug("WebSearchInterception: Converting stream=True to stream=False") + verbose_logger.debug( + "WebSearchInterception: Converting stream=True to stream=False" + ) kwargs["stream"] = False kwargs["_websearch_interception_converted_stream"] = True @@ -672,13 +731,20 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs=kwargs, ) - verbose_logger.debug("WebSearchInterception: Hook called! provider=%s, stream=%s", custom_llm_provider, stream) + verbose_logger.debug( + "WebSearchInterception: Hook called! provider=%s, stream=%s", + custom_llm_provider, + stream, + ) verbose_logger.debug("WebSearchInterception: Response type: %s", type(response)) # Check if provider should be intercepted # Note: custom_llm_provider is already normalized by get_llm_provider() # (e.g., "bedrock/invoke/..." -> "bedrock") - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( "WebSearchInterception: Skipping provider %s (not in enabled list: %s)", custom_llm_provider, @@ -700,11 +766,14 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug("WebSearchInterception: No WebSearch tool_use detected in response") + verbose_logger.debug( + "WebSearchInterception: No WebSearch tool_use detected in response" + ) return False, {} verbose_logger.debug( - "WebSearchInterception: Detected %s WebSearch tool call(s), executing agentic loop", len(tool_calls) + "WebSearchInterception: Detected %s WebSearch tool call(s), executing agentic loop", + len(tool_calls), ) # Extract thinking blocks from response content. @@ -732,14 +801,17 @@ class WebSearchInterceptionLogger(CustomLogger): thinking_block_dict: dict = {"type": block_type} if block_type == "thinking": thinking_block_dict["thinking"] = getattr(block, "thinking", "") - thinking_block_dict["signature"] = getattr(block, "signature", "") + thinking_block_dict["signature"] = getattr( + block, "signature", "" + ) else: # redacted_thinking thinking_block_dict["data"] = getattr(block, "data", "") thinking_blocks.append(thinking_block_dict) if thinking_blocks: verbose_logger.debug( - "WebSearchInterception: Extracted %s thinking block(s) from response", len(thinking_blocks) + "WebSearchInterception: Extracted %s thinking block(s) from response", + len(thinking_blocks), ) # Return tools dict with tool calls and thinking blocks @@ -769,12 +841,17 @@ class WebSearchInterceptionLogger(CustomLogger): """ verbose_logger.debug( - "WebSearchInterception: Chat completion hook called! provider=%s, stream=%s", custom_llm_provider, stream + "WebSearchInterception: Chat completion hook called! provider=%s, stream=%s", + custom_llm_provider, + stream, ) verbose_logger.debug("WebSearchInterception: Response type: %s", type(response)) # Check if provider should be intercepted - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( "WebSearchInterception: Skipping provider %s (not in enabled list: %s)", custom_llm_provider, @@ -783,9 +860,13 @@ class WebSearchInterceptionLogger(CustomLogger): return False, {} # Check if tools include any web search tool (strict check for chat completions) - has_websearch_tool: Final = any(is_web_search_tool_chat_completion(t) for t in (tools or [])) + has_websearch_tool: Final = any( + is_web_search_tool_chat_completion(t) for t in (tools or []) + ) if not has_websearch_tool: - verbose_logger.debug("WebSearchInterception: No litellm_web_search tool in request") + verbose_logger.debug( + "WebSearchInterception: No litellm_web_search tool in request" + ) return False, {} # Detect WebSearch tool_calls in response (OpenAI format) @@ -796,11 +877,14 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug("WebSearchInterception: No WebSearch tool_calls detected in response") + verbose_logger.debug( + "WebSearchInterception: No WebSearch tool_calls detected in response" + ) return False, {} verbose_logger.debug( - "WebSearchInterception: Detected %s WebSearch tool call(s), executing agentic loop", len(tool_calls) + "WebSearchInterception: Detected %s WebSearch tool call(s), executing agentic loop", + len(tool_calls), ) # Return tools dict with tool calls @@ -824,10 +908,15 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> tuple[bool, dict]: """Check if WebSearch interception is needed for the Responses API.""" verbose_logger.debug( - "WebSearchInterception: Responses hook called! provider=%s, stream=%s", custom_llm_provider, stream + "WebSearchInterception: Responses hook called! provider=%s, stream=%s", + custom_llm_provider, + stream, ) - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( "WebSearchInterception: Skipping provider %s (not in enabled list: %s)", custom_llm_provider, @@ -835,9 +924,13 @@ class WebSearchInterceptionLogger(CustomLogger): ) return False, {} - has_websearch_tool: Final = any(is_web_search_tool_responses(t) for t in (tools or [])) + has_websearch_tool: Final = any( + is_web_search_tool_responses(t) for t in (tools or []) + ) if not has_websearch_tool: - verbose_logger.debug("WebSearchInterception: No litellm_web_search tool in responses request") + verbose_logger.debug( + "WebSearchInterception: No litellm_web_search tool in responses request" + ) return False, {} should_intercept, tool_calls = WebSearchTransformation.transform_request( @@ -847,11 +940,14 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug("WebSearchInterception: No WebSearch function_call detected in responses output") + verbose_logger.debug( + "WebSearchInterception: No WebSearch function_call detected in responses output" + ) return False, {} verbose_logger.debug( - "WebSearchInterception: Detected %s WebSearch function_call(s), executing agentic loop", len(tool_calls) + "WebSearchInterception: Detected %s WebSearch function_call(s), executing agentic loop", + len(tool_calls), ) tools_dict: Final = { @@ -883,7 +979,10 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls: Final = tools["tool_calls"] thinking_blocks: Final = tools.get("thinking_blocks", []) - verbose_logger.debug("WebSearchInterception: Executing agentic loop for %s search(es)", len(tool_calls)) + verbose_logger.debug( + "WebSearchInterception: Executing agentic loop for %s search(es)", + len(tool_calls), + ) return await self._execute_agentic_loop( model=model, @@ -954,9 +1053,11 @@ class WebSearchInterceptionLogger(CustomLogger): # (while we still have the structured SearchResponse list) and stash # them on plan metadata for the post-hook to inject. if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): - metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( - tool_calls=tool_calls, - structured_results=structured_results, + metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = ( + self._build_native_result_blocks( + tool_calls=tool_calls, + structured_results=structured_results, + ) ) return AgenticLoopPlan( @@ -982,7 +1083,9 @@ class WebSearchInterceptionLogger(CustomLogger): render citations / sources alongside the model's textual reply. """ metadata_view: Final[_PlanMetadataView] = { - "websearch_native_blocks": plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) + "websearch_native_blocks": plan.metadata.get( + WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY + ) } native_blocks: Final = metadata_view["websearch_native_blocks"] if not native_blocks: @@ -1007,7 +1110,9 @@ class WebSearchInterceptionLogger(CustomLogger): for i, tool_call in enumerate(tool_calls) for block in WebSearchInterceptionLogger._native_result_pair( query=WebSearchInterceptionLogger._tool_call_query(tool_call), - search_response=structured_results[i] if i < len(structured_results) else None, + search_response=( + structured_results[i] if i < len(structured_results) else None + ), ) ) @@ -1026,7 +1131,9 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> tuple[Mapping[str, object], Mapping[str, object]]: tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" return ( - AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(), + AnthropicServerToolUseBlock( + id=tool_use_id, input=AnthropicSearchQuery(query=query) + ).model_dump(), WebSearchTransformation.build_web_search_tool_result_block( tool_use_id=tool_use_id, search_response=search_response, @@ -1034,7 +1141,9 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: + def _inject_native_blocks( + response: _ResponseT, native_blocks: Sequence[Mapping[str, object]] + ) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response @@ -1044,7 +1153,9 @@ class WebSearchInterceptionLogger(CustomLogger): return response existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or [] try: - setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing)) + setattr( + response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing) + ) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. @@ -1075,7 +1186,8 @@ class WebSearchInterceptionLogger(CustomLogger): response_format: Final = tools.get("response_format", "openai") verbose_logger.debug( - "WebSearchInterception: Executing chat completion agentic loop for %s search(es)", len(tool_calls) + "WebSearchInterception: Executing chat completion agentic loop for %s search(es)", + len(tool_calls), ) return await self._execute_chat_completion_agentic_loop( @@ -1152,17 +1264,29 @@ class WebSearchInterceptionLogger(CustomLogger): """Execute litellm.asearch() and build a Responses API rerun patch.""" search_tasks: Final = [ ( - self._execute_search(tool_call["input"]["query"], kwargs=kwargs) - if isinstance(tool_call.get("input"), dict) and tool_call["input"].get("query") + self._execute_search( + tool_call["input"]["query"], + kwargs=kwargs, + rich=self._rich_search_input(tool_call["input"]), + ) + if isinstance(tool_call.get("input"), dict) + and tool_call["input"].get("query") else self._create_empty_search_result() ) for tool_call in tool_calls ] - verbose_logger.debug("WebSearchInterception: Executing %s responses search(es) in parallel", len(search_tasks)) - search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) + verbose_logger.debug( + "WebSearchInterception: Executing %s responses search(es) in parallel", + len(search_tasks), + ) + search_results: Final = await asyncio.gather( + *search_tasks, return_exceptions=True + ) - search_texts: Final = [self._extract_search_text(result) for result in search_results] + search_texts: Final = [ + self._extract_search_text(result) for result in search_results + ] followup_items: Final = [ item @@ -1188,7 +1312,15 @@ class WebSearchInterceptionLogger(CustomLogger): optional_params_clean: Final = { k: v for k, v in optional_params.items() - if k not in {"tools", "tool_choice", "stream", "model_alias_map", "stream_response", "custom_prompt_dict"} + if k + not in { + "tools", + "tool_choice", + "stream", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } } kwargs_for_followup: Final = { @@ -1235,12 +1367,16 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _extract_search_text(result: object) -> str: if isinstance(result, Exception): - verbose_logger.error("WebSearchInterception: Responses search failed with error: %s", result) + verbose_logger.error( + "WebSearchInterception: Responses search failed with error: %s", result + ) return f"Search failed: {result}" if isinstance(result, tuple) and len(result) == 2: text_value, _ = result return text_value if isinstance(text_value, str) else str(text_value) - verbose_logger.debug("WebSearchInterception: Unexpected search result type %s", type(result)) + verbose_logger.debug( + "WebSearchInterception: Unexpected search result type %s", type(result) + ) return str(result) @staticmethod @@ -1291,7 +1427,9 @@ class WebSearchInterceptionLogger(CustomLogger): """ _internal_keys: Final = {"litellm_logging_obj"} return { - k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in _internal_keys + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") and k not in _internal_keys } async def _execute_agentic_loop( @@ -1311,7 +1449,9 @@ class WebSearchInterceptionLogger(CustomLogger): messages=messages, tool_calls=tool_calls, thinking_blocks=thinking_blocks, - anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params), + anthropic_messages_optional_request_params=dict[str, object]( + anthropic_messages_optional_request_params + ), logging_obj=logging_obj, kwargs=dict[str, object](kwargs), ) @@ -1329,13 +1469,15 @@ class WebSearchInterceptionLogger(CustomLogger): max_tokens = cast(int, kwargs.get("max_tokens", 1024)) patch_kwargs: Final = dict[str, object](request_patch.kwargs) - response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( - max_tokens=max_tokens, - messages=request_patch.messages, - model=request_patch.model or model, - **_NO_ACREATE_NAMED, - **optional_params, - **patch_kwargs, + response: AnthropicMessagesResponse | AsyncIterator[object] = ( + await anthropic_messages.acreate( + max_tokens=max_tokens, + messages=request_patch.messages, + model=request_patch.model or model, + **_NO_ACREATE_NAMED, + **optional_params, + **patch_kwargs, + ) ) # Legacy path: the new path goes through the typed plan + core @@ -1375,16 +1517,31 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: query = tool_call["input"].get("query") if query: - verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query) - search_tasks.append(self._execute_search(query, kwargs=kwargs)) + verbose_logger.debug( + "WebSearchInterception: Queuing search for query='%s'", query + ) + search_tasks.append( + self._execute_search( + query, + kwargs=kwargs, + rich=self._rich_search_input(tool_call["input"]), + ) + ) else: - verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call["id"]) + verbose_logger.debug( + "WebSearchInterception: Tool call %s has no query", tool_call["id"] + ) # Add empty result for tools without query search_tasks.append(self._create_empty_search_result()) # Execute searches in parallel - verbose_logger.debug("WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks)) - search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) + verbose_logger.debug( + "WebSearchInterception: Executing %s search(es) in parallel", + len(search_tasks), + ) + search_results: Final = await asyncio.gather( + *search_tasks, return_exceptions=True + ) # Split the gathered (text, structured) tuples into two parallel lists. # The text list feeds the follow-up model call; the structured list @@ -1393,17 +1550,31 @@ class WebSearchInterceptionLogger(CustomLogger): structured_results: Final[list[SearchResponse | None]] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result) + verbose_logger.error( + "WebSearchInterception: Search %s failed with error: %s", i, result + ) final_search_results.append(f"Search failed: {result}") structured_results.append(None) elif isinstance(result, tuple) and len(result) == 2: text_value, structured_value = result - final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) - structured_results.append(structured_value if isinstance(structured_value, SearchResponse) else None) + final_search_results.append( + cast(str, text_value) + if isinstance(text_value, str) + else str(text_value) + ) + structured_results.append( + structured_value + if isinstance(structured_value, SearchResponse) + else None + ) else: # Defensive: legacy callers / unexpected shape — preserve text, # drop structure. - verbose_logger.debug("WebSearchInterception: Unexpected result type %s at index %s", type(result), i) + verbose_logger.debug( + "WebSearchInterception: Unexpected result type %s at index %s", + type(result), + i, + ) final_search_results.append(str(result)) structured_results.append(None) @@ -1414,25 +1585,39 @@ class WebSearchInterceptionLogger(CustomLogger): thinking_blocks=thinking_blocks, ) - follow_up_messages: Final = messages + [assistant_message, cast(dict, user_message)] + follow_up_messages: Final = messages + [ + assistant_message, + cast(dict, user_message), + ] # Correlation context for structured logging - _call_id: Final = getattr(logging_obj, "litellm_call_id", None) or kwargs.get("litellm_call_id", "unknown") + _call_id: Final = getattr(logging_obj, "litellm_call_id", None) or kwargs.get( + "litellm_call_id", "unknown" + ) full_model_name = model # safe default before try block - max_tokens: Final = self._resolve_max_tokens(anthropic_messages_optional_request_params, kwargs) + max_tokens: Final = self._resolve_max_tokens( + anthropic_messages_optional_request_params, kwargs + ) - verbose_logger.debug("WebSearchInterception: Using max_tokens=%s for follow-up request", max_tokens) + verbose_logger.debug( + "WebSearchInterception: Using max_tokens=%s for follow-up request", + max_tokens, + ) optional_params_without_max_tokens: Final = { - k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" + k: v + for k, v in anthropic_messages_optional_request_params.items() + if k != "max_tokens" } kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: agentic_view: Final[_AgenticLoopParamsView] = { - "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {}) + "agentic_loop_params": logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) } full_model_name = agentic_view["agentic_loop_params"].get("model", model) verbose_logger.debug( @@ -1451,8 +1636,50 @@ class WebSearchInterceptionLogger(CustomLogger): ) return patch, structured_results + @staticmethod + def _rich_search_input(tool_input: object) -> RichWebSearchInput | None: + """ + Extract the optional objective/search_queries pair from a tool input. + + Returns None when the input carries neither, so callers can pass the + result straight through as ``_execute_search``'s ``rich`` argument. + """ + if not isinstance(tool_input, Mapping): + return None + rich: RichWebSearchInput = {} + objective = tool_input.get("objective") + if isinstance(objective, str) and objective.strip(): + rich["objective"] = objective + raw_queries = tool_input.get("search_queries") + if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str): + queries = [q for q in raw_queries if isinstance(q, str) and q.strip()] + if queries: + # Providers cap multi-query requests (Parallel drops queries + # past the fifth); trim here so nothing is silently ignored. + rich["search_queries"] = queries[:5] + return rich or None + + @staticmethod + def _provider_supports_rich_search(search_provider: str | None) -> bool: + """Whether the provider's search config accepts objective + multi-query input.""" + if not search_provider: + return False + try: + from litellm.utils import ProviderConfigManager + except ImportError: + return False + # SearchProviders is a str enum, so an unknown provider string simply + # misses the config map and returns None rather than raising. + config = ProviderConfigManager.get_provider_search_config( + search_provider + ) # pyright: ignore[reportArgumentType] + return config is not None and config.supports_rich_search_input() + async def _execute_search( - self, query: str, kwargs: Mapping[str, object] | None = None + self, + query: str, + kwargs: Mapping[str, object] | None = None, + rich: RichWebSearchInput | None = None, ) -> tuple[str, SearchResponse | None]: """ Execute a single web search using router's search tools. @@ -1475,13 +1702,21 @@ class WebSearchInterceptionLogger(CustomLogger): ) llm_router = None - search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) + search_tool: Final = self._select_search_tool_from_router( + llm_router=llm_router + ) search_provider: str | None = None search_litellm_params: Mapping[str, object] = {} - search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) + search_tool_name: Final = self._selected_search_tool_name( + search_tool=search_tool + ) if search_tool is not None: - await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) - tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {} + await self._authorize_search_tool( + search_tool=search_tool, kwargs=kwargs + ) + tool_params: Final[_SearchToolLitellmParams] = ( + search_tool.get("litellm_params", {}) or {} + ) search_litellm_params = dict[str, object](tool_params) search_provider = tool_params.get("search_provider") @@ -1494,7 +1729,9 @@ class WebSearchInterceptionLogger(CustomLogger): ) verbose_logger.debug( - "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider + "WebSearchInterception: Executing search for '%s' using provider '%s'", + query, + search_provider, ) user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs) search_metadata: Final = ( @@ -1510,13 +1747,27 @@ class WebSearchInterceptionLogger(CustomLogger): for key, value in search_litellm_params.items() if key != "search_provider" and value is not None } + # Forward the model's richer shape (objective + keyword queries) + # only to providers whose search API takes it natively; everyone + # else keeps the single query string the model also provided. + query_arg: str | list[str] = query + if rich and self._provider_supports_rich_search(search_provider): + rich_queries = rich.get("search_queries") + if rich_queries: + query_arg = rich_queries + rich_objective = rich.get("objective") + if rich_objective and "objective" not in search_kwargs: + search_kwargs["objective"] = rich_objective result: Final = ( await litellm.asearch( - query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs + query=query_arg, + search_provider=search_provider, + **_NO_ASEARCH_NAMED, + **search_kwargs, ) if search_metadata is None else await litellm.asearch( - query=query, + query=query_arg, search_provider=search_provider, litellm_metadata=search_metadata, **_NO_ASEARCH_NAMED, @@ -1525,14 +1776,20 @@ class WebSearchInterceptionLogger(CustomLogger): ) # Format using transformation function - search_result_text: Final = WebSearchTransformation.format_search_response(result) + search_result_text: Final = WebSearchTransformation.format_search_response( + result + ) verbose_logger.debug( - "WebSearchInterception: Search completed for '%s', got %s chars", query, len(search_result_text) + "WebSearchInterception: Search completed for '%s', got %s chars", + query, + len(search_result_text), ) return search_result_text, result except Exception as e: - verbose_logger.error("WebSearchInterception: Search failed for '%s': %s", query, e) + verbose_logger.error( + "WebSearchInterception: Search failed for '%s': %s", query, e + ) raise async def _authorize_search_tool( @@ -1592,7 +1849,9 @@ class WebSearchInterceptionLogger(CustomLogger): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=user_api_key_auth + ) ) return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches **user_api_key_metadata, @@ -1602,20 +1861,31 @@ class WebSearchInterceptionLogger(CustomLogger): } @staticmethod - def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None: + def _selected_search_tool_name( + search_tool: Mapping[str, object] | None, + ) -> str | None: if search_tool is None: return None search_tool_name: Final = search_tool.get("search_tool_name") - return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None + return ( + search_tool_name + if isinstance(search_tool_name, str) and search_tool_name + else None + ) @staticmethod - def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": + def _get_user_api_key_auth_from_kwargs( + kwargs: Mapping[str, object] | None, + ) -> "UserAPIKeyAuth | None": if not kwargs: return None for metadata_key in ("metadata", "litellm_metadata"): metadata = kwargs.get(metadata_key) - if isinstance(metadata, dict) and metadata.get("user_api_key_auth") is not None: + if ( + isinstance(metadata, dict) + and metadata.get("user_api_key_auth") is not None + ): return metadata["user_api_key_auth"] litellm_params: Final = kwargs.get("litellm_params") @@ -1624,16 +1894,23 @@ class WebSearchInterceptionLogger(CustomLogger): for metadata_key in ("metadata", "litellm_metadata"): metadata = litellm_params.get(metadata_key) - if isinstance(metadata, dict) and metadata.get("user_api_key_auth") is not None: + if ( + isinstance(metadata, dict) + and metadata.get("user_api_key_auth") is not None + ): return metadata["user_api_key_auth"] return None - def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": + def _select_search_tool_from_router( + self, llm_router: object + ) -> "_SearchToolConfig | None": if llm_router is None or not hasattr(llm_router, "search_tools"): return None search_tools: Final = tuple(getattr(llm_router, "search_tools", None) or ()) - return self._select_search_tool_from_list(search_tools=search_tools, source="router") + return self._select_search_tool_from_list( + search_tools=search_tools, source="router" + ) def _select_search_tool_from_list( self, @@ -1642,10 +1919,14 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> "_SearchToolConfig | None": if self.search_tool_name: matching_tools: Final = tuple( - tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name + tool + for tool in search_tools + if tool.get("search_tool_name") == self.search_tool_name ) if matching_tools: - search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider") + search_provider = ( + matching_tools[0].get("litellm_params", {}) or {} + ).get("search_provider") verbose_logger.debug( "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", self.search_tool_name, @@ -1661,7 +1942,9 @@ class WebSearchInterceptionLogger(CustomLogger): if search_tools: first_tool: Final = search_tools[0] - search_provider = (first_tool.get("litellm_params", {}) or {}).get("search_provider") + search_provider = (first_tool.get("litellm_params", {}) or {}).get( + "search_provider" + ) verbose_logger.debug( "WebSearchInterception: Using first available search tool from %s with provider '%s'", source, @@ -1721,39 +2004,66 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: # Handle both Anthropic-style input and OpenAI-style function.arguments query = None + tool_args: dict | None = None if "input" in tool_call and isinstance(tool_call["input"], dict): - query = tool_call["input"].get("query") + tool_args = tool_call["input"] + query = tool_args.get("query") elif "function" in tool_call: func = tool_call["function"] if isinstance(func, dict): args = func.get("arguments", {}) if isinstance(args, dict): + tool_args = args query = args.get("query") if query: - verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query) - search_tasks.append(self._execute_search(query, kwargs=kwargs)) + verbose_logger.debug( + "WebSearchInterception: Queuing search for query='%s'", query + ) + search_tasks.append( + self._execute_search( + query, kwargs=kwargs, rich=self._rich_search_input(tool_args) + ) + ) else: - verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call.get("id")) + verbose_logger.debug( + "WebSearchInterception: Tool call %s has no query", + tool_call.get("id"), + ) # Add empty result for tools without query search_tasks.append(self._create_empty_search_result()) # Execute searches in parallel - verbose_logger.debug("WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks)) - search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) + verbose_logger.debug( + "WebSearchInterception: Executing %s search(es) in parallel", + len(search_tasks), + ) + search_results: Final = await asyncio.gather( + *search_tasks, return_exceptions=True + ) # Chat-completion path only needs text — OpenAI tool_result format # has no equivalent of Anthropic's web_search_tool_result block. final_search_results: Final[list[str]] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result) + verbose_logger.error( + "WebSearchInterception: Search %s failed with error: %s", i, result + ) final_search_results.append(f"Search failed: {result}") elif isinstance(result, tuple) and len(result) == 2: text_value, _ = result - final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) + final_search_results.append( + cast(str, text_value) + if isinstance(text_value, str) + else str(text_value) + ) else: - verbose_logger.debug("WebSearchInterception: Unexpected result type %s at index %s", type(result), i) + verbose_logger.debug( + "WebSearchInterception: Unexpected result type %s at index %s", + type(result), + i, + ) final_search_results.append(str(result)) # Build assistant and tool messages using transformation @@ -1769,7 +2079,9 @@ class WebSearchInterceptionLogger(CustomLogger): # Make follow-up request with search results # For OpenAI format, tool_messages_or_user is a list of tool messages if response_format == "openai": - follow_up_messages = messages + [assistant_message] + cast(list[dict], tool_messages_or_user) + follow_up_messages = ( + messages + [assistant_message] + cast(list[dict], tool_messages_or_user) + ) else: # For Anthropic format (shouldn't happen in this method, but handle it) follow_up_messages = messages + [ @@ -1777,8 +2089,13 @@ class WebSearchInterceptionLogger(CustomLogger): cast(dict, tool_messages_or_user), ] - verbose_logger.debug("WebSearchInterception: Making follow-up chat completion request with search results") - verbose_logger.debug("WebSearchInterception: Follow-up messages count: %s", len(follow_up_messages)) + verbose_logger.debug( + "WebSearchInterception: Making follow-up chat completion request with search results" + ) + verbose_logger.debug( + "WebSearchInterception: Follow-up messages count: %s", + len(follow_up_messages), + ) # Remove internal parameters that shouldn't be passed to follow-up request internal_params: Final = { @@ -1791,7 +2108,9 @@ class WebSearchInterceptionLogger(CustomLogger): "custom_prompt_dict", } kwargs_for_followup: Final = { - k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in internal_params + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") and k not in internal_params } full_model_name = model @@ -1864,7 +2183,9 @@ class WebSearchInterceptionLogger(CustomLogger): websearch_params: WebSearchInterceptionConfig = {} if "websearch_interception_params" in litellm_settings: settings_view: Final[_WebSearchSettingsView] = { - "websearch_interception_params": litellm_settings["websearch_interception_params"] + "websearch_interception_params": litellm_settings[ + "websearch_interception_params" + ] } websearch_params = settings_view["websearch_interception_params"] elif "websearch_interception" in callback_specific_params and isinstance( diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 97c6c90d2ba..9e3d3fd91f3 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -11,6 +11,50 @@ from typing import Any, Final from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME +_WEB_SEARCH_TOOL_DESCRIPTION: Final = ( + "Search the web for information. Use this when you need current " + "information or answers to questions that require up-to-date data." +) + + +def _web_search_input_schema() -> dict[str, object]: + """ + JSON schema for the web search tool's input, shared by every tool format. + + ``query`` stays required so providers and callers that only understand a + single query string keep working unchanged. ``objective`` and + ``search_queries`` are optional richer inputs; they are forwarded only to + search providers that support them (see + ``BaseSearchConfig.supports_rich_search_input``). + """ + return { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to execute", + }, + "objective": { + "type": "string", + "description": ( + "Natural-language description of the goal behind the " + "search, including any source or freshness requirements." + ), + }, + "search_queries": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Two to five short keyword queries (3-6 words each) " + "covering different angles of the objective, e.g. varying " + "names, synonyms, or phrasings. Provide together with " + "objective for the best results." + ), + }, + }, + "required": ["query"], + } + def get_litellm_web_search_tool() -> dict[str, object]: """ @@ -33,20 +77,8 @@ def get_litellm_web_search_tool() -> dict[str, object]: """ return { "name": LITELLM_WEB_SEARCH_TOOL_NAME, - "description": ( - "Search the web for information. Use this when you need current " - "information or answers to questions that require up-to-date data." - ), - "input_schema": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query to execute", - } - }, - "required": ["query"], - }, + "description": _WEB_SEARCH_TOOL_DESCRIPTION, + "input_schema": _web_search_input_schema(), } @@ -65,20 +97,8 @@ def get_litellm_web_search_tool_openai() -> dict[str, object]: "type": "function", "function": { "name": LITELLM_WEB_SEARCH_TOOL_NAME, - "description": ( - "Search the web for information. Use this when you need current " - "information or answers to questions that require up-to-date data." - ), - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query to execute", - } - }, - "required": ["query"], - }, + "description": _WEB_SEARCH_TOOL_DESCRIPTION, + "parameters": _web_search_input_schema(), }, } @@ -98,20 +118,8 @@ def get_litellm_web_search_tool_responses() -> dict[str, object]: return { "type": "function", "name": LITELLM_WEB_SEARCH_TOOL_NAME, - "description": ( - "Search the web for information. Use this when you need current " - "information or answers to questions that require up-to-date data." - ), - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query to execute", - } - }, - "required": ["query"], - }, + "description": _WEB_SEARCH_TOOL_DESCRIPTION, + "parameters": _web_search_input_schema(), } diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 7668c6132d6..c183d538c01 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -95,6 +95,18 @@ class BaseSearchConfig: """ return "Unknown Search Provider" + def supports_rich_search_input(self) -> bool: + """ + Whether this provider's search API accepts a natural-language + objective plus multiple keyword queries in one request. + + Integrations that collect the richer shape (e.g. websearch + interception) forward ``query`` as a list plus an ``objective`` + optional param to providers that return True; every other provider + keeps receiving the single query string. + """ + return False + def get_http_method(self) -> Literal["GET", "POST"]: """ Get HTTP method for search requests. @@ -185,12 +197,20 @@ class BaseSearchConfig: def sign_request( self, - headers: dict[str, str], # mutable-ok: matches the request header dict every other hook on this base takes - optional_params: dict[str, object], # mutable-ok: matches every other hook on this base - request_data: dict[str, object] | list[dict[str, object]], # mutable-ok: transform_search_request's body + headers: dict[ + str, str + ], # mutable-ok: matches the request header dict every other hook on this base takes + optional_params: dict[ + str, object + ], # mutable-ok: matches every other hook on this base + request_data: ( + dict[str, object] | list[dict[str, object]] + ), # mutable-ok: transform_search_request's body api_base: str, api_key: str | None = None, - ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the handler passes these headers straight to httpx + ) -> tuple[ + dict[str, str], bytes | None + ]: # mutable-ok: the handler passes these headers straight to httpx """ OPTIONAL @@ -250,7 +270,9 @@ class BaseSearchConfig: Returns: Dict with request data """ - raise NotImplementedError("transform_search_request must be implemented by provider") + raise NotImplementedError( + "transform_search_request must be implemented by provider" + ) def transform_search_response( self, @@ -262,7 +284,9 @@ class BaseSearchConfig: Transform provider-specific Search response to standard format. Override in provider-specific implementations. """ - raise NotImplementedError("transform_search_response must be implemented by provider") + raise NotImplementedError( + "transform_search_response must be implemented by provider" + ) def get_error_class( self, diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index bde7b7b86db..4154a497d2c 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -90,6 +90,11 @@ class ParallelAISearchConfig(BaseSearchConfig): def ui_friendly_name() -> str: return "Parallel AI" + def supports_rich_search_input(self) -> bool: + # The v1 search API takes `objective` + multiple `search_queries` + # natively; sending both is the documented best practice. + return True + def validate_environment( self, headers: dict, @@ -105,7 +110,9 @@ class ParallelAISearchConfig(BaseSearchConfig): default_api_base=self.PARALLEL_AI_API_BASE, ) if not resolved_api_key: - raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") + raise ValueError( + "PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable." + ) headers["x-api-key"] = resolved_api_key headers["Content-Type"] = "application/json" return headers @@ -117,7 +124,11 @@ class ParallelAISearchConfig(BaseSearchConfig): data: dict | list[dict] | None = None, **kwargs, ) -> str: - resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE + resolved_api_base: Final = ( + api_base + or get_secret_str("PARALLEL_AI_API_BASE") + or self.PARALLEL_AI_API_BASE + ) trimmed: Final = resolved_api_base.rstrip("/") if trimmed.endswith("/v1/search"): @@ -184,7 +195,9 @@ class ParallelAISearchConfig(BaseSearchConfig): advanced_settings["location"] = params.pop("location") if "max_chars_per_result" in params: - advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")} + advanced_settings["excerpt_settings"] = { + "max_chars_per_result": params.pop("max_chars_per_result") + } if "fetch_policy" in params: advanced_settings["fetch_policy"] = params.pop("fetch_policy") @@ -277,4 +290,6 @@ class ParallelAISearchConfig(BaseSearchConfig): } ) - return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields})) + return SearchResponse.model_validate( + MappingProxyType({"results": results, "object": "search", **extra_fields}) + ) diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 7926b9eee0a..bf01340630e 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -27,6 +27,22 @@ class AnthropicServerToolUseBlock(BaseModel): input: AnthropicSearchQuery +class RichWebSearchInput(TypedDict, total=False): + """ + Optional richer search shape a model may emit alongside ``query``. + + Collected from the intercepted tool call and forwarded only to search + providers whose config reports ``supports_rich_search_input()``; every + other provider keeps receiving the single ``query`` string. + """ + + objective: str + """Natural-language description of the goal behind the search.""" + + search_queries: list[str] + """Two to five short keyword queries covering different angles.""" + + class WebSearchInterceptionConfig(TypedDict, total=False): """ Configuration parameters for WebSearchInterceptionLogger. diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py new file mode 100644 index 00000000000..f8d20a3d5fd --- /dev/null +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py @@ -0,0 +1,188 @@ +""" +Unit tests for the rich web-search input shape (objective + search_queries). + +The intercepted web search tool exposes optional `objective` and +`search_queries` fields alongside the required single `query` string. The +handler forwards the richer shape only to search providers whose config +reports supports_rich_search_input(); every other provider keeps receiving +the single query string the model also provided. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, +) +from litellm.integrations.websearch_interception.tools import ( + get_litellm_web_search_tool, + get_litellm_web_search_tool_openai, + get_litellm_web_search_tool_responses, +) +from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse +from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig + +RICH_INPUT = { + "query": "stripe node sdk v14 authentication", + "objective": "Find the current authentication flow for the Stripe Node SDK v14", + "search_queries": ["stripe node sdk v14 auth", "stripe api key rotation node"], +} + + +def _search_response() -> SearchResponse: + return SearchResponse(object="search", results=[]) + + +def _mock_router(search_provider: str) -> MagicMock: + """Router stub exposing one configured search tool.""" + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "test-search", + "litellm_params": { + "search_provider": search_provider, + "api_key": "sk-test", + }, + } + ] + return router + + +class TestToolSchema: + def test_all_formats_expose_rich_fields_and_keep_query_required(self): + anthropic_schema = get_litellm_web_search_tool()["input_schema"] + openai_schema = get_litellm_web_search_tool_openai()["function"]["parameters"] + responses_schema = get_litellm_web_search_tool_responses()["parameters"] + + for schema in (anthropic_schema, openai_schema, responses_schema): + assert schema["required"] == ["query"] + assert "objective" in schema["properties"] + assert "search_queries" in schema["properties"] + assert schema["properties"]["search_queries"]["type"] == "array" + + +class TestRichInputExtraction: + def test_extracts_objective_and_queries(self): + rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT) + assert rich == { + "objective": RICH_INPUT["objective"], + "search_queries": RICH_INPUT["search_queries"], + } + + def test_returns_none_when_only_query_present(self): + assert ( + WebSearchInterceptionLogger._rich_search_input({"query": "plain"}) is None + ) + + def test_returns_none_for_non_mapping_input(self): + assert WebSearchInterceptionLogger._rich_search_input(None) is None + assert WebSearchInterceptionLogger._rich_search_input("query") is None + + def test_drops_invalid_queries_and_caps_at_five(self): + rich = WebSearchInterceptionLogger._rich_search_input( + { + "query": "q", + "search_queries": ["a", "", 3, "b", "c", "d", "e", "f"], + } + ) + assert rich == {"search_queries": ["a", "b", "c", "d", "e"]} + + def test_ignores_string_valued_search_queries(self): + # A string is a Sequence; it must not be treated as a list of queries. + assert ( + WebSearchInterceptionLogger._rich_search_input( + {"query": "q", "search_queries": "not a list"} + ) + is None + ) + + +class TestProviderSupport: + def test_parallel_ai_supports_rich_input(self): + assert ParallelAISearchConfig().supports_rich_search_input() is True + + def test_base_config_defaults_to_unsupported(self): + assert BaseSearchConfig().supports_rich_search_input() is False + + def test_unknown_provider_is_unsupported(self): + assert WebSearchInterceptionLogger._provider_supports_rich_search(None) is False + assert ( + WebSearchInterceptionLogger._provider_supports_rich_search("not_a_provider") + is False + ) + + +class TestExecuteSearchShape: + @pytest.mark.asyncio + async def test_rich_shape_reaches_supporting_provider(self, monkeypatch): + """Parallel AI receives the query list plus objective.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + mock_asearch = AsyncMock(return_value=_search_response()) + monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai")) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT) + await logger._execute_search(RICH_INPUT["query"], rich=rich) + + call_kwargs = mock_asearch.await_args.kwargs + assert call_kwargs["query"] == RICH_INPUT["search_queries"] + assert call_kwargs["objective"] == RICH_INPUT["objective"] + assert call_kwargs["search_provider"] == "parallel_ai" + + @pytest.mark.asyncio + async def test_string_only_provider_keeps_single_query(self, monkeypatch): + """A provider without rich support receives the plain query string.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + mock_asearch = AsyncMock(return_value=_search_response()) + monkeypatch.setattr(proxy_server, "llm_router", _mock_router("perplexity")) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT) + await logger._execute_search(RICH_INPUT["query"], rich=rich) + + call_kwargs = mock_asearch.await_args.kwargs + assert call_kwargs["query"] == RICH_INPUT["query"] + assert "objective" not in call_kwargs + + @pytest.mark.asyncio + async def test_single_string_callers_unchanged(self, monkeypatch): + """No rich input: behavior is identical to before for any provider.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + mock_asearch = AsyncMock(return_value=_search_response()) + monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai")) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("plain query") + + call_kwargs = mock_asearch.await_args.kwargs + assert call_kwargs["query"] == "plain query" + assert "objective" not in call_kwargs + + @pytest.mark.asyncio + async def test_configured_objective_not_overwritten(self, monkeypatch): + """An objective set on the search tool's litellm_params wins over the model's.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + router = _mock_router("parallel_ai") + router.search_tools[0]["litellm_params"]["objective"] = "configured objective" + mock_asearch = AsyncMock(return_value=_search_response()) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + rich = WebSearchInterceptionLogger._rich_search_input(RICH_INPUT) + await logger._execute_search(RICH_INPUT["query"], rich=rich) + + call_kwargs = mock_asearch.await_args.kwargs + assert call_kwargs["objective"] == "configured objective" From 66ebc722d6751a7a8d0aef751e8d2e5f6a2efe49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:33:09 -0700 Subject: [PATCH 024/393] perf(azure): reuse the token refresh credential across image generation requests With enable_azure_ad_token_refresh, every keyless image request built a new DefaultAzureCredential and fetched a token. Cache the provider per scope like the Entra ID one. --- litellm/llms/azure/common_utils.py | 9 ++-- .../test_azure_image_generation_init.py | 47 +++++++++++++++++++ .../llms/azure/test_azure_common_utils.py | 3 ++ 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index f276d8b18d1..9f70761514b 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -95,6 +95,11 @@ def _cached_entra_id_token_provider( return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) +@lru_cache(maxsize=128) +def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=scope) + + def get_azure_ad_token_from_entra_id( tenant_id: str, client_id: str, @@ -649,9 +654,7 @@ class BaseAzureLLM(BaseOpenAILLM): "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth" ) try: - azure_ad_token_provider = get_azure_ad_token_provider( - azure_scope=scope, - ) + azure_ad_token_provider = _cached_azure_ad_token_refresh_provider(scope) except ValueError: verbose_logger.debug("Azure AD Token Provider could not be used.") if api_version is None: diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index a30aa277f3d..cfde1760389 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -11,6 +11,7 @@ import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.common_utils import ( + _cached_azure_ad_token_refresh_provider, _cached_entra_id_token_provider, get_azure_request_auth_headers, redact_azure_auth_headers, @@ -775,6 +776,52 @@ def test_azure_image_generation_with_api_key_keeps_api_key_header( assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***" +@pytest.fixture +def fake_default_azure_credential(monkeypatch: pytest.MonkeyPatch): + built_credentials = [] + + class FakeDefaultAzureCredential: + def __init__(self) -> None: + built_credentials.append(self) + + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", "AZURE_AD_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", FakeDefaultAzureCredential) + monkeypatch.setattr( + "azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "default-credential-token" + ) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + _cached_azure_ad_token_refresh_provider.cache_clear() + yield built_credentials + _cached_azure_ad_token_refresh_provider.cache_clear() + + +def test_azure_image_generation_token_refresh_reuses_credential_across_requests( + respx_mock: respx.MockRouter, fake_default_azure_credential: list +): + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + + for _ in range(3): + AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + ) + + assert route.call_count == 3 + assert all(call.request.headers["Authorization"] == "Bearer default-credential-token" for call in route.calls) + assert len(fake_default_azure_credential) == 1 + + @pytest.mark.parametrize( "caller_auth_header", [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}], diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f000abb4c9a..7189be7c052 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, + _cached_azure_ad_token_refresh_provider, _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, @@ -34,6 +35,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_TENANT_ID", raising=False) monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) + _cached_azure_ad_token_refresh_provider.cache_clear() with ( patch( @@ -78,6 +80,7 @@ def setup_mocks(monkeypatch): "logger": mock_logger, "select_url": mock_select_url, } + _cached_azure_ad_token_refresh_provider.cache_clear() def test_initialize_with_api_key(setup_mocks): From 05908bbe5767a020c2d4b3c88bc7941e3746fe71 Mon Sep 17 00:00:00 2001 From: Aidan Sinclair Date: Wed, 9 Sep 2026 08:41:04 -0400 Subject: [PATCH 025/393] fix(websearch): address review - ReadOnly TypedDict fields, suppression reason, call-site coverage - RichWebSearchInput fields are ReadOnly and constructed literally - the pyright suppression now states why the str provider name is safe - new tests drive _build_anthropic_request_patch and _build_chat_completion_request_patch end to end so the tool-call -> _rich_search_input wiring is covered, not just _execute_search Co-Authored-By: Claude Fable 5 --- .../websearch_interception/handler.py | 19 +++-- .../integrations/websearch_interception.py | 4 +- .../test_websearch_rich_query_shape.py | 70 +++++++++++++++++++ 3 files changed, 85 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 4fca0a36797..eda0413326e 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1646,18 +1646,25 @@ class WebSearchInterceptionLogger(CustomLogger): """ if not isinstance(tool_input, Mapping): return None - rich: RichWebSearchInput = {} objective = tool_input.get("objective") - if isinstance(objective, str) and objective.strip(): - rich["objective"] = objective + valid_objective = ( + objective if isinstance(objective, str) and objective.strip() else None + ) raw_queries = tool_input.get("search_queries") + valid_queries: list[str] | None = None if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str): queries = [q for q in raw_queries if isinstance(q, str) and q.strip()] if queries: # Providers cap multi-query requests (Parallel drops queries # past the fifth); trim here so nothing is silently ignored. - rich["search_queries"] = queries[:5] - return rich or None + valid_queries = queries[:5] + if valid_objective is not None and valid_queries is not None: + return {"objective": valid_objective, "search_queries": valid_queries} + if valid_objective is not None: + return {"objective": valid_objective} + if valid_queries is not None: + return {"search_queries": valid_queries} + return None @staticmethod def _provider_supports_rich_search(search_provider: str | None) -> bool: @@ -1672,7 +1679,7 @@ class WebSearchInterceptionLogger(CustomLogger): # misses the config map and returns None rather than raising. config = ProviderConfigManager.get_provider_search_config( search_provider - ) # pyright: ignore[reportArgumentType] + ) # pyright: ignore[reportArgumentType] -- SearchProviders is a str enum, so the router's provider string hashes to the matching member; unknown strings miss the map and yield None return config is not None and config.supports_rich_search_input() async def _execute_search( diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index bf01340630e..6b5b1519874 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -36,10 +36,10 @@ class RichWebSearchInput(TypedDict, total=False): other provider keeps receiving the single ``query`` string. """ - objective: str + objective: ReadOnly[str] """Natural-language description of the goal behind the search.""" - search_queries: list[str] + search_queries: ReadOnly[list[str]] """Two to five short keyword queries covering different angles.""" diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py index f8d20a3d5fd..e836a0d7062 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py @@ -186,3 +186,73 @@ class TestExecuteSearchShape: call_kwargs = mock_asearch.await_args.kwargs assert call_kwargs["objective"] == "configured objective" + + +class TestCallSiteWiring: + """Drive the patch builders end to end so regressions in the tool-call -> + _rich_search_input wiring are caught, not just _execute_search itself.""" + + @pytest.mark.asyncio + async def test_anthropic_tool_call_forwards_rich_shape(self, monkeypatch): + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + mock_asearch = AsyncMock(return_value=_search_response()) + monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai")) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + tool_calls = [ + {"id": "toolu_1", "name": "litellm_web_search", "input": dict(RICH_INPUT)} + ] + await logger._build_anthropic_request_patch( + model="claude", + messages=[{"role": "user", "content": "hi"}], + tool_calls=tool_calls, + thinking_blocks=[], + anthropic_messages_optional_request_params={}, + logging_obj=None, + kwargs={}, + ) + + call_kwargs = mock_asearch.await_args.kwargs + assert call_kwargs["query"] == RICH_INPUT["search_queries"] + assert call_kwargs["objective"] == RICH_INPUT["objective"] + + @pytest.mark.asyncio + async def test_chat_completion_tool_call_forwards_rich_shape(self, monkeypatch): + import json + + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger() + mock_asearch = AsyncMock(return_value=_search_response()) + monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai")) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + # The normalized shape transform_request produces for OpenAI responses: + # function.arguments (raw) plus top-level name/input (parsed). + tool_calls = [ + { + "id": "call_1", + "type": "function", + "name": "litellm_web_search", + "function": { + "name": "litellm_web_search", + "arguments": json.dumps(RICH_INPUT), + }, + "input": dict(RICH_INPUT), + } + ] + await logger._build_chat_completion_request_patch( + model="claude", + messages=[{"role": "user", "content": "hi"}], + tool_calls=tool_calls, + optional_params={}, + kwargs={}, + ) + + call_kwargs = mock_asearch.await_args.kwargs + assert call_kwargs["query"] == RICH_INPUT["search_queries"] + assert call_kwargs["objective"] == RICH_INPUT["objective"] From 038a4bb38429910298fc21db3d2a72d2729ee080 Mon Sep 17 00:00:00 2001 From: Aidan Sinclair Date: Wed, 9 Sep 2026 08:55:25 -0400 Subject: [PATCH 026/393] style(websearch): apply ruff format to changed files CI's lint gate checks ruff format, not black; black's output differs on a few line splits. No logic changes. Co-Authored-By: Claude Fable 5 --- .../websearch_interception/handler.py | 376 +++++------------- .../llms/base_llm/search/transformation.py | 24 +- .../llms/parallel_ai/search/transformation.py | 18 +- .../test_websearch_rich_query_shape.py | 20 +- 4 files changed, 109 insertions(+), 329 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index eda0413326e..f47751f4762 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -174,9 +174,7 @@ class _AcompletionNamedParams(TypedDict, total=False): logprobs: ReadOnly[bool | None] top_logprobs: ReadOnly[int | None] deployment_id: ReadOnly[str | None] - reasoning_effort: ReadOnly[ - Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None - ] + reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None] verbosity: ReadOnly[Literal["low", "medium", "high"] | None] safety_identifier: ReadOnly[str | None] service_tier: ReadOnly[str | None] @@ -234,9 +232,7 @@ class WebSearchInterceptionLogger(CustomLogger): if enabled_providers is None: self.enabled_providers = [LlmProviders.BEDROCK.value] else: - self.enabled_providers = [ - p.value if isinstance(p, LlmProviders) else p for p in enabled_providers - ] + self.enabled_providers = [p.value if isinstance(p, LlmProviders) else p for p in enabled_providers] self.search_tool_name = search_tool_name self.max_agentic_loops = self._validated_max_agentic_loops(max_agentic_loops) self._request_has_websearch = False # Track if current request has web search @@ -246,9 +242,7 @@ class WebSearchInterceptionLogger(CustomLogger): """ Reject loop ceilings the agentic loop cannot honor, at config load time. """ - return validated_max_agentic_loops( - max_agentic_loops, field="websearch_interception_params.max_agentic_loops" - ) + return validated_max_agentic_loops(max_agentic_loops, field="websearch_interception_params.max_agentic_loops") async def try_short_circuit_search( self, @@ -283,10 +277,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if provider is in enabled list provider_str: Final = custom_llm_provider or "" - if ( - self.enabled_providers is not None - and provider_str not in self.enabled_providers - ): + if self.enabled_providers is not None and provider_str not in self.enabled_providers: return None # Only short-circuit for providers whose Anthropic Messages agentic loop @@ -302,15 +293,10 @@ class WebSearchInterceptionLogger(CustomLogger): # web-search-only requests against it. try: provider_enum: Final = LlmProviders(provider_str) - anthropic_config: Final = ( - ProviderConfigManager.get_provider_anthropic_messages_config( - model=model, provider=provider_enum - ) + anthropic_config: Final = ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=provider_enum ) - if ( - anthropic_config is not None - and anthropic_config.handles_web_search_natively() - ): + if anthropic_config is not None and anthropic_config.handles_web_search_natively(): verbose_logger.debug( "WebSearchInterception: Skipping short-circuit for %s (provider handles web search natively via the agentic loop)", provider_str, @@ -355,13 +341,9 @@ class WebSearchInterceptionLogger(CustomLogger): if kwargs is None: search_result_text, structured = await self._execute_search(query) else: - search_result_text, structured = await self._execute_search( - query, kwargs=kwargs - ) + search_result_text, structured = await self._execute_search(query, kwargs=kwargs) except Exception as e: - verbose_logger.error( - "WebSearchInterception: Short-circuit search failed: %s", e - ) + verbose_logger.error("WebSearchInterception: Short-circuit search failed: %s", e) search_result_text, structured = f"Search failed: {e}", None content: Final[list[dict[str, object]]] = [] @@ -421,14 +403,12 @@ class WebSearchInterceptionLogger(CustomLogger): "litellm_params": kwargs.get("litellm_params", {}), "model": kwargs.get("model", ""), } - custom_llm_provider = call_kwargs_view[ - "custom_llm_provider" - ] or call_kwargs_view["litellm_params"].get("custom_llm_provider", "") + custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get( + "custom_llm_provider", "" + ) if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=call_kwargs_view["model"] - ) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"]) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -447,9 +427,7 @@ class WebSearchInterceptionLogger(CustomLogger): if not has_websearch: return None - verbose_logger.debug( - "WebSearchInterception: Converting native web_search tools to LiteLLM standard" - ) + verbose_logger.debug("WebSearchInterception: Converting native web_search tools to LiteLLM standard") # If the client sent an Anthropic-native web_search_* tool, mark the # request so the agentic loop emits native web_search_tool_result @@ -479,9 +457,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs["tools"] = converted_tools if kwargs.get("stream"): - verbose_logger.debug( - "WebSearchInterception: deployment hook converting stream=True to stream=False" - ) + verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") kwargs["stream"] = False kwargs["_websearch_interception_converted_stream"] = True @@ -494,34 +470,23 @@ class WebSearchInterceptionLogger(CustomLogger): if not any(is_web_search_tool_responses(tool) for tool in tools): return None - verbose_logger.debug( - "WebSearchInterception: Converting Responses web_search tools to LiteLLM standard" - ) + verbose_logger.debug("WebSearchInterception: Converting Responses web_search tools to LiteLLM standard") converted_tools: Final = [ - ( - get_litellm_web_search_tool_responses() - if is_web_search_tool_responses(tool) - else tool - ) - for tool in tools + (get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool) for tool in tools ] converted_kwargs: Final = {**kwargs, "tools": converted_tools} if kwargs.get("stream"): - verbose_logger.debug( - "WebSearchInterception: deployment hook converting stream=True to stream=False" - ) + verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") converted_kwargs["stream"] = False converted_kwargs["_websearch_interception_converted_stream"] = True return converted_kwargs @classmethod - def from_config_yaml( - cls, config: WebSearchInterceptionConfig - ) -> "WebSearchInterceptionLogger": + def from_config_yaml(cls, config: WebSearchInterceptionConfig) -> "WebSearchInterceptionLogger": """ Initialize WebSearchInterceptionLogger from proxy config.yaml parameters. @@ -576,9 +541,7 @@ class WebSearchInterceptionLogger(CustomLogger): return tool.get("name") @classmethod - def _sync_forced_tool_choice( - cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]] - ) -> object: + def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]]) -> object: """Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it names a web-search tool that was just converted away. @@ -595,9 +558,7 @@ class WebSearchInterceptionLogger(CustomLogger): return tool_choice return {**tool_choice, "name": LITELLM_WEB_SEARCH_TOOL_NAME} - async def async_pre_request_hook( - self, model: str, messages: list[dict], kwargs: dict - ) -> dict | None: + async def async_pre_request_hook(self, model: str, messages: list[dict], kwargs: dict) -> dict | None: """ Pre-request hook to convert native web search tools to LiteLLM standard. @@ -613,9 +574,7 @@ class WebSearchInterceptionLogger(CustomLogger): Modified kwargs dict with converted tools, or None if no modifications needed """ # Check if this request is for an enabled provider - custom_llm_provider: Final = kwargs.get("litellm_params", {}).get( - "custom_llm_provider", "" - ) + custom_llm_provider: Final = kwargs.get("litellm_params", {}).get("custom_llm_provider", "") verbose_logger.debug( "WebSearchInterception: Pre-request hook called - custom_llm_provider=%s - enabled_providers=%s", @@ -623,10 +582,7 @@ class WebSearchInterceptionLogger(CustomLogger): self.enabled_providers or "ALL", ) - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( "WebSearchInterception: Skipping - provider %s not in %s", custom_llm_provider, @@ -651,9 +607,7 @@ class WebSearchInterceptionLogger(CustomLogger): deployment_max_agentic_loops: Final = kwargs.get("max_agentic_loops") if self.max_agentic_loops is not None and deployment_max_agentic_loops is None: - kwargs["max_agentic_loops"] = ( - self.max_agentic_loops - ) # rebind-ok: this hook returns the kwargs it edits + kwargs["max_agentic_loops"] = self.max_agentic_loops # rebind-ok: this hook returns the kwargs it edits # If the client sent an Anthropic-native web_search_* tool, mark the # request so the agentic loop emits native web_search_tool_result @@ -685,15 +639,11 @@ class WebSearchInterceptionLogger(CustomLogger): ) if "tool_choice" in kwargs: - kwargs["tool_choice"] = self._sync_forced_tool_choice( - kwargs.get("tool_choice"), converted_tools - ) + kwargs["tool_choice"] = self._sync_forced_tool_choice(kwargs.get("tool_choice"), converted_tools) # Also convert here for direct callers that bypass the deployment hook. if kwargs.get("stream"): - verbose_logger.debug( - "WebSearchInterception: Converting stream=True to stream=False" - ) + verbose_logger.debug("WebSearchInterception: Converting stream=True to stream=False") kwargs["stream"] = False kwargs["_websearch_interception_converted_stream"] = True @@ -741,10 +691,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if provider should be intercepted # Note: custom_llm_provider is already normalized by get_llm_provider() # (e.g., "bedrock/invoke/..." -> "bedrock") - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( "WebSearchInterception: Skipping provider %s (not in enabled list: %s)", custom_llm_provider, @@ -766,9 +713,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug( - "WebSearchInterception: No WebSearch tool_use detected in response" - ) + verbose_logger.debug("WebSearchInterception: No WebSearch tool_use detected in response") return False, {} verbose_logger.debug( @@ -801,9 +746,7 @@ class WebSearchInterceptionLogger(CustomLogger): thinking_block_dict: dict = {"type": block_type} if block_type == "thinking": thinking_block_dict["thinking"] = getattr(block, "thinking", "") - thinking_block_dict["signature"] = getattr( - block, "signature", "" - ) + thinking_block_dict["signature"] = getattr(block, "signature", "") else: # redacted_thinking thinking_block_dict["data"] = getattr(block, "data", "") thinking_blocks.append(thinking_block_dict) @@ -848,10 +791,7 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug("WebSearchInterception: Response type: %s", type(response)) # Check if provider should be intercepted - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( "WebSearchInterception: Skipping provider %s (not in enabled list: %s)", custom_llm_provider, @@ -860,13 +800,9 @@ class WebSearchInterceptionLogger(CustomLogger): return False, {} # Check if tools include any web search tool (strict check for chat completions) - has_websearch_tool: Final = any( - is_web_search_tool_chat_completion(t) for t in (tools or []) - ) + has_websearch_tool: Final = any(is_web_search_tool_chat_completion(t) for t in (tools or [])) if not has_websearch_tool: - verbose_logger.debug( - "WebSearchInterception: No litellm_web_search tool in request" - ) + verbose_logger.debug("WebSearchInterception: No litellm_web_search tool in request") return False, {} # Detect WebSearch tool_calls in response (OpenAI format) @@ -877,9 +813,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug( - "WebSearchInterception: No WebSearch tool_calls detected in response" - ) + verbose_logger.debug("WebSearchInterception: No WebSearch tool_calls detected in response") return False, {} verbose_logger.debug( @@ -913,10 +847,7 @@ class WebSearchInterceptionLogger(CustomLogger): stream, ) - if ( - self.enabled_providers is not None - and custom_llm_provider not in self.enabled_providers - ): + if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: verbose_logger.debug( "WebSearchInterception: Skipping provider %s (not in enabled list: %s)", custom_llm_provider, @@ -924,13 +855,9 @@ class WebSearchInterceptionLogger(CustomLogger): ) return False, {} - has_websearch_tool: Final = any( - is_web_search_tool_responses(t) for t in (tools or []) - ) + has_websearch_tool: Final = any(is_web_search_tool_responses(t) for t in (tools or [])) if not has_websearch_tool: - verbose_logger.debug( - "WebSearchInterception: No litellm_web_search tool in responses request" - ) + verbose_logger.debug("WebSearchInterception: No litellm_web_search tool in responses request") return False, {} should_intercept, tool_calls = WebSearchTransformation.transform_request( @@ -940,9 +867,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) if not should_intercept: - verbose_logger.debug( - "WebSearchInterception: No WebSearch function_call detected in responses output" - ) + verbose_logger.debug("WebSearchInterception: No WebSearch function_call detected in responses output") return False, {} verbose_logger.debug( @@ -1053,11 +978,9 @@ class WebSearchInterceptionLogger(CustomLogger): # (while we still have the structured SearchResponse list) and stash # them on plan metadata for the post-hook to inject. if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): - metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = ( - self._build_native_result_blocks( - tool_calls=tool_calls, - structured_results=structured_results, - ) + metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( + tool_calls=tool_calls, + structured_results=structured_results, ) return AgenticLoopPlan( @@ -1083,9 +1006,7 @@ class WebSearchInterceptionLogger(CustomLogger): render citations / sources alongside the model's textual reply. """ metadata_view: Final[_PlanMetadataView] = { - "websearch_native_blocks": plan.metadata.get( - WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY - ) + "websearch_native_blocks": plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) } native_blocks: Final = metadata_view["websearch_native_blocks"] if not native_blocks: @@ -1110,9 +1031,7 @@ class WebSearchInterceptionLogger(CustomLogger): for i, tool_call in enumerate(tool_calls) for block in WebSearchInterceptionLogger._native_result_pair( query=WebSearchInterceptionLogger._tool_call_query(tool_call), - search_response=( - structured_results[i] if i < len(structured_results) else None - ), + search_response=(structured_results[i] if i < len(structured_results) else None), ) ) @@ -1131,9 +1050,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> tuple[Mapping[str, object], Mapping[str, object]]: tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" return ( - AnthropicServerToolUseBlock( - id=tool_use_id, input=AnthropicSearchQuery(query=query) - ).model_dump(), + AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(), WebSearchTransformation.build_web_search_tool_result_block( tool_use_id=tool_use_id, search_response=search_response, @@ -1141,9 +1058,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks( - response: _ResponseT, native_blocks: Sequence[Mapping[str, object]] - ) -> _ResponseT: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response @@ -1153,9 +1068,7 @@ class WebSearchInterceptionLogger(CustomLogger): return response existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or [] try: - setattr( - response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing) - ) + setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. @@ -1269,8 +1182,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs=kwargs, rich=self._rich_search_input(tool_call["input"]), ) - if isinstance(tool_call.get("input"), dict) - and tool_call["input"].get("query") + if isinstance(tool_call.get("input"), dict) and tool_call["input"].get("query") else self._create_empty_search_result() ) for tool_call in tool_calls @@ -1280,13 +1192,9 @@ class WebSearchInterceptionLogger(CustomLogger): "WebSearchInterception: Executing %s responses search(es) in parallel", len(search_tasks), ) - search_results: Final = await asyncio.gather( - *search_tasks, return_exceptions=True - ) + search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) - search_texts: Final = [ - self._extract_search_text(result) for result in search_results - ] + search_texts: Final = [self._extract_search_text(result) for result in search_results] followup_items: Final = [ item @@ -1367,16 +1275,12 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _extract_search_text(result: object) -> str: if isinstance(result, Exception): - verbose_logger.error( - "WebSearchInterception: Responses search failed with error: %s", result - ) + verbose_logger.error("WebSearchInterception: Responses search failed with error: %s", result) return f"Search failed: {result}" if isinstance(result, tuple) and len(result) == 2: text_value, _ = result return text_value if isinstance(text_value, str) else str(text_value) - verbose_logger.debug( - "WebSearchInterception: Unexpected search result type %s", type(result) - ) + verbose_logger.debug("WebSearchInterception: Unexpected search result type %s", type(result)) return str(result) @staticmethod @@ -1427,9 +1331,7 @@ class WebSearchInterceptionLogger(CustomLogger): """ _internal_keys: Final = {"litellm_logging_obj"} return { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") and k not in _internal_keys + k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in _internal_keys } async def _execute_agentic_loop( @@ -1449,9 +1351,7 @@ class WebSearchInterceptionLogger(CustomLogger): messages=messages, tool_calls=tool_calls, thinking_blocks=thinking_blocks, - anthropic_messages_optional_request_params=dict[str, object]( - anthropic_messages_optional_request_params - ), + anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params), logging_obj=logging_obj, kwargs=dict[str, object](kwargs), ) @@ -1469,15 +1369,13 @@ class WebSearchInterceptionLogger(CustomLogger): max_tokens = cast(int, kwargs.get("max_tokens", 1024)) patch_kwargs: Final = dict[str, object](request_patch.kwargs) - response: AnthropicMessagesResponse | AsyncIterator[object] = ( - await anthropic_messages.acreate( - max_tokens=max_tokens, - messages=request_patch.messages, - model=request_patch.model or model, - **_NO_ACREATE_NAMED, - **optional_params, - **patch_kwargs, - ) + response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( + max_tokens=max_tokens, + messages=request_patch.messages, + model=request_patch.model or model, + **_NO_ACREATE_NAMED, + **optional_params, + **patch_kwargs, ) # Legacy path: the new path goes through the typed plan + core @@ -1517,9 +1415,7 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: query = tool_call["input"].get("query") if query: - verbose_logger.debug( - "WebSearchInterception: Queuing search for query='%s'", query - ) + verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query) search_tasks.append( self._execute_search( query, @@ -1528,9 +1424,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) ) else: - verbose_logger.debug( - "WebSearchInterception: Tool call %s has no query", tool_call["id"] - ) + verbose_logger.debug("WebSearchInterception: Tool call %s has no query", tool_call["id"]) # Add empty result for tools without query search_tasks.append(self._create_empty_search_result()) @@ -1539,9 +1433,7 @@ class WebSearchInterceptionLogger(CustomLogger): "WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks), ) - search_results: Final = await asyncio.gather( - *search_tasks, return_exceptions=True - ) + search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) # Split the gathered (text, structured) tuples into two parallel lists. # The text list feeds the follow-up model call; the structured list @@ -1550,23 +1442,13 @@ class WebSearchInterceptionLogger(CustomLogger): structured_results: Final[list[SearchResponse | None]] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error( - "WebSearchInterception: Search %s failed with error: %s", i, result - ) + verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result) final_search_results.append(f"Search failed: {result}") structured_results.append(None) elif isinstance(result, tuple) and len(result) == 2: text_value, structured_value = result - final_search_results.append( - cast(str, text_value) - if isinstance(text_value, str) - else str(text_value) - ) - structured_results.append( - structured_value - if isinstance(structured_value, SearchResponse) - else None - ) + final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) + structured_results.append(structured_value if isinstance(structured_value, SearchResponse) else None) else: # Defensive: legacy callers / unexpected shape — preserve text, # drop structure. @@ -1591,15 +1473,11 @@ class WebSearchInterceptionLogger(CustomLogger): ] # Correlation context for structured logging - _call_id: Final = getattr(logging_obj, "litellm_call_id", None) or kwargs.get( - "litellm_call_id", "unknown" - ) + _call_id: Final = getattr(logging_obj, "litellm_call_id", None) or kwargs.get("litellm_call_id", "unknown") full_model_name = model # safe default before try block - max_tokens: Final = self._resolve_max_tokens( - anthropic_messages_optional_request_params, kwargs - ) + max_tokens: Final = self._resolve_max_tokens(anthropic_messages_optional_request_params, kwargs) verbose_logger.debug( "WebSearchInterception: Using max_tokens=%s for follow-up request", @@ -1607,17 +1485,13 @@ class WebSearchInterceptionLogger(CustomLogger): ) optional_params_without_max_tokens: Final = { - k: v - for k, v in anthropic_messages_optional_request_params.items() - if k != "max_tokens" + k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" } kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: agentic_view: Final[_AgenticLoopParamsView] = { - "agentic_loop_params": logging_obj.model_call_details.get( - "agentic_loop_params", {} - ) + "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {}) } full_model_name = agentic_view["agentic_loop_params"].get("model", model) verbose_logger.debug( @@ -1647,9 +1521,7 @@ class WebSearchInterceptionLogger(CustomLogger): if not isinstance(tool_input, Mapping): return None objective = tool_input.get("objective") - valid_objective = ( - objective if isinstance(objective, str) and objective.strip() else None - ) + valid_objective = objective if isinstance(objective, str) and objective.strip() else None raw_queries = tool_input.get("search_queries") valid_queries: list[str] | None = None if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str): @@ -1677,9 +1549,7 @@ class WebSearchInterceptionLogger(CustomLogger): return False # SearchProviders is a str enum, so an unknown provider string simply # misses the config map and returns None rather than raising. - config = ProviderConfigManager.get_provider_search_config( - search_provider - ) # pyright: ignore[reportArgumentType] -- SearchProviders is a str enum, so the router's provider string hashes to the matching member; unknown strings miss the map and yield None + config = ProviderConfigManager.get_provider_search_config(search_provider) # pyright: ignore[reportArgumentType] -- SearchProviders is a str enum, so the router's provider string hashes to the matching member; unknown strings miss the map and yield None return config is not None and config.supports_rich_search_input() async def _execute_search( @@ -1709,21 +1579,13 @@ class WebSearchInterceptionLogger(CustomLogger): ) llm_router = None - search_tool: Final = self._select_search_tool_from_router( - llm_router=llm_router - ) + search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None search_litellm_params: Mapping[str, object] = {} - search_tool_name: Final = self._selected_search_tool_name( - search_tool=search_tool - ) + search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: - await self._authorize_search_tool( - search_tool=search_tool, kwargs=kwargs - ) - tool_params: Final[_SearchToolLitellmParams] = ( - search_tool.get("litellm_params", {}) or {} - ) + await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) + tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {} search_litellm_params = dict[str, object](tool_params) search_provider = tool_params.get("search_provider") @@ -1783,9 +1645,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) # Format using transformation function - search_result_text: Final = WebSearchTransformation.format_search_response( - result - ) + search_result_text: Final = WebSearchTransformation.format_search_response(result) verbose_logger.debug( "WebSearchInterception: Search completed for '%s', got %s chars", @@ -1794,9 +1654,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) return search_result_text, result except Exception as e: - verbose_logger.error( - "WebSearchInterception: Search failed for '%s': %s", query, e - ) + verbose_logger.error("WebSearchInterception: Search failed for '%s': %s", query, e) raise async def _authorize_search_tool( @@ -1856,9 +1714,7 @@ class WebSearchInterceptionLogger(CustomLogger): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( - LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_auth - ) + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) ) return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches **user_api_key_metadata, @@ -1874,11 +1730,7 @@ class WebSearchInterceptionLogger(CustomLogger): if search_tool is None: return None search_tool_name: Final = search_tool.get("search_tool_name") - return ( - search_tool_name - if isinstance(search_tool_name, str) and search_tool_name - else None - ) + return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None @staticmethod def _get_user_api_key_auth_from_kwargs( @@ -1889,10 +1741,7 @@ class WebSearchInterceptionLogger(CustomLogger): for metadata_key in ("metadata", "litellm_metadata"): metadata = kwargs.get(metadata_key) - if ( - isinstance(metadata, dict) - and metadata.get("user_api_key_auth") is not None - ): + if isinstance(metadata, dict) and metadata.get("user_api_key_auth") is not None: return metadata["user_api_key_auth"] litellm_params: Final = kwargs.get("litellm_params") @@ -1901,23 +1750,16 @@ class WebSearchInterceptionLogger(CustomLogger): for metadata_key in ("metadata", "litellm_metadata"): metadata = litellm_params.get(metadata_key) - if ( - isinstance(metadata, dict) - and metadata.get("user_api_key_auth") is not None - ): + if isinstance(metadata, dict) and metadata.get("user_api_key_auth") is not None: return metadata["user_api_key_auth"] return None - def _select_search_tool_from_router( - self, llm_router: object - ) -> "_SearchToolConfig | None": + def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": if llm_router is None or not hasattr(llm_router, "search_tools"): return None search_tools: Final = tuple(getattr(llm_router, "search_tools", None) or ()) - return self._select_search_tool_from_list( - search_tools=search_tools, source="router" - ) + return self._select_search_tool_from_list(search_tools=search_tools, source="router") def _select_search_tool_from_list( self, @@ -1926,14 +1768,10 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> "_SearchToolConfig | None": if self.search_tool_name: matching_tools: Final = tuple( - tool - for tool in search_tools - if tool.get("search_tool_name") == self.search_tool_name + tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name ) if matching_tools: - search_provider = ( - matching_tools[0].get("litellm_params", {}) or {} - ).get("search_provider") + search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider") verbose_logger.debug( "WebSearchInterception: Found search tool '%s' from %s with provider '%s'", self.search_tool_name, @@ -1949,9 +1787,7 @@ class WebSearchInterceptionLogger(CustomLogger): if search_tools: first_tool: Final = search_tools[0] - search_provider = (first_tool.get("litellm_params", {}) or {}).get( - "search_provider" - ) + search_provider = (first_tool.get("litellm_params", {}) or {}).get("search_provider") verbose_logger.debug( "WebSearchInterception: Using first available search tool from %s with provider '%s'", source, @@ -2024,14 +1860,8 @@ class WebSearchInterceptionLogger(CustomLogger): query = args.get("query") if query: - verbose_logger.debug( - "WebSearchInterception: Queuing search for query='%s'", query - ) - search_tasks.append( - self._execute_search( - query, kwargs=kwargs, rich=self._rich_search_input(tool_args) - ) - ) + verbose_logger.debug("WebSearchInterception: Queuing search for query='%s'", query) + search_tasks.append(self._execute_search(query, kwargs=kwargs, rich=self._rich_search_input(tool_args))) else: verbose_logger.debug( "WebSearchInterception: Tool call %s has no query", @@ -2045,26 +1875,18 @@ class WebSearchInterceptionLogger(CustomLogger): "WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks), ) - search_results: Final = await asyncio.gather( - *search_tasks, return_exceptions=True - ) + search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) # Chat-completion path only needs text — OpenAI tool_result format # has no equivalent of Anthropic's web_search_tool_result block. final_search_results: Final[list[str]] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): - verbose_logger.error( - "WebSearchInterception: Search %s failed with error: %s", i, result - ) + verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result) final_search_results.append(f"Search failed: {result}") elif isinstance(result, tuple) and len(result) == 2: text_value, _ = result - final_search_results.append( - cast(str, text_value) - if isinstance(text_value, str) - else str(text_value) - ) + final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) else: verbose_logger.debug( "WebSearchInterception: Unexpected result type %s at index %s", @@ -2086,9 +1908,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Make follow-up request with search results # For OpenAI format, tool_messages_or_user is a list of tool messages if response_format == "openai": - follow_up_messages = ( - messages + [assistant_message] + cast(list[dict], tool_messages_or_user) - ) + follow_up_messages = messages + [assistant_message] + cast(list[dict], tool_messages_or_user) else: # For Anthropic format (shouldn't happen in this method, but handle it) follow_up_messages = messages + [ @@ -2096,9 +1916,7 @@ class WebSearchInterceptionLogger(CustomLogger): cast(dict, tool_messages_or_user), ] - verbose_logger.debug( - "WebSearchInterception: Making follow-up chat completion request with search results" - ) + verbose_logger.debug("WebSearchInterception: Making follow-up chat completion request with search results") verbose_logger.debug( "WebSearchInterception: Follow-up messages count: %s", len(follow_up_messages), @@ -2115,9 +1933,7 @@ class WebSearchInterceptionLogger(CustomLogger): "custom_prompt_dict", } kwargs_for_followup: Final = { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") and k not in internal_params + k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in internal_params } full_model_name = model @@ -2190,9 +2006,7 @@ class WebSearchInterceptionLogger(CustomLogger): websearch_params: WebSearchInterceptionConfig = {} if "websearch_interception_params" in litellm_settings: settings_view: Final[_WebSearchSettingsView] = { - "websearch_interception_params": litellm_settings[ - "websearch_interception_params" - ] + "websearch_interception_params": litellm_settings["websearch_interception_params"] } websearch_params = settings_view["websearch_interception_params"] elif "websearch_interception" in callback_specific_params and isinstance( diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index c183d538c01..4794fdd0d74 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -197,20 +197,12 @@ class BaseSearchConfig: def sign_request( self, - headers: dict[ - str, str - ], # mutable-ok: matches the request header dict every other hook on this base takes - optional_params: dict[ - str, object - ], # mutable-ok: matches every other hook on this base - request_data: ( - dict[str, object] | list[dict[str, object]] - ), # mutable-ok: transform_search_request's body + headers: dict[str, str], # mutable-ok: matches the request header dict every other hook on this base takes + optional_params: dict[str, object], # mutable-ok: matches every other hook on this base + request_data: (dict[str, object] | list[dict[str, object]]), # mutable-ok: transform_search_request's body api_base: str, api_key: str | None = None, - ) -> tuple[ - dict[str, str], bytes | None - ]: # mutable-ok: the handler passes these headers straight to httpx + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: the handler passes these headers straight to httpx """ OPTIONAL @@ -270,9 +262,7 @@ class BaseSearchConfig: Returns: Dict with request data """ - raise NotImplementedError( - "transform_search_request must be implemented by provider" - ) + raise NotImplementedError("transform_search_request must be implemented by provider") def transform_search_response( self, @@ -284,9 +274,7 @@ class BaseSearchConfig: Transform provider-specific Search response to standard format. Override in provider-specific implementations. """ - raise NotImplementedError( - "transform_search_response must be implemented by provider" - ) + raise NotImplementedError("transform_search_response must be implemented by provider") def get_error_class( self, diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index 4154a497d2c..d91e532a2cf 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -110,9 +110,7 @@ class ParallelAISearchConfig(BaseSearchConfig): default_api_base=self.PARALLEL_AI_API_BASE, ) if not resolved_api_key: - raise ValueError( - "PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable." - ) + raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") headers["x-api-key"] = resolved_api_key headers["Content-Type"] = "application/json" return headers @@ -124,11 +122,7 @@ class ParallelAISearchConfig(BaseSearchConfig): data: dict | list[dict] | None = None, **kwargs, ) -> str: - resolved_api_base: Final = ( - api_base - or get_secret_str("PARALLEL_AI_API_BASE") - or self.PARALLEL_AI_API_BASE - ) + resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE trimmed: Final = resolved_api_base.rstrip("/") if trimmed.endswith("/v1/search"): @@ -195,9 +189,7 @@ class ParallelAISearchConfig(BaseSearchConfig): advanced_settings["location"] = params.pop("location") if "max_chars_per_result" in params: - advanced_settings["excerpt_settings"] = { - "max_chars_per_result": params.pop("max_chars_per_result") - } + advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")} if "fetch_policy" in params: advanced_settings["fetch_policy"] = params.pop("fetch_policy") @@ -290,6 +282,4 @@ class ParallelAISearchConfig(BaseSearchConfig): } ) - return SearchResponse.model_validate( - MappingProxyType({"results": results, "object": "search", **extra_fields}) - ) + return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields})) diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py index e836a0d7062..72149e8a435 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_rich_query_shape.py @@ -71,9 +71,7 @@ class TestRichInputExtraction: } def test_returns_none_when_only_query_present(self): - assert ( - WebSearchInterceptionLogger._rich_search_input({"query": "plain"}) is None - ) + assert WebSearchInterceptionLogger._rich_search_input({"query": "plain"}) is None def test_returns_none_for_non_mapping_input(self): assert WebSearchInterceptionLogger._rich_search_input(None) is None @@ -90,12 +88,7 @@ class TestRichInputExtraction: def test_ignores_string_valued_search_queries(self): # A string is a Sequence; it must not be treated as a list of queries. - assert ( - WebSearchInterceptionLogger._rich_search_input( - {"query": "q", "search_queries": "not a list"} - ) - is None - ) + assert WebSearchInterceptionLogger._rich_search_input({"query": "q", "search_queries": "not a list"}) is None class TestProviderSupport: @@ -107,10 +100,7 @@ class TestProviderSupport: def test_unknown_provider_is_unsupported(self): assert WebSearchInterceptionLogger._provider_supports_rich_search(None) is False - assert ( - WebSearchInterceptionLogger._provider_supports_rich_search("not_a_provider") - is False - ) + assert WebSearchInterceptionLogger._provider_supports_rich_search("not_a_provider") is False class TestExecuteSearchShape: @@ -202,9 +192,7 @@ class TestCallSiteWiring: monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai")) monkeypatch.setattr(litellm, "asearch", mock_asearch) - tool_calls = [ - {"id": "toolu_1", "name": "litellm_web_search", "input": dict(RICH_INPUT)} - ] + tool_calls = [{"id": "toolu_1", "name": "litellm_web_search", "input": dict(RICH_INPUT)}] await logger._build_anthropic_request_patch( model="claude", messages=[{"role": "user", "content": "hi"}], From 6cc3a6022193ae54cb04c73989c5337e5fe0db75 Mon Sep 17 00:00:00 2001 From: Aidan Sinclair Date: Wed, 9 Sep 2026 09:27:41 -0400 Subject: [PATCH 027/393] chore(websearch): justify new mutable annotations for the type-discipline gate Adds the required mutable-ok reasons to the five annotations this change introduced; no logic changes. Co-Authored-By: Claude Fable 5 --- litellm/integrations/websearch_interception/handler.py | 6 +++--- litellm/integrations/websearch_interception/tools.py | 2 +- litellm/types/integrations/websearch_interception.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index f47751f4762..093e0351c70 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1523,7 +1523,7 @@ class WebSearchInterceptionLogger(CustomLogger): objective = tool_input.get("objective") valid_objective = objective if isinstance(objective, str) and objective.strip() else None raw_queries = tool_input.get("search_queries") - valid_queries: list[str] | None = None + valid_queries: list[str] | None = None # mutable-ok: matches litellm.asearch's list[str] query parameter if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str): queries = [q for q in raw_queries if isinstance(q, str) and q.strip()] if queries: @@ -1619,7 +1619,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Forward the model's richer shape (objective + keyword queries) # only to providers whose search API takes it natively; everyone # else keeps the single query string the model also provided. - query_arg: str | list[str] = query + query_arg: str | list[str] = query # mutable-ok: litellm.asearch declares query as str | list[str] if rich and self._provider_supports_rich_search(search_provider): rich_queries = rich.get("search_queries") if rich_queries: @@ -1847,7 +1847,7 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: # Handle both Anthropic-style input and OpenAI-style function.arguments query = None - tool_args: dict | None = None + tool_args: dict | None = None # mutable-ok: the tool call's own arguments dict if "input" in tool_call and isinstance(tool_call["input"], dict): tool_args = tool_call["input"] query = tool_args.get("query") diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 9e3d3fd91f3..2e1ae07eb68 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -17,7 +17,7 @@ _WEB_SEARCH_TOOL_DESCRIPTION: Final = ( ) -def _web_search_input_schema() -> dict[str, object]: +def _web_search_input_schema() -> dict[str, object]: # mutable-ok: plain-dict tool shape, as the get_* builders """ JSON schema for the web search tool's input, shared by every tool format. diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 6b5b1519874..ea5e6d51749 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -39,7 +39,7 @@ class RichWebSearchInput(TypedDict, total=False): objective: ReadOnly[str] """Natural-language description of the goal behind the search.""" - search_queries: ReadOnly[list[str]] + search_queries: ReadOnly[list[str]] # mutable-ok: forwarded verbatim as litellm.asearch's list[str] query argument """Two to five short keyword queries covering different angles.""" From 233337628f4b6f1c9ec0527d5d442cfe512ac08b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:06:25 -0400 Subject: [PATCH 028/393] feat(batches): support Mistral files/batches and per-page OCR batch cost tracking Adds MistralFilesConfig and MistralBatchesConfig so Mistral can be used as a Files and Batches provider through the shared BaseLLMHTTPHandler path, the same way Bedrock plugs in. /v1/ocr is now an accepted batch endpoint, and completed OCR batches are billed per page (ocr_cost_per_page_batches, half the synchronous rate) instead of per token. Resolves #29914 --- litellm/batches/batch_utils.py | 41 ++- litellm/batches/main.py | 20 +- litellm/cost_calculator.py | 61 ++++ litellm/files/main.py | 9 +- litellm/files/types.py | 2 +- litellm/llms/mistral/batches/__init__.py | 0 .../llms/mistral/batches/transformation.py | 186 +++++++++++++ litellm/llms/mistral/common_utils.py | 36 +++ litellm/llms/mistral/files/__init__.py | 0 litellm/llms/mistral/files/transformation.py | 226 +++++++++++++++ ...odel_prices_and_context_window_backup.json | 40 ++- litellm/types/llms/openai.py | 2 +- litellm/types/utils.py | 4 + litellm/utils.py | 10 + model_prices_and_context_window.json | 40 ++- .../test_litellm/batches/test_batch_utils.py | 83 ++++++ tests/test_litellm/batches/test_main.py | 42 +++ .../llms/mistral/batches/__init__.py | 0 .../test_mistral_batches_transformation.py | 260 ++++++++++++++++++ .../llms/mistral/files/__init__.py | 0 .../test_mistral_files_transformation.py | 189 +++++++++++++ .../llms/mistral/ocr/test_mistral_ocr_cost.py | 4 +- tests/test_litellm/test_utils.py | 2 + 23 files changed, 1216 insertions(+), 41 deletions(-) create mode 100644 litellm/llms/mistral/batches/__init__.py create mode 100644 litellm/llms/mistral/batches/transformation.py create mode 100644 litellm/llms/mistral/common_utils.py create mode 100644 litellm/llms/mistral/files/__init__.py create mode 100644 litellm/llms/mistral/files/transformation.py create mode 100644 tests/test_litellm/llms/mistral/batches/__init__.py create mode 100644 tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py create mode 100644 tests/test_litellm/llms/mistral/files/__init__.py create mode 100644 tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 87f8fd3946e..077d6e72fd7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details +from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -50,7 +51,7 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -80,7 +81,7 @@ async def calculate_batch_cost_and_usage( async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, litellm_params: dict | None = None, model_info: ModelInfo | None = None, @@ -166,7 +167,7 @@ class _BatchOutputLineStats: def _classify_output_line_stats( entries: Iterable[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats | _LineOutcome]: @@ -185,7 +186,7 @@ def _classify_output_line_stats( def _safe_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats | None: @@ -207,7 +208,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats: @@ -218,6 +219,7 @@ def _compute_output_line_stats( response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details line_prompt_cost, line_completion_cost = _output_line_cost( + response_body=response_body, usage=usage, custom_llm_provider=custom_llm_provider, model_name=model_name, @@ -237,19 +239,36 @@ def _compute_output_line_stats( ) +def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> OCRUsageInfo | None: + """OCR results report ``usage_info`` (pages) instead of ``usage`` (tokens); None for non-OCR lines.""" + raw_usage_info: Final = response_body.get("usage_info") + if not isinstance(raw_usage_info, Mapping): + return None + return OCRUsageInfo.model_validate(raw_usage_info) + + def _output_line_cost( + response_body: Mapping[str, object], usage: Usage, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, ) -> tuple[float, float]: """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" - from litellm.cost_calculator import batch_cost_calculator + from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost cost_model: Final = ( model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" ) + ocr_usage: Final = _ocr_usage_info_from_response_body(response_body) + if ocr_usage is not None: + return ocr_batch_cost( + model=cost_model, + custom_llm_provider=custom_llm_provider, + usage_info=ocr_usage, + model_info=model_info, + ) return batch_cost_calculator( usage=usage, model=cost_model, @@ -260,7 +279,7 @@ def _output_line_cost( def _aggregate_batch_cost_usage_models( entries: Iterable[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -427,7 +446,7 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -457,7 +476,7 @@ async def _fetch_batch_managed_file_content( async def _fetch_batch_output_file_content( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -479,7 +498,7 @@ async def _fetch_batch_output_file_content( async def count_error_file_failed_requests( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], litellm_params: dict | None, ) -> int: """Count failed requests reported only in the batch's separate error file. diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 77a4fdebf16..76b6c73b375 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -105,9 +105,11 @@ def _resolve_timeout( @client async def acreate_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -155,9 +157,11 @@ async def acreate_batch( @client def create_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -341,7 +345,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -389,7 +393,7 @@ def _handle_retrieve_batch_providers_without_provider_config( _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", logging_obj: LiteLLMLoggingObj | None = None, ): @@ -497,7 +501,7 @@ def _handle_retrieve_batch_providers_without_provider_config( message=( f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. " "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. " - "'bedrock' is supported but requires `model` to be passed so the provider config can be loaded." + "'bedrock' and 'mistral' are supported but require `model` to be passed so the provider config can be loaded." ), model="n/a", llm_provider=custom_llm_provider, @@ -514,7 +518,7 @@ def _handle_retrieve_batch_providers_without_provider_config( def retrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..0c0c6b05df8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -139,6 +139,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LitellmLoggingObject, ) + from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo else: LitellmLoggingObject = Any @@ -1982,6 +1983,66 @@ def ocr_cost( return ocr_pages_cost + annotation_pages_cost, 0.0 +_OCR_PRICING_KEYS: Final = ( + "ocr_cost_per_page", + "ocr_cost_per_page_batches", + "annotation_cost_per_page", + "annotation_cost_per_page_batches", +) + + +def ocr_batch_cost( + model: str, + custom_llm_provider: str | None, + usage_info: "OCRUsageInfo", + model_info: ModelInfo | None = None, +) -> tuple[float, float]: + """Per-page cost of one OCR result inside a batch output file. + + Batch OCR is billed per page at the ``*_batches`` rate, falling back to the + synchronous per-page rate when a model has no batch price recorded, the same + fallback ``batch_cost_calculator`` applies to per-token batch pricing. Returns + ``(prompt_cost, completion_cost)`` with the whole cost in the first slot, like + ``ocr_cost``. + """ + has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) + if has_ocr_pricing: + resolved_info: ModelInfo | None = model_info + else: + try: + resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + resolved_info = None + if resolved_info is None: + verbose_logger.warning( + "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", + model, + custom_llm_provider, + ) + return 0.0, 0.0 + + page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") + annotation_rate: Final = _first_price( + resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page" + ) + pages_processed: Final = usage_info.pages_processed or 0 + annotation_pages: Final = usage_info.pages_processed_annotation or 0 + if page_rate is None and pages_processed > 0: + verbose_logger.warning( + "OCR batch cost: model=%s custom_llm_provider=%s reported pages_processed=%s but no " + "ocr_cost_per_page is configured; returning 0.0 cost for those pages.", + model, + custom_llm_provider, + pages_processed, + ) + effective_annotation_rate: Final = annotation_rate if annotation_rate is not None else page_rate + return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0 + + +def _first_price(model_info: ModelInfo, *keys: str) -> float | None: + return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) + + def vector_store_search_cost( model: str | None, custom_llm_provider: str, diff --git a/litellm/files/main.py b/litellm/files/main.py index 218518eb3cd..3d90bf4f299 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -27,12 +27,15 @@ FileCreateProvider = Literal[ "litellm_proxy", "manus", "anthropic", + "mistral", ] FileRetrieveProvider = Literal[ - "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"] -FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] +FileDeleteProvider = Literal[ + "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral" +] +FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse diff --git a/litellm/files/types.py b/litellm/files/types.py index b4ec9996f37..01c7970144b 100644 --- a/litellm/files/types.py +++ b/litellm/files/types.py @@ -2,7 +2,7 @@ from collections.abc import AsyncIterator, Iterator from typing import Literal, NamedTuple FileContentProvider = Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus", "mistral" ] diff --git a/litellm/llms/mistral/batches/__init__.py b/litellm/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py new file mode 100644 index 00000000000..399319e590b --- /dev/null +++ b/litellm/llms/mistral/batches/transformation.py @@ -0,0 +1,186 @@ +""" +Mistral Batch API. Reference: https://docs.mistral.ai/api/#tag/batch + +Mistral runs one model per job (set on the job, not per input line) and accepts +``/v1/ocr`` as a batch endpoint, which is how OCR gets its 50% batch discount. +Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_code, body}}``), +so the shared batch cost accounting reads them without a provider branch. +""" + +from types import MappingProxyType +from typing import Final, Literal + +import httpx +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Errors as BatchErrors +from openai.types.batch_error import BatchError +from pydantic import BaseModel, ConfigDict + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders + +from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error + +MistralBatchStatus = Literal[ + "QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED" +] +OpenAIBatchStatus = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + +_STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType( + { + "QUEUED": "validating", + "RUNNING": "in_progress", + "SUCCESS": "completed", + "FAILED": "failed", + "TIMEOUT_EXCEEDED": "expired", + "CANCELLATION_REQUESTED": "cancelling", + "CANCELLED": "cancelled", + } +) + + +class MistralBatchError(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + message: str + count: int = 1 + + +class MistralBatchJob(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + input_files: tuple[str, ...] = () + endpoint: str + model: str | None = None + status: MistralBatchStatus + created_at: int + started_at: int | None = None + completed_at: int | None = None + total_requests: int = 0 + completed_requests: int = 0 + succeeded_requests: int = 0 + failed_requests: int = 0 + output_file: str | None = None + error_file: str | None = None + errors: tuple[MistralBatchError, ...] = () + metadata: dict[str, str] | None = None + + +def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: + status: Final = _STATUS_MAP[job.status] + terminal_at: Final = job.completed_at + return LiteLLMBatch( + id=job.id, + object="batch", + endpoint=job.endpoint, + input_file_id=job.input_files[0] if job.input_files else "", + completion_window="24h", + status=status, + created_at=job.created_at, + in_progress_at=job.started_at, + completed_at=terminal_at if status == "completed" else None, + failed_at=terminal_at if status == "failed" else None, + expired_at=terminal_at if status == "expired" else None, + cancelled_at=terminal_at if status == "cancelled" else None, + output_file_id=job.output_file, + error_file_id=job.error_file, + errors=( + BatchErrors( + object="list", + data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors], + ) + if job.errors + else None + ), + request_counts=BatchRequestCounts( + total=job.total_requests, + completed=job.succeeded_requests, + failed=job.failed_requests, + ), + metadata=job.metadata, + ) + + +class MistralBatchesConfig(BaseBatchesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.MISTRAL + + def validate_environment( + self, + headers: dict, + model: str, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + return get_mistral_auth_headers(headers, api_key) + + def get_complete_batch_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + data: CreateBatchRequest, + ) -> str: + return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" + + def transform_create_batch_request( + self, + model: str, + create_batch_data: CreateBatchRequest, + optional_params: dict, + litellm_params: dict, + ) -> dict[str, object]: + metadata: Final = create_batch_data.get("metadata") + return { + "input_files": [create_batch_data["input_file_id"]], + "endpoint": create_batch_data["endpoint"], + "model": model, + **({"metadata": metadata} if metadata else {}), + **(create_batch_data.get("extra_body") or {}), + } + + def transform_create_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: dict, + ) -> LiteLLMBatch: + return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) + + def transform_retrieve_batch_request( + self, + batch_id: str, + optional_params: dict, + litellm_params: dict, + ) -> dict[str, object]: + encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") + return { + "method": "GET", + "url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}", + "headers": get_mistral_auth_headers({}, litellm_params.get("api_key")), + } + + def transform_retrieve_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: dict, + ) -> LiteLLMBatch: + return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) + + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + return mistral_error(error_message, status_code, headers) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py new file mode 100644 index 00000000000..9ea501c860d --- /dev/null +++ b/litellm/llms/mistral/common_utils.py @@ -0,0 +1,36 @@ +from typing import Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str + +MISTRAL_API_BASE: Final = "https://api.mistral.ai" +MISTRAL_API_KEY_ENV_VAR: Final = "MISTRAL_API_KEY" + + +class MistralError(BaseLLMException): + pass + + +def get_mistral_api_base(api_base: str | None) -> str: + """Return the Mistral origin without a trailing ``/v1``, so callers can append ``/v1/``.""" + resolved: Final = (api_base or get_secret_str("MISTRAL_API_BASE") or MISTRAL_API_BASE).rstrip("/") + return resolved.removesuffix("/v1") + + +def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict: + resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR) + if resolved_key is None: + raise ValueError( + "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" + ) + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + +def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError: + return MistralError( + status_code=status_code, + message=error_message, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), + ) diff --git a/litellm/llms/mistral/files/__init__.py b/litellm/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py new file mode 100644 index 00000000000..071b6f58569 --- /dev/null +++ b/litellm/llms/mistral/files/transformation.py @@ -0,0 +1,226 @@ +""" +Mistral Files API. Reference: https://docs.mistral.ai/api/#tag/files + +Mistral's file objects already carry the OpenAI field names (id, bytes, created_at, +filename, purpose), so this config is URL routing, auth, and a purpose mapping: +Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes. +""" + +import time +from typing import Final, Literal + +import httpx +from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, ConfigDict + +from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj +from litellm.types.llms.openai import ( + CreateFileRequest, + FileContentRequest, + HttpxBinaryResponseContent, + OpenAICreateFileRequestOptionalParams, + OpenAIFileObject, + OpenAIFilesPurpose, +) +from litellm.types.utils import LlmProviders + +from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error + +MistralFilePurpose = Literal["fine-tune", "batch", "ocr"] + + +class MistralFile(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + bytes: int = 0 + created_at: int | None = None + filename: str = "" + purpose: MistralFilePurpose = "batch" + expires_at: int | None = None + + +class MistralFileList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[MistralFile, ...] = () + + +class MistralFileDeleted(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + deleted: bool = True + + +def _to_openai_file_object(file: MistralFile) -> OpenAIFileObject: + return OpenAIFileObject( + id=file.id, + bytes=file.bytes, + created_at=file.created_at if file.created_at is not None else int(time.time()), + filename=file.filename, + object="file", + purpose=_to_openai_purpose(file.purpose), + status="uploaded", + expires_at=file.expires_at, + ) + + +def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose: + match purpose: + case "fine-tune" | "batch": + return purpose + case "ocr": + return "user_data" + + +def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: + match purpose: + case "fine-tune" | "ocr": + return purpose + case _: + return "batch" + + +class MistralFilesConfig(BaseFilesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.MISTRAL + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + stream: bool | None = None, + ) -> str: + return f"{get_mistral_api_base(api_base)}/v1/files" + + def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str: + encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") + return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}" + + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + return mistral_error(error_message, status_code, headers) + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + return get_mistral_auth_headers(headers, api_key) + + def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: + return ["purpose"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + return optional_params + + def transform_create_file_request( + self, + model: str, + create_file_data: CreateFileRequest, + optional_params: dict, + litellm_params: dict, + ) -> dict: + file_data: Final = create_file_data.get("file") + if file_data is None: + raise ValueError("File data is required") + extracted: Final = extract_file_data(file_data) + filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" + content_type: Final = extracted.get("content_type") or "application/octet-stream" + return { + "file": (filename, extracted["content"], content_type), + "purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + } + + def transform_create_file_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) + + def transform_retrieve_file_request( + self, + file_id: str, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_id, litellm_params), {} + + def transform_retrieve_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) + + def transform_delete_file_request( + self, + file_id: str, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_id, litellm_params), {} + + def transform_delete_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> FileDeleted: + deleted: Final = MistralFileDeleted.model_validate(raw_response.json()) + return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") + + def transform_list_files_request( + self, + purpose: str | None, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {} + return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params + + def transform_list_files_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> list[OpenAIFileObject]: + return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data] + + def transform_file_content_request( + self, + file_content_request: FileContentRequest, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {} + + def transform_file_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 54ebdc85be9..5a934301edd 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -35262,51 +35262,66 @@ "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-1": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, + "ocr_cost_per_page_batches": 0.0005, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-2512": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, @@ -59822,31 +59837,40 @@ "mistral/mistral-ocr-3": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-3-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/voxtral-mini-latest": { diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..defd59f2be8 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -498,7 +498,7 @@ class CreateBatchRequest(TypedDict, total=False): """ completion_window: Literal["24h"] - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"] + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"] input_file_id: str metadata: dict[str, str] | None output_expires_after: FileExpiresAfter diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d62f00f3676..ea23e00d2bb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -320,8 +320,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models + ocr_cost_per_page_batches: ReadOnly[float | None] ocr_cost_per_credit: float | None # for OCR models priced by credit annotation_cost_per_page: float | None # for OCR models + annotation_cost_per_page_batches: ReadOnly[float | None] search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None @@ -3598,8 +3600,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_512k_tokens: float | None = None output_vector_size: int | None = None ocr_cost_per_page: float | None = None + ocr_cost_per_page_batches: float | None = None ocr_cost_per_credit: float | None = None annotation_cost_per_page: float | None = None + annotation_cost_per_page_batches: float | None = None regional_processing_uplift_multiplier_eu: float | None = None regional_processing_uplift_multiplier_us: float | None = None regional_endpoint_uplift_multiplier: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 8df28870544..3ca81605a95 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5963,8 +5963,10 @@ def _get_model_info_helper( tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None), + ocr_cost_per_page_batches=_model_info.get("ocr_cost_per_page_batches", None), ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None), annotation_cost_per_page=_model_info.get("annotation_cost_per_page", None), + annotation_cost_per_page_batches=_model_info.get("annotation_cost_per_page_batches", None), provider_specific_entry=_model_info.get("provider_specific_entry", None), uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), @@ -8909,6 +8911,10 @@ class ProviderConfigManager: from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig return AnthropicFilesConfig() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.files.transformation import MistralFilesConfig + + return MistralFilesConfig() return None @staticmethod @@ -8920,6 +8926,10 @@ class ProviderConfigManager: from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig return BedrockBatchesConfig() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.batches.transformation import MistralBatchesConfig + + return MistralBatchesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 54ebdc85be9..5a934301edd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -35262,51 +35262,66 @@ "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-1": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, + "ocr_cost_per_page_batches": 0.0005, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-2512": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, @@ -59822,31 +59837,40 @@ "mistral/mistral-ocr-3": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-3-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/voxtral-mini-latest": { diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 768ea332677..85f267c2db0 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1787,3 +1787,86 @@ class TestBatchCostIsFinal: @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) def test_other_terminal_statuses_are_final(self, status): assert bu.batch_cost_is_final(_retrieved_batch(status)) is True + + +# =========================================================================== # +# OCR batch output lines (Mistral /v1/ocr batches) are billed per page, not per token +# =========================================================================== # + + +def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): + usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} + if annotation_pages is not None: + usage_info["pages_processed_annotation"] = annotation_pages + return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info) + + +def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002}, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(3), _ocr_row(5), _failed_row(model="mistral-ocr-latest")], + custom_llm_provider="mistral", + model_name="mistral/mistral-ocr-latest", + ) + assert result.cost == pytest.approx(8 * 0.002) + assert result.prompt_cost == pytest.approx(8 * 0.002) + assert result.completion_cost == 0.0 + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert result.usage.total_tokens == 0 + assert result.models == ["mistral/mistral-ocr-latest"] + + +def test_ocr_rows_fall_back_to_sync_page_rate_without_batch_price(monkeypatch): + monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004}) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(2)], custom_llm_provider="mistral") + assert result.cost == pytest.approx(2 * 0.004) + + +def test_ocr_rows_bill_annotation_pages_separately(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: { + "ocr_cost_per_page_batches": 0.002, + "annotation_cost_per_page_batches": 0.0025, + }, + ) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral") + assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025) + + +def test_ocr_rows_use_deployment_model_info_pricing_over_cost_map(monkeypatch): + monkeypatch.setattr( + litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted") + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(10)], + custom_llm_provider="mistral", + model_info={"ocr_cost_per_page_batches": 0.001}, + ) + assert result.cost == pytest.approx(0.01) + + +def test_ocr_rows_without_pricing_bill_zero_but_count_as_successful(monkeypatch): + monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"mode": "ocr"}) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(3)], custom_llm_provider="mistral") + assert result.cost == 0.0 + assert (result.successful_requests, result.failed_requests) == (1, 0) + + +def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_success_row(model="mistral-small-latest", usage=_usage(10, 5))], + custom_llm_provider="mistral", + ) + assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2) + assert result.usage.total_tokens == 15 diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index b87f9489250..c5a33dd6508 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -778,3 +778,45 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] assert "_litellm_internal_model_credentials" not in litellm_params + + +# =========================================================================== # +# mistral - a provider-config provider, like bedrock, so it requires `model` +# =========================================================================== # + + +def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams): + with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg: + result = bm.create_batch( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-abc", + custom_llm_provider="mistral", + model="mistral/mistral-ocr-latest", + ) + + assert result is seams.base_http.create_batch.return_value + _assert_only(seams.base_http.create_batch, seams, "create_batch") + get_cfg.assert_called_once() + forwarded = seams.base_http.create_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" + assert forwarded["model"] == "mistral-ocr-latest" + assert forwarded["create_batch_data"]["endpoint"] == "/v1/ocr" + + +def test_create__mistral_without_model_raises_badrequest(seams): + with pytest.raises(litellm.exceptions.BadRequestError): + bm.create_batch(**CREATE_KW, custom_llm_provider="mistral") + + for m in _all_seam_methods(seams, "create_batch"): + m.assert_not_called() + + +def test_retrieve__mistral_routes_to_base_http_handler_with_mistral_config(seams): + result = bm.retrieve_batch(batch_id="job-1", custom_llm_provider="mistral", model="mistral/mistral-ocr-latest") + + assert result is seams.base_http.retrieve_batch.return_value + _assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch") + forwarded = seams.base_http.retrieve_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" + assert forwarded["batch_id"] == "job-1" diff --git a/tests/test_litellm/llms/mistral/batches/__init__.py b/tests/test_litellm/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py new file mode 100644 index 00000000000..03e9c351a30 --- /dev/null +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -0,0 +1,260 @@ +""" +Regression tests for ``MistralBatchesConfig``, the BaseBatchesConfig implementation +behind ``custom_llm_provider="mistral"`` on /v1/batches. + +Locks the request shape Mistral's ``POST /v1/batch/jobs`` accepts (input_files list, +model set on the job, endpoint passed through untouched so ``/v1/ocr`` batches work), +the Mistral -> OpenAI status mapping, request-count and file-id mapping, and auth. +Everything runs for real against canned httpx responses; only the API key env var is +set. +""" + +import json + +import httpx +import pytest + +from litellm.llms.mistral.batches.transformation import MistralBatchesConfig +from litellm.llms.mistral.common_utils import MistralError +from litellm.types.llms.openai import CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders + +STATUS_MAP = { + "QUEUED": "validating", + "RUNNING": "in_progress", + "SUCCESS": "completed", + "FAILED": "failed", + "TIMEOUT_EXCEEDED": "expired", + "CANCELLATION_REQUESTED": "cancelling", + "CANCELLED": "cancelled", +} + + +def _job(**overrides): + base = { + "id": "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b", + "object": "batch", + "input_files": ["c1a2b3d4-0000-4000-8000-000000000001"], + "endpoint": "/v1/ocr", + "model": "mistral-ocr-latest", + "status": "SUCCESS", + "created_at": 1_757_400_000, + "started_at": 1_757_400_010, + "completed_at": 1_757_400_500, + "total_requests": 3, + "completed_requests": 3, + "succeeded_requests": 2, + "failed_requests": 1, + "output_file": "out-0000-4000-8000-000000000002", + "error_file": "err-0000-4000-8000-000000000003", + "errors": [], + "metadata": {"job_type": "testing"}, + } + return {**base, **overrides} + + +def _response(payload: dict, status_code: int = 200) -> httpx.Response: + return httpx.Response( + status_code=status_code, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://api.mistral.ai/v1/batch/jobs/x"), + ) + + +@pytest.fixture +def config() -> MistralBatchesConfig: + return MistralBatchesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + return "sk-mistral-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.MISTRAL + + +# --------------------------------------------------------------------------- # +# create +# --------------------------------------------------------------------------- # + + +def test_create_request_maps_openai_fields_onto_mistral_job(config): + data = CreateBatchRequest( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-123", + metadata={"team": "docs"}, + ) + body = config.transform_create_batch_request( + model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={} + ) + assert body == { + "input_files": ["file-123"], + "endpoint": "/v1/ocr", + "model": "mistral-ocr-latest", + "metadata": {"team": "docs"}, + } + + +def test_create_request_omits_empty_metadata_and_forwards_extra_body(config): + data = CreateBatchRequest( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-123", + metadata=None, + extra_body={"timeout_hours": 48}, + ) + body = config.transform_create_batch_request( + model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={} + ) + assert "metadata" not in body + assert body["timeout_hours"] == 48 + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.mistral.ai/v1/batch/jobs"), + ("https://api.mistral.ai/v1", "https://api.mistral.ai/v1/batch/jobs"), + ("https://proxy.example.com/", "https://proxy.example.com/v1/batch/jobs"), + ], +) +def test_create_url(config, api_base, expected): + url = config.get_complete_batch_url( + api_base=api_base, api_key="k", model="m", optional_params={}, litellm_params={}, data={} + ) + assert url == expected + + +def test_validate_environment_uses_bearer_auth(config, api_key): + headers = config.validate_environment( + headers={"x-extra": "1"}, model="m", messages=[], optional_params={}, litellm_params={} + ) + assert headers == {"x-extra": "1", "Authorization": f"Bearer {api_key}"} + + +def test_validate_environment_explicit_key_wins(config, api_key): + headers = config.validate_environment( + headers={}, model="m", messages=[], optional_params={}, litellm_params={}, api_key="sk-explicit" + ) + assert headers["Authorization"] == "Bearer sk-explicit" + + +def test_validate_environment_without_key_raises(config, monkeypatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + with pytest.raises(ValueError, match="Missing Mistral API Key"): + config.validate_environment(headers={}, model="m", messages=[], optional_params={}, litellm_params={}) + + +def test_create_response_maps_job_onto_openai_batch(config): + batch = config.transform_create_batch_response( + model="mistral-ocr-latest", + raw_response=_response(_job(status="QUEUED", started_at=None, completed_at=None)), + logging_obj=None, + litellm_params={}, + ) + assert isinstance(batch, LiteLLMBatch) + assert batch.id == "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b" + assert batch.endpoint == "/v1/ocr" + assert batch.input_file_id == "c1a2b3d4-0000-4000-8000-000000000001" + assert batch.status == "validating" + assert batch.created_at == 1_757_400_000 + assert batch.in_progress_at is None + assert batch.completed_at is None + assert batch.metadata == {"job_type": "testing"} + + +# --------------------------------------------------------------------------- # +# retrieve +# --------------------------------------------------------------------------- # + + +def test_retrieve_request_is_presigned_get_with_auth(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"} + ) + assert req["method"] == "GET" + assert req["url"] == "https://api.mistral.ai/v1/batch/jobs/job%2Fwith%20slash" + assert req["headers"] == {"Authorization": f"Bearer {api_key}"} + + +def test_retrieve_request_prefers_litellm_params_api_key(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="job-1", optional_params={}, litellm_params={"api_key": "sk-from-deployment"} + ) + assert req["headers"]["Authorization"] == "Bearer sk-from-deployment" + + +@pytest.mark.parametrize("mistral_status,openai_status", sorted(STATUS_MAP.items())) +def test_retrieve_response_status_mapping(config, mistral_status, openai_status): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={} + ) + assert batch.status == openai_status + + +@pytest.mark.parametrize( + "mistral_status,populated_field", + [ + ("SUCCESS", "completed_at"), + ("FAILED", "failed_at"), + ("TIMEOUT_EXCEEDED", "expired_at"), + ("CANCELLED", "cancelled_at"), + ], +) +def test_retrieve_response_terminal_timestamp_lands_on_matching_field(config, mistral_status, populated_field): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={} + ) + terminal_fields = {"completed_at", "failed_at", "expired_at", "cancelled_at"} + assert getattr(batch, populated_field) == 1_757_400_500 + for other in terminal_fields - {populated_field}: + assert getattr(batch, other) is None + assert batch.in_progress_at == 1_757_400_010 + + +def test_retrieve_response_maps_counts_and_files(config): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job()), logging_obj=None, litellm_params={} + ) + assert batch.request_counts.total == 3 + assert batch.request_counts.completed == 2 + assert batch.request_counts.failed == 1 + assert batch.output_file_id == "out-0000-4000-8000-000000000002" + assert batch.error_file_id == "err-0000-4000-8000-000000000003" + assert batch.errors is None + + +def test_retrieve_response_surfaces_job_errors(config): + batch = config.transform_retrieve_batch_response( + model=None, + raw_response=_response( + _job(status="FAILED", errors=[{"message": "invalid document", "count": 2}, {"message": "timeout"}]) + ), + logging_obj=None, + litellm_params={}, + ) + assert [e.message for e in batch.errors.data] == ["invalid document (x2)", "timeout"] + + +def test_retrieve_response_without_files_or_input(config): + batch = config.transform_retrieve_batch_response( + model=None, + raw_response=_response(_job(input_files=[], output_file=None, error_file=None, metadata=None)), + logging_obj=None, + litellm_params={}, + ) + assert batch.input_file_id == "" + assert batch.output_file_id is None + assert batch.error_file_id is None + assert batch.metadata is None + + +def test_get_error_class(config): + err = config.get_error_class("nope", 401, {"x-request-id": "r1"}) + assert isinstance(err, MistralError) + assert err.status_code == 401 + assert err.message == "nope" diff --git a/tests/test_litellm/llms/mistral/files/__init__.py b/tests/test_litellm/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py new file mode 100644 index 00000000000..d6ad8a34b35 --- /dev/null +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -0,0 +1,189 @@ +""" +Regression tests for ``MistralFilesConfig``, the BaseFilesConfig implementation behind +``custom_llm_provider="mistral"`` on /v1/files. + +Locks the URL routing for each file operation, the multipart upload shape Mistral's +``POST /v1/files`` accepts (purpose restricted to fine-tune/batch/ocr), and the +Mistral -> OpenAI file object mapping. Runs against canned httpx responses. +""" + +import json + +import httpx +import pytest +from openai.types.file_deleted import FileDeleted + +from litellm.llms.mistral.files.transformation import MistralFilesConfig +from litellm.types.llms.openai import CreateFileRequest, FileContentRequest, OpenAIFileObject +from litellm.types.utils import LlmProviders + +FILE_ID = "497f6eca-6276-4993-bfeb-53cbbbba6f09" + + +def _file(**overrides): + base = { + "id": FILE_ID, + "object": "file", + "bytes": 13000, + "created_at": 1_716_963_433, + "filename": "batch_input.jsonl", + "purpose": "batch", + "sample_type": "batch_request", + "num_lines": 3, + "source": "upload", + } + return {**base, **overrides} + + +def _response(payload) -> httpx.Response: + return httpx.Response( + status_code=200, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://api.mistral.ai/v1/files"), + ) + + +@pytest.fixture +def config() -> MistralFilesConfig: + return MistralFilesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + return "sk-mistral-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.MISTRAL + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.mistral.ai/v1/files"), + ("https://api.mistral.ai/v1/", "https://api.mistral.ai/v1/files"), + ("https://proxy.example.com", "https://proxy.example.com/v1/files"), + ], +) +def test_upload_url(config, api_base, expected): + url = config.get_complete_url(api_base=api_base, api_key="k", model="", optional_params={}, litellm_params={}) + assert url == expected + + +def test_validate_environment_uses_bearer_auth(config, api_key): + headers = config.validate_environment(headers={}, model="", messages=[], optional_params={}, litellm_params={}) + assert headers == {"Authorization": f"Bearer {api_key}"} + + +def test_upload_request_is_multipart_with_batch_purpose(config): + body = config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"), + optional_params={}, + litellm_params={}, + ) + assert body == { + "file": ("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), + "purpose": (None, "batch"), + } + + +@pytest.mark.parametrize( + "openai_purpose,mistral_purpose", + [("batch", "batch"), ("fine-tune", "fine-tune"), ("ocr", "ocr"), ("assistants", "batch"), ("user_data", "batch")], +) +def test_upload_request_maps_purpose_onto_mistral_enum(config, openai_purpose, mistral_purpose): + body = config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=openai_purpose), + optional_params={}, + litellm_params={}, + ) + assert body["purpose"] == (None, mistral_purpose) + + +def test_upload_request_requires_file(config): + with pytest.raises(ValueError, match="File data is required"): + config.transform_create_file_request( + model="", create_file_data=CreateFileRequest(purpose="batch"), optional_params={}, litellm_params={} + ) + + +def test_upload_response_maps_onto_openai_file_object(config): + obj = config.transform_create_file_response( + model=None, raw_response=_response(_file()), logging_obj=None, litellm_params={} + ) + assert obj == OpenAIFileObject( + id=FILE_ID, + bytes=13000, + created_at=1_716_963_433, + filename="batch_input.jsonl", + object="file", + purpose="batch", + status="uploaded", + ) + + +def test_file_response_with_ocr_purpose_maps_onto_user_data(config): + obj = config.transform_retrieve_file_response( + raw_response=_response(_file(purpose="ocr", expires_at=1_800_000_000)), logging_obj=None, litellm_params={} + ) + assert obj.purpose == "user_data" + assert obj.expires_at == 1_800_000_000 + + +@pytest.mark.parametrize( + "method,suffix", + [ + ("transform_retrieve_file_request", ""), + ("transform_delete_file_request", ""), + ], +) +def test_single_file_urls_encode_id_and_honor_api_base(config, method, suffix): + url, params = getattr(config, method)( + file_id="id/with slash", optional_params={}, litellm_params={"api_base": "https://mistral.internal/v1"} + ) + assert url == f"https://mistral.internal/v1/files/id%2Fwith%20slash{suffix}" + assert params == {} + + +def test_file_content_url(config): + url, params = config.transform_file_content_request( + file_content_request=FileContentRequest(file_id=FILE_ID), optional_params={}, litellm_params={} + ) + assert url == f"https://api.mistral.ai/v1/files/{FILE_ID}/content" + assert params == {} + + +def test_file_content_response_is_binary_passthrough(config): + raw = httpx.Response( + 200, content=b'{"custom_id":"0","response":{"status_code":200}}\n', request=httpx.Request("GET", "https://x") + ) + out = config.transform_file_content_response(raw_response=raw, logging_obj=None, litellm_params={}) + assert out.content == b'{"custom_id":"0","response":{"status_code":200}}\n' + + +def test_delete_response(config): + out = config.transform_delete_file_response( + raw_response=_response({"id": FILE_ID, "object": "file", "deleted": True}), logging_obj=None, litellm_params={} + ) + assert out == FileDeleted(id=FILE_ID, deleted=True, object="file") + + +def test_list_request_filters_by_mapped_purpose(config): + url, params = config.transform_list_files_request(purpose="batch", optional_params={}, litellm_params={}) + assert url == "https://api.mistral.ai/v1/files" + assert params == {"purpose": "batch"} + _, no_params = config.transform_list_files_request(purpose=None, optional_params={}, litellm_params={}) + assert no_params == {} + + +def test_list_response(config): + out = config.transform_list_files_response( + raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}), + logging_obj=None, + litellm_params={}, + ) + assert [f.id for f in out] == [FILE_ID, "second"] + assert out[1].filename == "b.jsonl" diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 40e54f71eeb..9fe6f38003f 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -72,9 +72,11 @@ def test_ocr3_pricing_entry(cost_map_path: Path) -> None: assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}" assert info["litellm_provider"] == "mistral" assert info["mode"] == "ocr" - assert info["supported_endpoints"] == ["/v1/ocr"] + assert info["supported_endpoints"] == ["/v1/ocr", "/v1/batch"] assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE + assert info["ocr_cost_per_page_batches"] == OCR3_COST_PER_PAGE / 2 + assert info["annotation_cost_per_page_batches"] == OCR3_ANNOTATION_COST_PER_PAGE / 2 def test_ocr3_model_info_price(local_model_cost_map) -> None: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index e42608c9904..61739846706 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -992,7 +992,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, "input_dbu_cost_per_token": {"type": "number"}, "annotation_cost_per_page": {"type": "number"}, + "annotation_cost_per_page_batches": {"type": "number"}, "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_page_batches": {"type": "number"}, "ocr_cost_per_credit": {"type": "number"}, "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, From 2e5f5a95c813b940dc7f654c10b2ce2036c6fb2a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:14:16 -0400 Subject: [PATCH 029/393] fix(proxy): retrieve model-routed file ids from the deployment's provider GET /v1/files/{id} for an id encoded with a non-OpenAI deployment forwarded the deployment credentials but let custom_llm_provider default to openai, so a Mistral file was fetched from api.openai.com with the Mistral key and 401'd. Delete and content already passed the provider through; retrieve now does too. --- .../openai_files_endpoints/files_endpoints.py | 5 +- .../test_files_endpoint.py | 59 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c315d30b8f3..49a01495d65 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1139,7 +1139,10 @@ async def get_file( include_internal_credentials=True, ) - response = await litellm.afile_retrieve(**data) + response = await litellm.afile_retrieve( + custom_llm_provider=credentials["custom_llm_provider"], + **data, + ) # Keep the encoded ID in response if it was originally encoded if original_file_id and response and hasattr(response, "id") and response.id: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5f1e7e1fe0c..5faae166fca 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4819,3 +4819,62 @@ def test_get_file_content_keeps_the_status_of_a_rejection_raised_inside_the_rout error = response.json()["error"] assert error["message"].startswith("Storage backend error") assert (error["type"], error["param"], error["code"]) == ("invalid_request_error", "file_id", "400") + + +def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFixture, monkeypatch): + """ + Regression: a file id encoded with a non-OpenAI deployment (here Mistral) must be + retrieved from that deployment's provider. Before the fix the retrieve path only + forwarded the credentials and let ``custom_llm_provider`` default to openai, so a + Mistral file id was sent to api.openai.com with the Mistral key and 401'd. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"}, + "model_info": {"id": "mistral-ocr-id"}, + } + ] + ) + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_retrieve(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", + object="file", + bytes=2, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["custom_llm_provider"] == "mistral" + assert captured_kwargs["api_key"] == "mistral-key" + assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df" + assert response.json()["id"] == encoded_id From c246f75e3ec93192f95e0cb2fc3f50485bf13ebc Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:37:01 -0400 Subject: [PATCH 030/393] refactor(mistral): satisfy type-discipline and basedpyright gates for files/batches configs --- litellm/cost_calculator.py | 23 ++-- litellm/files/main.py | 4 +- .../llms/mistral/batches/transformation.py | 114 ++++++++++------ litellm/llms/mistral/common_utils.py | 13 +- litellm/llms/mistral/files/transformation.py | 129 +++++++++++------- .../openai_files_endpoints/files_endpoints.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 46 +++++-- tests/test_litellm/batches/test_main.py | 76 +++-------- .../test_mistral_batches_transformation.py | 16 ++- .../test_mistral_files_transformation.py | 8 +- 10 files changed, 251 insertions(+), 180 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0c0c6b05df8..7c95941d77d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2006,13 +2006,11 @@ def ocr_batch_cost( ``ocr_cost``. """ has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) - if has_ocr_pricing: - resolved_info: ModelInfo | None = model_info - else: - try: - resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - resolved_info = None + resolved_info: Final = ( + model_info + if has_ocr_pricing + else _lookup_model_info_or_none(model=model, custom_llm_provider=custom_llm_provider) + ) if resolved_info is None: verbose_logger.warning( "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", @@ -2022,9 +2020,7 @@ def ocr_batch_cost( return 0.0, 0.0 page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") - annotation_rate: Final = _first_price( - resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page" - ) + annotation_rate: Final = _first_price(resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page") pages_processed: Final = usage_info.pages_processed or 0 annotation_pages: Final = usage_info.pages_processed_annotation or 0 if page_rate is None and pages_processed > 0: @@ -2039,6 +2035,13 @@ def ocr_batch_cost( return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0 +def _lookup_model_info_or_none(model: str, custom_llm_provider: str | None) -> ModelInfo | None: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + return None + + def _first_price(model_info: ModelInfo, *keys: str) -> float | None: return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) diff --git a/litellm/files/main.py b/litellm/files/main.py index 3d90bf4f299..5cf0f8e576a 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -32,9 +32,7 @@ FileCreateProvider = Literal[ FileRetrieveProvider = Literal[ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" ] -FileDeleteProvider = Literal[ - "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral" -] +FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"] FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] import litellm from litellm import get_secret_str diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py index 399319e590b..ef9ee5ff503 100644 --- a/litellm/llms/mistral/batches/transformation.py +++ b/litellm/llms/mistral/batches/transformation.py @@ -7,14 +7,16 @@ Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_ so the shared batch cost accounting reads them without a provider branch. """ +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, Literal +from typing import Final, Literal, TypeAlias import httpx from openai.types.batch import BatchRequestCounts from openai.types.batch import Errors as BatchErrors from openai.types.batch_error import BatchError from pydantic import BaseModel, ConfigDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig @@ -24,13 +26,14 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error -MistralBatchStatus = Literal[ +MistralBatchStatus: TypeAlias = Literal[ "QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED" ] -OpenAIBatchStatus = Literal[ +OpenAIBatchStatus: TypeAlias = Literal[ "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" ] +_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) # mutable-ok: frozen at module scope _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType( { "QUEUED": "validating", @@ -44,6 +47,23 @@ _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = Ma ) +class MistralCreateBatchJobRequest(TypedDict): + """Body of ``POST /v1/batch/jobs``.""" + + input_files: ReadOnly[tuple[str, ...]] + endpoint: ReadOnly[str] + model: ReadOnly[str] + metadata: NotRequired[ReadOnly[Mapping[str, str]]] + + +class MistralPresignedRequest(TypedDict): + """A fully-formed request the shared HTTP handler sends as-is (its ``method`` branch).""" + + method: ReadOnly[Literal["GET"]] + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + + class MistralBatchError(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -69,7 +89,18 @@ class MistralBatchJob(BaseModel): output_file: str | None = None error_file: str | None = None errors: tuple[MistralBatchError, ...] = () - metadata: dict[str, str] | None = None + metadata: dict[str, str] | None = None # mutable-ok: LiteLLMBatch.metadata is typed as dict + + +def _to_batch_errors(errors: Sequence[MistralBatchError]) -> BatchErrors | None: + if not errors: + return None + return BatchErrors( + object="list", + data=[ # mutable-ok: openai Batch.Errors.data is typed as list + BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in errors + ], + ) def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: @@ -90,14 +121,7 @@ def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: cancelled_at=terminal_at if status == "cancelled" else None, output_file_id=job.output_file, error_file_id=job.error_file, - errors=( - BatchErrors( - object="list", - data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors], - ) - if job.errors - else None - ), + errors=_to_batch_errors(job.errors), request_counts=BatchRequestCounts( total=job.total_requests, completed=job.succeeded_requests, @@ -114,14 +138,14 @@ class MistralBatchesConfig(BaseBatchesConfig): def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: # mutable-ok: BaseBatchesConfig signature return get_mistral_auth_headers(headers, api_key) def get_complete_batch_url( @@ -129,8 +153,8 @@ class MistralBatchesConfig(BaseBatchesConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], data: CreateBatchRequest, ) -> str: return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" @@ -139,48 +163,58 @@ class MistralBatchesConfig(BaseBatchesConfig): self, model: str, create_batch_data: CreateBatchRequest, - optional_params: dict, - litellm_params: dict, - ) -> dict[str, object]: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature + input_file_id: Final = create_batch_data.get("input_file_id") + endpoint: Final = create_batch_data.get("endpoint") + if input_file_id is None or endpoint is None: + raise ValueError("input_file_id and endpoint are required to create a Mistral batch job") metadata: Final = create_batch_data.get("metadata") - return { - "input_files": [create_batch_data["input_file_id"]], - "endpoint": create_batch_data["endpoint"], - "model": model, - **({"metadata": metadata} if metadata else {}), - **(create_batch_data.get("extra_body") or {}), - } + body: Final = ( + MistralCreateBatchJobRequest( + input_files=(input_file_id,), endpoint=endpoint, model=model, metadata=metadata + ) + if metadata + else MistralCreateBatchJobRequest(input_files=(input_file_id,), endpoint=endpoint, model=model) + ) + return dict(body) # mutable-ok: BaseBatchesConfig signature def transform_create_batch_response( self, model: str | None, raw_response: httpx.Response, logging_obj: object, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> LiteLLMBatch: return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) def transform_retrieve_batch_request( self, batch_id: str, - optional_params: dict, - litellm_params: dict, - ) -> dict[str, object]: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - return { - "method": "GET", - "url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}", - "headers": get_mistral_auth_headers({}, litellm_params.get("api_key")), - } + api_base: Final = litellm_params.get("api_base") + api_key: Final = litellm_params.get("api_key") + request: Final = MistralPresignedRequest( + method="GET", + url=f"{get_mistral_api_base(api_base if isinstance(api_base, str) else None)}/v1/batch/jobs/{encoded_batch_id}", + headers=get_mistral_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None), + ) + return dict(request) # mutable-ok: BaseBatchesConfig signature def transform_retrieve_batch_response( self, model: str | None, raw_response: httpx.Response, logging_obj: object, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> LiteLLMBatch: return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: return mistral_error(error_message, status_code, headers) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py index 9ea501c860d..2f14328afdf 100644 --- a/litellm/llms/mistral/common_utils.py +++ b/litellm/llms/mistral/common_utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -19,18 +20,22 @@ def get_mistral_api_base(api_base: str | None) -> str: return resolved.removesuffix("/v1") -def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict: +def get_mistral_auth_headers( + headers: Mapping[str, str], api_key: str | None +) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR) if resolved_key is None: raise ValueError( "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" ) - return {**headers, "Authorization": f"Bearer {resolved_key}"} + return dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict -def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError: +def mistral_error(error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers) -> MistralError: return MistralError( status_code=status_code, message=error_message, - headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), + headers=headers + if isinstance(headers, httpx.Headers) + else httpx.Headers(dict(headers)), # mutable-ok: httpx.Headers takes a dict ) diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 071b6f58569..6d58311813c 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -7,11 +7,13 @@ Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes. """ import time -from typing import Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, TypeAlias import httpx from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -29,7 +31,16 @@ from litellm.types.utils import LlmProviders from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error -MistralFilePurpose = Literal["fine-tune", "batch", "ocr"] +MistralFilePurpose: TypeAlias = Literal["fine-tune", "batch", "ocr"] + +_NO_QUERY_PARAMS: Final[dict[str, str]] = {} # mutable-ok: BaseFilesConfig request transforms return tuple[str, dict] + + +class MistralMultipartUpload(TypedDict): + """``files=`` payload for ``POST /v1/files``: each value is an httpx multipart tuple.""" + + file: ReadOnly[tuple[str, object, str]] + purpose: ReadOnly[tuple[None, MistralFilePurpose]] class MistralFile(BaseModel): @@ -85,6 +96,11 @@ def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: return "batch" +def _api_base_from(litellm_params: Mapping[str, object]) -> str: + api_base: Final = litellm_params.get("api_base") + return get_mistral_api_base(api_base if isinstance(api_base, str) else None) + + class MistralFilesConfig(BaseFilesConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -95,99 +111,103 @@ class MistralFilesConfig(BaseFilesConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: return f"{get_mistral_api_base(api_base)}/v1/files" - def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str: + def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}" + return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: return mistral_error(error_message, status_code, headers) def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list, - optional_params: dict, - litellm_params: dict, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: # mutable-ok: BaseFilesConfig signature return get_mistral_auth_headers(headers, api_key) - def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: - return ["purpose"] + def get_supported_openai_params( + self, model: str + ) -> list[OpenAICreateFileRequestOptionalParams]: # mutable-ok: BaseFilesConfig signature + return ["purpose"] # mutable-ok: BaseFilesConfig signature def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: Mapping[str, object], + optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: # mutable-ok: BaseConfig signature return optional_params def transform_create_file_request( self, model: str, create_file_data: CreateFileRequest, - optional_params: dict, - litellm_params: dict, - ) -> dict: - file_data: Final = create_file_data.get("file") - if file_data is None: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseFilesConfig signature + if "file" not in create_file_data: raise ValueError("File data is required") - extracted: Final = extract_file_data(file_data) + extracted: Final = extract_file_data(create_file_data["file"]) filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" content_type: Final = extracted.get("content_type") or "application/octet-stream" - return { - "file": (filename, extracted["content"], content_type), - "purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), - } + upload: Final = MistralMultipartUpload( + file=(filename, extracted["content"], content_type), + purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + ) + return dict(upload) # mutable-ok: BaseFilesConfig signature def transform_create_file_response( self, model: str | None, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) def transform_retrieve_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_id, litellm_params), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS def transform_retrieve_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) def transform_delete_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_id, litellm_params), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS def transform_delete_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> FileDeleted: deleted: Final = MistralFileDeleted.model_validate(raw_response.json()) return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") @@ -195,32 +215,39 @@ class MistralFilesConfig(BaseFilesConfig): def transform_list_files_request( self, purpose: str | None, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {} - return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + url: Final = f"{_api_base_from(litellm_params)}/v1/files" + if not purpose: + return url, _NO_QUERY_PARAMS + return url, {"purpose": _to_mistral_purpose(purpose)} # mutable-ok: BaseFilesConfig signature returns dict def transform_list_files_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - ) -> list[OpenAIFileObject]: - return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data] + litellm_params: Mapping[str, object], + ) -> list[OpenAIFileObject]: # mutable-ok: BaseFilesConfig signature + return [ # mutable-ok: BaseFilesConfig signature + _to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data + ] def transform_file_content_request( self, file_content_request: FileContentRequest, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + file_id: Final = file_content_request.get("file_id") + if file_id is None: + raise ValueError("file_id is required to download file content") + return self._file_url(file_id, litellm_params, suffix="/content"), _NO_QUERY_PARAMS def transform_file_content_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> HttpxBinaryResponseContent: return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 49a01495d65..b3bb1fa9a01 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1130,7 +1130,7 @@ async def get_file( check_file_id_encoding=True, ) - if should_route: + if should_route and credentials is not None: # Use model-based routing with credentials from config prepare_data_with_credentials( data=data, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 85f267c2db0..56cd3298db6 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -645,9 +645,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): lambda content, model: pytest.fail("raw vertex path should not run"), ) - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[], custom_llm_provider="vertex_ai" - ) + result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") assert result.cost == 0.0 assert result.usage.total_tokens == 0 assert result.models == [] @@ -1250,6 +1248,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): result set - zero cost, zero usage, no models - instead of letting the file fetch raise "Output file id is None" on every aretrieve_batch logging poll. """ + # The output-file fetch must not even be attempted when there is no output file. async def _must_not_fetch(*args, **kwargs): pytest.fail("_fetch_batch_output_file_content should not be called") @@ -1376,7 +1375,10 @@ def test_anthropic_response_body_is_result_message(): def test_anthropic_usage_conversion_includes_cache_tokens(): - body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)} + body = { + "model": "claude-sonnet-4-5-20250929", + "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000), + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic") assert usage.prompt_tokens == 11000 assert usage.completion_tokens == 200 @@ -1391,7 +1393,9 @@ def test_bedrock_model_output_line_success_check(): "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True - assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + assert ( + bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + ) def test_bedrock_cost_uses_deployment_model_name(): @@ -1445,7 +1449,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): rows = [ { "custom_id": "req-1", - "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, + "response": { + "status_code": 200, + "body": { + "model": "gpt-5.2", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, } ] result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") @@ -1487,7 +1497,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"), ) - result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic") + result = bu._aggregate_batch_cost_usage_models( + entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic" + ) assert result.cost == pytest.approx(0.3) assert seen[0]["model"] == "claude-sonnet-4-5-20250929" @@ -1522,7 +1534,11 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): ) assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 11000, + 200, + 11200, + ) assert result.models == ["claude-sonnet-4-5"] @@ -1689,7 +1705,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> def test_bedrock_converse_shaped_batch_usage_is_parsed(): - body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}} + body = { + "model": "us.amazon.nova-lite-v1:0", + "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}, + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742) @@ -1739,6 +1758,7 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): # batch_cost_is_final # --------------------------------------------------------------------------- # + def _retrieved_batch( status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None ) -> LiteLLMBatch: @@ -1798,7 +1818,9 @@ def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest") usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: usage_info["pages_processed_annotation"] = annotation_pages - return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info) + return _success_row( + model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info + ) def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch): @@ -1835,7 +1857,9 @@ def test_ocr_rows_bill_annotation_pages_separately(monkeypatch): "annotation_cost_per_page_batches": 0.0025, }, ) - result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral") + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral" + ) assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025) diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c5a33dd6508..26dc4083b0b 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -66,9 +66,7 @@ def seams(): stack.enter_context(patch.object(bm, "openai_batches_instance", openai_i)) stack.enter_context(patch.object(bm, "azure_batches_instance", azure_i)) stack.enter_context(patch.object(bm, "vertex_ai_batches_instance", vertex_i)) - stack.enter_context( - patch.object(bm, "anthropic_batches_instance", anthropic_i) - ) + stack.enter_context(patch.object(bm, "anthropic_batches_instance", anthropic_i)) stack.enter_context(patch.object(bm, "base_llm_http_handler", base_http)) stack.enter_context(patch.object(bm, "BedrockBatchesHandler", bedrock_arn)) yield Seams( @@ -174,9 +172,7 @@ def test_create__provider_config_routes_to_base_http_handler(seams): "get_provider_batches_config", return_value=MagicMock(name="provider_config"), ): - result = bm.create_batch( - **CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model" - ) + result = bm.create_batch(**CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model") assert result is seams.base_http.create_batch.return_value _assert_only(seams.base_http.create_batch, seams, "create_batch") @@ -281,9 +277,7 @@ def test_retrieve__bedrock_model_invocation_job_arn(seams): result = bm.retrieve_batch(batch_id=arn, custom_llm_provider="bedrock") seams.bedrock_arn._handle_model_invocation_job_status.assert_called_once() - assert ( - result is seams.bedrock_arn._handle_model_invocation_job_status.return_value - ) + assert result is seams.bedrock_arn._handle_model_invocation_job_status.return_value seams.bedrock_arn._handle_async_invoke_status.assert_not_called() @@ -385,9 +379,7 @@ def test_cancel__unsupported_provider_raises_badrequest(seams): def test_cancel__async_flag_propagates_is_async(seams): - bm.cancel_batch( - batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True - ) + bm.cancel_batch(batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True) assert seams.openai.cancel_batch.call_args.kwargs["_is_async"] is True @@ -415,9 +407,7 @@ async def test_acreate_batch_delegates_to_create_batch(): @pytest.mark.asyncio async def test_aretrieve_batch_delegates_to_retrieve_batch(): with patch.object(bm, "retrieve_batch", MagicMock(return_value="SENTINEL")) as m: - result = await bm.aretrieve_batch( - batch_id="batch-1", custom_llm_provider="azure" - ) + result = await bm.aretrieve_batch(batch_id="batch-1", custom_llm_provider="azure") assert result == "SENTINEL" assert m.call_count == 1 @@ -429,9 +419,7 @@ async def test_aretrieve_batch_delegates_to_retrieve_batch(): @pytest.mark.asyncio async def test_alist_batches_delegates_to_list_batches(): with patch.object(bm, "list_batches", MagicMock(return_value="SENTINEL")) as m: - result = await bm.alist_batches( - after="cur", limit=3, custom_llm_provider="vertex_ai" - ) + result = await bm.alist_batches(after="cur", limit=3, custom_llm_provider="vertex_ai") assert result == "SENTINEL" assert m.call_count == 1 @@ -444,9 +432,7 @@ async def test_alist_batches_delegates_to_list_batches(): @pytest.mark.asyncio async def test_acancel_batch_delegates_to_cancel_batch(): with patch.object(bm, "cancel_batch", MagicMock(return_value="SENTINEL")) as m: - result = await bm.acancel_batch( - batch_id="batch-1", custom_llm_provider="openai" - ) + result = await bm.acancel_batch(batch_id="batch-1", custom_llm_provider="openai") assert result == "SENTINEL" assert m.call_count == 1 @@ -499,9 +485,7 @@ def _sent(mock_method, *keys): def test_create__openai_credentials_passthrough(seams): bm.create_batch(**CREATE_KW, custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries" - ) == { + assert _sent(seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -512,9 +496,7 @@ def test_create__openai_credentials_passthrough(seams): def test_create__azure_credentials_passthrough(seams): bm.create_batch(**CREATE_KW, custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.create_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.create_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -564,9 +546,7 @@ def test_create__provider_config_credentials_passthrough(seams): def test_retrieve__openai_credentials_passthrough(seams): bm.retrieve_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.retrieve_batch, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.retrieve_batch, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -576,9 +556,7 @@ def test_retrieve__openai_credentials_passthrough(seams): def test_retrieve__azure_credentials_passthrough(seams): bm.retrieve_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.retrieve_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.retrieve_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -640,9 +618,7 @@ def test_retrieve__provider_config_credentials_passthrough(seams): def test_list__openai_credentials_passthrough(seams): bm.list_batches(custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.list_batches, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.list_batches, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -652,9 +628,7 @@ def test_list__openai_credentials_passthrough(seams): def test_list__azure_credentials_passthrough(seams): bm.list_batches(custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.list_batches, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.list_batches, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -682,9 +656,7 @@ def test_list__vertex_credentials_passthrough(seams): def test_cancel__openai_credentials_passthrough(seams): bm.cancel_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.cancel_batch, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.cancel_batch, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -694,9 +666,7 @@ def test_cancel__openai_credentials_passthrough(seams): def test_cancel__azure_credentials_passthrough(seams): bm.cancel_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.cancel_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.cancel_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -786,18 +756,16 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams): - with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg: - result = bm.create_batch( - completion_window="24h", - endpoint="/v1/ocr", - input_file_id="file-abc", - custom_llm_provider="mistral", - model="mistral/mistral-ocr-latest", - ) + result = bm.create_batch( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-abc", + custom_llm_provider="mistral", + model="mistral/mistral-ocr-latest", + ) assert result is seams.base_http.create_batch.return_value _assert_only(seams.base_http.create_batch, seams, "create_batch") - get_cfg.assert_called_once() forwarded = seams.base_http.create_batch.call_args.kwargs assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" assert forwarded["model"] == "mistral-ocr-latest" diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py index 03e9c351a30..03cfeedece2 100644 --- a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -92,26 +92,34 @@ def test_create_request_maps_openai_fields_onto_mistral_job(config): model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={} ) assert body == { - "input_files": ["file-123"], + "input_files": ("file-123",), "endpoint": "/v1/ocr", "model": "mistral-ocr-latest", "metadata": {"team": "docs"}, } -def test_create_request_omits_empty_metadata_and_forwards_extra_body(config): +def test_create_request_omits_empty_metadata(config): data = CreateBatchRequest( completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-123", metadata=None, - extra_body={"timeout_hours": 48}, ) body = config.transform_create_batch_request( model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={} ) assert "metadata" not in body - assert body["timeout_hours"] == 48 + + +def test_create_request_requires_input_file_and_endpoint(config): + with pytest.raises(ValueError, match="input_file_id and endpoint are required"): + config.transform_create_batch_request( + model="m", + create_batch_data=CreateBatchRequest(completion_window="24h"), + optional_params={}, + litellm_params={}, + ) @pytest.mark.parametrize( diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index d6ad8a34b35..f62645be7ee 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -79,7 +79,9 @@ def test_validate_environment_uses_bearer_auth(config, api_key): def test_upload_request_is_multipart_with_batch_purpose(config): body = config.transform_create_file_request( model="", - create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"), + create_file_data=CreateFileRequest( + file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch" + ), optional_params={}, litellm_params={}, ) @@ -181,7 +183,9 @@ def test_list_request_filters_by_mapped_purpose(config): def test_list_response(config): out = config.transform_list_files_response( - raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}), + raw_response=_response( + {"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2} + ), logging_obj=None, litellm_params={}, ) From 91e7df3d8fe83c37268b61080a37e48ef0e63e3f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 13:35:35 -0400 Subject: [PATCH 031/393] chore: regenerate cost-map schema and UI API types, drop test banner comments --- model_prices_and_context_window.schema.json | 8 ++++++++ tests/test_litellm/batches/test_batch_utils.py | 5 ----- .../batches/test_mistral_batches_transformation.py | 10 ---------- .../llms/mistral/ocr/test_mistral_ocr_cost.py | 1 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++++++++ 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 47a1934a703..09385db728b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -53,6 +53,10 @@ "type": "number", "minimum": 0 }, + "annotation_cost_per_page_batches": { + "type": "number", + "minimum": 0 + }, "audio_transcription_config": { "type": "string" }, @@ -432,6 +436,10 @@ "type": "number", "minimum": 0 }, + "ocr_cost_per_page_batches": { + "type": "number", + "minimum": 0 + }, "output_cost_per_audio_token": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 56cd3298db6..0d66c0eb5ec 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1809,11 +1809,6 @@ class TestBatchCostIsFinal: assert bu.batch_cost_is_final(_retrieved_batch(status)) is True -# =========================================================================== # -# OCR batch output lines (Mistral /v1/ocr batches) are billed per page, not per token -# =========================================================================== # - - def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py index 03cfeedece2..4073879e3b8 100644 --- a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -76,11 +76,6 @@ def test_custom_llm_provider(config): assert config.custom_llm_provider == LlmProviders.MISTRAL -# --------------------------------------------------------------------------- # -# create -# --------------------------------------------------------------------------- # - - def test_create_request_maps_openai_fields_onto_mistral_job(config): data = CreateBatchRequest( completion_window="24h", @@ -175,11 +170,6 @@ def test_create_response_maps_job_onto_openai_batch(config): assert batch.metadata == {"job_type": "testing"} -# --------------------------------------------------------------------------- # -# retrieve -# --------------------------------------------------------------------------- # - - def test_retrieve_request_is_presigned_get_with_auth(config, api_key): req = config.transform_retrieve_batch_request( batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"} diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 9fe6f38003f..d72e866949f 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -63,7 +63,6 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) - @pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) def test_ocr3_pricing_entry(cost_map_path: Path) -> None: with open(cost_map_path) as f: diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0f0c1fc9af4..5d1ed79098e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29493,6 +29493,8 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Annotation Cost Per Page Batches */ + annotation_cost_per_page_batches?: number | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -29704,6 +29706,8 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Ocr Cost Per Page Batches */ + ocr_cost_per_page_batches?: number | null; /** Organization */ organization?: string | null; /** Otpm */ @@ -39684,6 +39688,8 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Annotation Cost Per Page Batches */ + annotation_cost_per_page_batches?: number | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -39895,6 +39901,8 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Ocr Cost Per Page Batches */ + ocr_cost_per_page_batches?: number | null; /** Organization */ organization?: string | null; /** Otpm */ From edd5727f3c9077f449eec37367ace43945e649fa Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 17:28:27 -0400 Subject: [PATCH 032/393] fix(proxy): enforce key/team/org/project model grants on model-routed file and batch credentials Files and batches routes take their model from a header, query param or a model-encoded resource id, which the auth layer never sees, so any key could name any deployment and act on that provider account with its server-side key. Every caller-supplied model now goes through can_key_call_resolved_model before deployment credentials are resolved, covering file create/retrieve/content/ delete/list, batch create/retrieve/list/cancel, and vector store files. --- litellm/proxy/batches_endpoints/endpoints.py | 17 +- .../openai_files_endpoints/common_utils.py | 61 +++++- .../openai_files_endpoints/files_endpoints.py | 20 +- .../vector_store_files_endpoints/endpoints.py | 6 +- .../proxy/batches_endpoints/test_endpoints.py | 58 +++++- .../test_files_endpoint.py | 184 ++++++++++++++++-- .../test_batch_x_litellm_model_encoding.py | 53 ++--- 7 files changed, 322 insertions(+), 77 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..c99f66d032e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -34,9 +34,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, + get_authorized_credentials_for_model, get_batch_from_database, get_batch_id_from_unified_batch_id, - get_credentials_for_model, get_model_id_from_unified_batch_id, get_models_from_unified_file_id, get_original_file_id, @@ -218,9 +218,10 @@ async def create_batch( # SCENARIO 1: File ID is encoded with model info if model_from_file_id is not None and input_file_id: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_file_id, + user_api_key_dict=user_api_key_dict, operation_context="batch creation (file created with model)", ) @@ -310,9 +311,10 @@ async def create_batch( # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback if model_param: # SCENARIO 2: Use model-based routing from header/query/body - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, + user_api_key_dict=user_api_key_dict, operation_context="batch creation", ) @@ -540,9 +542,10 @@ async def retrieve_batch( # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context="batch retrieval (batch created with model)", ) @@ -764,9 +767,10 @@ async def list_batches( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ): # SCENARIO 2: Use model-based routing from header/query/body - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, + user_api_key_dict=user_api_key_dict, operation_context="batch listing", ) @@ -952,9 +956,10 @@ async def cancel_batch( # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context="batch cancellation (batch created with model)", ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b1f282a0978..4202a6d1689 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -351,6 +351,10 @@ def get_credentials_for_model( """ Retrieve API credentials for a model from the LLM Router. + Does not check whether the caller may use ``model_id``; use + ``get_authorized_credentials_for_model`` for anything driven by a caller-supplied + model name (request body, header, query param, or a model-encoded resource id). + Args: llm_router: LiteLLM Router instance model_id: Model name or deployment ID @@ -381,6 +385,48 @@ def get_credentials_for_model( return credentials +async def authorize_model_for_key( + model_id: str, + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", +) -> None: + """ + Enforce the caller's model grants on a model name the auth layer never saw. + + The files and batches routes carry their model in a header, query param, or a + model-encoded resource id rather than the request body, so ``user_api_key_auth`` + cannot check it. Run the same key, team (incl. team-member and access-group + fallbacks), org and project allowlist checks a chat request would get, so a + restricted key cannot borrow another deployment's server-side credentials. + + Raises: + ProxyException (403): the caller is not allowed to use ``model_id`` + """ + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + await can_key_call_resolved_model( + model=model_id, + llm_model_list=None, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + + +async def get_authorized_credentials_for_model( + llm_router: Optional["Router"], + model_id: str, + user_api_key_dict: "UserAPIKeyAuth", + operation_context: str = "file operation", +) -> dict: # mutable-ok: same contract as get_credentials_for_model, callers merge it into request data + """``get_credentials_for_model`` gated by ``authorize_model_for_key``.""" + await authorize_model_for_key(model_id=model_id, llm_router=llm_router, user_api_key_dict=user_api_key_dict) + return get_credentials_for_model( + llm_router=llm_router, + model_id=model_id, + operation_context=operation_context, + ) + + def get_team_provider_credentials( llm_router: Optional["Router"], user_api_key_dict: "UserAPIKeyAuth", @@ -573,21 +619,27 @@ def prepare_data_with_credentials( data["file_id"] = file_id -def handle_model_based_routing( +async def handle_model_based_routing( file_id: str, request, # FastAPI Request object llm_router, # Router instance data: dict, + user_api_key_dict: "UserAPIKeyAuth", check_file_id_encoding: bool = True, ) -> tuple[bool, str | None, str | None, dict | None]: """ Orchestrate model-based credential routing for file operations. + The model name comes from the caller (embedded in the file id, or a header, query + param or body field), so it is authorized against the caller's key, team, org and + project grants before any deployment credentials are resolved. + Args: file_id: File ID (may contain embedded model info) request: FastAPI request object llm_router: LiteLLM Router instance data: Request data dictionary + user_api_key_dict: The authenticated caller check_file_id_encoding: Whether to check for embedded model in file_id Returns: @@ -599,6 +651,7 @@ def handle_model_based_routing( Raises: HTTPException: If router unavailable or model not found + ProxyException: If the caller is not allowed to use the model """ model_from_id, model_from_param = extract_model_from_sources( file_id=file_id, @@ -608,9 +661,10 @@ def handle_model_based_routing( # Priority 1: Model embedded in file_id if check_file_id_encoding and model_from_id is not None: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context=f"file operation (file created with model '{model_from_id}')", ) original_file_id: Final = get_original_file_id(file_id) @@ -618,9 +672,10 @@ def handle_model_based_routing( # Priority 2: Model from header/query/body elif model_from_param is not None: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_param, + user_api_key_dict=user_api_key_dict, operation_context="file operation", ) return True, model_from_param, None, credentials diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b3bb1fa9a01..0efd618e171 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -68,7 +68,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( apply_team_provider_credentials, encode_file_id_with_model, extract_file_creation_params, - get_credentials_for_model, + get_authorized_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, validate_file_list_limit, @@ -267,9 +267,10 @@ async def route_create_file( # NEW: Handle model-based routing (no DB required) if model is not None: # Get credentials from model_list via router - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model, + user_api_key_dict=user_api_key_dict, operation_context="file upload", ) @@ -907,11 +908,12 @@ async def get_file_content( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1122,11 +1124,12 @@ async def get_file( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1330,11 +1333,12 @@ async def delete_file( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1517,11 +1521,12 @@ async def list_files( response: Any | None = None # Check for model-based credential routing (no file_id encoding check for list) - should_route, model_used, _, credentials = handle_model_based_routing( + should_route, model_used, _, credentials = await handle_model_based_routing( file_id="", # No file_id for list endpoint request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=False, ) @@ -1548,9 +1553,10 @@ async def list_files( status_code=500, detail="LLM Router not initialized. Ensure models added to proxy.", ) - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=target_model_names_list[0], + user_api_key_dict=user_api_key_dict, operation_context="file list", ) prepare_data_with_credentials(data=data, credentials=credentials) diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 957ed9fd0b9..11ef8efb598 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -144,11 +144,12 @@ async def _update_request_data_with_managed_file_id( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -273,11 +274,12 @@ async def _update_request_data_with_model_routing_hint( _model_used, _original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id="", request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=False, ) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..57e42e79a42 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -177,6 +177,8 @@ def harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1161,6 +1163,8 @@ def retrieve_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1616,6 +1620,8 @@ def list_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2012,6 +2018,8 @@ def cancel_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2733,8 +2741,6 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc assert cancel_harness.router_acancel.call_count == 1 - - @pytest.mark.asyncio async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness): with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): @@ -2762,3 +2768,51 @@ async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retriev metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {} assert metadata.get("batch_ignore_default_logging") is None + + +def _key_restricted_to(*models: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-restricted", team_id="team-a", team_models=list(models), models=list(models)) + + +@pytest.mark.asyncio +async def test_create__header_model_rejects_key_without_model_grant(harness): + """A key not granted the model named in x-litellm-model must not receive that deployment's credentials.""" + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness, user=_key_restricted_to("azure/gpt-4o"), headers={"x-litellm-model": "vertex-model"}) + + assert exc_info.value.code == "403" + harness.creds_resolver.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__header_model_allows_key_with_model_grant(harness): + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + await call_create(harness, user=_key_restricted_to("vertex-model"), headers={"x-litellm-model": "vertex-model"}) + + harness.creds_resolver.assert_called_once_with(model_id="vertex-model") + assert harness.acreate_kwargs()["custom_llm_provider"] == "vertex_ai" + + +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id_rejects_key_without_model_grant(retrieve_harness): + """The model embedded in a batch id is caller-controlled, so it is checked against the key's grants too.""" + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + retrieve_harness.creds_resolver.assert_not_called() + retrieve_harness.litellm_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_harness): + with pytest.raises(ProxyException) as exc_info: + await call_cancel(cancel_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + cancel_harness.creds_resolver.assert_not_called() + cancel_harness.litellm_acancel.assert_not_called() diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5faae166fca..548c0eb0d91 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1877,7 +1877,7 @@ def test_get_file_content_streams_openai_direct_path( monkeypatch.setattr(litellm, "afile_content", _mock_afile_content) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: (False, None, None, None), + AsyncMock(return_value=(False, None, None, None)), ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -1942,15 +1942,17 @@ def test_get_file_content_routed_provider_skips_streaming_when_resolved_provider ) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: ( - True, - "azure-gpt-3-5-turbo", - "file-original-123", - { - "custom_llm_provider": "azure", - "api_key": "azure-key", - "api_base": "https://azure.example.com", - }, + AsyncMock( + return_value=( + True, + "azure-gpt-3-5-turbo", + "file-original-123", + { + "custom_llm_provider": "azure", + "api_key": "azure-key", + "api_base": "https://azure.example.com", + }, + ) ), ) @@ -2015,7 +2017,7 @@ def test_get_file_content_non_openai_provider_skips_streaming_handler( ) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: (False, None, None, None), + AsyncMock(return_value=(False, None, None, None)), ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -2463,14 +2465,16 @@ def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice( monkeypatch.setattr(litellm, "afile_list", _mock_afile_list) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: ( - True, - "azure-gpt-4o", - None, - { - "custom_llm_provider": "azure", - "api_key": "azure-key", - }, + AsyncMock( + return_value=( + True, + "azure-gpt-4o", + None, + { + "custom_llm_provider": "azure", + "api_key": "azure-key", + }, + ) ), ) @@ -4878,3 +4882,145 @@ def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFix assert captured_kwargs["api_key"] == "mistral-key" assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df" assert response.json()["id"] == encoded_id + + +def _mistral_plus_anthropic_router() -> Router: + return Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"}, + "model_info": {"id": "mistral-ocr-id"}, + }, + { + "model_name": "claude-opus-4-6", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "anthropic-key"}, + "model_info": {"id": "claude-id"}, + }, + ] + ) + + +def _restricted_key(key_models: list[str]) -> UserAPIKeyAuth: + from litellm.proxy._types import LitellmUserRoles + + return UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + team_id="team-a", + team_models=["claude-opus-4-6", "mistral-ocr"], + models=key_models, + ) + + +@pytest.mark.parametrize( + "http_method, path_suffix, litellm_fn", + [ + ("get", "", "afile_retrieve"), + ("get", "/content", "afile_content"), + ("delete", "", "afile_delete"), + ], +) +def test_model_routed_file_ops_reject_key_without_model_grant( + mocker: MockerFixture, monkeypatch, http_method: str, path_suffix: str, litellm_fn: str +): + """ + Regression: a key whose allowlist does not include the deployment named in a + model-encoded file id must be refused before that deployment's server-side + credentials are resolved. Previously any key could name any deployment via the + id (or the x-litellm-model header) and act on that provider account's files. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called")) + monkeypatch.setattr(litellm, litellm_fn, upstream) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"]) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = getattr(client, http_method)( + f"/v1/files/{encoded_id}{path_suffix}", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + assert "not allowed to access model" in response.text + upstream.assert_not_called() + + +def test_list_files_header_model_rejects_key_without_model_grant(mocker: MockerFixture, monkeypatch): + import litellm.proxy.proxy_server as ps + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called")) + monkeypatch.setattr(litellm, "afile_list", upstream) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"]) + + try: + response = client.get( + "/v1/files", headers={"Authorization": "Bearer test-key", "x-litellm-model": "mistral-ocr"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + upstream.assert_not_called() + + +def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFixture, monkeypatch): + """The grant check must not break the happy path: a key allowed the deployment still resolves its credentials.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_retrieve(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", + object="file", + bytes=2, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["mistral-ocr"]) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["api_key"] == "mistral-key" + assert captured_kwargs["custom_llm_provider"] == "mistral" diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index 3d1831bb4cd..fe4903b547c 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -58,10 +58,7 @@ def _make_batch_response( def test_get_batch_id_from_unified_batch_id_handles_appended_fields(): - decoded_id = ( - "litellm_proxy;model_id:deployment-123;" - "llm_batch_id:batch_openai_123;llm_output_file_id:file-output" - ) + decoded_id = "litellm_proxy;model_id:deployment-123;llm_batch_id:batch_openai_123;llm_output_file_id:file-output" assert get_batch_id_from_unified_batch_id(decoded_id) == "batch_openai_123" @@ -107,12 +104,10 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id(): } ), ), + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, + "litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model", + new=AsyncMock(return_value=mock_credentials), ), patch( "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", @@ -165,23 +160,15 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id(): ) # The batch_id should be encoded with model info - assert ( - response.id != raw_batch_id - ), f"Expected batch_id to be encoded, but got raw ID: {response.id}" - assert response.id.startswith( - "batch_" - ), f"Encoded batch_id should keep batch_ prefix, got: {response.id}" + assert response.id != raw_batch_id, f"Expected batch_id to be encoded, but got raw ID: {response.id}" + assert response.id.startswith("batch_"), f"Encoded batch_id should keep batch_ prefix, got: {response.id}" # Should be decodable back to the original decoded_model = decode_model_from_file_id(response.id) - assert ( - decoded_model == model_name - ), f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" + assert decoded_model == model_name, f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" original_id = get_original_file_id(response.id) - assert ( - original_id == raw_batch_id - ), f"Expected original ID '{raw_batch_id}', got: {original_id}" + assert original_id == raw_batch_id, f"Expected original ID '{raw_batch_id}', got: {original_id}" assert mock_create_batch.call_args.kwargs["metadata"] == {"customer_id": "cust-123"} @@ -227,12 +214,10 @@ async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_i } ), ), + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, + "litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model", + new=AsyncMock(return_value=mock_credentials), ), patch( "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", @@ -316,9 +301,7 @@ async def test_create_batch_without_x_litellm_model_returns_raw_ids(monkeypatch) } ), ), - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( "litellm.acreate_batch", new=AsyncMock(return_value=mock_response), @@ -383,9 +366,7 @@ class TestBatchIdRoundTripWithRetrieve: raw_batch_id = "batch_vllm_12345" # What create_batch does: - encoded_id = encode_file_id_with_model( - file_id=raw_batch_id, model=model_name, id_type="batch" - ) + encoded_id = encode_file_id_with_model(file_id=raw_batch_id, model=model_name, id_type="batch") # What retrieve_batch does: decoded_model = decode_model_from_file_id(encoded_id) @@ -410,9 +391,7 @@ class TestBatchIdRoundTripWithRetrieve: ] for raw_id, model in test_cases: - encoded = encode_file_id_with_model( - file_id=raw_id, model=model, id_type="batch" - ) + encoded = encode_file_id_with_model(file_id=raw_id, model=model, id_type="batch") assert encoded.startswith("batch_") assert decode_model_from_file_id(encoded) == model assert get_original_file_id(encoded) == raw_id @@ -440,9 +419,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ mock_user_api_key_dict.team_metadata = {} with ( - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( "litellm.proxy.batches_endpoints.endpoints.update_batch_in_database", new=AsyncMock(), From bae731ddfc394b23d3c6f44a85cfaa58472897e7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 17:54:57 -0400 Subject: [PATCH 033/393] fix(proxy): apply model grants to unified file and batch ids on batch routes Unified ids carry the deployment model inside the id, so a restricted key could create, retrieve or cancel a batch on a deployment it is not granted. The model parsed from a unified id now goes through the same grant check as header, query and model-encoded id sources before the router is called. --- litellm/proxy/batches_endpoints/endpoints.py | 16 ++++- .../proxy/batches_endpoints/test_endpoints.py | 58 +++++++++++++++++++ .../test_batch_x_litellm_model_encoding.py | 7 +-- 3 files changed, 75 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index c99f66d032e..c2489ce52ea 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, + authorize_model_for_key, batch_cost_poller_is_active, decode_model_from_file_id, encode_batch_response_ids, @@ -286,6 +287,7 @@ async def create_batch( detail={"error": f"Expected 1 model, got {len(target_model_names)}"}, ) model: Final = target_model_names[0] + await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict) _create_batch_data["model"] = model resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id) @@ -582,10 +584,17 @@ async def retrieve_batch( ) if unified_batch_id: + unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) + if unified_model_id is not None: + await authorize_model_for_key( + model_id=llm_router.resolve_model_name_from_model_id(unified_model_id) or unified_model_id, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) add_internal_model_credentials( data=data, llm_router=llm_router, - model_id=get_model_id_from_unified_batch_id(unified_batch_id), + model_id=unified_model_id, ) response = await llm_router.aretrieve_batch(**data) @@ -998,6 +1007,11 @@ async def cancel_batch( status_code=400, detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."}, ) + await authorize_model_for_key( + model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) data["model"] = model_id_from_batch data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id) response = await llm_router.acancel_batch(**data) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 57e42e79a42..5be64ca0847 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -179,6 +179,8 @@ def harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1165,6 +1167,8 @@ def retrieve_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1622,6 +1626,8 @@ def list_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2020,6 +2026,8 @@ def cancel_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2816,3 +2824,53 @@ async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_h assert exc_info.value.code == "403" cancel_harness.creds_resolver.assert_not_called() cancel_harness.litellm_acancel.assert_not_called() + + +def _b64_unified_id(decoded: str) -> str: + return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=") + + +UNIFIED_FILE_ID_FOR_GPT4O_MINI = _b64_unified_id( + "litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;" + "target_model_names,gpt-4o-mini;llm_output_file_id,file-provider;llm_output_file_model_id,dep-1" +) +UNIFIED_BATCH_ID_FOR_GPT4O_MINI = _b64_unified_id(UNIFIED_BATCH_ID) + + +@pytest.mark.asyncio +async def test_create__unified_file_id_rejects_key_without_model_grant(harness): + """The model carried inside a unified file id is caller-controlled too, so it is checked against the key's grants.""" + set_body( + harness, + { + "input_file_id": UNIFIED_FILE_ID_FOR_GPT4O_MINI, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_rejects_key_without_model_grant(retrieve_harness): + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + retrieve_harness.router_aretrieve.assert_not_called() + retrieve_harness.creds_resolver.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_harness): + with pytest.raises(ProxyException) as exc_info: + await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + cancel_harness.router_acancel.assert_not_called() diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index fe4903b547c..3161fe99e68 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, get_batch_id_from_unified_batch_id, @@ -412,11 +413,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ mock_request.url.path = f"/v1/batches/{unified_batch_id}/cancel" mock_fastapi_response = MagicMock() mock_fastapi_response.headers = {} - mock_user_api_key_dict = MagicMock() - mock_user_api_key_dict.parent_otel_span = None - mock_user_api_key_dict.user_id = "test_user" - mock_user_api_key_dict.allowed_model_region = None - mock_user_api_key_dict.team_metadata = {} + mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="test_user", team_metadata={}) with ( patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, From 95fdefa390af6586affb5ff825955b0ccb3bce17 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 19:14:45 -0400 Subject: [PATCH 034/393] fix(logging): tolerate a missing api_base in pre_call for presigned batch retrieves Provider batch configs that build their own request URL (Mistral, Bedrock) hand pre_call api_base=None, and mask_api_base_credentials raised TypeError on it, so every such retrieve logged a non-blocking LoggingError and lost its pre-call logging. --- litellm/litellm_core_utils/litellm_logging.py | 4 +- .../test_litellm_logging.py | 63 ++++++++++--------- 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cb9209be267..38e88493986 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1199,8 +1199,8 @@ class Logging(LiteLLMLoggingBaseClass): return {"error": f"Unable to parse raw request body. Got - {data}"} return data - def _get_masked_api_base(self, api_base: str) -> str: - return str(mask_api_base_credentials(api_base)) + def _get_masked_api_base(self, api_base: str | None) -> str: + return str(mask_api_base_credentials(api_base or "")) def _pre_call(self, input, api_key, model=None, additional_args={}): """ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 6aa77745e3d..570f4339cd4 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -58,6 +58,16 @@ def test_get_masked_api_base(logging_obj): assert type(masked_api_base) == str +def test_pre_call_tolerates_missing_api_base(logging_obj): + """Presigned batch retrieves (Mistral, Bedrock) build their own URL and pass api_base=None + to pre_call; masking must not raise or the request's pre-call logging is silently lost.""" + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="", additional_args={"api_base": None, "headers": {}}) + + assert logging_obj.model_call_details["litellm_params"]["api_base"] == "" + + def test_post_call_serializes_dict_with_datetime(logging_obj): import datetime @@ -3976,9 +3986,7 @@ def test_get_standard_logging_object_payload_carries_matched_access_groups(loggi "model": "gpt-4o", "messages": [], "litellm_params": { - "metadata": { - "user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"] - }, + "metadata": {"user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"]}, "proxy_server_request": {"body": {}}, }, }, @@ -4062,9 +4070,7 @@ def _model_router_response(selected_model: str, stamp: bool): from litellm.types.utils import ModelResponse response = ModelResponse(model=selected_model) - response._hidden_params = ( - {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} - ) + response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} return response @@ -4088,9 +4094,7 @@ def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response( - "azure_ai/grok-4-1-fast-reasoning", stamp=True - ), + init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True), start_time=now, end_time=now, logging_obj=logging_obj, @@ -4122,9 +4126,7 @@ def test_standard_logging_payload_keeps_requested_model_without_router_stamp( "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response( - "azure_ai/grok-4-1-fast-reasoning", stamp=False - ), + init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False), start_time=now, end_time=now, logging_obj=logging_obj, @@ -5536,9 +5538,7 @@ class TestNonInferenceCallTypesAreNotBilled: init_response_obj=self._retrieved_response(), start_time=now, end_time=now, - logging_obj=self._logging_obj( - "aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA - ), + logging_obj=self._logging_obj("aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA), status="success", ) @@ -5784,9 +5784,7 @@ async def test_streaming_success_callbacks_survive_cost_calculation_failure(): releasing.async_log_success_event = AsyncMock() patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) - with patcher, patch.object( - logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block") - ): + with patcher, patch.object(logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block")): await logging_obj.async_success_handler(result=_assembled_stream_result()) assert logging_obj.model_call_details["response_cost"] is None @@ -5799,8 +5797,9 @@ async def test_streaming_success_callbacks_survive_standard_logging_payload_fail releasing.async_log_success_event = AsyncMock() patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) - with patcher, patch.object( - logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream") + with ( + patcher, + patch.object(logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream")), ): await logging_obj.async_success_handler(result=_assembled_stream_result()) @@ -6073,6 +6072,8 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa ) for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook) + + def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch): """With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic" callback builds the OTel v2 logger (per-team credential routing); with the @@ -6228,7 +6229,9 @@ def test_get_error_information_skips_traceback_for_budget_rejection_with_provide from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup assert litellm.log_client_error_tracebacks is False - over_budget = _raise_and_catch(litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic")) + over_budget = _raise_and_catch( + litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic") + ) result = StandardLoggingPayloadSetup.get_error_information(over_budget) assert result["error_code"] == "429" assert result["llm_provider"] == "anthropic" @@ -6745,9 +6748,7 @@ def test_passthrough_embeddings_result_swapped_for_callbacks(): ], "model": "EmbeddingsGigaR", }, - request=httpx.Request( - "POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings" - ), + request=httpx.Request("POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings"), ) _, _, swapped_result = logging_obj._success_handler_helper_fn( @@ -6766,12 +6767,14 @@ def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervene request-level guardrail_status but never mask an intervention.""" flagged = {"guardrail_status": "guardrail_flagged"} - assert _get_status_fields( - "success", [{"guardrail_status": "success"}, flagged], None - )["guardrail_status"] == "guardrail_flagged" - assert _get_status_fields( - "success", [flagged, {"guardrail_status": "guardrail_intervened"}], None - )["guardrail_status"] == "guardrail_intervened" + assert ( + _get_status_fields("success", [{"guardrail_status": "success"}, flagged], None)["guardrail_status"] + == "guardrail_flagged" + ) + assert ( + _get_status_fields("success", [flagged, {"guardrail_status": "guardrail_intervened"}], None)["guardrail_status"] + == "guardrail_intervened" + ) def test_get_error_information_redacts_provider_key_from_upstream_url(): From e6bc4e47c7a63e8540a856f1b2f66710a4b47142 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 19:46:55 -0400 Subject: [PATCH 035/393] fix(mistral): reject file purposes Mistral lacks instead of mapping them to batch The proxy runs batch-file validation and guardrails only for purpose=batch, so a purpose such as assistants that was silently rewritten to batch on the way to Mistral let an upload skip both. Only batch, fine-tune and ocr pass through now; anything else is a 400. --- litellm/llms/mistral/files/transformation.py | 9 ++++-- .../test_mistral_files_transformation.py | 29 ++++++++++++++----- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 6d58311813c..bf1ef7cb69f 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -89,11 +89,14 @@ def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose: def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: + """Only Mistral's own purposes pass through. Silently mapping anything else to ``batch`` + would let an upload skip the proxy's batch-file validation and guardrails, which only + run when the caller says ``purpose=batch``.""" match purpose: - case "fine-tune" | "ocr": + case "batch" | "fine-tune" | "ocr": return purpose case _: - return "batch" + raise ValueError(f"Mistral does not support purpose={purpose!r}. Use one of: batch, fine-tune, ocr") def _api_base_from(litellm_params: Mapping[str, object]) -> str: @@ -166,7 +169,7 @@ class MistralFilesConfig(BaseFilesConfig): content_type: Final = extracted.get("content_type") or "application/octet-stream" upload: Final = MistralMultipartUpload( file=(filename, extracted["content"], content_type), - purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + purpose=(None, _to_mistral_purpose(create_file_data.get("purpose") or "batch")), ) return dict(upload) # mutable-ok: BaseFilesConfig signature diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index f62645be7ee..b81740c0429 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -91,18 +91,28 @@ def test_upload_request_is_multipart_with_batch_purpose(config): } -@pytest.mark.parametrize( - "openai_purpose,mistral_purpose", - [("batch", "batch"), ("fine-tune", "fine-tune"), ("ocr", "ocr"), ("assistants", "batch"), ("user_data", "batch")], -) -def test_upload_request_maps_purpose_onto_mistral_enum(config, openai_purpose, mistral_purpose): +@pytest.mark.parametrize("purpose", ["batch", "fine-tune", "ocr"]) +def test_upload_request_passes_mistral_purposes_through(config, purpose): body = config.transform_create_file_request( model="", - create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=openai_purpose), + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), optional_params={}, litellm_params={}, ) - assert body["purpose"] == (None, mistral_purpose) + assert body["purpose"] == (None, purpose) + + +@pytest.mark.parametrize("purpose", ["assistants", "user_data", "vision", "evals"]) +def test_upload_request_rejects_purposes_mistral_lacks(config, purpose): + """Regression: these used to be silently rewritten to ``batch``, so an upload that skipped the + proxy's batch-only validation and guardrails still landed on Mistral as a batch input file.""" + with pytest.raises(ValueError, match=f"purpose={purpose!r}"): + config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), + optional_params={}, + litellm_params={}, + ) def test_upload_request_requires_file(config): @@ -181,6 +191,11 @@ def test_list_request_filters_by_mapped_purpose(config): assert no_params == {} +def test_list_request_rejects_purposes_mistral_lacks(config): + with pytest.raises(ValueError, match="purpose='assistants'"): + config.transform_list_files_request(purpose="assistants", optional_params={}, litellm_params={}) + + def test_list_response(config): out = config.transform_list_files_response( raw_response=_response( From 0f3c4ccfbba894c0a66d92ed44ca18f312bf75f0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 11 Sep 2026 00:00:17 +0000 Subject: [PATCH 036/393] fix(auth): inherit organization_alias from the org for JWT and team-linked keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 42 ++++++- .../proxy/auth/test_user_api_key_auth.py | 104 +++++++++++++++++- 2 files changed, 143 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 20ab9904f46..110c524ecdf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( get_end_user_object, get_jwt_key_mapping_object, get_object_permission, + get_org_object, get_project_object, get_team_object, get_user_object, @@ -2398,6 +2399,37 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() +async def _inherit_org_identity( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTableCachedObj | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> None: + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + if ( + user_api_key_auth_obj.org_id is None + or user_api_key_auth_obj.organization_alias is not None + or prisma_client is None + ): + return + try: + org_object: Final = await get_org_object( + org_id=user_api_key_auth_obj.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.debug("org alias lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) + return + if org_object is not None: + user_api_key_auth_obj.organization_alias = org_object.organization_alias + + @tracer.wrap() async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, @@ -2622,8 +2654,14 @@ async def _run_centralized_common_checks( ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result - if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: - user_api_key_auth_obj.org_id = team_object.organization_id + await _inherit_org_identity( + user_api_key_auth_obj=user_api_key_auth_obj, + team_object=cast(LiteLLM_TeamTableCachedObj | None, team_object), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6cce6d0316b..03efbfa7185 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -31,7 +32,7 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object +from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, @@ -5293,6 +5294,107 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_id,team_org_id,existing_alias,lookup_mode,expected_org_id,expected_alias", + [ + (None, "t1", "org-from-team", None, "success", "org-from-team", "acme-org"), + ("org-jwt", None, None, None, "success", "org-jwt", "acme-org"), + ("org-pinned", None, None, "preset", "success", "org-pinned", "preset"), + ("org-missing", None, None, None, "missing", "org-missing", None), + ], +) +async def test_centralized_common_checks_inherits_org_alias( + key_org_id, + team_id, + team_org_id, + existing_alias, + lookup_mode, + expected_org_id, + expected_alias, +): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id=team_id, + org_id=key_org_id, + organization_alias=existing_alias, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = ( + LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) if team_id is not None else None + ) + organization = LiteLLM_OrganizationTable( + organization_id=expected_org_id, + organization_alias="acme-org", + budget_id="budget-id", + models=[], + created_by="test", + updated_by="test", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["prisma_client"] = MagicMock() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + identity_seen_by_common_checks = [] + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ) as mock_get_team_object, + patch( + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + return_value=organization, + ) as mock_get_org_object, + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + side_effect=lambda **kw: identity_seen_by_common_checks.append( + (kw["valid_token"].org_id, kw["valid_token"].organization_alias) + ), + ) as mock_checks, + ): + if lookup_mode == "missing": + mock_get_org_object.side_effect = OrganizationNotFoundError("x") + + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id == expected_org_id + assert token.organization_alias == expected_alias + assert identity_seen_by_common_checks == [(expected_org_id, expected_alias)] + if team_id is None: + mock_get_team_object.assert_not_awaited() + else: + mock_get_team_object.assert_awaited_once() + if existing_alias is not None: + mock_get_org_object.assert_not_awaited() + else: + mock_get_org_object.assert_awaited_once() + assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_cli_session_token_org_backfilled_from_team(monkeypatch): """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted From 20d80b5420508c73391cca91be232b7f74041d1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:45:07 -0700 Subject: [PATCH 037/393] fix(guardrails): scan each choice's tool-call arguments apart on n>1 streams and log why a rewrite was discarded The rebuilt streamed response keyed tool-call fragments by tool index alone, so on n>1 chat streams the two choices' argument fragments were concatenated into one string and post_call guardrails scanned garbled JSON. Fragments are now keyed by (choice index, tool index). When a guardrail's rewrite cannot be written back to the stream (multi-choice streams, a rewrite that adds or drops a tool call, legacy-hook shapes the translation cannot rescan), the pipeline now logs a warning naming the guardrail and the exact reason before releasing the original stream. Also commits the regenerated dashboard API types that make check produced. --- .../streaming_chunk_builder_utils.py | 38 +++++---- .../chat/guardrail_translation/handler.py | 11 ++- .../chat/guardrail_translation/handler.py | 26 +++++-- .../guardrail_translation/handler.py | 11 ++- .../proxy/policy_engine/pipeline_executor.py | 78 ++++++++++++++----- .../test_streaming_chunk_builder_utils.py | 55 +++++++++++++ .../test_openai_guardrail_handler.py | 57 +++++++++++++- .../policy_engine/test_pipeline_executor.py | 42 ++++++---- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 9 files changed, 254 insertions(+), 66 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 90698296142..5ffe36573d5 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -138,9 +138,13 @@ class _ToolCallDelta(TypedDict, total=False): class _ToolCallChoice(TypedDict, total=False): + index: ReadOnly[int] delta: ReadOnly[_ToolCallDelta] +_ToolCallKey: TypeAlias = tuple[int, int] + + class _ToolCallChunk(TypedDict): choices: ReadOnly[Sequence[_ToolCallChoice]] @@ -416,40 +420,41 @@ class ChunkProcessor: @staticmethod def _iter_tool_call_fragments( tool_call_chunks: Sequence["_ToolCallChunk"], - ) -> Iterator[tuple[int, str, str]]: + ) -> Iterator[tuple[_ToolCallKey, str, str]]: for chunk in tool_call_chunks: for choice in chunk["choices"]: delta = choice.get("delta") if not delta: continue + choice_index = choice.get("index", 0) for tool_call in delta.get("tool_calls", ()): if not tool_call: continue if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + key = (choice_index, tool_call.get("index", 0)) function = tool_call.get("function") if isinstance(function, dict): if fragment_arguments := function.get("arguments"): - yield index, "arguments", fragment_arguments + yield key, "arguments", fragment_arguments elif function_arguments := getattr(function, "arguments", None): - yield index, "arguments", function_arguments + yield key, "arguments", function_arguments custom = tool_call.get("custom") if isinstance(custom, dict) and (custom_input := custom.get("input")): - yield index, "custom_input", custom_input + yield key, "custom_input", custom_input else: - index = getattr(tool_call, "index", 0) + key = (choice_index, getattr(tool_call, "index", 0)) function = getattr(tool_call, "function", None) if object_arguments := getattr(function, "arguments", None): - yield index, "arguments", object_arguments + yield key, "arguments", object_arguments custom = getattr(tool_call, "custom", None) if object_custom_input := getattr(custom, "input", None): - yield index, "custom_input", object_custom_input + yield key, "custom_input", object_custom_input @staticmethod - def _join_fragments_by_index_and_field( - fragment_records: Iterator[tuple[int, str, str]], - ) -> Mapping[tuple[int, str], str]: - def group_key(record: tuple[int, str, str]) -> tuple[int, str]: + def _join_fragments_by_key_and_field( + fragment_records: Iterator[tuple[_ToolCallKey, str, str]], + ) -> Mapping[tuple[_ToolCallKey, str], str]: + def group_key(record: tuple[_ToolCallKey, str, str]) -> tuple[_ToolCallKey, str]: return record[0], record[1] return MappingProxyType( @@ -467,13 +472,14 @@ class ChunkProcessor: tool_calls_list: list[ ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall ] = [] # mutable-ok: see return type - tool_call_map: Final[dict[int, dict[str, Any]]] = {} # Map to store tool calls by index + tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {} # Map to store tool calls by choice and index for chunk in tool_call_chunks: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) tool_calls = delta.get("tool_calls", []) + choice_index = choice.get("index", 0) for tool_call in tool_calls: # Handle both dict and object formats @@ -495,9 +501,9 @@ class ChunkProcessor: # Get index (handle both dict and object) if isinstance(tool_call, dict): - index = tool_call.get("index", 0) + index = (choice_index, tool_call.get("index", 0)) else: - index = getattr(tool_call, "index", 0) + index = (choice_index, getattr(tool_call, "index", 0)) if index not in tool_call_map: tool_call_map[index] = { @@ -572,7 +578,7 @@ class ChunkProcessor: if isinstance(provider_fields, dict): merged_provider_fields.update(provider_fields) - joined_fragments: Final = self._join_fragments_by_index_and_field( + joined_fragments: Final = self._join_fragments_by_key_and_field( self._iter_tool_call_fragments(tool_call_chunks) ) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..c7ba5daec56 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1172,7 +1172,10 @@ class AnthropicMessagesHandler(BaseTranslation): if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + "the stream never reported a stop_reason, so the text rewrite has no assembled response to land on", + ) return responses_so_far def _prepare_request_data( @@ -1318,7 +1321,11 @@ class AnthropicMessagesHandler(BaseTranslation): if len(block_indices) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " + f"{len(block_indices)} tool_use blocks", + ) rewrites_by_block: Final = MappingProxyType( { index: after diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 58ff03e6a0d..e4943690639 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1041,13 +1041,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice.index for response in responses_so_far for choice in response.choices ) if len(stream_choice_indices) != 1: - # stream_chunk_builder collapses every choice into one index-0 - # choice, so a rewrite of the rebuilt response cannot be attributed - # back to a single choice on an n>1 stream: report it undeliverable - # rather than deliver the rewrite on the wrong choice from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and the rebuilt response's text rewrite " + "cannot be attributed to one of them", + ) target_choice_index: Final = next(iter(stream_choice_indices)) await self._apply_guardrail_responses_to_output_streaming( responses=responses_so_far, @@ -1105,10 +1105,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation): choice.index for response in responses_so_far for choice in response.choices ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) - if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + if len(stream_choice_indices) != 1: from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " + "back on single-choice streams", + ) + if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " + f"{len(fragments_by_tool_call)}", + ) for before, (name, arguments), fragments in zip( pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call ): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..2be36a826f7 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -957,7 +957,10 @@ class OpenAIResponsesHandler(BaseTranslation): if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_to_apply.guardrail_name or "unknown", + "the stream carried no terminal response envelope to write the text rewrite back into", + ) return responses_so_far @staticmethod @@ -1070,7 +1073,11 @@ class OpenAIResponsesHandler(BaseTranslation): ): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - raise UndeliverableStreamRewrite(guardrail_name) + raise UndeliverableStreamRewrite( + guardrail_name, + f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls and the stream's " + f"{len(tool_call_items)} function_call items could not be lined up with them by call_id", + ) for output_item, rewrite in ( (output_item, rewrites_by_call_id[call_id]) for output_item, call_id in zip(tool_call_items, call_ids) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ad45781d5d2..26dd806b3e5 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -50,12 +50,13 @@ except ImportError: class UndeliverableStreamRewrite(Exception): - def __init__(self, guardrail_name: str) -> None: + def __init__(self, guardrail_name: str, reason: str) -> None: super().__init__( - f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's " - "streaming pipeline cannot deliver" + f"Guardrail '{guardrail_name}' rewrote the streamed response but the rewrite cannot be written " + f"back to the stream: {reason}" ) self.guardrail_name: Final = guardrail_name + self.reason: Final = reason class UnappliableRequestRewrite(Exception): @@ -91,8 +92,22 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent -def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: - return sent is not None and returned is not None and len(returned) != len(sent) +def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> tuple[int, int] | None: + if sent is None or returned is None or len(returned) == len(sent): + return None + return (len(sent), len(returned)) + + +def _tool_call_mismatch_reason( + sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None +) -> str | None: + if sent == returned: + return None + sent_count: Final = len(sent or ()) + returned_count: Final = len(returned or ()) + if sent_count == returned_count: + return "the legacy hook changed a tool call's name or arguments, which this path cannot write back" + return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}" _GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) @@ -119,7 +134,7 @@ class _StreamRewriteObserver(CustomGuardrail): self.inner: Final = inner self.rewrote_texts = False self.rewrote_tool_calls = False - self.changed_tool_call_count = False + self.tool_call_count_change: tuple[int, int] | None = None def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -140,11 +155,22 @@ class _StreamRewriteObserver(CustomGuardrail): returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls")) self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes) - self.changed_tool_call_count = self.changed_tool_call_count or _changed_count( + self.tool_call_count_change = self.tool_call_count_change or _count_change( sent_tool_shapes, returned_tool_shapes ) return outputs + def discard_reason(self, deliver_rewrites: bool) -> str | None: + if self.tool_call_count_change is not None: + sent, returned = self.tool_call_count_change + return ( + f"the guardrail returned {returned} tool calls for a stream that carried {sent}, and a rewrite " + "that drops or adds a tool call cannot be written back" + ) + if not deliver_rewrites and (self.rewrote_texts or self.rewrote_tool_calls): + return "this endpoint's streaming pipeline does not write ended-stream rewrites back yet" + return None + class _ScannedTextRecorder(CustomGuardrail): def __init__(self, guardrail_name: str) -> None: @@ -209,13 +235,24 @@ class _LegacyHookStreamAdapter(CustomGuardrail): if rewrite is None: return inputs rescanned: Final = await self._rescan(rewrite, logging_obj) + guardrail_name: Final = self.guardrail_name or "unknown" if rescanned is None: - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + raise UndeliverableStreamRewrite( + guardrail_name, "the legacy hook's response could not be rescanned by this endpoint's translation" + ) rewritten: Final = rescanned.get("texts") - if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") - if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): - raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + returned_text_count: Final = len(_scanned_texts(rewritten)) + sent_text_count: Final = len(_scanned_texts(inputs.get("texts"))) + if returned_text_count != sent_text_count: + raise UndeliverableStreamRewrite( + guardrail_name, + f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}", + ) + tool_call_mismatch: Final = _tool_call_mismatch_reason( + _tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls")) + ) + if tool_call_mismatch is not None: + raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch) if not rewritten: return inputs rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} @@ -262,14 +299,16 @@ def _prepare_hook_input( def _release_original_chunks( guardrail_name: str, + reason: str, streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place originals: Sequence[object], ) -> None: streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives verbose_proxy_logger.warning( - "Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming " - "pipeline cannot deliver yet; the rewrite was discarded and the original stream released", + "Pipeline: guardrail '%s' rewrote the streamed response but the rewrite could not be written back to " + "the stream: %s. The whole rewrite, text rewrites included, was discarded and the original stream released", guardrail_name, + reason, ) @@ -442,13 +481,12 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, request_data=hook_input, ) - except UndeliverableStreamRewrite: - _release_original_chunks(step.guardrail, streaming_chunks, originals) + except UndeliverableStreamRewrite as undeliverable: + _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return - if observer.changed_tool_call_count or ( - not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls) - ): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + discard_reason: Final = observer.discard_reason(deliver_rewrites) + if discard_reason is not None: + _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals) return if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index efe4209c1c9..2266258bf20 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1288,6 +1288,61 @@ def _tool_call_delta_chunk(tool_call: dict[str, object] | ChatCompletionDeltaToo return {"choices": [{"delta": {"tool_calls": [tool_call]}}]} +def _choice_tool_call_delta_chunk(choice_index: int, tool_call: dict[str, object]) -> dict[str, object]: + return {"choices": [{"index": choice_index, "delta": {"tool_calls": [tool_call]}}]} + + +def test_get_combined_tool_content_keeps_each_choices_arguments_apart_when_choices_share_a_tool_index(): + processor = ChunkProcessor.__new__(ChunkProcessor) + chunks = [ + _choice_tool_call_delta_chunk(0, {"index": 0, "id": "call_a", "type": "function", "function": {"name": "f"}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "id": "call_b", "type": "function", "function": {"name": "f"}}), + _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": '{"fruit": "pers'}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": '{"fruit": "dur'}}), + _choice_tool_call_delta_chunk(0, {"index": 0, "function": {"arguments": 'immon"}'}}), + _choice_tool_call_delta_chunk(1, {"index": 0, "function": {"arguments": 'ian"}'}}), + ] + + combined = processor.get_combined_tool_content(chunks) + + assert [(tool_call.id, tool_call.function.arguments) for tool_call in combined] == [ + ("call_a", '{"fruit": "persimmon"}'), + ("call_b", '{"fruit": "durian"}'), + ] + + +def test_stream_chunk_builder_keeps_each_choices_tool_call_arguments_apart(): + def chunk(choice_index: int, tool_call: ChatCompletionDeltaToolCall) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + object="chat.completion.chunk", + created=1234567890, + model="gpt-4.1-mini", + choices=[StreamingChoices(index=choice_index, delta=Delta(tool_calls=[tool_call]), finish_reason=None)], + ) + + def fragment(arguments: str, name: str | None = None, call_id: str | None = None) -> ChatCompletionDeltaToolCall: + return ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments) + ) + + response = stream_chunk_builder( + chunks=[ + chunk(0, fragment("", name="lookup_fruit", call_id="call_a")), + chunk(1, fragment("", name="lookup_fruit", call_id="call_b")), + chunk(0, fragment('{"fruit": "pers')), + chunk(1, fragment('{"fruit": "dur')), + chunk(0, fragment('immon"}')), + chunk(1, fragment('ian"}')), + ] + ) + + assert [(tool_call.id, tool_call.function.arguments) for tool_call in response.choices[0].message.tool_calls] == [ + ("call_a", '{"fruit": "persimmon"}'), + ("call_b", '{"fruit": "durian"}'), + ] + + def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_order(): processor = ChunkProcessor.__new__(ChunkProcessor) first_fragments = [f"a{i};" for i in range(300)] diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5a29a96829f..60a5752e83a 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1267,7 +1267,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: handler = OpenAIChatCompletionsHandler() chunks = self._two_choice_stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: await handler.process_output_streaming_response( responses_so_far=chunks, guardrail_to_apply=self._world_masking_guardrail(), @@ -1275,6 +1275,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + assert raised.value.guardrail_name == "test-mask" + assert raised.value.reason == ( + "the stream carries 2 choices and the rebuilt response's text rewrite cannot be attributed to one of them" + ) + @staticmethod def _two_choice_tool_call_stream_chunks() -> list: from litellm.types.utils import ( @@ -1310,12 +1315,51 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: return [ chunk(0, fragment("", name="lookup_fruit", call_id="call_1")), chunk(1, fragment("", name="lookup_fruit", call_id="call_2")), - chunk(0, fragment('{"fruit": "persimmon"}')), - chunk(1, fragment('{"fruit": "durian"}')), + chunk(0, fragment('{"fruit": "pers')), + chunk(1, fragment('{"fruit": "dur')), + chunk(0, fragment('immon"}')), + chunk(1, fragment('ian"}')), chunk(0, None, finish_reason="tool_calls"), chunk(1, None, finish_reason="tool_calls"), ] + @staticmethod + def _recording_guardrail() -> CustomGuardrail: + class Recorder(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="recorder") + self.seen_inputs: list[GenericGuardrailAPIInputs] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.seen_inputs.append(inputs) + return inputs + + return Recorder() + + @pytest.mark.asyncio + async def test_ended_multi_choice_stream_scans_each_choices_tool_call_arguments_apart(self): + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() + guardrail = self._recording_guardrail() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert [ + (tool_call["id"], tool_call["function"]["arguments"]) + for tool_call in guardrail.seen_inputs[-1]["tool_calls"] + ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1323,7 +1367,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: handler = OpenAIChatCompletionsHandler() chunks = self._two_choice_tool_call_stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: await handler.process_output_streaming_response( responses_so_far=chunks, guardrail_to_apply=MockGuardrail(guardrail_name="test"), @@ -1331,6 +1375,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) + assert raised.value.guardrail_name == "test" + assert raised.value.reason == ( + "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" + ) + @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 0a2641082dc..e7689cc7d0c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1122,7 +1122,7 @@ class _RefusingTranslation: deliver_ended_stream_rewrites=False, ): responses_so_far[0]["text"] = "half-written" - raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name) + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name, "the translation refused it") def _chunk(): @@ -1143,10 +1143,22 @@ async def _run_streaming_step(translation, streaming_chunks=None): ) -def _assert_passed_with_discard_warning(result, caplog): +NO_WRITE_BACK_REASON = "this endpoint's streaming pipeline does not write ended-stream rewrites back yet" + + +def _assert_passed_with_discard_warning(result, caplog, reason): assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] - assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + discard_warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + and "'masker'" in record.getMessage() + and "discarded" in record.getMessage() + ] + assert len(discard_warnings) == 1 + assert reason in discard_warnings[0] + assert "text rewrites included" in discard_warnings[0] assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) @@ -1159,7 +1171,7 @@ async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(translation, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] assert translation.seen_guardrail_names == ["masker"] @@ -1196,7 +1208,7 @@ async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(m with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_TextTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] @@ -1258,7 +1270,9 @@ async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_WritingTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning( + result, caplog, "the guardrail returned 0 tool calls for a stream that carried 1" + ) assert chunks == [_chunk()] @@ -1270,7 +1284,7 @@ async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_ with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_TextTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, NO_WRITE_BACK_REASON) assert chunks == [_chunk()] @@ -1326,7 +1340,7 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_streaming_step(_RefusingTranslation(), chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the translation refused it") assert chunks == [_chunk()] @@ -1561,7 +1575,7 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook returned 2 texts for a stream that carried 1") assert chunks == [_chunk()] @@ -1576,7 +1590,7 @@ async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(m with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") assert chunks == [_chunk()] @@ -1588,7 +1602,9 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls( with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning( + result, caplog, "the legacy hook returned 0 tool calls for a stream that carried 1" + ) assert chunks == [_chunk()] @@ -1620,7 +1636,7 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() ) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments") assert chunks == [_tool_only_chunk()] @@ -1658,7 +1674,7 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) - _assert_passed_with_discard_warning(result, caplog) + _assert_passed_with_discard_warning(result, caplog, "the legacy hook's response could not be rescanned") assert chunks == [_chunk()] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 65160a97c54da63f24bf674f4f03551b22ac97c3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:55:03 -0700 Subject: [PATCH 038/393] fix(guardrails): keep tool calls carried by a later choice of a packed multi-choice chunk The rebuild's tool-call selection and its text-only fast path only looked at choice 0 of each chunk, so a chunk that packs several choices (Gemini with candidateCount above 1) lost a tool call carried by a later candidate, and a chunk whose later choice had no tool calls at all made the rebuild raise. Both now consider every choice in the chunk. --- .../streaming_chunk_builder_utils.py | 4 +- litellm/main.py | 66 +++++++++++-------- tests/test_litellm/test_main.py | 62 +++++++++++++++++ 3 files changed, 104 insertions(+), 28 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 5ffe36573d5..f5b723755aa 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -427,7 +427,7 @@ class ChunkProcessor: if not delta: continue choice_index = choice.get("index", 0) - for tool_call in delta.get("tool_calls", ()): + for tool_call in delta.get("tool_calls") or (): if not tool_call: continue if isinstance(tool_call, dict): @@ -478,7 +478,7 @@ class ChunkProcessor: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) - tool_calls = delta.get("tool_calls", []) + tool_calls = delta.get("tool_calls") or () choice_index = choice.get("index", 0) for tool_call in tool_calls: diff --git a/litellm/main.py b/litellm/main.py index 17edafcdfca..a4a648acd4e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8749,6 +8749,39 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o setattr(usage, "cost", computed_cost) +_NON_TEXT_DELTA_FIELDS: Final = ( + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "annotations", + "audio", + "images", + "provider_specific_fields", +) + + +def _stream_choice_delta(choice: object) -> Mapping[str, object]: + delta: Final = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) + if isinstance(delta, Mapping): + return delta + if isinstance(delta, BaseModel): + return delta.model_dump() + return {} + + +def _delta_carries_more_than_text(delta: Mapping[str, object]) -> bool: + return any(delta.get(field) is not None for field in _NON_TEXT_DELTA_FIELDS) + + +def _simple_text_part(choices: Sequence[object]) -> str | None: + deltas: Final = tuple(_stream_choice_delta(choice) for choice in choices) + if any(_delta_carries_more_than_text(delta) for delta in deltas): + return None + content: Final = deltas[0].get("content") + return content if isinstance(content, str) else "" + + def stream_chunk_builder( chunks: list, messages: Sequence | None = None, @@ -8793,31 +8826,11 @@ def stream_chunk_builder( if not chunk.get("choices"): continue - choice = chunk["choices"][0] - delta_obj = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {}) - if isinstance(delta_obj, dict): - delta = delta_obj - elif hasattr(delta_obj, "model_dump"): - delta = cast(dict[str, Any], delta_obj.model_dump()) - else: - delta = {} - - if ( - delta.get("tool_calls") is not None - or delta.get("function_call") is not None - or delta.get("reasoning_content") is not None - or delta.get("thinking_blocks") is not None - or delta.get("annotations") is not None - or delta.get("audio") is not None - or delta.get("images") is not None - or delta.get("provider_specific_fields") is not None - ): + if (part := _simple_text_part(chunk["choices"])) is None: is_simple_text_stream = False break - - content = delta.get("content") - if isinstance(content, str) and content: - simple_content_parts.append(content) + if part: + simple_content_parts.append(part) if is_simple_text_stream: if simple_content_parts: @@ -8854,9 +8867,10 @@ def stream_chunk_builder( tool_call_chunks: Final = [ chunk for chunk in chunks - if chunk.get("choices") - and "tool_calls" in chunk["choices"][0]["delta"] - and chunk["choices"][0]["delta"]["tool_calls"] is not None + if any( + "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None + for choice in chunk.get("choices") or () + ) ] if len(tool_call_chunks) > 0: diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4f7a51eb531..42599f45ade 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1659,6 +1659,68 @@ async def test_async_mock_delay(): assert delay >= 0.01 +def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of_a_multi_choice_chunk(): + from litellm import stream_chunk_builder + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + def chunk(choices: list[StreamingChoices]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-multi-choice", + created=1751934860, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=choices, + ) + + chunks = [ + chunk( + [ + StreamingChoices(index=0, delta=Delta(role="assistant", content="hello")), + StreamingChoices( + index=1, + delta=Delta( + role="assistant", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + index=0, + type="function", + function=Function(name="lookup_fruit", arguments='{"fruit":'), + ) + ], + ), + ), + ] + ), + chunk( + [ + StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop"), + StreamingChoices( + index=1, + delta=Delta( + tool_calls=[ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"kiwi"}'))] + ), + finish_reason="tool_calls", + ), + ] + ), + ] + + response = stream_chunk_builder(chunks=chunks) + + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None + assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ + ("call_1", "lookup_fruit", '{"fruit":"kiwi"}') + ] + + def test_stream_chunk_builder_thinking_blocks(): from litellm import stream_chunk_builder from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices From 785c6cffc4826ef44c73a797981e28a294a45668 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Thu, 13 Aug 2026 23:34:40 -0400 Subject: [PATCH 039/393] fix(cost): carry image and video input tokens through the Responses usage bridge Realtime cost is computed from *_tokens_details after the usage round-trips through the Responses shape, and the input half of that shape carried audio only, so image and video prompt tokens stopped being billable as themselves. Vertex splits prompt tokens by modality, so a session sending camera frames arrives with image_tokens set. Those were folded into text_tokens and lost their attribution. The amount happens not to move today, because the calculator falls back to input_cost_per_token when no per-modality rate is set, but the tokens have to survive before any such rate can ever apply. InputTokensDetails now declares image_tokens and video_tokens instead of leaning on pydantic extras, the repeated per-field copying is a loop over the modality names so adding a modality no longer adds a branch, and the read-back in ResponseAPILoggingUtils picks up video_tokens, which PromptTokensDetailsWrapper already declared. The output half of the original change is dropped: 449c091391 landed the same OutputTokensDetails.audio_tokens fix upstream, with its own coverage in test_gemini_realtime_transformation.py, and it always sets output_tokens_details rather than only when non-empty. That structure is kept as upstream wrote it. --- .../transformation.py | 2 ++ litellm/responses/utils.py | 1 + litellm/types/llms/openai.py | 2 ++ .../test_litellm_completion_responses.py | 33 +++++++++++++++++++ 4 files changed, 38 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 01fb6cb483d..ffd7ce491b1 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2851,6 +2851,8 @@ class LiteLLMCompletionResponsesConfig: cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, text_tokens=prompt_details.text_tokens, audio_tokens=prompt_details.audio_tokens, + image_tokens=prompt_details.image_tokens, + video_tokens=prompt_details.video_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 41a3ded7022..f50c17aff85 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1182,6 +1182,7 @@ class ResponseAPILoggingUtils: cached_tokens_details=getattr( response_api_usage.input_tokens_details, "cached_tokens_details", None ), + video_tokens=getattr(response_api_usage.input_tokens_details, "video_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..98548705979 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1291,7 +1291,9 @@ class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 cached_tokens_details: CachedTokensDetails | None = None + image_tokens: int | None = None text_tokens: int | None = None + video_tokens: int | None = None model_config = {"extra": "allow"} diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..2f9f7adfcf1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2885,6 +2885,39 @@ class TestUsageTransformation: assert getattr(response_usage.input_tokens_details, "cache_write_tokens", None) == 800 assert response_usage.input_tokens_details.model_dump()["cache_write_tokens"] == 800 + def test_transform_usage_preserves_input_modality_tokens(self): + """Regression: the bridge dropped image and video input tokens. + + Vertex reports prompt tokens split by modality, so a Live session that sends + camera frames arrives with image_tokens set. InputTokensDetails declared only + audio/cached/text, so those tokens were folded into text and lost their + attribution, and any per-modality rate could never apply to them. + """ + usage = Usage( + prompt_tokens=300, + completion_tokens=10, + total_tokens=310, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, audio_tokens=80, image_tokens=150, video_tokens=50, cached_tokens=0 + ), + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), + ) + + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=usage + ) + details = response_usage.input_tokens_details + assert details is not None + assert getattr(details, "image_tokens", None) == 150 + assert getattr(details, "video_tokens", None) == 50 + assert getattr(details, "audio_tokens", None) == 80 + + from litellm.responses.utils import ResponseAPILoggingUtils + + back = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_usage.model_dump()) + assert back.prompt_tokens_details.image_tokens == 150 + assert back.prompt_tokens_details.video_tokens == 50 + def test_transform_usage_with_reasoning_tokens_gemini(self): """Test that reasoning_tokens from Gemini are properly transformed to output_tokens_details""" # Setup: Simulate Gemini usage with thoughtsTokenCount From 76488beaf8d1a44e0f07b6a4a66c06b9b4390222 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:55:25 -0700 Subject: [PATCH 040/393] fix(utils): reject an untranslatable tool_choice with a 400 instead of a 500 --- litellm/main.py | 2 +- litellm/utils.py | 16 +++- tests/litellm_utils_tests/test_utils.py | 6 +- .../test_validate_tool_choice.py | 74 ++++++++++--------- .../test_litellm_completion_responses.py | 15 ++++ tests/test_litellm/test_main.py | 14 ++++ 6 files changed, 85 insertions(+), 42 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index f6f4ec1bf63..9eea6779abf 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5102,7 +5102,7 @@ def completion( messages = validate_and_fix_openai_messages(messages=messages) tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice - tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice) + tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice, model=model) # validate optional params stop = validate_openai_optional_params(stop=stop) thinking = validate_and_fix_thinking_param(thinking=thinking) diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..17088f0475b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8038,6 +8038,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, + model: str, ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. @@ -8053,12 +8054,19 @@ def validate_chat_completion_tool_choice( # Standard OpenAI format: {"type": "function", "function": {...}} if tool_choice.get("type") is None or tool_choice.get("function") is None: - raise Exception( - f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec" + raise BadRequestError( + message=f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec", + model=model, + llm_provider="", ) return tool_choice - raise Exception( - f"Invalid tool choice, tool_choice={tool_choice}. Got={type(tool_choice)}. Expecting str, or dict. Please ensure tool_choice follows the OpenAI tool_choice spec" + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. Got={type(tool_choice)}. Expecting str, or dict. " + "Please ensure tool_choice follows the OpenAI tool_choice spec" + ), + model=model, + llm_provider="", ) diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 0ccfae55290..11d089719c6 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1334,10 +1334,10 @@ def test_validate_chat_completion_tool_choice(tool_choice, expected_bool): from litellm.utils import validate_chat_completion_tool_choice if expected_bool: - validate_chat_completion_tool_choice(tool_choice=tool_choice) + validate_chat_completion_tool_choice(tool_choice=tool_choice, model="gpt-5.6-sol") else: - with pytest.raises(Exception, match="Invalid tool choice"): - validate_chat_completion_tool_choice(tool_choice=tool_choice) + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"): + validate_chat_completion_tool_choice(tool_choice=tool_choice, model="gpt-5.6-sol") def test_models_by_provider(): diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index b8246fe0deb..b4272af7b90 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -1,60 +1,66 @@ +import re +from typing import Final + import pytest - +import litellm from litellm.utils import validate_chat_completion_tool_choice +MODEL: Final = "anthropic/claude-haiku-4-5" + def test_validate_tool_choice_none(): """Test that None is returned as-is.""" - result = validate_chat_completion_tool_choice(None) + result = validate_chat_completion_tool_choice(None, model=MODEL) assert result is None def test_validate_tool_choice_string(): """Test that string values are returned as-is.""" - assert validate_chat_completion_tool_choice("auto") == "auto" - assert validate_chat_completion_tool_choice("none") == "none" - assert validate_chat_completion_tool_choice("required") == "required" + assert validate_chat_completion_tool_choice("auto", model=MODEL) == "auto" + assert validate_chat_completion_tool_choice("none", model=MODEL) == "none" + assert validate_chat_completion_tool_choice("required", model=MODEL) == "required" def test_validate_tool_choice_standard_dict(): """Test standard OpenAI format with function.""" tool_choice = {"type": "function", "function": {"name": "my_function"}} - result = validate_chat_completion_tool_choice(tool_choice) + result = validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert result == tool_choice def test_validate_tool_choice_cursor_format(): """Cursor IDE format {"type": "auto"} is unwrapped to the bare string.""" - assert validate_chat_completion_tool_choice({"type": "auto"}) == "auto" - assert validate_chat_completion_tool_choice({"type": "none"}) == "none" - assert validate_chat_completion_tool_choice({"type": "required"}) == "required" + assert validate_chat_completion_tool_choice({"type": "auto"}, model=MODEL) == "auto" + assert validate_chat_completion_tool_choice({"type": "none"}, model=MODEL) == "none" + assert validate_chat_completion_tool_choice({"type": "required"}, model=MODEL) == "required" -def test_validate_tool_choice_invalid_dict(): - """Test that invalid dict formats raise exceptions.""" - # Missing both type and function - with pytest.raises(Exception, match='Invalid tool choice, tool_choice=\\{\\}\\. Please ensure') as exc_info: - validate_chat_completion_tool_choice({}) - assert "Invalid tool choice" in str(exc_info.value) - - # Invalid type value - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'invalid'\\}\\.") as exc_info: - validate_chat_completion_tool_choice({"type": "invalid"}) - assert "Invalid tool choice" in str(exc_info.value) - - # Has type but missing function when type is "function" - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'function'\\}\\.") as exc_info: - validate_chat_completion_tool_choice({"type": "function"}) - assert "Invalid tool choice" in str(exc_info.value) +@pytest.mark.parametrize( + "tool_choice", + [ + {}, + {"type": "invalid"}, + {"type": "function"}, + {"name": "lookup_fruit"}, + {"type": "file_search"}, + ], +) +def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): + """A dict shape chat completions cannot carry is the caller's mistake: a 400 that names the field, never a 500.""" + with pytest.raises( + litellm.BadRequestError, match=f"Invalid tool choice, tool_choice={re.escape(str(tool_choice))}\\. Please ensure" + ) as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == MODEL -def test_validate_tool_choice_invalid_type(): - """Test that invalid types raise exceptions.""" - with pytest.raises(Exception, match="\\. Expecting str, or dict\\. Please ensure") as exc_info: - validate_chat_completion_tool_choice(123) - assert "Got=" in str(exc_info.value) - - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\[\\]\\. Got=\\.") as exc_info: - validate_chat_completion_tool_choice([]) - assert "Got=" in str(exc_info.value) +@pytest.mark.parametrize("tool_choice", [123, []]) +def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): + """A non-str, non-dict tool_choice is rejected as a 400 that names the type it got.""" + with pytest.raises( + litellm.BadRequestError, match=f"Got={re.escape(str(type(tool_choice)))}\\. Expecting str, or dict\\." + ) as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..3950fd549ef 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4906,3 +4906,18 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +@pytest.mark.parametrize("stream", [True, False]) +async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool): + with pytest.raises(litellm.BadRequestError) as exc_info: + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Which fruit is red?", + tools=[{"type": "function", "name": "lookup_fruit", "parameters": {"type": "object"}}], + tool_choice={"type": "file_search"}, + stream=stream, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3dccb2b35bf..81ad161772e 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3850,3 +3850,17 @@ def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_ assert "extra_headers" not in body assert body["model"] == "gpt-5.4" assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS + + +@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) +def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Which fruit is red?"}], + tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert f"tool_choice={tool_choice}" in str(exc_info.value) From 2bbf34c6520ef3266a62121510470868f73e499e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:00:33 -0700 Subject: [PATCH 041/393] fix(utils): keep the tool_choice validator's model argument optional --- litellm/utils.py | 2 +- tests/litellm_utils_tests/test_validate_tool_choice.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 17088f0475b..c2dbbda2b68 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8038,7 +8038,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, - model: str, + model: str = "", ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index b4272af7b90..a9dacf9fa15 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -64,3 +64,11 @@ def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): ) as exc_info: validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert exc_info.value.status_code == 400 + + +def test_validate_tool_choice_without_model_is_still_a_400(): + """Callers that predate the model argument keep getting a 400, with an empty model on the error.""" + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc_info: + validate_chat_completion_tool_choice({"type": "bogus"}) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == "" From 61ac4f57394e022382459235baa598ac89ce00d1 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:03:56 +0000 Subject: [PATCH 042/393] test(e2e): add scripted-provider cost calculation suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 3 +- tests/e2e/conftest.py | 9 +- tests/e2e/cost_calculation/conftest.py | 139 ++++ tests/e2e/cost_calculation/cost_matrix.py | 458 +++++++++++++ tests/e2e/cost_calculation/scripted_client.py | 70 ++ .../e2e/cost_calculation/scripted_provider.py | 631 ++++++++++++++++++ .../test_token_pricing_e2e.py | 115 ++++ .../cost_calculation/test_wire_formats_e2e.py | 186 ++++++ tests/e2e/cost_map.json | 352 ++++++++++ .../coverage_registry/quota_management.yaml | 2 + tests/e2e/e2e_config.py | 16 + tests/e2e/pytest.ini | 1 + 12 files changed, 1980 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/cost_calculation/conftest.py create mode 100644 tests/e2e/cost_calculation/cost_matrix.py create mode 100644 tests/e2e/cost_calculation/scripted_client.py create mode 100644 tests/e2e/cost_calculation/scripted_provider.py create mode 100644 tests/e2e/cost_calculation/test_token_pricing_e2e.py create mode 100644 tests/e2e/cost_calculation/test_wire_formats_e2e.py create mode 100644 tests/e2e/cost_map.json diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0541ce25d4b..b6c3840f626 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,6 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` @@ -221,7 +222,7 @@ other... ``` ## Hard Rules -- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description +- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests; the one carve-out is a scripted upstream served through a real HTTP sidecar (the cost_calculation suite's scripted provider), allowed because provider-response-shape coverage needs a controlled usage payload and every hop from the proxy's upstream call to the spend row still executes for real. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description - use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index b1a75d5f862..7ab41b8ff68 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -22,9 +22,9 @@ from typing import Final import pytest import requests - from e2e_config import ( CONTROL_PLANE_BASE_URL, + COST_MAP_OPT_IN_ENV, FIXTURE_DIR, FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, @@ -53,6 +53,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "cost_map_stack": COST_MAP_OPT_IN_ENV, } ) @@ -120,6 +121,12 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set", ) + config.addinivalue_line( + "markers", + "cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json " + "(LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless " + "E2E_COST_MAP_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py new file mode 100644 index 00000000000..1bba3d50e1d --- /dev/null +++ b/tests/e2e/cost_calculation/conftest.py @@ -0,0 +1,139 @@ +"""Cost-calculation suite fixtures. + +Runs against a dedicated proxy whose whole model cost map is the test-owned +``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL), so every deployment +bills at rates the test asserts literal arithmetic on. Provider calls are +answered by the scripted-provider sidecar (``scripted_provider.py``), registered +per scenario over its control API. + +Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). +""" + +from __future__ import annotations + +import importlib.util +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Final, Protocol, cast + +import pytest + +from cost_matrix import Case, FrontierModel +from e2e_config import COST_MAP_PROXY_URL +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody +from proxy_client import ProxyClient, build_proxy_client +from scripted_client import ScenarioHandle, delete_scenario, register_scenario +from scripted_provider import Scenario + + +def _load_cost_rows() -> ModuleType: + """Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree + has no package layout), the same trick the mcp suite uses for + logging/datadog_reader.py.""" + path = ( + Path(__file__).resolve().parent.parent + / "quota_management" + / "spend_tracking" + / "cost_rows.py" + ) + name = "e2e_spend_tracking_cost_rows" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +class SpendCostBreakdown(Protocol): + input_cost: float | None + output_cost: float | None + cache_read_cost: float | None + cache_creation_cost: float | None + reasoning_cost: float | None + tool_usage_cost: float | None + total_cost: float | None + service_tier: str | None + + def model_dump(self) -> dict[str, object]: ... + + +class SpendRowMetadata(Protocol): + cost_breakdown: SpendCostBreakdown | None + + +class SpendCostRow(Protocol): + """The slice of spend_tracking.cost_rows.CostRow this suite reads.""" + + spend: float | None + prompt_tokens: int | None + completion_tokens: int | None + metadata: SpendRowMetadata | None + + @property + def breakdown(self) -> SpendCostBreakdown: ... + + +class CostRowsModule(Protocol): + """cost_rows.py loaded by path has no importable name for basedpyright, so + its surface is declared here and reached through a single cast.""" + + approx_equal: Callable[[float, float], bool] + assert_total_is_sum_of_components: Callable[[SpendCostRow], None] + poll_cost_row_where: Callable[ + [ProxyClient, str, Callable[[SpendCostRow], bool]], SpendCostRow | None + ] + + +cost_rows: Final[CostRowsModule] = cast(CostRowsModule, _load_cost_rows()) + + +@dataclass(frozen=True, slots=True) +class CostCalcClient: + """The suite's client: a ProxyClient pointed at the cost-map proxy pod.""" + + proxy: ProxyClient + + +@pytest.fixture(scope="session") +def client() -> CostCalcClient: + proxy = build_proxy_client( + base_url=COST_MAP_PROXY_URL, + control_plane_base_url=COST_MAP_PROXY_URL, + replica_urls=(COST_MAP_PROXY_URL,), + ) + return CostCalcClient(proxy=proxy) + + +def register_scenario_deployment( + client: CostCalcClient, + resources: ResourceManager, + model: FrontierModel, + case: Case, + marker: str, +) -> tuple[str, ScenarioHandle]: + """Register the case's scenario on the sidecar plus a deployment pointed at + it; both are torn down by ``resources``. Returns the callable model_name.""" + scenario: Scenario = case.scenario( + scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}" + ) + handle = register_scenario(scenario) + resources.defer(lambda: delete_scenario(handle)) + model_name = f"{model.model_name}-{marker}" + model_id = client.proxy.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody( + model=model.litellm_model, + api_key="sk-scripted-provider", + api_base=handle.api_base(), + ), + model_info=ModelInfoBody(), + ) + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model_name, handle diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py new file mode 100644 index 00000000000..bc466d7d823 --- /dev/null +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -0,0 +1,458 @@ +"""The cost-calculation matrix: frontier model set, the pricing-component cases +each model runs, and the expected-cost arithmetic. + +Rates come from ``tests/e2e/cost_map.json``, which the proxy under test loads as +its ENTIRE model cost map (LITELLM_MODEL_COST_MAP_URL), so an entry's rates are +exactly what the proxy bills and nothing in the suite depends on the bundled +map. Each model's rates are a distinct multiple of a shared base set, so a +component billed at the wrong model's rate (or the wrong case's rate) can never +coincidentally match. + +Case applicability is pricing-field-gated AND wire-gated: a case runs for a +model only when the entry carries the rate the case exercises and the wire can +report the token kind that rate prices. When the wire cannot report a kind +(e.g. Anthropic has no reasoning-token field, Responses reports no cache +creation), the case is absent from the matrix rather than silently zero. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from scripted_provider import Scenario, ScriptedOutput, ScriptedUsage, Wire + +COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" + + +class SearchContextCostPerQuery(BaseModel): + model_config = ConfigDict(frozen=True) + + search_context_size_low: float | None = None + search_context_size_medium: float | None = None + search_context_size_high: float | None = None + + +class CostMapEntry(BaseModel): + """The pricing fields of a cost-map entry the matrix reads. Shaped like a + ``model_prices_and_context_window.json`` entry; unmodelled keys are ignored.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_provider: str + mode: str + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + cache_creation_input_token_cost_above_1hr: float | None = None + output_cost_per_reasoning_token: float | None = None + input_cost_per_audio_token: float | None = None + output_cost_per_audio_token: float | None = None + input_cost_per_token_above_200k_tokens: float | None = None + output_cost_per_token_above_200k_tokens: float | None = None + input_cost_per_token_flex: float | None = None + output_cost_per_token_flex: float | None = None + input_cost_per_token_priority: float | None = None + output_cost_per_token_priority: float | None = None + search_context_cost_per_query: SearchContextCostPerQuery | None = None + web_search_billing_unit: str | None = None + + +_COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) +_COST_MAP: Final[dict[str, CostMapEntry]] = _COST_MAP_ADAPTER.validate_python( + json.loads(COST_MAP_PATH.read_text()) +) + +TIER_THRESHOLD_TOKENS: Final = 200_000 + + +@dataclass(frozen=True, slots=True) +class FrontierModel: + """One deployment under test: the model_name the suite registers, the + provider-prefixed litellm model string, the wire the scripted upstream + speaks, its cost-map key, and the sibling map model the response_model + override case reports.""" + + model_name: str + litellm_model: str + wire: Wire + map_key: str + override_model: str + + @property + def rates(self) -> CostMapEntry: + return _COST_MAP[self.map_key] + + @property + def override_rates(self) -> CostMapEntry: + return _COST_MAP[self.override_map_key] + + @property + def override_map_key(self) -> str: + return _OVERRIDE_MAP_KEYS[self.override_model] + + @property + def provider(self) -> str: + return self.rates.litellm_provider + + @property + def api_key(self) -> str: + # The scripted upstream ignores auth; a fixed bogus key proves the suite + # spends zero real provider calls. + return "sk-scripted-provider" + + +# Response-model override targets: emit a sibling's bare provider-facing name so +# the biller's provider-prefixed lookup lands on that sibling's map key. +_OVERRIDE_MODELS: Final[dict[str, str]] = { + "gpt-5.6": "gpt-5.4-mini", + "gpt-5.5-pro": "gpt-5.3-codex", + "gpt-5.3-codex": "gpt-5.5-pro", + "gpt-5.4-mini": "gpt-5.6", + "claude-opus-5": "claude-sonnet-5", + "claude-sonnet-5": "claude-opus-5", + "claude-haiku-4-5": "claude-sonnet-5", + "gemini/gemini-3.8-flash": "gemini-3.1-pro-preview", + "gemini/gemini-3.1-pro-preview": "gemini-3.8-flash", + "together_ai/moonshotai/Kimi-K3": "zai-org/GLM-5.3", + "together_ai/zai-org/GLM-5.3": "moonshotai/Kimi-K3", + "fireworks_ai/kimi-k3": "qwen3p8-max", + "fireworks_ai/qwen3p8-max": "kimi-k3", + "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", +} + +_OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { + "gpt-5.4-mini": "gpt-5.4-mini", + "gpt-5.6": "gpt-5.6", + "gpt-5.3-codex": "gpt-5.3-codex", + "gpt-5.5-pro": "gpt-5.5-pro", + "claude-sonnet-5": "claude-sonnet-5", + "claude-opus-5": "claude-opus-5", + "gemini-3.1-pro-preview": "gemini/gemini-3.1-pro-preview", + "gemini-3.8-flash": "gemini/gemini-3.8-flash", + "zai-org/GLM-5.3": "together_ai/zai-org/GLM-5.3", + "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", + "qwen3p8-max": "fireworks_ai/qwen3p8-max", + "kimi-k3": "fireworks_ai/kimi-k3", +} + + +_FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( + ("gpt-5.6", "openai/gpt-5.6", "openai_chat"), + ("gpt-5.5-pro", "openai/gpt-5.5-pro", "openai_responses"), + ("gpt-5.3-codex", "openai/gpt-5.3-codex", "openai_responses"), + ("gpt-5.4-mini", "openai/gpt-5.4-mini", "openai_chat"), + ("claude-opus-5", "anthropic/claude-opus-5", "anthropic_messages"), + ("claude-sonnet-5", "anthropic/claude-sonnet-5", "anthropic_messages"), + ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "anthropic_messages"), + ("gemini/gemini-3.8-flash", "gemini/gemini-3.8-flash", "gemini_generate"), + ("gemini/gemini-3.1-pro-preview", "gemini/gemini-3.1-pro-preview", "gemini_generate"), + ("together_ai/moonshotai/Kimi-K3", "together_ai/moonshotai/Kimi-K3", "together_chat"), + ("together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3", "together_chat"), + ("fireworks_ai/kimi-k3", "fireworks_ai/kimi-k3", "fireworks_chat"), + ("fireworks_ai/qwen3p8-max", "fireworks_ai/qwen3p8-max", "fireworks_chat"), + ("fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/deepseek-v4p1-flash", "fireworks_chat"), +) + + +def _frontier() -> tuple[FrontierModel, ...]: + return tuple( + FrontierModel( + model_name=f"cc-{map_key.replace('/', '-').lower()}", + litellm_model=litellm_model, + wire=wire, + map_key=map_key, + override_model=_OVERRIDE_MODELS[map_key], + ) + for map_key, litellm_model, wire in _FRONTIER_SPECS + ) + + +FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() + +# Token kinds each wire can report, gating which pricing cases apply. +_WIRE_CAPS: Final[dict[str, frozenset[str]]] = { + "openai_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), + "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), + # Product gap: litellm hard-indexes message_delta["usage"] in + # anthropic/chat/handler.py, so a usage-absent anthropic stream raises + # KeyError; the real wire always carries it, so the case cannot be + # represented. + "anthropic_messages": frozenset({"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model"}), + # Product gap: the gemini transform sets ModelResponse.model from the + # request and drops the provider's modelVersion, so a response-model + # override can never be priced on this wire. + "gemini_generate": frozenset({"cache_read", "reasoning", "audio", "web_search", "absent_usage"}), + "together_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), + "fireworks_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), +} + +CaseName = Literal[ + "basic", + "cache_read", + "cache_write_5m", + "cache_write_1h", + "reasoning", + "audio", + "tiered", + "service_tier_flex", + "service_tier_priority", + "web_search", + "stream", + "stream_no_usage", + "response_model_override", +] + + +@dataclass(frozen=True, slots=True) +class Case: + name: CaseName + usage: ScriptedUsage + stream: bool = False + stream_usage: Literal["final_chunk", "absent"] = "final_chunk" + service_tier: Literal["flex", "priority"] | None = None + # For web_search the wire's reported call count is not always what gets + # billed: chat-completions surfaces only expose url_citation annotations, so + # the biller floors to one call; responses/messages/gemini report a real + # count. + billed_web_search_calls: int = 0 + response_model_override: bool = False + exact_spend: bool = True + # stream_usage=absent on a wire with no proxy-side token recount means the + # bill is exactly zero; asserted as such rather than skipped. + expect_zero_bill: bool = False + + def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: + return Scenario( + scenario_id=scenario_id, + wire=model.wire, + usage=self.usage, + output=ScriptedOutput( + text=text, + response_model=model.override_model if self.response_model_override else None, + ), + stream_usage=self.stream_usage, + service_tier=self.service_tier, + ) + + +_BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) + + +def _web_search_case(model: FrontierModel) -> Case: + counts_exactly = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + return Case( + name="web_search", + usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), + billed_web_search_calls=3 if counts_exactly else 1, + ) + + +def cases_for(model: FrontierModel) -> tuple[Case, ...]: + rates = model.rates + caps = _WIRE_CAPS[model.wire] + cases: list[Case] = [Case(name="basic", usage=_BASIC_USAGE)] + if rates.cache_read_input_token_cost is not None and "cache_read" in caps: + cases.append( + Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) + ) + if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps: + cases.append( + Case( + name="cache_write_5m", + usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), + ) + ) + if ( + rates.cache_creation_input_token_cost_above_1hr is not None + and rates.cache_creation_input_token_cost is not None + and "cache_write_1h" in caps + ): + cases.append( + Case( + name="cache_write_1h", + usage=ScriptedUsage( + fresh_input_tokens=90, + cache_write_5m_tokens=20, + cache_write_1h_tokens=40, + output_tokens=30, + ), + ) + ) + if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps: + cases.append( + Case( + name="reasoning", + usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), + ) + ) + if ( + rates.input_cost_per_audio_token is not None + and rates.output_cost_per_audio_token is not None + and "audio" in caps + ): + cases.append( + Case( + name="audio", + usage=ScriptedUsage( + fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 + ), + ) + ) + if ( + rates.input_cost_per_token_above_200k_tokens is not None + and rates.output_cost_per_token_above_200k_tokens is not None + ): + cases.append( + Case( + name="tiered", + usage=ScriptedUsage( + fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 + ), + ) + ) + if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None: + cases.append( + Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") + ) + if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None: + cases.append( + Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") + ) + if rates.search_context_cost_per_query is not None and "web_search" in caps: + cases.append(_web_search_case(model)) + cases.append(Case(name="stream", usage=_BASIC_USAGE, stream=True)) + if "absent_usage" in caps: + cases.append( + Case( + name="stream_no_usage", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + exact_spend=False, + # The responses surface bills only provider-reported usage; + # with no usage in the stream the spend row is zero. Other + # wires recount tokens proxy-side and bill a nonzero amount. + expect_zero_bill=model.wire == "openai_responses", + ) + ) + if "response_model" in caps: + cases.append(Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True)) + return tuple(cases) + + +@dataclass(frozen=True, slots=True) +class ExpectedCost: + """The expected bill split the way the spend row's cost_breakdown reports + it: the gross input component (cache reads/writes folded in), the output + component, and the tool-usage component.""" + + input_cost: float + output_cost: float + tool_cost: float + + @property + def total(self) -> float: + return self.input_cost + self.output_cost + self.tool_cost + + +def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: + """Literal arithmetic on the test-map rates over the scripted token counts. + + Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; + output = text*out + reasoning*reasoning + audio_out*audio_out; plus the + billed web-search calls at the medium search-context rate. Above-threshold + swaps every input/output rate to its ``_above_200k_tokens`` variant when + total prompt tokens exceed the threshold; a service tier swaps input/output + to the tier's variants, falling back to the base rate when a variant is + unset -- mirroring _get_token_base_cost in litellm's cost calculator. + """ + rates = model.override_rates if case.response_model_override else model.rates + u = case.usage + prompt_tokens = ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + u.audio_input_tokens + ) + tiered = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate = rates.input_cost_per_token or 0.0 + out_rate = rates.output_cost_per_token or 0.0 + if case.service_tier == "flex": + in_rate = rates.input_cost_per_token_flex or in_rate + out_rate = rates.output_cost_per_token_flex or out_rate + if case.service_tier == "priority": + in_rate = rates.input_cost_per_token_priority or in_rate + out_rate = rates.output_cost_per_token_priority or out_rate + if tiered: + in_rate = rates.input_cost_per_token_above_200k_tokens or in_rate + out_rate = rates.output_cost_per_token_above_200k_tokens or out_rate + input_cost = ( + u.fresh_input_tokens * in_rate + + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + ) + output_cost = ( + u.output_tokens * out_rate + + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + ) + search = rates.search_context_cost_per_query + tool_cost = case.billed_web_search_calls * ( + search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + ) + return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) + + +def expected_cost(model: FrontierModel, case: Case) -> float: + return expected_breakdown(model, case).total + + +def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: + """(prompt_tokens, completion_tokens) the spend row should carry, per the + wire's normalization: Anthropic folds cache read/write into prompt_tokens, + everyone else reports the totals the wire emitted.""" + u = case.usage + if model.wire == "anthropic_messages": + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, + u.output_tokens, + ) + if model.wire == "gemini_generate": + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + if model.wire == "openai_responses": + return ( + u.fresh_input_tokens + u.cache_read_tokens, + u.output_tokens + u.reasoning_tokens, + ) + return ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) diff --git a/tests/e2e/cost_calculation/scripted_client.py b/tests/e2e/cost_calculation/scripted_client.py new file mode 100644 index 00000000000..dceec02630a --- /dev/null +++ b/tests/e2e/cost_calculation/scripted_client.py @@ -0,0 +1,70 @@ +"""Client side of the scripted-provider sidecar: register scenarios over its +control API through the shared transport helpers and get back a handle whose +``api_base`` is what a /model/new deployment should register for the proxy to +reach the scripted wire.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BASE +from e2e_http import URL, NoBody, unwrap, post +from e2e_http import delete as http_delete +from scripted_provider import ( + Scenario, + ScenarioDeleted, + ScenarioRegistered, + Wire, +) + + +@dataclass(frozen=True, slots=True) +class ScenarioHandle: + scenario_id: str + wire: Wire + proxy_base: str + + def api_base(self) -> str: + return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}" + + def _mount(self) -> str: + return { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", + }[self.wire] + + +def register_scenario(scenario: Scenario) -> ScenarioHandle: + """POST the scenario to the sidecar's control API and return its handle.""" + result = unwrap( + post( + URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"), + headers=NoBody(), + json=scenario, + response_type=ScenarioRegistered, + ) + ) + return ScenarioHandle( + scenario_id=result.scenario_id, + wire=scenario.wire, + proxy_base=SCRIPTED_PROVIDER_PROXY_BASE, + ) + + +def delete_scenario(handle: ScenarioHandle) -> None: + unwrap( + http_delete( + URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios/{handle.scenario_id}"), + headers=NoBody(), + json=NoBody(), + response_type=ScenarioDeleted, + ) + ) + + +CONTROL_URL: Final = SCRIPTED_PROVIDER_CONTROL_URL diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py new file mode 100644 index 00000000000..93a6f49ec25 --- /dev/null +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -0,0 +1,631 @@ +"""Scripted provider sidecar for the cost-calculation e2e suite. + +A standalone process (``python -m cost_calculation.scripted_provider``) that +pretends to be an LLM provider for the proxy under test. The suite registers a +Scenario over a small control API; the provider wire routes then answer the +proxy's upstream calls with the scripted usage figures, in the exact wire shape +the real provider would emit (OpenAI chat completions, OpenAI Responses, +Anthropic Messages, Gemini generateContent, or the OpenAI-compatible Together / +Fireworks surfaces). Because the usage is scripted, expected spend is literal +arithmetic on the test cost map's rates, with no dependency on what a real +provider would report. + +Layout on one port: + +- ``GET /health`` liveness +- ``POST /_scenarios`` register a Scenario JSON, returns its id +- ``DELETE /_scenarios/`` remove it +- ``POST ///`` provider wire; mount is one of + ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks`` and the + remainder is whatever path the provider client appends (``chat/completions``, + ``responses``, ``v1/messages``, ``models/:generateContent`` ...) + +A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini +verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the +final stream chunk carries usage or the provider reports none. +""" + +from __future__ import annotations + +import json +import sys +import threading +import time +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, Literal +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +Wire = Literal[ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate", + "together_chat", + "fireworks_chat", +] + +_WIRE_MOUNTS: Final[dict[str, str]] = { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", +} + +StreamUsage = Literal["final_chunk", "absent"] +ServiceTier = Literal["flex", "priority"] + + +class ScriptedUsage(BaseModel): + """Physical token counts the scripted response reports. ``fresh_input_tokens`` + is the uncached, never-written, non-audio input count; ``output_tokens`` is + the non-reasoning, non-audio output count. Renderers add the cached, written, + audio, and reasoning counts into the wire's total fields the way the real + provider does (inside prompt_tokens for OpenAI/Gemini, as uncached-only + input_tokens for Anthropic).""" + + model_config = ConfigDict(frozen=True) + + fresh_input_tokens: int = 0 + output_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_5m_tokens: int = 0 + cache_write_1h_tokens: int = 0 + reasoning_tokens: int = 0 + audio_input_tokens: int = 0 + audio_output_tokens: int = 0 + web_search_calls: int = 0 + + +class ScriptedOutput(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str + finish_reason: str = "stop" + # When set, emitted verbatim as the response's model field, letting a test + # prove the biller prices the provider-reported model. + response_model: str | None = None + # OpenAI-compatible providers can report a provider-computed cost; emitted as + # the top-level "cost" field on the together/fireworks wire. + provider_cost: float | None = None + + +class Scenario(BaseModel): + model_config = ConfigDict(frozen=True) + + scenario_id: str + wire: Wire + usage: ScriptedUsage + output: ScriptedOutput + stream_usage: StreamUsage = "final_chunk" + service_tier: ServiceTier | None = None + + @property + def mount(self) -> str: + return _WIRE_MOUNTS[self.wire] + + +class ScenarioRegistered(BaseModel): + scenario_id: str + + +class ScenarioDeleted(BaseModel): + deleted: bool + + +class HealthStatus(BaseModel): + status: str + + +@dataclass(frozen=True, slots=True) +class RenderedResponse: + status_code: int + content_type: str + body: bytes + + +def _json_bytes(payload: dict[str, object]) -> bytes: + return json.dumps(payload).encode("utf-8") + + +def _sse(events: tuple[tuple[str | None, dict[str, object] | str], ...]) -> bytes: + frames: list[str] = [] + for event_name, data in events: + head = f"event: {event_name}\n" if event_name is not None else "" + payload = data if isinstance(data, str) else json.dumps(data) + frames.append(f"{head}data: {payload}\n\n") + return "".join(frames).encode("utf-8") + + +# ---------- per-wire usage shapes ---------- + + +def _openai_usage(u: ScriptedUsage) -> dict[str, object]: + prompt_tokens = ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens + ) + completion_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + prompt_details: dict[str, object] = {} + if u.cache_read_tokens: + prompt_details["cached_tokens"] = u.cache_read_tokens + if u.cache_write_5m_tokens or u.cache_write_1h_tokens: + prompt_details["cache_write_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens + prompt_details["cache_creation_token_details"] = { + "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, + "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, + } + if u.audio_input_tokens: + prompt_details["audio_tokens"] = u.audio_input_tokens + completion_details: dict[str, object] = {} + if u.reasoning_tokens: + completion_details["reasoning_tokens"] = u.reasoning_tokens + if u.audio_output_tokens: + completion_details["audio_tokens"] = u.audio_output_tokens + usage: dict[str, object] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + if prompt_details: + usage["prompt_tokens_details"] = prompt_details + if completion_details: + usage["completion_tokens_details"] = completion_details + return usage + + +def _anthropic_usage(u: ScriptedUsage) -> dict[str, object]: + # Anthropic reports uncached-only input_tokens; cache reads and writes ride + # top-level fields, with the 5m/1h write split under cache_creation. + usage: dict[str, object] = { + "input_tokens": u.fresh_input_tokens, + "output_tokens": u.output_tokens, + } + if u.cache_read_tokens: + usage["cache_read_input_tokens"] = u.cache_read_tokens + if u.cache_write_5m_tokens or u.cache_write_1h_tokens: + usage["cache_creation_input_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens + usage["cache_creation"] = { + "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, + "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, + } + if u.web_search_calls: + usage["server_tool_use"] = {"web_search_requests": u.web_search_calls} + return usage + + +def _gemini_usage(u: ScriptedUsage) -> dict[str, object]: + # promptTokenCount carries the cached count inside it; TEXT modality is the + # cached-inclusive text count so litellm's implicit-caching subtraction lands + # on the fresh figure. candidatesTokenCount includes reasoning + audio. + prompt_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + candidates = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + usage: dict[str, object] = { + "promptTokenCount": prompt_tokens, + "candidatesTokenCount": candidates, + "totalTokenCount": prompt_tokens + candidates, + } + if u.cache_read_tokens: + usage["cachedContentTokenCount"] = u.cache_read_tokens + if u.reasoning_tokens: + usage["thoughtsTokenCount"] = u.reasoning_tokens + prompt_details = [{"modality": "TEXT", "tokenCount": u.fresh_input_tokens + u.cache_read_tokens}] + if u.audio_input_tokens: + prompt_details.append({"modality": "AUDIO", "tokenCount": u.audio_input_tokens}) + usage["promptTokensDetails"] = prompt_details + if u.audio_output_tokens: + usage["candidatesTokensDetails"] = [ + {"modality": "TEXT", "tokenCount": u.output_tokens + u.reasoning_tokens}, + {"modality": "AUDIO", "tokenCount": u.audio_output_tokens}, + ] + return usage + + +def _responses_usage(u: ScriptedUsage) -> dict[str, object]: + input_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + output_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + usage: dict[str, object] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + input_details: dict[str, object] = {} + if u.cache_read_tokens: + input_details["cached_tokens"] = u.cache_read_tokens + if input_details: + usage["input_tokens_details"] = input_details + if u.reasoning_tokens: + usage["output_tokens_details"] = {"reasoning_tokens": u.reasoning_tokens} + return usage + + +# ---------- per-wire responses ---------- + + +def _openai_message(scenario: Scenario) -> dict[str, object]: + message: dict[str, object] = {"role": "assistant", "content": scenario.output.text} + if scenario.usage.web_search_calls: + message["annotations"] = [ + { + "type": "url_citation", + "url_citation": { + "url": "https://scripted.example/source", + "title": "scripted source", + "start_index": 0, + "end_index": 1, + }, + } + for _ in range(scenario.usage.web_search_calls) + ] + return message + + +def _openai_chat_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + body: dict[str, object] = { + "id": f"chatcmpl-{scenario.scenario_id}", + "object": "chat.completion", + "created": int(time.time()), + "model": scenario.output.response_model or requested_model, + "choices": [ + { + "index": 0, + "message": _openai_message(scenario), + "finish_reason": scenario.output.finish_reason, + } + ], + "usage": _openai_usage(scenario.usage), + } + if scenario.service_tier is not None: + body["service_tier"] = scenario.service_tier + if scenario.output.provider_cost is not None: + body["cost"] = scenario.output.provider_cost + return body + + +def _openai_chunk(scenario: Scenario, requested_model: str, **kw: object) -> dict[str, object]: + chunk: dict[str, object] = { + "id": f"chatcmpl-{scenario.scenario_id}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": scenario.output.response_model or requested_model, + } + chunk.update(kw) + return chunk + + +def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: + _EMPTY_DELTA: Final[dict[str, object]] = {} + delta: dict[str, object] = {"role": "assistant", "content": scenario.output.text} + if scenario.usage.web_search_calls: + delta["annotations"] = _openai_message(scenario)["annotations"] + events: list[tuple[str | None, dict[str, object] | str]] = [ + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + ), + ), + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ), + ), + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[ + { + "index": 0, + "delta": _EMPTY_DELTA, + "finish_reason": scenario.output.finish_reason, + } + ], + ), + ), + ] + if scenario.stream_usage == "final_chunk": + events.append( + (None, _openai_chunk(scenario, requested_model, choices=(), usage=_openai_usage(scenario.usage))) + ) + events.append((None, "[DONE]")) + return _sse(tuple(events)) + + +def _anthropic_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + return { + "id": f"msg_{scenario.scenario_id}", + "type": "message", + "role": "assistant", + "model": scenario.output.response_model or requested_model, + "content": [{"type": "text", "text": scenario.output.text}], + "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + "usage": _anthropic_usage(scenario.usage), + } + + +def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: + emit_usage = scenario.stream_usage == "final_chunk" + input_usage = {k: v for k, v in _anthropic_usage(scenario.usage).items() if k != "output_tokens"} + message_start: dict[str, object] = { + "type": "message_start", + "message": { + "id": f"msg_{scenario.scenario_id}", + "type": "message", + "role": "assistant", + "model": scenario.output.response_model or requested_model, + "content": [], + "stop_reason": None, + **({"usage": input_usage} if emit_usage else {}), + }, + } + message_delta: dict[str, object] = { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + }, + **({"usage": {"output_tokens": scenario.usage.output_tokens}} if emit_usage else {}), + } + return _sse( + ( + ("message_start", message_start), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": scenario.output.text}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", message_delta), + ("message_stop", {"type": "message_stop"}), + ) + ) + + +def _gemini_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + candidate: dict[str, object] = { + "content": {"parts": [{"text": scenario.output.text}], "role": "model"}, + "finishReason": "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), + "index": 0, + } + if scenario.usage.web_search_calls: + candidate["groundingMetadata"] = { + "webSearchQueries": [f"query {i}" for i in range(scenario.usage.web_search_calls)] + } + return { + "candidates": [candidate], + "usageMetadata": _gemini_usage(scenario.usage), + "modelVersion": scenario.output.response_model or requested_model, + } + + +def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: + first = _gemini_body(scenario, requested_model) + if scenario.stream_usage == "absent": + first = {k: v for k, v in first.items() if k != "usageMetadata"} + events: list[tuple[str | None, dict[str, object] | str]] = [(None, first)] + if scenario.stream_usage == "final_chunk": + events.append( + ( + None, + { + "candidates": [], + "usageMetadata": _gemini_usage(scenario.usage), + "modelVersion": scenario.output.response_model or requested_model, + }, + ) + ) + return _sse(tuple(events)) + + +def _responses_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + output: list[dict[str, object]] = [ + {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} + for i in range(scenario.usage.web_search_calls) + ] + output.append( + { + "type": "message", + "id": f"msg_{scenario.scenario_id}", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": scenario.output.text, + "annotations": [], + } + ], + } + ) + return { + "id": f"resp_{scenario.scenario_id}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": scenario.output.response_model or requested_model, + "output": output, + "usage": _responses_usage(scenario.usage), + } + + +def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: + completed = _responses_body(scenario, requested_model) + if scenario.stream_usage == "absent": + completed = {k: v for k, v in completed.items() if k != "usage"} + created = {**completed, "status": "in_progress", "usage": None} + return _sse( + ( + ("response.created", {"type": "response.created", "response": created}), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": f"msg_{scenario.scenario_id}", + "output_index": scenario.usage.web_search_calls, + "content_index": 0, + "delta": scenario.output.text, + }, + ), + ("response.completed", {"type": "response.completed", "response": completed}), + ) + ) + + +def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> RenderedResponse: + if scenario.wire == "anthropic_messages": + if stream: + return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_anthropic_body(scenario, requested_model))) + if scenario.wire == "gemini_generate": + if stream: + return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model))) + if scenario.wire == "openai_responses": + if stream: + return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model))) + # openai_chat, together_chat, fireworks_chat share the OpenAI chat shape. + if stream: + return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model))) + + +# ---------- registry + request routing ---------- + + +class _ScenarioStore: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self._scenarios: dict[str, Scenario] = {} # mutable-ok: server state, guarded by _lock + + def put(self, scenario: Scenario) -> None: + with self._lock: + self._scenarios[scenario.scenario_id] = scenario + + def drop(self, scenario_id: str) -> bool: + with self._lock: + return self._scenarios.pop(scenario_id, None) is not None + + def get(self, scenario_id: str) -> Scenario | None: + with self._lock: + return self._scenarios.get(scenario_id) + + +_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +def _request_body(body: bytes) -> dict[str, object]: + try: + return _REQUEST_BODY.validate_json(body) + except ValueError: + return {} + + +def _request_wants_stream(path_tail: str, body: bytes) -> bool: + if ":streamGenerateContent" in path_tail: + return True + if not body: + return False + return _request_body(body).get("stream") is True + + +def _request_model(body: bytes) -> str: + model = _request_body(body).get("model") + return model if isinstance(model, str) else "unknown" + + +def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: + path = urlsplit(raw_path).path + segments = [segment for segment in path.split("/") if segment] + if method == "GET" and segments == ["health"]: + return RenderedResponse(200, "application/json", _json_bytes({"status": "ok"})) + if segments and segments[0] == "_scenarios": + if method == "POST" and len(segments) == 1: + try: + scenario = Scenario.model_validate_json(body) + except ValidationError as exc: + return RenderedResponse(400, "application/json", _json_bytes({"error": str(exc)})) + store.put(scenario) + return RenderedResponse(200, "application/json", _json_bytes({"scenario_id": scenario.scenario_id})) + if method == "DELETE" and len(segments) == 2: + deleted = store.drop(segments[1]) + return RenderedResponse( + 200 if deleted else 404, "application/json", _json_bytes({"deleted": deleted}) + ) + return RenderedResponse(404, "application/json", _json_bytes({"error": "unknown control route"})) + if len(segments) < 2 or method != "POST": + return RenderedResponse(404, "application/json", _json_bytes({"error": f"no route for {method} {path}"})) + scenario_id, mount = segments[0], segments[1] + scenario = store.get(scenario_id) + if scenario is None: + return RenderedResponse(404, "application/json", _json_bytes({"error": f"unknown scenario {scenario_id}"})) + if scenario.mount != mount: + return RenderedResponse( + 400, + "application/json", + _json_bytes({"error": f"scenario {scenario_id} is wire {scenario.wire}, not mount {mount}"}), + ) + tail = "/".join(segments[2:]) + return _render(scenario, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + + +class _ScriptedHandler(BaseHTTPRequestHandler): + store: Final[_ScenarioStore] = _ScenarioStore() + + def _dispatch(self, method: str) -> None: + length = int(self.headers.get("content-length") or 0) + body = self.rfile.read(length) if length else b"" + rendered = handle_request(self.store, method, self.path, body) + self.send_response(rendered.status_code) + self.send_header("content-type", rendered.content_type) + self.send_header("content-length", str(len(rendered.body))) + self.end_headers() + self.wfile.write(rendered.body) + + def do_GET(self) -> None: + self._dispatch("GET") + + def do_POST(self) -> None: + self._dispatch("POST") + + def do_DELETE(self) -> None: + self._dispatch("DELETE") + + + +DEFAULT_PORT: Final = 9100 + + +def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None: + server = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) + sys.stderr.write(f"scripted-provider listening on http://{bind_host}:{port}\n") + server.serve_forever() + + +if __name__ == "__main__": + port_arg = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT + serve(port=port_arg) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py new file mode 100644 index 00000000000..8d7678cf9ca --- /dev/null +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -0,0 +1,115 @@ +"""Token-pricing e2e: every (frontier model, pricing-component case) cell runs a +scripted-usage call through a deployment registered on the cost-map proxy, and +the spend row plus response-cost header must equal literal arithmetic on the +test map's rates. + +Nothing here touches a real provider or the bundled cost map: the proxy's +upstream is the scripted-provider sidecar and its entire cost map is +tests/e2e/cost_map.json. +""" + +from __future__ import annotations + +import pytest + +from conftest import CostCalcClient, cost_rows, register_scenario_deployment +from cost_matrix import ( + FRONTIER_MODELS, + Case, + FrontierModel, + cases_for, + expected_cost, + expected_token_columns, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatStreamOptions + +pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] + +_MATRIX: list[tuple[FrontierModel, Case]] = [ + (model, case) for model in FRONTIER_MODELS for case in cases_for(model) +] + + +def _case_id(param: tuple[FrontierModel, Case]) -> str: + model, case = param + return f"{model.map_key.replace('/', '-')}-{case.name}" + + +def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: + return ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted pricing call")], + stream=case.stream, + stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, + service_tier=case.service_tier, + ) + + +class TestTokenPricing: + @pytest.mark.parametrize("model_case", _MATRIX, ids=_case_id) + @pytest.mark.covers("quota_management.spend_tracking.cost_matrix.logs_cost") + def test_scripted_usage_bills_at_map_rates( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + model_case: tuple[FrontierModel, Case], + ) -> None: + model, case = model_case + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=_chat_body(model_name, marker, case), + stream=case.stream, + ) + assert response.ok, ( + f"{model.map_key}/{case.name}: proxy returned {response.status_code}: {response.body[:400]}" + ) + assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" + + expected = expected_cost(model, case) + if case.exact_spend and not case.stream: + # Streamed responses commit headers before the bill is computed, so + # the x-litellm-response-cost header is asserted only on non-stream + # calls. + assert response.response_cost is not None and cost_rows.approx_equal( + response.response_cost, expected + ), ( + f"x-litellm-response-cost {response.response_cost} != expected {expected}" + ) + + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}" + + if not case.exact_spend and case.expect_zero_bill: + # The provider reported no usage and this wire has no proxy-side + # recount, so the bill is exactly zero. + assert row.spend is not None and row.spend == 0, f"no-usage stream billed {row.spend}: {row}" + return + if not case.exact_spend: + # stream_usage=absent: the provider reported no usage, so the row's + # token counts are the proxy's own recount; only assert a bill landed. + assert row.spend is not None and row.spend > 0, f"no-usage stream billed nothing: {row}" + return + + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( + f"{model.map_key}/{case.name}: spend {row.spend} != expected {expected} " + f"(breakdown {row.breakdown.model_dump()})" + ) + + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py new file mode 100644 index 00000000000..b1ef675d9ef --- /dev/null +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -0,0 +1,186 @@ +"""Wire-format e2e: one scripted upstream per provider wire, answering with a +usage payload where every token kind the wire can report is nonzero. The spend +row's gross input cost must equal fresh tokens at the input rate plus each cache +and audio component at its own rate -- proving the wire's usage shape landed the +cached tokens inside the total (OpenAI/Gemini) or as separate fields +(Anthropic), and that the biller subtracted them before billing fresh tokens. + +Also covers the Responses API wire (an openai/gpt-5.5-pro deployment bridged by +the proxy to POST /responses) and a streamed Anthropic-messages case. +""" + +from __future__ import annotations + +import pytest + +from conftest import CostCalcClient, cost_rows, register_scenario_deployment +from cost_matrix import ( + FRONTIER_MODELS, + Case, + FrontierModel, + expected_breakdown, + expected_token_columns, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatStreamOptions +from scripted_provider import ScriptedUsage + +pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] + +_MODELS: dict[str, FrontierModel] = {model.map_key: model for model in FRONTIER_MODELS} + +# One scripted usage per wire, every reportable token kind nonzero. +_WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { + "openai_chat": ( + "gpt-5.6", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "openai_responses": ( + "gpt-5.5-pro", + ScriptedUsage( + fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25, reasoning_tokens=15 + ), + ), + "anthropic_messages": ( + "claude-sonnet-5", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + ), + ), + "gemini_generate": ( + "gemini/gemini-3.8-flash", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "together_chat": ( + "together_ai/moonshotai/Kimi-K3", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "fireworks_chat": ( + "fireworks_ai/kimi-k3", + ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), + ), +} + + +class TestWireFormats: + @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_wire_usage_shape_bills_each_component( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + wire: str, + ) -> None: + map_key, usage = _WIRE_USAGE[wire] + model = _MODELS[map_key] + case = Case(name="basic", usage=usage) + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted wire call")], + ), + ) + assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" + + expected = expected_breakdown(model, case) + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"{wire}: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"{wire}: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + breakdown = row.breakdown + assert breakdown.input_cost is not None and cost_rows.approx_equal( + breakdown.input_cost, expected.input_cost + ), ( + f"{wire}: gross input_cost {breakdown.input_cost} != expected {expected.input_cost}; " + "cached/written tokens billed at the input rate" + ) + assert breakdown.output_cost is not None and cost_rows.approx_equal( + breakdown.output_cost, expected.output_cost + ), f"{wire}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" + + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"{wire}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"{wire}: completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) + + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_anthropic_streamed_usage_bills_each_component( + self, client: CostCalcClient, resources: ResourceManager, scoped_key: str + ) -> None: + map_key, usage = _WIRE_USAGE["anthropic_messages"] + model = _MODELS[map_key] + case = Case(name="stream", usage=usage, stream=True) + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted anthropic stream")], + stream=True, + stream_options=ChatStreamOptions(include_usage=True), + ), + stream=True, + ) + assert response.ok, f"anthropic stream: proxy returned {response.status_code}: {response.body[:400]}" + assert response.stream_done, "anthropic stream did not reach its terminal event" + assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" + + expected = expected_breakdown(model, case) + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, "anthropic stream: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"anthropic stream: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json new file mode 100644 index 00000000000..b761710bae3 --- /dev/null +++ b/tests/e2e/cost_map.json @@ -0,0 +1,352 @@ +{ + "claude-haiku-4-5": { + "cache_creation_input_token_cost": 0.00021, + "cache_creation_input_token_cost_above_1hr": 0.00028000000000000003, + "cache_read_input_token_cost": 7e-06, + "input_cost_per_token": 7.000000000000001e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00014000000000000001, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "claude-opus-5": { + "cache_creation_input_token_cost": 0.00015000000000000001, + "cache_creation_input_token_cost_above_1hr": 0.0002, + "cache_read_input_token_cost": 4.9999999999999996e-06, + "input_cost_per_token": 5e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "claude-sonnet-5": { + "cache_creation_input_token_cost": 0.00018, + "cache_creation_input_token_cost_above_1hr": 0.00024000000000000003, + "cache_read_input_token_cost": 6e-06, + "input_cost_per_token": 6.000000000000001e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00012000000000000002, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "fireworks_ai/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 1.4e-05, + "input_cost_per_token": 0.00014000000000000001, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00028000000000000003, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "fireworks_ai/kimi-k3": { + "cache_read_input_token_cost": 1.2e-05, + "input_cost_per_token": 0.00012000000000000002, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00024000000000000003, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "fireworks_ai/qwen3p8-max": { + "cache_read_input_token_cost": 1.3e-05, + "input_cost_per_token": 0.00013000000000000002, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00026000000000000003, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gemini/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 9e-06, + "input_cost_per_audio_token": 0.00054, + "input_cost_per_token": 9e-05, + "input_cost_per_token_above_200k_tokens": 0.00072, + "input_cost_per_token_flex": 0.000135, + "input_cost_per_token_priority": 0.000153, + "litellm_provider": "gemini", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0006299999999999999, + "output_cost_per_reasoning_token": 0.00045000000000000004, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_200k_tokens": 0.0008100000000000001, + "output_cost_per_token_flex": 0.00022500000000000002, + "output_cost_per_token_priority": 0.000243, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-3.8-flash": { + "cache_read_input_token_cost": 8e-06, + "input_cost_per_audio_token": 0.00048, + "input_cost_per_token": 8e-05, + "input_cost_per_token_above_200k_tokens": 0.00064, + "input_cost_per_token_flex": 0.00012, + "input_cost_per_token_priority": 0.000136, + "litellm_provider": "gemini", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00056, + "output_cost_per_reasoning_token": 0.0004, + "output_cost_per_token": 0.00016, + "output_cost_per_token_above_200k_tokens": 0.00072, + "output_cost_per_token_flex": 0.0002, + "output_cost_per_token_priority": 0.000216, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 3e-06, + "input_cost_per_token": 3.0000000000000004e-05, + "input_cost_per_token_above_200k_tokens": 0.00024000000000000003, + "input_cost_per_token_flex": 4.5e-05, + "input_cost_per_token_priority": 5.1e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_reasoning_token": 0.00015000000000000001, + "output_cost_per_token": 6.000000000000001e-05, + "output_cost_per_token_above_200k_tokens": 0.00027, + "output_cost_per_token_flex": 7.500000000000001e-05, + "output_cost_per_token_priority": 8.099999999999999e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00012, + "cache_creation_input_token_cost_above_1hr": 0.00016, + "cache_read_input_token_cost": 4e-06, + "input_cost_per_audio_token": 0.00024, + "input_cost_per_token": 4e-05, + "input_cost_per_token_above_200k_tokens": 0.00032, + "input_cost_per_token_flex": 6e-05, + "input_cost_per_token_priority": 6.8e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00028, + "output_cost_per_reasoning_token": 0.0002, + "output_cost_per_token": 8e-05, + "output_cost_per_token_above_200k_tokens": 0.00036, + "output_cost_per_token_flex": 0.0001, + "output_cost_per_token_priority": 0.000108, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "input_cost_per_token_above_200k_tokens": 0.00016, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_priority": 3.4e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_reasoning_token": 0.0001, + "output_cost_per_token": 4e-05, + "output_cost_per_token_above_200k_tokens": 0.00018, + "output_cost_per_token_flex": 5e-05, + "output_cost_per_token_priority": 5.4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gpt-5.6": { + "cache_creation_input_token_cost": 3e-05, + "cache_creation_input_token_cost_above_1hr": 4e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_audio_token": 6e-05, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_200k_tokens": 8e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_priority": 1.7e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 7e-05, + "output_cost_per_reasoning_token": 5e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_200k_tokens": 9e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 2.7e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "together_ai/moonshotai/Kimi-K3": { + "cache_creation_input_token_cost": 0.00030000000000000003, + "cache_creation_input_token_cost_above_1hr": 0.0004, + "cache_read_input_token_cost": 9.999999999999999e-06, + "input_cost_per_audio_token": 0.0006000000000000001, + "input_cost_per_token": 0.0001, + "input_cost_per_token_above_200k_tokens": 0.0008, + "input_cost_per_token_flex": 0.00015000000000000001, + "input_cost_per_token_priority": 0.00017, + "litellm_provider": "together_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0006999999999999999, + "output_cost_per_reasoning_token": 0.0005, + "output_cost_per_token": 0.0002, + "output_cost_per_token_above_200k_tokens": 0.0009000000000000001, + "output_cost_per_token_flex": 0.00025, + "output_cost_per_token_priority": 0.00027, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "together_ai/zai-org/GLM-5.3": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, + "cache_read_input_token_cost": 1.1e-05, + "input_cost_per_audio_token": 0.00066, + "input_cost_per_token": 0.00011, + "input_cost_per_token_above_200k_tokens": 0.00088, + "input_cost_per_token_flex": 0.000165, + "input_cost_per_token_priority": 0.000187, + "litellm_provider": "together_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, + "output_cost_per_token": 0.00022, + "output_cost_per_token_above_200k_tokens": 0.00099, + "output_cost_per_token_flex": 0.000275, + "output_cost_per_token_priority": 0.000297, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + } +} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index ad0914d455b..6b40e70125c 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -63,3 +63,5 @@ - {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} - {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} - {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} +- {id: quota_management.spend_tracking.cost_matrix.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_matrix, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "A scripted-usage call through the cost-map proxy bills every reported token kind at the deployment's test-map rate (input, output, cache read, 5m/1h cache write, reasoning, audio, above-threshold tiers, flex/priority service tiers, web search, response-model override) and lands on the row's cost_breakdown, streamed or not"} +- {id: quota_management.spend_tracking.scripted_wire.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: scripted_wire, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "Each provider wire shape (openai chat, responses, anthropic messages, gemini generateContent, together, fireworks) parses usage into the same spend components: the gross input cost is fresh tokens at the input rate plus each cache/audio component at its own rate, streamed anthropic included"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..a891d9dcba2 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,6 +145,22 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +# The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL +# pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a +# scripted-provider sidecar; deselected unless the opt-in env var is set. +COST_MAP_OPT_IN_ENV = "E2E_COST_MAP_STACK" +# Base URL of the proxy running the test cost map. Defaults to the shared proxy +# so a local run only has to set the opt-in and boot the proxy accordingly. +COST_MAP_PROXY_URL = os.environ.get("E2E_COST_MAP_PROXY_URL", PROXY_BASE_URL).rstrip("/") +# Where the test runner reaches the scripted-provider sidecar's control API. +SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get( + "E2E_SCRIPTED_PROVIDER_CONTROL_URL", "http://127.0.0.1:9100" +).rstrip("/") +# The api_base root deployments register with: how the proxy (possibly in +# another container) reaches the sidecar's provider wire. +SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get( + "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL +).rstrip("/") ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..7d37bcc6d3e 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -11,3 +11,4 @@ markers = managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set + cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json (LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless E2E_COST_MAP_STACK is set From 269afbe382df06d33780571a40a55e527afea2b7 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:28:44 +0000 Subject: [PATCH 043/393] test(e2e): apply review nits to cost calculation suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 28 +- tests/e2e/cost_calculation/cost_matrix.py | 174 ++-- tests/e2e/cost_calculation/scripted_client.py | 12 +- .../e2e/cost_calculation/scripted_provider.py | 803 +++++++++++------- .../test_token_pricing_e2e.py | 17 +- .../cost_calculation/test_wire_formats_e2e.py | 43 +- 6 files changed, 620 insertions(+), 457 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 1bba3d50e1d..345ca26f7e3 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -13,7 +13,7 @@ from __future__ import annotations import importlib.util import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path from types import ModuleType @@ -34,16 +34,16 @@ def _load_cost_rows() -> ModuleType: """Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree has no package layout), the same trick the mcp suite uses for logging/datadog_reader.py.""" - path = ( + path: Final = ( Path(__file__).resolve().parent.parent / "quota_management" / "spend_tracking" / "cost_rows.py" ) - name = "e2e_spend_tracking_cost_rows" - spec = importlib.util.spec_from_file_location(name, path) + name: Final = "e2e_spend_tracking_cost_rows" + spec: Final = importlib.util.spec_from_file_location(name, path) assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) + module: Final = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) return module @@ -59,7 +59,7 @@ class SpendCostBreakdown(Protocol): total_cost: float | None service_tier: str | None - def model_dump(self) -> dict[str, object]: ... + def model_dump(self) -> Mapping[str, object]: ... class SpendRowMetadata(Protocol): @@ -89,7 +89,9 @@ class CostRowsModule(Protocol): ] -cost_rows: Final[CostRowsModule] = cast(CostRowsModule, _load_cost_rows()) +cost_rows: Final[CostRowsModule] = cast( # cast-ok: cost_rows.py is loaded by path, so basedpyright has no importable name for it; its surface is declared in CostRowsModule + CostRowsModule, _load_cost_rows() +) @dataclass(frozen=True, slots=True) @@ -101,7 +103,7 @@ class CostCalcClient: @pytest.fixture(scope="session") def client() -> CostCalcClient: - proxy = build_proxy_client( + proxy: Final = build_proxy_client( base_url=COST_MAP_PROXY_URL, control_plane_base_url=COST_MAP_PROXY_URL, replica_urls=(COST_MAP_PROXY_URL,), @@ -118,18 +120,18 @@ def register_scenario_deployment( ) -> tuple[str, ScenarioHandle]: """Register the case's scenario on the sidecar plus a deployment pointed at it; both are torn down by ``resources``. Returns the callable model_name.""" - scenario: Scenario = case.scenario( + scenario: Final[Scenario] = case.scenario( scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}" ) - handle = register_scenario(scenario) + handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) - model_name = f"{model.model_name}-{marker}" - model_id = client.proxy.register_model( + model_name: Final = f"{model.model_name}-{marker}" + model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, litellm_params=LiteLLMParamsBody( model=model.litellm_model, - api_key="sk-scripted-provider", + api_key=model.api_key, api_base=handle.api_base(), ), model_info=ModelInfoBody(), diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index bc466d7d823..e8b1d249559 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -18,9 +18,11 @@ creation), the case is absent from the matrix rather than silently zero. from __future__ import annotations import json +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -64,8 +66,8 @@ class CostMapEntry(BaseModel): _COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) -_COST_MAP: Final[dict[str, CostMapEntry]] = _COST_MAP_ADAPTER.validate_python( - json.loads(COST_MAP_PATH.read_text()) +_COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( + _COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) ) TIER_THRESHOLD_TOKENS: Final = 200_000 @@ -109,7 +111,7 @@ class FrontierModel: # Response-model override targets: emit a sibling's bare provider-facing name so # the biller's provider-prefixed lookup lands on that sibling's map key. -_OVERRIDE_MODELS: Final[dict[str, str]] = { +_OVERRIDE_MODELS: Final[Mapping[str, str]] = MappingProxyType({ "gpt-5.6": "gpt-5.4-mini", "gpt-5.5-pro": "gpt-5.3-codex", "gpt-5.3-codex": "gpt-5.5-pro", @@ -124,9 +126,9 @@ _OVERRIDE_MODELS: Final[dict[str, str]] = { "fireworks_ai/kimi-k3": "qwen3p8-max", "fireworks_ai/qwen3p8-max": "kimi-k3", "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", -} +}) -_OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { +_OVERRIDE_MAP_KEYS: Final[Mapping[str, str]] = MappingProxyType({ "gpt-5.4-mini": "gpt-5.4-mini", "gpt-5.6": "gpt-5.6", "gpt-5.3-codex": "gpt-5.3-codex", @@ -139,7 +141,7 @@ _OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", "qwen3p8-max": "fireworks_ai/qwen3p8-max", "kimi-k3": "fireworks_ai/kimi-k3", -} +}) _FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( @@ -176,7 +178,7 @@ def _frontier() -> tuple[FrontierModel, ...]: FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() # Token kinds each wire can report, gating which pricing cases apply. -_WIRE_CAPS: Final[dict[str, frozenset[str]]] = { +_WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "openai_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", @@ -205,9 +207,9 @@ _WIRE_CAPS: Final[dict[str, frozenset[str]]] = { "web_search", "response_model", "absent_usage", } ), -} +}) -CaseName = Literal[ +CaseName: TypeAlias = Literal[ "basic", "cache_read", "cache_write_5m", @@ -260,7 +262,7 @@ _BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) def _web_search_case(model: FrontierModel) -> Case: - counts_exactly = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") return Case( name="web_search", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), @@ -269,26 +271,24 @@ def _web_search_case(model: FrontierModel) -> Case: def cases_for(model: FrontierModel) -> tuple[Case, ...]: - rates = model.rates - caps = _WIRE_CAPS[model.wire] - cases: list[Case] = [Case(name="basic", usage=_BASIC_USAGE)] - if rates.cache_read_input_token_cost is not None and "cache_read" in caps: - cases.append( + rates: Final = model.rates + caps: Final = _WIRE_CAPS[model.wire] + candidates: Final[tuple[Case | None, ...]] = ( + Case(name="basic", usage=_BASIC_USAGE), + ( Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) - ) - if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps: - cases.append( + if rates.cache_read_input_token_cost is not None and "cache_read" in caps + else None + ), + ( Case( name="cache_write_5m", usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), ) - ) - if ( - rates.cache_creation_input_token_cost_above_1hr is not None - and rates.cache_creation_input_token_cost is not None - and "cache_write_1h" in caps - ): - cases.append( + if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps + else None + ), + ( Case( name="cache_write_1h", usage=ScriptedUsage( @@ -298,52 +298,61 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: output_tokens=30, ), ) - ) - if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps: - cases.append( + if ( + rates.cache_creation_input_token_cost_above_1hr is not None + and rates.cache_creation_input_token_cost is not None + and "cache_write_1h" in caps + ) + else None + ), + ( Case( name="reasoning", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), ) - ) - if ( - rates.input_cost_per_audio_token is not None - and rates.output_cost_per_audio_token is not None - and "audio" in caps - ): - cases.append( + if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps + else None + ), + ( Case( name="audio", usage=ScriptedUsage( fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 ), ) - ) - if ( - rates.input_cost_per_token_above_200k_tokens is not None - and rates.output_cost_per_token_above_200k_tokens is not None - ): - cases.append( + if ( + rates.input_cost_per_audio_token is not None + and rates.output_cost_per_audio_token is not None + and "audio" in caps + ) + else None + ), + ( Case( name="tiered", usage=ScriptedUsage( fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 ), ) - ) - if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None: - cases.append( + if ( + rates.input_cost_per_token_above_200k_tokens is not None + and rates.output_cost_per_token_above_200k_tokens is not None + ) + else None + ), + ( Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") - ) - if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None: - cases.append( + if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None + else None + ), + ( Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") - ) - if rates.search_context_cost_per_query is not None and "web_search" in caps: - cases.append(_web_search_case(model)) - cases.append(Case(name="stream", usage=_BASIC_USAGE, stream=True)) - if "absent_usage" in caps: - cases.append( + if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None + else None + ), + _web_search_case(model) if rates.search_context_cost_per_query is not None and "web_search" in caps else None, + Case(name="stream", usage=_BASIC_USAGE, stream=True), + ( Case( name="stream_no_usage", usage=_BASIC_USAGE, @@ -355,10 +364,16 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: # wires recount tokens proxy-side and bill a nonzero amount. expect_zero_bill=model.wire == "openai_responses", ) - ) - if "response_model" in caps: - cases.append(Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True)) - return tuple(cases) + if "absent_usage" in caps + else None + ), + ( + Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True) + if "response_model" in caps + else None + ), + ) + return tuple(case for case in candidates if case is not None) @dataclass(frozen=True, slots=True) @@ -387,38 +402,41 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: to the tier's variants, falling back to the base rate when a variant is unset -- mirroring _get_token_base_cost in litellm's cost calculator. """ - rates = model.override_rates if case.response_model_override else model.rates - u = case.usage - prompt_tokens = ( + rates: Final = model.override_rates if case.response_model_override else model.rates + u: Final = case.usage + prompt_tokens: Final = ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens + u.audio_input_tokens ) - tiered = prompt_tokens > TIER_THRESHOLD_TOKENS - in_rate = rates.input_cost_per_token or 0.0 - out_rate = rates.output_cost_per_token or 0.0 - if case.service_tier == "flex": - in_rate = rates.input_cost_per_token_flex or in_rate - out_rate = rates.output_cost_per_token_flex or out_rate - if case.service_tier == "priority": - in_rate = rates.input_cost_per_token_priority or in_rate - out_rate = rates.output_cost_per_token_priority or out_rate - if tiered: - in_rate = rates.input_cost_per_token_above_200k_tokens or in_rate - out_rate = rates.output_cost_per_token_above_200k_tokens or out_rate - input_cost = ( + tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate: Final = ( + (rates.input_cost_per_token_above_200k_tokens if tiered else None) + or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.input_cost_per_token + or 0.0 + ) + out_rate: Final = ( + (rates.output_cost_per_token_above_200k_tokens if tiered else None) + or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.output_cost_per_token + or 0.0 + ) + input_cost: Final = ( u.fresh_input_tokens * in_rate + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) - output_cost = ( + output_cost: Final = ( u.output_tokens * out_rate + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) ) - search = rates.search_context_cost_per_query - tool_cost = case.billed_web_search_calls * ( + search: Final = rates.search_context_cost_per_query + tool_cost: Final = case.billed_web_search_calls * ( search.search_context_size_medium if search and search.search_context_size_medium else 0.0 ) return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) @@ -432,7 +450,7 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: """(prompt_tokens, completion_tokens) the spend row should carry, per the wire's normalization: Anthropic folds cache read/write into prompt_tokens, everyone else reports the totals the wire emitted.""" - u = case.usage + u: Final = case.usage if model.wire == "anthropic_messages": return ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, diff --git a/tests/e2e/cost_calculation/scripted_client.py b/tests/e2e/cost_calculation/scripted_client.py index dceec02630a..9dbf9c98986 100644 --- a/tests/e2e/cost_calculation/scripted_client.py +++ b/tests/e2e/cost_calculation/scripted_client.py @@ -12,6 +12,7 @@ from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BA from e2e_http import URL, NoBody, unwrap, post from e2e_http import delete as http_delete from scripted_provider import ( + WIRE_MOUNTS, Scenario, ScenarioDeleted, ScenarioRegistered, @@ -29,19 +30,12 @@ class ScenarioHandle: return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}" def _mount(self) -> str: - return { - "openai_chat": "openai", - "openai_responses": "openai", - "anthropic_messages": "anthropic", - "gemini_generate": "gemini", - "together_chat": "together", - "fireworks_chat": "fireworks", - }[self.wire] + return WIRE_MOUNTS[self.wire] def register_scenario(scenario: Scenario) -> ScenarioHandle: """POST the scenario to the sidecar's control API and return its handle.""" - result = unwrap( + result: Final = unwrap( post( URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"), headers=NoBody(), diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 93a6f49ec25..f1deafd1bc5 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -31,14 +31,16 @@ import json import sys import threading import time +from collections.abc import Mapping from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError -Wire = Literal[ +Wire: TypeAlias = Literal[ "openai_chat", "openai_responses", "anthropic_messages", @@ -47,17 +49,19 @@ Wire = Literal[ "fireworks_chat", ] -_WIRE_MOUNTS: Final[dict[str, str]] = { - "openai_chat": "openai", - "openai_responses": "openai", - "anthropic_messages": "anthropic", - "gemini_generate": "gemini", - "together_chat": "together", - "fireworks_chat": "fireworks", -} +WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( + { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", + } +) -StreamUsage = Literal["final_chunk", "absent"] -ServiceTier = Literal["flex", "priority"] +StreamUsage: TypeAlias = Literal["final_chunk", "absent"] +ServiceTier: TypeAlias = Literal["flex", "priority"] class ScriptedUsage(BaseModel): @@ -106,7 +110,7 @@ class Scenario(BaseModel): @property def mount(self) -> str: - return _WIRE_MOUNTS[self.wire] + return WIRE_MOUNTS[self.wire] class ScenarioRegistered(BaseModel): @@ -128,369 +132,494 @@ class RenderedResponse: body: bytes -def _json_bytes(payload: dict[str, object]) -> bytes: - return json.dumps(payload).encode("utf-8") +def _jobj(*pairs: tuple[str, object]) -> Mapping[str, object]: + """A JSON object payload built in one shot and frozen.""" + return MappingProxyType(dict(pairs)) -def _sse(events: tuple[tuple[str | None, dict[str, object] | str], ...]) -> bytes: - frames: list[str] = [] - for event_name, data in events: - head = f"event: {event_name}\n" if event_name is not None else "" - payload = data if isinstance(data, str) else json.dumps(data) - frames.append(f"{head}data: {payload}\n\n") - return "".join(frames).encode("utf-8") +def _jobj_opt(*pairs: tuple[str, object] | None) -> Mapping[str, object]: + """``_jobj`` where a ``None`` pair means the field is absent.""" + return MappingProxyType(dict(pair for pair in pairs if pair is not None)) + + +def _json_bytes(payload: Mapping[str, object]) -> bytes: + return json.dumps(payload, default=dict).encode("utf-8") + + +def _sse_frame(event_name: str | None, data: Mapping[str, object] | str) -> str: + head: Final = f"event: {event_name}\n" if event_name is not None else "" + payload: Final = data if isinstance(data, str) else json.dumps(data, default=dict) + return f"{head}data: {payload}\n\n" + + +def _sse(events: tuple[tuple[str | None, Mapping[str, object] | str], ...]) -> bytes: + return "".join(_sse_frame(event_name, data) for event_name, data in events).encode("utf-8") # ---------- per-wire usage shapes ---------- -def _openai_usage(u: ScriptedUsage) -> dict[str, object]: - prompt_tokens = ( +def _openai_usage(u: ScriptedUsage) -> Mapping[str, object]: + prompt_tokens: Final = ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens + u.audio_input_tokens ) - completion_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - prompt_details: dict[str, object] = {} - if u.cache_read_tokens: - prompt_details["cached_tokens"] = u.cache_read_tokens - if u.cache_write_5m_tokens or u.cache_write_1h_tokens: - prompt_details["cache_write_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens - prompt_details["cache_creation_token_details"] = { - "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, - "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, - } - if u.audio_input_tokens: - prompt_details["audio_tokens"] = u.audio_input_tokens - completion_details: dict[str, object] = {} - if u.reasoning_tokens: - completion_details["reasoning_tokens"] = u.reasoning_tokens - if u.audio_output_tokens: - completion_details["audio_tokens"] = u.audio_output_tokens - usage: dict[str, object] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } - if prompt_details: - usage["prompt_tokens_details"] = prompt_details - if completion_details: - usage["completion_tokens_details"] = completion_details - return usage + completion_tokens: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + prompt_details: Final = _jobj_opt( + ("cached_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ( + ("cache_write_tokens", u.cache_write_5m_tokens + u.cache_write_1h_tokens) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ( + "cache_creation_token_details", + _jobj( + ("ephemeral_5m_input_tokens", u.cache_write_5m_tokens), + ("ephemeral_1h_input_tokens", u.cache_write_1h_tokens), + ), + ) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ("audio_tokens", u.audio_input_tokens) if u.audio_input_tokens else None, + ) + completion_details: Final = _jobj_opt( + ("reasoning_tokens", u.reasoning_tokens) if u.reasoning_tokens else None, + ("audio_tokens", u.audio_output_tokens) if u.audio_output_tokens else None, + ) + return _jobj_opt( + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("total_tokens", prompt_tokens + completion_tokens), + ("prompt_tokens_details", prompt_details) if prompt_details else None, + ("completion_tokens_details", completion_details) if completion_details else None, + ) -def _anthropic_usage(u: ScriptedUsage) -> dict[str, object]: +def _anthropic_usage(u: ScriptedUsage) -> Mapping[str, object]: # Anthropic reports uncached-only input_tokens; cache reads and writes ride # top-level fields, with the 5m/1h write split under cache_creation. - usage: dict[str, object] = { - "input_tokens": u.fresh_input_tokens, - "output_tokens": u.output_tokens, - } - if u.cache_read_tokens: - usage["cache_read_input_tokens"] = u.cache_read_tokens - if u.cache_write_5m_tokens or u.cache_write_1h_tokens: - usage["cache_creation_input_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens - usage["cache_creation"] = { - "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, - "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, - } - if u.web_search_calls: - usage["server_tool_use"] = {"web_search_requests": u.web_search_calls} - return usage + return _jobj_opt( + ("input_tokens", u.fresh_input_tokens), + ("output_tokens", u.output_tokens), + ("cache_read_input_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ( + ("cache_creation_input_tokens", u.cache_write_5m_tokens + u.cache_write_1h_tokens) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ( + "cache_creation", + _jobj( + ("ephemeral_5m_input_tokens", u.cache_write_5m_tokens), + ("ephemeral_1h_input_tokens", u.cache_write_1h_tokens), + ), + ) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ("server_tool_use", _jobj(("web_search_requests", u.web_search_calls))) + if u.web_search_calls + else None + ), + ) -def _gemini_usage(u: ScriptedUsage) -> dict[str, object]: +def _gemini_usage(u: ScriptedUsage) -> Mapping[str, object]: # promptTokenCount carries the cached count inside it; TEXT modality is the # cached-inclusive text count so litellm's implicit-caching subtraction lands # on the fresh figure. candidatesTokenCount includes reasoning + audio. - prompt_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens - candidates = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - usage: dict[str, object] = { - "promptTokenCount": prompt_tokens, - "candidatesTokenCount": candidates, - "totalTokenCount": prompt_tokens + candidates, - } - if u.cache_read_tokens: - usage["cachedContentTokenCount"] = u.cache_read_tokens - if u.reasoning_tokens: - usage["thoughtsTokenCount"] = u.reasoning_tokens - prompt_details = [{"modality": "TEXT", "tokenCount": u.fresh_input_tokens + u.cache_read_tokens}] - if u.audio_input_tokens: - prompt_details.append({"modality": "AUDIO", "tokenCount": u.audio_input_tokens}) - usage["promptTokensDetails"] = prompt_details - if u.audio_output_tokens: - usage["candidatesTokensDetails"] = [ - {"modality": "TEXT", "tokenCount": u.output_tokens + u.reasoning_tokens}, - {"modality": "AUDIO", "tokenCount": u.audio_output_tokens}, - ] - return usage + prompt_tokens: Final = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + candidates: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + return _jobj_opt( + ("promptTokenCount", prompt_tokens), + ("candidatesTokenCount", candidates), + ("totalTokenCount", prompt_tokens + candidates), + ("cachedContentTokenCount", u.cache_read_tokens) if u.cache_read_tokens else None, + ("thoughtsTokenCount", u.reasoning_tokens) if u.reasoning_tokens else None, + ( + "promptTokensDetails", + ( + _jobj(("modality", "TEXT"), ("tokenCount", u.fresh_input_tokens + u.cache_read_tokens)), + *( + (_jobj(("modality", "AUDIO"), ("tokenCount", u.audio_input_tokens)),) + if u.audio_input_tokens + else () + ), + ), + ), + ( + ( + "candidatesTokensDetails", + ( + _jobj(("modality", "TEXT"), ("tokenCount", u.output_tokens + u.reasoning_tokens)), + _jobj(("modality", "AUDIO"), ("tokenCount", u.audio_output_tokens)), + ), + ) + if u.audio_output_tokens + else None + ), + ) -def _responses_usage(u: ScriptedUsage) -> dict[str, object]: - input_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens - output_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - usage: dict[str, object] = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - input_details: dict[str, object] = {} - if u.cache_read_tokens: - input_details["cached_tokens"] = u.cache_read_tokens - if input_details: - usage["input_tokens_details"] = input_details - if u.reasoning_tokens: - usage["output_tokens_details"] = {"reasoning_tokens": u.reasoning_tokens} - return usage +def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: + input_tokens: Final = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + output_tokens: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + input_details: Final = _jobj_opt( + ("cached_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ) + return _jobj_opt( + ("input_tokens", input_tokens), + ("output_tokens", output_tokens), + ("total_tokens", input_tokens + output_tokens), + ("input_tokens_details", input_details) if input_details else None, + ( + ("output_tokens_details", _jobj(("reasoning_tokens", u.reasoning_tokens))) + if u.reasoning_tokens + else None + ), + ) # ---------- per-wire responses ---------- -def _openai_message(scenario: Scenario) -> dict[str, object]: - message: dict[str, object] = {"role": "assistant", "content": scenario.output.text} - if scenario.usage.web_search_calls: - message["annotations"] = [ - { - "type": "url_citation", - "url_citation": { - "url": "https://scripted.example/source", - "title": "scripted source", - "start_index": 0, - "end_index": 1, - }, - } - for _ in range(scenario.usage.web_search_calls) - ] - return message +def _openai_message(scenario: Scenario) -> Mapping[str, object]: + return _jobj_opt( + ("role", "assistant"), + ("content", scenario.output.text), + ( + ( + "annotations", + tuple( + _jobj( + ("type", "url_citation"), + ( + "url_citation", + _jobj( + ("url", "https://scripted.example/source"), + ("title", "scripted source"), + ("start_index", 0), + ("end_index", 1), + ), + ), + ) + for _ in range(scenario.usage.web_search_calls) + ), + ) + if scenario.usage.web_search_calls + else None + ), + ) -def _openai_chat_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - body: dict[str, object] = { - "id": f"chatcmpl-{scenario.scenario_id}", - "object": "chat.completion", - "created": int(time.time()), - "model": scenario.output.response_model or requested_model, - "choices": [ - { - "index": 0, - "message": _openai_message(scenario), - "finish_reason": scenario.output.finish_reason, - } - ], - "usage": _openai_usage(scenario.usage), - } - if scenario.service_tier is not None: - body["service_tier"] = scenario.service_tier - if scenario.output.provider_cost is not None: - body["cost"] = scenario.output.provider_cost - return body +def _openai_chat_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj_opt( + ("id", f"chatcmpl-{scenario.scenario_id}"), + ("object", "chat.completion"), + ("created", int(time.time())), + ("model", scenario.output.response_model or requested_model), + ( + "choices", + ( + _jobj( + ("index", 0), + ("message", _openai_message(scenario)), + ("finish_reason", scenario.output.finish_reason), + ), + ), + ), + ("usage", _openai_usage(scenario.usage)), + ("service_tier", scenario.service_tier) if scenario.service_tier is not None else None, + ("cost", scenario.output.provider_cost) if scenario.output.provider_cost is not None else None, + ) -def _openai_chunk(scenario: Scenario, requested_model: str, **kw: object) -> dict[str, object]: - chunk: dict[str, object] = { - "id": f"chatcmpl-{scenario.scenario_id}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": scenario.output.response_model or requested_model, - } - chunk.update(kw) - return chunk +def _openai_chunk( + scenario: Scenario, + requested_model: str, + choices: tuple[Mapping[str, object], ...] = (), + usage: Mapping[str, object] | None = None, +) -> Mapping[str, object]: + return _jobj_opt( + ("id", f"chatcmpl-{scenario.scenario_id}"), + ("object", "chat.completion.chunk"), + ("created", int(time.time())), + ("model", scenario.output.response_model or requested_model), + ("choices", choices), + ("usage", usage), + ) def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: - _EMPTY_DELTA: Final[dict[str, object]] = {} - delta: dict[str, object] = {"role": "assistant", "content": scenario.output.text} - if scenario.usage.web_search_calls: - delta["annotations"] = _openai_message(scenario)["annotations"] - events: list[tuple[str | None, dict[str, object] | str]] = [ + delta: Final = _jobj_opt( + ("role", "assistant"), + ("content", scenario.output.text), ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], - ), + ("annotations", _openai_message(scenario)["annotations"]) + if scenario.usage.web_search_calls + else None ), + ) + return _sse( ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", _jobj(("role", "assistant"))), ("finish_reason", None)),), + ), ), - ), - ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[ - { - "index": 0, - "delta": _EMPTY_DELTA, - "finish_reason": scenario.output.finish_reason, - } - ], + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", delta), ("finish_reason", None)),), + ), ), - ), - ] - if scenario.stream_usage == "final_chunk": - events.append( - (None, _openai_chunk(scenario, requested_model, choices=(), usage=_openai_usage(scenario.usage))) + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=( + _jobj( + ("index", 0), + ("delta", _jobj()), + ("finish_reason", scenario.output.finish_reason), + ), + ), + ), + ), + *( + ((None, _openai_chunk(scenario, requested_model, usage=_openai_usage(scenario.usage))),) + if scenario.stream_usage == "final_chunk" + else () + ), + (None, "[DONE]"), ) - events.append((None, "[DONE]")) - return _sse(tuple(events)) + ) -def _anthropic_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - return { - "id": f"msg_{scenario.scenario_id}", - "type": "message", - "role": "assistant", - "model": scenario.output.response_model or requested_model, - "content": [{"type": "text", "text": scenario.output.text}], - "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - "usage": _anthropic_usage(scenario.usage), - } +def _anthropic_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ("id", f"msg_{scenario.scenario_id}"), + ("type", "message"), + ("role", "assistant"), + ("model", scenario.output.response_model or requested_model), + ("content", (_jobj(("type", "text"), ("text", scenario.output.text)),)), + ( + "stop_reason", + "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + ), + ("usage", _anthropic_usage(scenario.usage)), + ) def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: - emit_usage = scenario.stream_usage == "final_chunk" - input_usage = {k: v for k, v in _anthropic_usage(scenario.usage).items() if k != "output_tokens"} - message_start: dict[str, object] = { - "type": "message_start", - "message": { - "id": f"msg_{scenario.scenario_id}", - "type": "message", - "role": "assistant", - "model": scenario.output.response_model or requested_model, - "content": [], - "stop_reason": None, - **({"usage": input_usage} if emit_usage else {}), - }, - } - message_delta: dict[str, object] = { - "type": "message_delta", - "delta": { - "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason - }, - **({"usage": {"output_tokens": scenario.usage.output_tokens}} if emit_usage else {}), - } + emit_usage: Final = scenario.stream_usage == "final_chunk" + input_usage: Final = _jobj( + *( + (key, value) + for key, value in _anthropic_usage(scenario.usage).items() + if key != "output_tokens" + ) + ) + message_start: Final = _jobj( + ("type", "message_start"), + ( + "message", + _jobj_opt( + ("id", f"msg_{scenario.scenario_id}"), + ("type", "message"), + ("role", "assistant"), + ("model", scenario.output.response_model or requested_model), + ("content", ()), + ("stop_reason", None), + ("usage", input_usage) if emit_usage else None, + ), + ), + ) + message_delta: Final = _jobj_opt( + ("type", "message_delta"), + ( + "delta", + _jobj( + ( + "stop_reason", + "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + ) + ), + ), + ( + ("usage", _jobj(("output_tokens", scenario.usage.output_tokens))) + if emit_usage + else None + ), + ) return _sse( ( ("message_start", message_start), ( "content_block_start", - { - "type": "content_block_start", - "index": 0, - "content_block": {"type": "text", "text": ""}, - }, + _jobj( + ("type", "content_block_start"), + ("index", 0), + ("content_block", _jobj(("type", "text"), ("text", ""))), + ), ), ( "content_block_delta", - { - "type": "content_block_delta", - "index": 0, - "delta": {"type": "text_delta", "text": scenario.output.text}, - }, + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), + ), ), - ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("content_block_stop", _jobj(("type", "content_block_stop"), ("index", 0))), ("message_delta", message_delta), - ("message_stop", {"type": "message_stop"}), + ("message_stop", _jobj(("type", "message_stop"))), ) ) -def _gemini_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - candidate: dict[str, object] = { - "content": {"parts": [{"text": scenario.output.text}], "role": "model"}, - "finishReason": "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), - "index": 0, - } - if scenario.usage.web_search_calls: - candidate["groundingMetadata"] = { - "webSearchQueries": [f"query {i}" for i in range(scenario.usage.web_search_calls)] - } - return { - "candidates": [candidate], - "usageMetadata": _gemini_usage(scenario.usage), - "modelVersion": scenario.output.response_model or requested_model, - } +def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ( + "candidates", + ( + _jobj_opt( + ( + "content", + _jobj( + ("parts", (_jobj(("text", scenario.output.text)),)), + ("role", "model"), + ), + ), + ( + "finishReason", + "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), + ), + ("index", 0), + ( + ( + "groundingMetadata", + _jobj( + ( + "webSearchQueries", + tuple(f"query {i}" for i in range(scenario.usage.web_search_calls)), + ) + ), + ) + if scenario.usage.web_search_calls + else None + ), + ), + ), + ), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ) def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: - first = _gemini_body(scenario, requested_model) - if scenario.stream_usage == "absent": - first = {k: v for k, v in first.items() if k != "usageMetadata"} - events: list[tuple[str | None, dict[str, object] | str]] = [(None, first)] - if scenario.stream_usage == "final_chunk": - events.append( - ( - None, - { - "candidates": [], - "usageMetadata": _gemini_usage(scenario.usage), - "modelVersion": scenario.output.response_model or requested_model, - }, - ) - ) - return _sse(tuple(events)) - - -def _responses_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - output: list[dict[str, object]] = [ - {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} - for i in range(scenario.usage.web_search_calls) - ] - output.append( - { - "type": "message", - "id": f"msg_{scenario.scenario_id}", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": scenario.output.text, - "annotations": [], - } - ], - } + emit_usage: Final = scenario.stream_usage == "final_chunk" + first: Final = ( + _jobj(*((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata")) + if scenario.stream_usage == "absent" + else _gemini_body(scenario, requested_model) + ) + return _sse( + ( + (None, first), + *( + ( + ( + None, + _jobj( + ("candidates", ()), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ), + ), + ) + if emit_usage + else () + ), + ) + ) + + +def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ("id", f"resp_{scenario.scenario_id}"), + ("object", "response"), + ("created_at", int(time.time())), + ("status", "completed"), + ("model", scenario.output.response_model or requested_model), + ( + "output", + ( + *( + _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) + for i in range(scenario.usage.web_search_calls) + ), + _jobj( + ("type", "message"), + ("id", f"msg_{scenario.scenario_id}"), + ("status", "completed"), + ("role", "assistant"), + ( + "content", + ( + _jobj( + ("type", "output_text"), + ("text", scenario.output.text), + ("annotations", ()), + ), + ), + ), + ), + ), + ), + ("usage", _responses_usage(scenario.usage)), ) - return { - "id": f"resp_{scenario.scenario_id}", - "object": "response", - "created_at": int(time.time()), - "status": "completed", - "model": scenario.output.response_model or requested_model, - "output": output, - "usage": _responses_usage(scenario.usage), - } def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: - completed = _responses_body(scenario, requested_model) - if scenario.stream_usage == "absent": - completed = {k: v for k, v in completed.items() if k != "usage"} - created = {**completed, "status": "in_progress", "usage": None} + completed: Final = ( + _jobj(*((key, value) for key, value in _responses_body(scenario, requested_model).items() if key != "usage")) + if scenario.stream_usage == "absent" + else _responses_body(scenario, requested_model) + ) + created: Final = _jobj( + *((key, value) for key, value in completed.items() if key not in ("status", "usage")), + ("status", "in_progress"), + ("usage", None), + ) return _sse( ( - ("response.created", {"type": "response.created", "response": created}), + ("response.created", _jobj(("type", "response.created"), ("response", created))), ( "response.output_text.delta", - { - "type": "response.output_text.delta", - "item_id": f"msg_{scenario.scenario_id}", - "output_index": scenario.usage.web_search_calls, - "content_index": 0, - "delta": scenario.output.text, - }, + _jobj( + ("type", "response.output_text.delta"), + ("item_id", f"msg_{scenario.scenario_id}"), + ("output_index", scenario.usage.web_search_calls), + ("content_index", 0), + ("delta", scenario.output.text), + ), ), - ("response.completed", {"type": "response.completed", "response": completed}), + ("response.completed", _jobj(("type", "response.completed"), ("response", completed))), ) ) @@ -538,11 +667,11 @@ class _ScenarioStore: _REQUEST_BODY: Final = TypeAdapter(dict[str, object]) -def _request_body(body: bytes) -> dict[str, object]: +def _request_body(body: bytes) -> Mapping[str, object]: try: return _REQUEST_BODY.validate_json(body) except ValueError: - return {} + return MappingProxyType({}) def _request_wants_stream(path_tail: str, body: bytes) -> bool: @@ -554,52 +683,66 @@ def _request_wants_stream(path_tail: str, body: bytes) -> bool: def _request_model(body: bytes) -> str: - model = _request_body(body).get("model") + model: Final = _request_body(body).get("model") return model if isinstance(model, str) else "unknown" def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: - path = urlsplit(raw_path).path - segments = [segment for segment in path.split("/") if segment] - if method == "GET" and segments == ["health"]: - return RenderedResponse(200, "application/json", _json_bytes({"status": "ok"})) + path: Final = urlsplit(raw_path).path + segments: Final = tuple(segment for segment in path.split("/") if segment) + if method == "GET" and segments == ("health",): + return RenderedResponse(200, "application/json", _json_bytes(_jobj(("status", "ok")))) if segments and segments[0] == "_scenarios": if method == "POST" and len(segments) == 1: try: - scenario = Scenario.model_validate_json(body) + scenario: Final = Scenario.model_validate_json(body) except ValidationError as exc: - return RenderedResponse(400, "application/json", _json_bytes({"error": str(exc)})) + return RenderedResponse( + 400, "application/json", _json_bytes(_jobj(("error", str(exc)))) + ) store.put(scenario) - return RenderedResponse(200, "application/json", _json_bytes({"scenario_id": scenario.scenario_id})) - if method == "DELETE" and len(segments) == 2: - deleted = store.drop(segments[1]) return RenderedResponse( - 200 if deleted else 404, "application/json", _json_bytes({"deleted": deleted}) + 200, "application/json", _json_bytes(_jobj(("scenario_id", scenario.scenario_id))) ) - return RenderedResponse(404, "application/json", _json_bytes({"error": "unknown control route"})) + if method == "DELETE" and len(segments) == 2: + deleted: Final = store.drop(segments[1]) + return RenderedResponse( + 200 if deleted else 404, + "application/json", + _json_bytes(_jobj(("deleted", deleted))), + ) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", "unknown control route"))) + ) if len(segments) < 2 or method != "POST": - return RenderedResponse(404, "application/json", _json_bytes({"error": f"no route for {method} {path}"})) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) + ) scenario_id, mount = segments[0], segments[1] - scenario = store.get(scenario_id) - if scenario is None: - return RenderedResponse(404, "application/json", _json_bytes({"error": f"unknown scenario {scenario_id}"})) - if scenario.mount != mount: + found: Final = store.get(scenario_id) + if found is None: + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", f"unknown scenario {scenario_id}"))) + ) + if found.mount != mount: return RenderedResponse( 400, "application/json", - _json_bytes({"error": f"scenario {scenario_id} is wire {scenario.wire}, not mount {mount}"}), + _json_bytes( + _jobj(("error", f"scenario {scenario_id} is wire {found.wire}, not mount {mount}")) + ), ) - tail = "/".join(segments[2:]) - return _render(scenario, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + tail: Final = "/".join(segments[2:]) + return _render(found, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) class _ScriptedHandler(BaseHTTPRequestHandler): store: Final[_ScenarioStore] = _ScenarioStore() def _dispatch(self, method: str) -> None: - length = int(self.headers.get("content-length") or 0) - body = self.rfile.read(length) if length else b"" - rendered = handle_request(self.store, method, self.path, body) + length: Final = int(self.headers.get("content-length") or 0) + body: Final = self.rfile.read(length) if length else b"" + rendered: Final = handle_request(self.store, method, self.path, body) self.send_response(rendered.status_code) self.send_header("content-type", rendered.content_type) self.send_header("content-length", str(len(rendered.body))) @@ -621,11 +764,11 @@ DEFAULT_PORT: Final = 9100 def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None: - server = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) + server: Final = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) sys.stderr.write(f"scripted-provider listening on http://{bind_host}:{port}\n") server.serve_forever() if __name__ == "__main__": - port_arg = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT + port_arg: Final = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT serve(port=port_arg) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 8d7678cf9ca..e210dad94b1 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -11,6 +11,7 @@ tests/e2e/cost_map.json. from __future__ import annotations import pytest +from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( @@ -25,11 +26,11 @@ from e2e_config import unique_marker from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatStreamOptions -pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] +pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark -_MATRIX: list[tuple[FrontierModel, Case]] = [ +_MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple( (model, case) for model in FRONTIER_MODELS for case in cases_for(model) -] +) def _case_id(param: tuple[FrontierModel, Case]) -> str: @@ -40,7 +41,7 @@ def _case_id(param: tuple[FrontierModel, Case]) -> str: def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: return ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted pricing call")], + messages=(ChatMessage(role="user", content=f"{marker} scripted pricing call"),), stream=case.stream, stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, service_tier=case.service_tier, @@ -58,9 +59,9 @@ class TestTokenPricing: model_case: tuple[FrontierModel, Case], ) -> None: model, case = model_case - marker = unique_marker() + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=_chat_body(model_name, marker, case), @@ -71,7 +72,7 @@ class TestTokenPricing: ) assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected = expected_cost(model, case) + expected: Final = expected_cost(model, case) if case.exact_spend and not case.stream: # Streamed responses commit headers before the bill is computed, so # the x-litellm-response-cost header is asserted only on non-stream @@ -82,7 +83,7 @@ class TestTokenPricing: f"x-litellm-response-cost {response.response_cost} != expected {expected}" ) - row = cost_rows.poll_cost_row_where( + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index b1ef675d9ef..c0276cf370c 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -12,6 +12,9 @@ the proxy to POST /responses) and a streamed Anthropic-messages case. from __future__ import annotations import pytest +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( @@ -26,12 +29,14 @@ from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatStreamOptions from scripted_provider import ScriptedUsage -pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] +pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark -_MODELS: dict[str, FrontierModel] = {model.map_key: model for model in FRONTIER_MODELS} +_MODELS: Final[Mapping[str, FrontierModel]] = MappingProxyType( + {model.map_key: model for model in FRONTIER_MODELS} +) # One scripted usage per wire, every reportable token kind nonzero. -_WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { +_WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ "openai_chat": ( "gpt-5.6", ScriptedUsage( @@ -89,7 +94,7 @@ _WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { "fireworks_ai/kimi-k3", ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), ), -} +}) class TestWireFormats: @@ -103,22 +108,22 @@ class TestWireFormats: wire: str, ) -> None: map_key, usage = _WIRE_USAGE[wire] - model = _MODELS[map_key] - case = Case(name="basic", usage=usage) - marker = unique_marker() + model: Final = _MODELS[map_key] + case: Final = Case(name="basic", usage=usage) + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted wire call")], + messages=(ChatMessage(role="user", content=f"{marker} scripted wire call"),), ), ) assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" - expected = expected_breakdown(model, case) - row = cost_rows.poll_cost_row_where( + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, @@ -128,7 +133,7 @@ class TestWireFormats: f"{wire}: spend {row.spend} != expected {expected.total} " f"(breakdown {row.breakdown.model_dump()})" ) - breakdown = row.breakdown + breakdown: Final = row.breakdown assert breakdown.input_cost is not None and cost_rows.approx_equal( breakdown.input_cost, expected.input_cost ), ( @@ -153,16 +158,16 @@ class TestWireFormats: self, client: CostCalcClient, resources: ResourceManager, scoped_key: str ) -> None: map_key, usage = _WIRE_USAGE["anthropic_messages"] - model = _MODELS[map_key] - case = Case(name="stream", usage=usage, stream=True) - marker = unique_marker() + model: Final = _MODELS[map_key] + case: Final = Case(name="stream", usage=usage, stream=True) + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted anthropic stream")], + messages=(ChatMessage(role="user", content=f"{marker} scripted anthropic stream"),), stream=True, stream_options=ChatStreamOptions(include_usage=True), ), @@ -172,8 +177,8 @@ class TestWireFormats: assert response.stream_done, "anthropic stream did not reach its terminal event" assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected = expected_breakdown(model, case) - row = cost_rows.poll_cost_row_where( + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, From 99bf8e9b2ffb6c647813029debba23c788e86b41 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:32:39 +0000 Subject: [PATCH 044/393] test(e2e): add cost calculation CI proxy config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/gateway/cost_calculation_ci_config.yml | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/gateway/cost_calculation_ci_config.yml diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index b6c3840f626..49cfc29aa17 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/gateway/cost_calculation_ci_config.yml b/tests/e2e/gateway/cost_calculation_ci_config.yml new file mode 100644 index 00000000000..ac0603fa7c1 --- /dev/null +++ b/tests/e2e/gateway/cost_calculation_ci_config.yml @@ -0,0 +1,7 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_batch_write_at: 5 + +model_list: [] From 415b06f5ff6d5959e2144ea826922694b2c8a60b Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 04:42:07 +0000 Subject: [PATCH 045/393] test(e2e): assert the real bill for the four fixed cost gaps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 22 +++++-------------- .../test_token_pricing_e2e.py | 5 ----- tests/e2e/cost_map.json | 15 +++++++++++++ 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index e8b1d249559..8f39e89a358 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -186,15 +186,12 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ } ), "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), - # Product gap: litellm hard-indexes message_delta["usage"] in - # anthropic/chat/handler.py, so a usage-absent anthropic stream raises - # KeyError; the real wire always carries it, so the case cannot be - # represented. - "anthropic_messages": frozenset({"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model"}), - # Product gap: the gemini transform sets ModelResponse.model from the - # request and drops the provider's modelVersion, so a response-model - # override can never be priced on this wire. - "gemini_generate": frozenset({"cache_read", "reasoning", "audio", "web_search", "absent_usage"}), + "anthropic_messages": frozenset( + {"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model", "absent_usage"} + ), + "gemini_generate": frozenset( + {"cache_read", "reasoning", "audio", "web_search", "response_model", "absent_usage"} + ), "together_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", @@ -240,9 +237,6 @@ class Case: billed_web_search_calls: int = 0 response_model_override: bool = False exact_spend: bool = True - # stream_usage=absent on a wire with no proxy-side token recount means the - # bill is exactly zero; asserted as such rather than skipped. - expect_zero_bill: bool = False def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( @@ -359,10 +353,6 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: stream=True, stream_usage="absent", exact_spend=False, - # The responses surface bills only provider-reported usage; - # with no usage in the stream the spend row is zero. Other - # wires recount tokens proxy-side and bill a nonzero amount. - expect_zero_bill=model.wire == "openai_responses", ) if "absent_usage" in caps else None diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index e210dad94b1..ead86931424 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -90,11 +90,6 @@ class TestTokenPricing: ) assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}" - if not case.exact_spend and case.expect_zero_bill: - # The provider reported no usage and this wire has no proxy-side - # recount, so the bill is exactly zero. - assert row.spend is not None and row.spend == 0, f"no-usage stream billed {row.spend}: {row}" - return if not case.exact_spend: # stream_usage=absent: the provider reported no usage, so the row's # token counts are the proxy's own recount; only assert a bill landed. diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index b761710bae3..68d840870c9 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -63,13 +63,18 @@ "supports_web_search": true }, "fireworks_ai/deepseek-v4p1-flash": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.4e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00014000000000000001, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00028000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -82,13 +87,18 @@ "supports_web_search": true }, "fireworks_ai/kimi-k3": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.2e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00012000000000000002, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00024000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -101,13 +111,18 @@ "supports_web_search": true }, "fireworks_ai/qwen3p8-max": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.3e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00013000000000000002, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00026000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, From 67778cfe2625eb97fd3d4733f1fae41aed7599fb Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:18:08 +0000 Subject: [PATCH 046/393] fix(cost): resolve dated openai/azure snapshots to their undated cost map entry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 9 +++++++- tests/test_litellm/test_cost_calculator.py | 24 ++++++++++++++++++++++ tests/test_litellm/test_utils.py | 19 +++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 734522c0c6a..33fdfcc36ba 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5286,6 +5286,13 @@ def _strip_stable_vertex_version(model_name) -> str: return re.sub(r"-\d+$", "", model_name) +_DATED_SNAPSHOT_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$") + + +def _strip_dated_snapshot_suffix(model_name: str) -> str: + return _DATED_SNAPSHOT_SUFFIX.sub("", model_name) + + def _get_base_bedrock_model(model_name) -> str: """ Get the base model from the given model name. @@ -5333,7 +5340,7 @@ def _strip_model_name(model: str, custom_llm_provider: str | None) -> str: strip_finetune: Final = _strip_openai_finetune_model_name(model_name=model) return strip_finetune else: - return model + return _strip_dated_snapshot_suffix(model_name=model) # Global case-insensitive lookup map for model_cost (built eagerly at module import) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7b53d3a58df..c5bc8d80fed 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -21,7 +21,9 @@ from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, + Choices, LiteLLMRealtimeStreamLoggingObject, + Message, ModelInfo, ModelResponse, PromptTokensDetailsWrapper, @@ -110,6 +112,28 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: + dated_response = ModelResponse( + model="gpt-5.6-luna-2026-07-09", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + dated_response._hidden_params = {"custom_llm_provider": "azure"} + + undated_response = ModelResponse( + model="gpt-5.6-luna", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + undated_response._hidden_params = {"custom_llm_provider": "azure"} + + dated_cost = litellm.completion_cost(completion_response=dated_response) + undated_cost = litellm.completion_cost(completion_response=undated_response) + + assert dated_cost == undated_cost + assert dated_cost > 0 + + def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): _hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 1000}} diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 46149589371..2f9e27af797 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -186,6 +186,25 @@ def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local assert info["key"] == "ft:gpt-4o-2024-08-06" +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_key"), + [ + ("gpt-5.6-luna-2026-07-09", "openai", "gpt-5.6-luna"), + ("gpt-5.6-luna-2026-07-09", "azure", "azure/gpt-5.6-luna"), + ], +) +def test_get_model_info_falls_back_from_dated_snapshot_to_undated_entry( + local_model_cost_map, model, custom_llm_provider, expected_key +): + info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + assert info["key"] == expected_key + + +def test_get_model_info_prefers_exact_dated_key_over_stripped(local_model_cost_map): + info = litellm.get_model_info(model="gpt-4o-2024-08-06", custom_llm_provider="openai") + assert info["key"] == "gpt-4o-2024-08-06" + + def test_check_provider_match_azure_ai_allows_openai_and_azure(): """ Test that azure_ai provider can match openai and azure models. From feb69c5f789d44a65dbbfa348ce39eaa3874b37f Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 16:38:15 +0000 Subject: [PATCH 047/393] test(e2e): add tool-call, terminal, and image-input shapes to the cost matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 184 +++++++- .../e2e/cost_calculation/scripted_provider.py | 408 +++++++++++++++--- .../test_token_pricing_e2e.py | 61 ++- .../cost_calculation/test_wire_formats_e2e.py | 111 ++++- 4 files changed, 688 insertions(+), 76 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 8f39e89a358..5f634712778 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -17,7 +17,11 @@ creation), the case is absent from the matrix rather than silently zero. from __future__ import annotations +import base64 import json +import random +import struct +import zlib from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -26,7 +30,7 @@ from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter -from scripted_provider import Scenario, ScriptedOutput, ScriptedUsage, Wire +from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" @@ -182,26 +186,37 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "openai_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", + } + ), + "openai_responses": frozenset( + { + "cache_read", "reasoning", "web_search", "response_model", "absent_usage", + "tool_call", "image_input", "responses_terminal", } ), - "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), "anthropic_messages": frozenset( - {"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model", "absent_usage"} + { + "cache_read", "cache_write_5m", "cache_write_1h", "web_search", + "response_model", "absent_usage", "tool_call", "image_input", + } ), "gemini_generate": frozenset( - {"cache_read", "reasoning", "audio", "web_search", "response_model", "absent_usage"} + { + "cache_read", "reasoning", "audio", "web_search", "response_model", + "absent_usage", "tool_call", "image_input", "prompt_blocked", + } ), "together_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), "fireworks_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), }) @@ -220,6 +235,15 @@ CaseName: TypeAlias = Literal[ "stream", "stream_no_usage", "response_model_override", + "stream_response_model_override", + "tool_call", + "stream_no_usage_tool_call", + "stream_no_usage_image_input", + "stream_no_usage_incomplete", + "stream_unvalidated", + "stream_no_usage_unvalidated", + "prompt_blocked", + "stream_prompt_blocked", ] @@ -237,6 +261,9 @@ class Case: billed_web_search_calls: int = 0 response_model_override: bool = False exact_spend: bool = True + tool_call: bool = False + image_input: bool = False + terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( @@ -246,6 +273,10 @@ class Case: output=ScriptedOutput( text=text, response_model=model.override_model if self.response_model_override else None, + tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) + if self.tool_call + else None, + terminal=self.terminal, ), stream_usage=self.stream_usage, service_tier=self.service_tier, @@ -254,6 +285,15 @@ class Case: _BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) +TOOL_CALL_ARGUMENTS: Final = json.dumps({ + "city": "Berlin", + "days": 7, + "units": "metric", + "notes": "filler " * 30, +}) + +_PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tokens=0) + def _web_search_case(model: FrontierModel) -> Case: counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") @@ -362,6 +402,100 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: if "response_model" in caps else None ), + ( + Case( + name="stream_response_model_override", + usage=_BASIC_USAGE, + stream=True, + response_model_override=True, + ) + if "response_model" in caps + else None + ), + ( + Case(name="tool_call", usage=_BASIC_USAGE, tool_call=True) + if "tool_call" in caps + else None + ), + ( + Case( + name="stream_no_usage_tool_call", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + tool_call=True, + exact_spend=False, + ) + if "absent_usage" in caps and "tool_call" in caps + else None + ), + ( + Case( + name="stream_no_usage_image_input", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + image_input=True, + exact_spend=False, + ) + if "absent_usage" in caps and "image_input" in caps + else None + ), + ( + Case( + name="stream_no_usage_incomplete", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + terminal="incomplete", + exact_spend=False, + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="stream_unvalidated", + usage=_BASIC_USAGE, + stream=True, + terminal="unvalidated", + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="stream_no_usage_unvalidated", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + terminal="unvalidated", + exact_spend=False, + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="prompt_blocked", + usage=_PROMPT_BLOCKED_USAGE, + terminal="prompt_blocked", + response_model_override=True, + ) + if "prompt_blocked" in caps + else None + ), + ( + Case( + name="stream_prompt_blocked", + usage=_PROMPT_BLOCKED_USAGE, + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ) + if "prompt_blocked" in caps + else None + ), ) return tuple(case for case in candidates if case is not None) @@ -436,6 +570,42 @@ def expected_cost(model: FrontierModel, case: Case) -> float: return expected_breakdown(model, case).total +def recount_cost( + model: FrontierModel, case: Case, prompt_tokens: int, completion_tokens: int +) -> float: + """What the proxy's own token recount should cost at the case's rates, + without pinning the tokenizer's exact counts.""" + rates: Final = model.override_rates if case.response_model_override else model.rates + return prompt_tokens * (rates.input_cost_per_token or 0.0) + completion_tokens * ( + rates.output_cost_per_token or 0.0 + ) + + +def _png_chunk(tag: bytes, payload: bytes) -> bytes: + return struct.pack(">I", len(payload)) + tag + payload + struct.pack(">I", zlib.crc32(tag + payload)) + + +def image_input_data_url() -> str: + """A deterministic 256x256 RGB noise PNG as a data URL; noise compresses + poorly on purpose so the base64 payload stays well above 100 KB and would + blow up the prompt recount if the URL were ever tokenized as text.""" + rng: Final = random.Random(0) + side: Final = 256 + raw: Final = b"".join( + b"\x00" + rng.randbytes(side * 3) for _ in range(side) + ) + png: Final = ( + b"\x89PNG\r\n\x1a\n" + + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", side, side, 8, 2, 0, 0, 0)) + + _png_chunk(b"IDAT", zlib.compress(raw)) + + _png_chunk(b"IEND", b"") + ) + return "data:image/png;base64," + base64.b64encode(png).decode() + + +IMAGE_INPUT_DATA_URL: Final = image_input_data_url() + + def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: """(prompt_tokens, completion_tokens) the spend row should carry, per the wire's normalization: Anthropic folds cache read/write into prompt_tokens, diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index f1deafd1bc5..e1a6c430307 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -38,7 +38,7 @@ from types import MappingProxyType from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator Wire: TypeAlias = Literal[ "openai_chat", @@ -62,6 +62,26 @@ WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( StreamUsage: TypeAlias = Literal["final_chunk", "absent"] ServiceTier: TypeAlias = Literal["flex", "priority"] +TerminalKind: TypeAlias = Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] + +# Which terminal variant each wire can represent. +_TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "openai_responses": frozenset({"incomplete", "unvalidated"}), + "gemini_generate": frozenset({"prompt_blocked"}), + } +) + + +class ScriptedToolCall(BaseModel): + """A single function call the scripted output emits instead of text. + ``arguments`` is the wire's JSON string (~250 chars), sliced into deltas + for streams.""" + + model_config = ConfigDict(frozen=True) + + name: str + arguments: str class ScriptedUsage(BaseModel): @@ -96,6 +116,12 @@ class ScriptedOutput(BaseModel): # OpenAI-compatible providers can report a provider-computed cost; emitted as # the top-level "cost" field on the together/fireworks wire. provider_cost: float | None = None + # When set, the response is a tool call only: no text content on any wire. + tool_call: ScriptedToolCall | None = None + # Terminal shape: "unvalidated" makes the Responses terminal response fail + # pydantic validation so the proxy takes its model_construct dict path; + # "prompt_blocked" is a Gemini promptFeedback-only body. + terminal: TerminalKind = "completed" class Scenario(BaseModel): @@ -108,6 +134,17 @@ class Scenario(BaseModel): stream_usage: StreamUsage = "final_chunk" service_tier: ServiceTier | None = None + @model_validator(mode="after") + def _check_terminal_supported(self) -> Scenario: + if ( + self.output.terminal != "completed" + and self.output.terminal not in _TERMINAL_CAPS.get(self.wire, frozenset()) + ): + raise ValueError( + f"wire {self.wire} cannot emit terminal={self.output.terminal}" + ) + return self + @property def mount(self) -> str: return WIRE_MOUNTS[self.wire] @@ -291,10 +328,38 @@ def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: # ---------- per-wire responses ---------- +def _split_arguments(arguments: str) -> tuple[str, ...]: + """Slice a tool-call arguments JSON string into 2-3 streamed deltas.""" + third: Final = max(1, len(arguments) // 3) + return tuple( + slice_ + for slice_ in (arguments[:third], arguments[third : 2 * third], arguments[2 * third :]) + if slice_ + ) + + def _openai_message(scenario: Scenario) -> Mapping[str, object]: + tool_call: Final = scenario.output.tool_call return _jobj_opt( ("role", "assistant"), - ("content", scenario.output.text), + ("content", None if tool_call is not None else scenario.output.text), + ( + ( + "tool_calls", + ( + _jobj( + ("id", f"call_{scenario.scenario_id}"), + ("type", "function"), + ( + "function", + _jobj(("name", tool_call.name), ("arguments", tool_call.arguments)), + ), + ), + ), + ) + if tool_call is not None + else None + ), ( ( "annotations", @@ -332,7 +397,12 @@ def _openai_chat_body(scenario: Scenario, requested_model: str) -> Mapping[str, _jobj( ("index", 0), ("message", _openai_message(scenario)), - ("finish_reason", scenario.output.finish_reason), + ( + "finish_reason", + "tool_calls" + if scenario.output.tool_call is not None + else scenario.output.finish_reason, + ), ), ), ), @@ -359,6 +429,7 @@ def _openai_chunk( def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: + tool_call: Final = scenario.output.tool_call delta: Final = _jobj_opt( ("role", "assistant"), ("content", scenario.output.text), @@ -368,6 +439,43 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: else None ), ) + body_deltas: Final[tuple[Mapping[str, object], ...]] = ( + ( + _jobj( + ("role", "assistant"), + ( + "tool_calls", + ( + _jobj( + ("index", 0), + ("id", f"call_{scenario.scenario_id}"), + ("type", "function"), + ( + "function", + _jobj(("name", tool_call.name), ("arguments", "")), + ), + ), + ), + ), + ), + *( + _jobj( + ( + "tool_calls", + ( + _jobj( + ("index", 0), + ("function", _jobj(("arguments", arguments_slice))), + ), + ), + ) + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ), + ) + if tool_call is not None + else (delta,) + ) return _sse( ( ( @@ -378,13 +486,16 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: choices=(_jobj(("index", 0), ("delta", _jobj(("role", "assistant"))), ("finish_reason", None)),), ), ), - ( - None, - _openai_chunk( - scenario, - requested_model, - choices=(_jobj(("index", 0), ("delta", delta), ("finish_reason", None)),), - ), + *( + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", body_delta), ("finish_reason", None)),), + ), + ) + for body_delta in body_deltas ), ( None, @@ -395,7 +506,12 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: _jobj( ("index", 0), ("delta", _jobj()), - ("finish_reason", scenario.output.finish_reason), + ( + "finish_reason", + "tool_calls" + if tool_call is not None + else scenario.output.finish_reason, + ), ), ), ), @@ -410,17 +526,34 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: ) +def _anthropic_content(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ("type", "tool_use"), + ("id", f"toolu_{scenario.scenario_id}"), + ("name", tool_call.name), + ("input", json.loads(tool_call.arguments)), + ), + ) + return (_jobj(("type", "text"), ("text", scenario.output.text)),) + + +def _anthropic_stop_reason(scenario: Scenario) -> str: + if scenario.output.tool_call is not None: + return "tool_use" + return "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + + def _anthropic_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: return _jobj( ("id", f"msg_{scenario.scenario_id}"), ("type", "message"), ("role", "assistant"), ("model", scenario.output.response_model or requested_model), - ("content", (_jobj(("type", "text"), ("text", scenario.output.text)),)), - ( - "stop_reason", - "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - ), + ("content", _anthropic_content(scenario)), + ("stop_reason", _anthropic_stop_reason(scenario)), ("usage", _anthropic_usage(scenario.usage)), ) @@ -453,12 +586,7 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: ("type", "message_delta"), ( "delta", - _jobj( - ( - "stop_reason", - "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - ) - ), + _jobj(("stop_reason", _anthropic_stop_reason(scenario))), ), ( ("usage", _jobj(("output_tokens", scenario.usage.output_tokens))) @@ -474,16 +602,45 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: _jobj( ("type", "content_block_start"), ("index", 0), - ("content_block", _jobj(("type", "text"), ("text", ""))), + ( + "content_block", + _jobj( + ("type", "tool_use"), + ("id", f"toolu_{scenario.scenario_id}"), + ("name", scenario.output.tool_call.name), + ("input", _jobj()), + ) + if scenario.output.tool_call is not None + else _jobj(("type", "text"), ("text", "")), + ), ), ), - ( - "content_block_delta", - _jobj( - ("type", "content_block_delta"), - ("index", 0), - ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), - ), + *( + tuple( + ( + "content_block_delta", + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ( + "delta", + _jobj(("type", "input_json_delta"), ("partial_json", arguments_slice)), + ), + ), + ) + for arguments_slice in _split_arguments(scenario.output.tool_call.arguments) + ) + if scenario.output.tool_call is not None + else ( + ( + "content_block_delta", + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), + ), + ), + ) ), ("content_block_stop", _jobj(("type", "content_block_stop"), ("index", 0))), ("message_delta", message_delta), @@ -492,7 +649,49 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: ) +def _gemini_prompt_blocked_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ( + "promptFeedback", + _jobj( + ("blockReason", "SAFETY"), + ( + "safetyRatings", + ( + _jobj( + ("category", "HARM_CATEGORY_HARASSMENT"), + ("probability", "HIGH"), + ("blocked", True), + ), + ), + ), + ), + ), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ) + + +def _gemini_parts(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ( + "functionCall", + _jobj( + ("name", tool_call.name), + ("args", json.loads(tool_call.arguments)), + ), + ) + ), + ) + return (_jobj(("text", scenario.output.text)),) + + def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + if scenario.output.terminal == "prompt_blocked": + return _gemini_prompt_blocked_body(scenario, requested_model) return _jobj( ( "candidates", @@ -501,7 +700,7 @@ def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, objec ( "content", _jobj( - ("parts", (_jobj(("text", scenario.output.text)),)), + ("parts", _gemini_parts(scenario)), ("role", "model"), ), ), @@ -559,67 +758,148 @@ def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: ) -def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: - return _jobj( - ("id", f"resp_{scenario.scenario_id}"), - ("object", "response"), - ("created_at", int(time.time())), - ("status", "completed"), - ("model", scenario.output.response_model or requested_model), - ( - "output", +def _responses_output(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + return ( + *( ( - *( - _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) - for i in range(scenario.usage.web_search_calls) - ), - _jobj( - ("type", "message"), - ("id", f"msg_{scenario.scenario_id}"), - ("status", "completed"), - ("role", "assistant"), - ( - "content", - ( - _jobj( - ("type", "output_text"), - ("text", scenario.output.text), - ("annotations", ()), - ), - ), + _jobj(("type", "scripted_future_item"), ("id", f"fut_{scenario.scenario_id}"), ("status", "completed")), + ) + if scenario.output.terminal == "unvalidated" + else () + ), + *( + _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) + for i in range(scenario.usage.web_search_calls) + ), + _jobj( + ("type", "function_call"), + ("id", f"fc_{scenario.scenario_id}"), + ("call_id", f"call_{scenario.scenario_id}"), + ("name", tool_call.name), + ("arguments", tool_call.arguments), + ("status", "completed"), + ) + if tool_call is not None + else _jobj( + ("type", "message"), + ("id", f"msg_{scenario.scenario_id}"), + ("status", "completed"), + ("role", "assistant"), + ( + "content", + ( + _jobj( + ("type", "output_text"), + ("text", scenario.output.text), + ("annotations", ()), ), ), ), ), + ) + + +def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + incomplete: Final = scenario.output.terminal == "incomplete" + return _jobj_opt( + ("id", f"resp_{scenario.scenario_id}"), + ("object", "response"), + ( + "created_at", + "not-a-number" if scenario.output.terminal == "unvalidated" else int(time.time()), + ), + ("status", "incomplete" if incomplete else "completed"), + ( + ("incomplete_details", _jobj(("reason", "max_output_tokens"))) + if incomplete + else None + ), + ("model", scenario.output.response_model or requested_model), + ("output", _responses_output(scenario)), ("usage", _responses_usage(scenario.usage)), ) def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: - completed: Final = ( + tool_call: Final = scenario.output.tool_call + terminal: Final = ( _jobj(*((key, value) for key, value in _responses_body(scenario, requested_model).items() if key != "usage")) if scenario.stream_usage == "absent" else _responses_body(scenario, requested_model) ) created: Final = _jobj( - *((key, value) for key, value in completed.items() if key not in ("status", "usage")), + *((key, value) for key, value in terminal.items() if key not in ("status", "usage")), ("status", "in_progress"), ("usage", None), ) - return _sse( + terminal_event: Final = ( + "response.incomplete" if scenario.output.terminal == "incomplete" else "response.completed" + ) + output_index: Final = ( + scenario.usage.web_search_calls + (1 if scenario.output.terminal == "unvalidated" else 0) + ) + middle_events: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( ( - ("response.created", _jobj(("type", "response.created"), ("response", created))), + ( + "response.output_item.added", + _jobj( + ("type", "response.output_item.added"), + ("output_index", output_index), + ( + "item", + _jobj( + ("type", "function_call"), + ("id", f"fc_{scenario.scenario_id}"), + ("call_id", f"call_{scenario.scenario_id}"), + ("name", tool_call.name), + ("arguments", ""), + ("status", "in_progress"), + ), + ), + ), + ), + *( + ( + "response.function_call_arguments.delta", + _jobj( + ("type", "response.function_call_arguments.delta"), + ("item_id", f"fc_{scenario.scenario_id}"), + ("output_index", output_index), + ("delta", arguments_slice), + ), + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ), + ( + "response.function_call_arguments.done", + _jobj( + ("type", "response.function_call_arguments.done"), + ("item_id", f"fc_{scenario.scenario_id}"), + ("output_index", output_index), + ("arguments", tool_call.arguments), + ), + ), + ) + if tool_call is not None + else ( ( "response.output_text.delta", _jobj( ("type", "response.output_text.delta"), ("item_id", f"msg_{scenario.scenario_id}"), - ("output_index", scenario.usage.web_search_calls), + ("output_index", output_index), ("content_index", 0), ("delta", scenario.output.text), ), ), - ("response.completed", _jobj(("type", "response.completed"), ("response", completed))), + ) + ) + return _sse( + ( + ("response.created", _jobj(("type", "response.created"), ("response", created))), + *middle_events, + (terminal_event, _jobj(("type", terminal_event), ("response", terminal))), ) ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index ead86931424..0b4f3e1fd37 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -16,15 +16,26 @@ from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( FRONTIER_MODELS, + IMAGE_INPUT_DATA_URL, Case, FrontierModel, cases_for, expected_cost, expected_token_columns, + recount_cost, ) from e2e_config import unique_marker from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions +from models import ( + ChatBody, + ChatMessage, + ChatStreamOptions, + ChatTool, + ChatToolFunction, + ImageContentPart, + ImageUrl, + TextContentPart, +) pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark @@ -41,10 +52,37 @@ def _case_id(param: tuple[FrontierModel, Case]) -> str: def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: return ChatBody( model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted pricing call"),), + messages=( + ChatMessage( + role="user", + content=( + [ + TextContentPart(text=f"{marker} scripted pricing call"), + ImageContentPart(image_url=ImageUrl(url=IMAGE_INPUT_DATA_URL)), + ] + if case.image_input + else f"{marker} scripted pricing call" + ), + ), + ), stream=case.stream, stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, service_tier=case.service_tier, + tools=( + ( + ChatTool( + function=ChatToolFunction( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ) + ), + ) + if case.tool_call + else None + ), ) @@ -92,8 +130,23 @@ class TestTokenPricing: if not case.exact_spend: # stream_usage=absent: the provider reported no usage, so the row's - # token counts are the proxy's own recount; only assert a bill landed. - assert row.spend is not None and row.spend > 0, f"no-usage stream billed nothing: {row}" + # token counts are the proxy's own recount; assert the recount + # billed both directions at the case's rates. + assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( + f"no-usage stream counted no input tokens: {row}" + ) + assert row.completion_tokens is not None and row.completion_tokens > 0, ( + f"no-usage stream counted no output tokens: {row}" + ) + if case.image_input: + assert row.prompt_tokens < 4000, ( + f"image data URL looks tokenized as text: prompt_tokens={row.prompt_tokens}" + ) + assert row.spend is not None and cost_rows.approx_equal( + row.spend, + recount_cost(model, case, row.prompt_tokens, row.completion_tokens), + ), f"no-usage stream spend {row.spend} != recount at map rates: {row}" + cost_rows.assert_total_is_sum_of_components(row) return assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index c0276cf370c..3c6c34b24fb 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -26,7 +26,7 @@ from cost_matrix import ( ) from e2e_config import unique_marker from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions +from models import ChatBody, ChatMessage, ChatStreamOptions, ChatTool, ChatToolFunction from scripted_provider import ScriptedUsage pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark @@ -96,6 +96,53 @@ _WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ ), }) +_SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) + +# Renderer-level shapes the pricing matrix gates per cap, pinned here once per +# wire so the sidecar emits prove they survive the proxy end to end. +_SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( + *( + ( + f"tool_call_{'stream' if stream else 'sync'}", + wire, + Case(name="tool_call", usage=_SHAPE_USAGE, stream=stream, tool_call=True), + ) + for wire in _WIRE_USAGE + for stream in (False, True) + ), + ( + "responses_incomplete", + "openai_responses", + Case(name="stream_no_usage_incomplete", usage=_SHAPE_USAGE, stream=True, terminal="incomplete"), + ), + ( + "responses_unvalidated", + "openai_responses", + Case(name="stream_unvalidated", usage=_SHAPE_USAGE, stream=True, terminal="unvalidated"), + ), + ( + "gemini_prompt_blocked", + "gemini_generate", + Case( + name="prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "gemini_prompt_blocked_stream", + "gemini_generate", + Case( + name="stream_prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ), + ), +) + class TestWireFormats: @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) @@ -189,3 +236,65 @@ class TestWireFormats: f"(breakdown {row.breakdown.model_dump()})" ) cost_rows.assert_total_is_sum_of_components(row) + + @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=lambda entry: entry[0]) + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_response_shape_bills_reported_usage( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + shape_wire_case: tuple[str, str, Case], + ) -> None: + shape, wire, case = shape_wire_case + map_key, _usage = _WIRE_USAGE[wire] + model: Final = _MODELS[map_key] + marker: Final = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response: Final = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=(ChatMessage(role="user", content=f"{marker} scripted {shape}"),), + stream=case.stream, + stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, + tools=( + ( + ChatTool( + function=ChatToolFunction( + name="get_weather", + parameters={"type": "object", "properties": {"city": {"type": "string"}}}, + ) + ), + ) + if case.tool_call + else None + ), + ), + stream=case.stream, + ) + assert response.ok, f"{shape}: proxy returned {response.status_code}: {response.body[:400]}" + if case.stream: + assert response.stream_done, f"{shape}: stream did not reach its terminal event" + assert response.stream_error is None, f"{shape}: stream error: {response.stream_error}" + + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"{shape}: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"{shape}: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"{shape}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"{shape}: completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) From 5507de326e3e98f9069af5f9d1315c89bb3c3e25 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 16:45:09 +0000 Subject: [PATCH 048/393] test(e2e): type the wire-shape parametrize ids callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/test_wire_formats_e2e.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index 3c6c34b24fb..4da7b31a6ef 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -144,6 +144,10 @@ _SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( ) +def _shape_id(entry: tuple[str, str, Case]) -> str: + return entry[0] + + class TestWireFormats: @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") @@ -237,7 +241,7 @@ class TestWireFormats: ) cost_rows.assert_total_is_sum_of_components(row) - @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=lambda entry: entry[0]) + @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=_shape_id) @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") def test_response_shape_bills_reported_usage( self, From 489a3ecf95bd354da16b4d50433552fd3c22f100 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:19:45 -0700 Subject: [PATCH 049/393] feat(proxy): add RFC 8693 token exchange for IdP JWTs on the gateway token endpoint A registered gateway DCR client can now POST /token with grant_type=urn:ietf:params:oauth:grant-type:token-exchange and an IdP JWT as subject_token. The gateway proves the JWT the way its JWT auth does, resolves the user and team, and answers with the proxy-API credential and a refresh token, so a fresh laptop with only an IdP login gets a gateway key without a browser round trip. "/token" joins mcp_inference_routes so the default JWT team allowlist reaches the exchange, and the JWT auth builder accepts any header mapping so the request headers pass through unchanged. --- .../mcp_server/discoverable_endpoints.py | 13 +- .../mcp_server/gateway_dcr_flow.py | 148 ++++++++++++-- .../mcp_server/idp_token_exchange.py | 105 ++++++++++ litellm/proxy/_lazy_openapi_snapshot.json | 68 ++++++- litellm/proxy/_types.py | 1 + litellm/proxy/auth/handle_jwt.py | 8 +- .../mcp_server/test_gateway_dcr_flow.py | 186 +++++++++++++++++- .../mcp_server/test_idp_token_exchange.py | 115 +++++++++++ .../proxy/auth/test_auth_checks.py | 18 ++ .../proxy/auth/test_route_checks.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 ++ 11 files changed, 654 insertions(+), 21 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/idp_token_exchange.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ffb27d5f92e..ebd5f43bf02 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + TOKEN_EXCHANGE_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -60,6 +61,9 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( relative_request_url, revoke_refresh_token, ) +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + exchange_idp_subject_token, +) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, RefreshTokenPresented, @@ -1980,6 +1984,9 @@ async def token_endpoint( refresh_token: str | None = Form(None), scope: str | None = Form(None), resource: str | None = Form(None), + subject_token: str | None = Form(None), + subject_token_type: str | None = Form(None), + requested_token_type: str | None = Form(None), mcp_server_name: str | None = None, ): """ @@ -2010,6 +2017,10 @@ async def token_endpoint( cache=user_api_key_cache, resource=resource, mint_proxy_credential=mint_proxy_credential, + subject_token=subject_token, + subject_token_type=subject_token_type, + requested_token_type=requested_token_type, + exchange_subject_token=exchange_idp_subject_token, ) lookup_name: Final = mcp_server_name or client_id @@ -2638,7 +2649,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], } diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index f3fdd54b39d..9bdde3c5edc 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -51,7 +51,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, ValidationError -from typing_extensions import ReadOnly, TypedDict, assert_never +from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -187,6 +187,41 @@ class MintProxyCredential(Protocol): ) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ... +TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" +"""RFC 8693: a native client that already holds a token from the customer's identity +provider trades it for the proxy-API credential without a browser round trip.""" + +_IssuedTokenType = Literal["urn:ietf:params:oauth:token-type:access_token"] +ACCESS_TOKEN_TOKEN_TYPE: Final[_IssuedTokenType] = "urn:ietf:params:oauth:token-type:access_token" +SUBJECT_TOKEN_TYPES: Final = frozenset( + { + "urn:ietf:params:oauth:token-type:jwt", + "urn:ietf:params:oauth:token-type:id_token", + ACCESS_TOKEN_TOKEN_TYPE, + } +) + + +class SubjectIdentity(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str = Field(min_length=1) + team_id: str | None = None + + +class SubjectTokenRefusal(BaseModel): + model_config = ConfigDict(frozen=True) + error: Literal["unsupported_grant_type", "invalid_request"] + description: str = Field(min_length=1) + + +class ExchangeSubjectToken(Protocol): + """Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the + IdP token the way the proxy's own JWT auth does and names the litellm user and team it + stands for, or says why this gateway will not take it.""" + + def __call__(self, subject_token: str, request: Request, /) -> Awaitable[SubjectIdentity | SubjectTokenRefusal]: ... + + class ConsentTeam(BaseModel): model_config = ConfigDict(frozen=True) team_id: str = Field(min_length=1) @@ -213,6 +248,12 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _refuse_subject_token(subject_token: str, request: Request) -> SubjectTokenRefusal: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="this gateway is not configured to exchange IdP tokens" + ) + + async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: return "unavailable" @@ -382,7 +423,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str, "client_id_issued_at": int(now.timestamp()), "redirect_uris": list(raw_uris), "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token"], + "grant_types": ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE], "response_types": ["code"], }, ) @@ -595,7 +636,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract: "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": ("authorization_code", "refresh_token"), + "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), @@ -1033,20 +1074,26 @@ class _ProxyCredentialTokenResponse(TypedDict): refresh_token: ReadOnly[str] user_id: ReadOnly[str] team_id: ReadOnly[str | None] + issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]] def _proxy_credential_response( - minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime + minted: MintedProxyCredential, + principal: SessionPrincipal, + keys: SessionSigningKeys, + now: datetime, + issued_token_type: _IssuedTokenType | None = None, ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and the refresh token is a gateway-sealed rotating token bound to the team the credential - was minted for, so a renewal keeps the team the user consented to.""" + was minted for, so a renewal keeps the team the user consented to. A token exchange + also states ``issued_token_type``, which RFC 8693 section 2.2.1 requires.""" bound_principal: Final = principal.model_copy(update=MappingProxyType({"team_id": minted.team_id})) refresh: Final = mint_session_refresh_token(bound_principal, keys, now) if not isinstance(refresh, MintedSessionToken): return _oauth_error(500, "server_error", "failed to mint the session credential") - body: Final[_ProxyCredentialTokenResponse] = { + credential: Final[_ProxyCredentialTokenResponse] = { "access_token": minted.key, "token_type": "Bearer", "expires_in": minted.expires_in, @@ -1054,7 +1101,10 @@ def _proxy_credential_response( "user_id": minted.user_id, "team_id": minted.team_id, } - return JSONResponse(status_code=200, content=body, headers=TOKEN_NO_CACHE_HEADERS) + if issued_token_type is None: + return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS) + exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type} + return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS) def _reload_failure_response(failure: ReloadUserFailure) -> Response: @@ -1116,11 +1166,16 @@ async def aggregate_token( cache: DualCache, resource: str | None = None, mint_proxy_credential: MintProxyCredential = _refuse_proxy_credential, + subject_token: str | None = None, + subject_token_type: str | None = None, + requested_token_type: str | None = None, + exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token, ) -> Response: """The aggregate token verb: authorization_code and refresh_token grants for the identity-only session pair, or for the proxy-API credential when the grant was issued - with that audience. Every path re-validates the litellm user live before minting, so a - deactivated user cannot obtain or renew a session.""" + with that audience, and the RFC 8693 token exchange that turns an IdP token straight + into the proxy-API credential. Every path re-validates the litellm user live before + minting, so a deactivated user cannot obtain or renew a session.""" if master_key is None: verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") @@ -1159,7 +1214,20 @@ async def aggregate_token( now=now, issue=issue, ) - return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token") + if grant_type == TOKEN_EXCHANGE_GRANT_TYPE: + return await _token_exchange_grant( + subject_token=subject_token, + subject_token_type=subject_token_type, + requested_token_type=requested_token_type, + client_id=client_id, + exchange_subject_token=exchange_subject_token, + issue=issue, + ) + return _oauth_error( + 400, + "unsupported_grant_type", + f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}", + ) class _GrantIssuer: @@ -1211,10 +1279,9 @@ class _GrantIssuer: async def _issue_proxy_credential( self, principal: SessionPrincipal, claim_key: str, claim_ttl_seconds: int, replayed: str ) -> Response: - if self._resource is not None and not is_proxy_api_resource(self._request, self._resource): - return _oauth_error( - 400, "invalid_target", "resource does not match the proxy API this grant was issued for" - ) + target_refusal: Final = self._proxy_api_target_refusal() + if target_refusal is not None: + return target_refusal minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id) if not isinstance(minted, MintedProxyCredential): return _mint_failure_response(minted) @@ -1223,6 +1290,33 @@ class _GrantIssuer: return refusal return _proxy_credential_response(minted, principal, self._keys, self._now) + async def exchange( + self, subject_token: str, client_id: str, exchange_subject_token: ExchangeSubjectToken + ) -> Response: + """The RFC 8693 tail: prove the IdP token, then mint. No single-use marker, because + the subject token stays a valid proof for as long as the IdP says it is and every + exchange mints a fresh credential and refresh token of its own.""" + target_refusal: Final = self._proxy_api_target_refusal() + if target_refusal is not None: + return target_refusal + identity: Final = await exchange_subject_token(subject_token, self._request) + if isinstance(identity, SubjectTokenRefusal): + return _oauth_error(400, identity.error, identity.description) + principal: Final = SessionPrincipal( + user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id + ) + minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id) + if not isinstance(minted, MintedProxyCredential): + return _mint_failure_response(minted) + return _proxy_credential_response( + minted, principal, self._keys, self._now, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE + ) + + def _proxy_api_target_refusal(self) -> Response | None: + if self._resource is None or is_proxy_api_resource(self._request, self._resource): + return None + return _oauth_error(400, "invalid_target", "resource does not match the proxy API this grant was issued for") + async def _claim_refusal(self, claim_key: str, claim_ttl_seconds: int, replayed: str) -> Response | None: return _claim_refusal( await self._guard.claim(claim_key, claim_ttl_seconds), replayed=_oauth_error(400, "invalid_grant", replayed) @@ -1297,6 +1391,32 @@ async def _refresh_token_grant( ) +async def _token_exchange_grant( + subject_token: str | None, + subject_token_type: str | None, + requested_token_type: str | None, + client_id: str, + exchange_subject_token: ExchangeSubjectToken, + issue: _GrantIssuer, +) -> Response: + """RFC 8693 token exchange for a registered native client that already holds an IdP + token: the gateway proves the token the way its JWT auth does and answers with the + proxy-API credential, so a fresh laptop with only an IdP login gets a gateway key + without a browser round trip. The client must be registered because the refresh token + in the answer is bound to it.""" + if not is_gateway_dcr_client_id(client_id) or open_gateway_dcr_client(client_id) is None: + return _oauth_error(401, "invalid_client", "unknown or malformed client_id") + if not subject_token or not subject_token_type: + return _oauth_error(400, "invalid_request", "subject_token and subject_token_type are required") + if subject_token_type not in SUBJECT_TOKEN_TYPES: + return _oauth_error( + 400, "invalid_request", f"subject_token_type must be one of {', '.join(sorted(SUBJECT_TOKEN_TYPES))}" + ) + if requested_token_type is not None and requested_token_type != ACCESS_TOKEN_TOKEN_TYPE: + return _oauth_error(400, "invalid_request", f"requested_token_type must be {ACCESS_TOKEN_TOKEN_TYPE}") + return await issue.exchange(subject_token, client_id, exchange_subject_token) + + async def revoke_refresh_token(token: str, client_id: str, master_key: str | None, cache: DualCache) -> Response: """RFC 7009 revocation for the gateway's refresh tokens: burn the presented token's ``jti`` so neither the holder nor a thief can rotate it again. Access tokens are diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py new file mode 100644 index 00000000000..69b81713810 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -0,0 +1,105 @@ +"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native +client that already holds a JWT from the customer's IdP trades it for the same proxy-API +credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims, +and the user and team sync it performs), so no browser round trip is needed.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, Protocol + +from fastapi import HTTPException, Request + +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal +from litellm.proxy._types import JWTAuthBuilderResult, ProxyException +from litellm.proxy.auth.handle_jwt import JWTAuthManager + +EXCHANGE_ROUTE: Final = "/token" + + +class AuthorizeSubjectToken(Protocol): + """Injected JWT authorization ``(subject_token, request_headers)``: the proxy's + ``JWTAuthManager.auth_builder`` in production, which raises when the token is not + acceptable and otherwise names the user and team it resolved.""" + + def __call__( + self, subject_token: str, request_headers: Mapping[str, str], / + ) -> Awaitable[JWTAuthBuilderResult]: ... + + +async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + general_settings, + jwt_handler, + premium_user, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult: + return await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=jwt_handler, + request_data={}, + general_settings=general_settings, + route=EXCHANGE_ROUTE, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + request_headers=request_headers, + request_method="POST", + ) + + return await identity_from_subject_token( + subject_token, + request_headers=request.headers, + jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, + has_database=prisma_client is not None, + licensed=premium_user is True, + is_jwt=jwt_handler.is_jwt, + authorize=authorize, + ) + + +async def identity_from_subject_token( + subject_token: str, + request_headers: Mapping[str, str], + jwt_auth_enabled: bool, + has_database: bool, + licensed: bool, + is_jwt: Callable[[str], bool], + authorize: AuthorizeSubjectToken, +) -> SubjectIdentity | SubjectTokenRefusal: + """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the + proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token.""" + if not jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", + ) + if not has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot exchange IdP tokens", + ) + if not is_jwt(subject_token): + return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT") + if not licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set" + ) + try: + result: Final = await authorize(subject_token, request_headers) + except HTTPException as denied: + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.detail}") + except ProxyException as denied: + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.message}") + except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures + return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied}") + user_id: Final = result["user_id"] + if user_id is None: + return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows") + return SubjectIdentity(user_id=user_id, team_id=result["team_id"]) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 74f38b3ca6d..6893c06d2d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { @@ -23115,6 +23115,17 @@ ], "title": "Refresh Token" }, + "requested_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Requested Token Type" + }, "resource": { "anyOf": [ { @@ -23136,6 +23147,28 @@ } ], "title": "Scope" + }, + "subject_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" } }, "required": [ @@ -23189,6 +23222,17 @@ ], "title": "Refresh Token" }, + "requested_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Requested Token Type" + }, "resource": { "anyOf": [ { @@ -23210,6 +23254,28 @@ } ], "title": "Scope" + }, + "subject_token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token" + }, + "subject_token_type": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subject Token Type" } }, "required": [ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 63d0bfcc5b8..5f7365cedd5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -525,6 +525,7 @@ class LiteLLMRoutes(enum.Enum): "/mcp-rest/tools/call", "/v1/mcp/tools", "/introspect", + "/token", ] # MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS. diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6a28cd7ff99..803093ff93a 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1867,7 +1867,7 @@ class JWTAuthManager: @staticmethod def get_team_id_from_header( - request_headers: dict | None, + request_headers: Mapping[str, str] | None, allowed_team_ids: set[str], fallback_to_db_teams: bool = False, ) -> str | None: @@ -2037,7 +2037,7 @@ class JWTAuthManager: async def _attach_team_from_header_for_admin( admin_result: JWTAuthBuilderResult, route: str, - request_headers: dict | None, + request_headers: Mapping[str, str] | None, jwt_handler: JWTHandler, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, @@ -2293,7 +2293,7 @@ class JWTAuthManager: user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, - request_headers: dict | None = None, + request_headers: Mapping[str, str] | None = None, request_method: str | None = None, ) -> JWTAuthBuilderResult: return await JWTAuthManager.authorize_jwt( @@ -2390,7 +2390,7 @@ class JWTAuthManager: user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging, - request_headers: dict[str, str] | None = None, + request_headers: Mapping[str, str] | None = None, request_method: str | None = None, provisioning: _JWTProvisioning | None = None, ) -> JWTAuthBuilderResult: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 7c80ee77cd7..1be6ebb6e22 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -15,13 +15,18 @@ from starlette.requests import Request from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( _AUTH_CODE_DEBUG_KEY, + ACCESS_TOKEN_TOKEN_TYPE, CONNECT_FLOW_COOKIE_PREFIX, GATEWAY_AUTH_CODE_PREFIX, GATEWAY_AUTH_CODE_TTL_SECONDS, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, MAX_CLIENT_ID_LENGTH, + SUBJECT_TOKEN_TYPES, + TOKEN_EXCHANGE_GRANT_TYPE, ConsentTeam, MintedProxyCredential, + SubjectIdentity, + SubjectTokenRefusal, _GatewayAuthCode, _open_sealed, _seal, @@ -105,6 +110,7 @@ async def _reload_user_active(user_id: str): async def test_register_mints_stateless_public_client(): body = await _register([REDIRECT_URI]) assert body["token_endpoint_auth_method"] == "none" + assert body["grant_types"] == ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE] assert "client_secret" not in body assert body["redirect_uris"] == [REDIRECT_URI] assert is_gateway_dcr_client_id(body["client_id"]) @@ -1957,7 +1963,11 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): "revocation_endpoint": "https://llm.example.com/revoke", "resource": "https://llm.example.com", "response_types_supported": ["code"], - "grant_types_supported": ["authorization_code", "refresh_token"], + "grant_types_supported": [ + "authorization_code", + "refresh_token", + "urn:ietf:params:oauth:grant-type:token-exchange", + ], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none"], "revocation_endpoint_auth_methods_supported": ["none"], @@ -2148,3 +2158,177 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a ) assert renewed.status_code == 200 assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id" + + +JWT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt" +IDP_TOKEN = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" + + +class _Exchanger: + def __init__(self, result=None): + self.calls = [] + self.result = result + + async def __call__(self, subject_token, request): + self.calls.append((subject_token, request.url.path)) + if self.result is not None: + return self.result + return SubjectIdentity(user_id="u1", team_id="team-b") + + +async def _exchange_native(client_id, minter, exchanger, cache=None, **overrides): + arguments = { + "grant_type": TOKEN_EXCHANGE_GRANT_TYPE, + "subject_token": IDP_TOKEN, + "subject_token_type": JWT_SUBJECT_TOKEN_TYPE, + "exchange_subject_token": exchanger, + } + return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_token_exchange_mints_the_proxy_credential_for_the_idp_subject(): + """RFC 8693: a registered native client trades the IdP token it already holds for the + same credential the consent flow mints, attributed to the user and team the gateway's + JWT auth resolved, with a rotating refresh token bound to that team and the client. + The exchange can be repeated while the IdP token lives; nothing is burned.""" + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, exchanger, cache = _Minter(), _Exchanger(), DualCache() + response = await _exchange_native(client_id, minter, exchanger, cache=cache) + assert response.status_code == 200 + assert response.headers["cache-control"] == "no-store" + body = json.loads(response.body) + assert exchanger.calls == [(IDP_TOKEN, "/token")] + assert minter.calls == [("u1", "team-b")] + assert body["issued_token_type"] == ACCESS_TOKEN_TOKEN_TYPE + assert body["access_token"] == "sk-cli-u1" + assert body["token_type"] == "Bearer" + assert body["expires_in"] == 3600 + assert (body["user_id"], body["team_id"]) == ("u1", "team-b") + principal = _opened_refresh(body["refresh_token"], client_id) + assert (principal.user_id, principal.client_id, principal.audience, principal.team_id) == ( + "u1", + client_id, + "proxy_api", + "team-b", + ) + again = await _exchange_native(client_id, minter, exchanger, cache=cache) + assert again.status_code == 200 + assert json.loads(again.body)["refresh_token"] != body["refresh_token"] + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + + +@pytest.mark.asyncio +async def test_exchanged_credential_refreshes_and_rotates_like_a_consented_one(): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, cache = _Minter(), DualCache() + exchanged = json.loads((await _exchange_native(client_id, minter, _Exchanger(), cache=cache)).body) + refreshed = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache) + assert refreshed.status_code == 200 + body = json.loads(refreshed.body) + assert "issued_token_type" not in body + assert (body["access_token"], body["user_id"], body["team_id"]) == ("sk-cli-u1", "u1", "team-b") + assert body["refresh_token"] != exchanged["refresh_token"] + assert minter.calls == [("u1", "team-b"), ("u1", "team-b")] + replay = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache) + assert replay.status_code == 400 + assert json.loads(replay.body)["error"] == "invalid_grant" + + +@pytest.mark.asyncio +async def test_token_exchange_for_a_teamless_subject_mints_a_teamless_credential(): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + response = await _exchange_native(client_id, minter, _Exchanger(SubjectIdentity(user_id="u2"))) + assert response.status_code == 200 + body = json.loads(response.body) + assert minter.calls == [("u2", None)] + assert (body["user_id"], body["team_id"]) == ("u2", None) + assert _opened_refresh(body["refresh_token"], client_id).team_id is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("subject_token_type", sorted(SUBJECT_TOKEN_TYPES)) +async def test_token_exchange_accepts_every_advertised_subject_token_type(subject_token_type): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + response = await _exchange_native(client_id, _Minter(), _Exchanger(), subject_token_type=subject_token_type) + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_token_exchange_without_an_idp_exchanger_is_unsupported(): + """A gateway that wires no IdP verifier into the endpoint answers the way it always + answered an unknown grant, and never reaches the minter.""" + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + response = await _redeem_native( + None, + client_id, + minter, + grant_type=TOKEN_EXCHANGE_GRANT_TYPE, + subject_token=IDP_TOKEN, + subject_token_type=JWT_SUBJECT_TOKEN_TYPE, + ) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "unsupported_grant_type" + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, status, error", + [ + ({"subject_token": None}, 400, "invalid_request"), + ({"subject_token": ""}, 400, "invalid_request"), + ({"subject_token_type": None}, 400, "invalid_request"), + ({"subject_token_type": "urn:ietf:params:oauth:token-type:saml2"}, 400, "invalid_request"), + ({"requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token"}, 400, "invalid_request"), + ({"resource": "https://other.example.com"}, 400, "invalid_target"), + ({"resource": "https://llm.example.com/mcp"}, 400, "invalid_target"), + ({"client_id": "llm_dcrc_forged"}, 401, "invalid_client"), + ({"client_id": "not-a-gateway-client"}, 401, "invalid_client"), + ], +) +async def test_token_exchange_refuses_a_malformed_request_before_touching_the_idp_token(overrides, status, error): + registered = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter, exchanger = _Minter(), _Exchanger() + response = await _exchange_native( + overrides.get("client_id", registered), + minter, + exchanger, + **{name: value for name, value in overrides.items() if name != "client_id"}, + ) + assert response.status_code == status + assert json.loads(response.body)["error"] == error + assert exchanger.calls == [] + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", ["unsupported_grant_type", "invalid_request"]) +async def test_token_exchange_relays_the_idp_refusal_and_never_mints(error): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + minter = _Minter() + exchanger = _Exchanger(SubjectTokenRefusal(error=error, description="subject_token was rejected: bad signature")) + response = await _exchange_native(client_id, minter, exchanger) + assert response.status_code == 400 + body = json.loads(response.body) + assert (body["error"], body["error_description"]) == (error, "subject_token was rejected: bad signature") + assert minter.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure, status, error", + [ + ("not_a_member", 400, "invalid_grant"), + ("team_required", 400, "invalid_grant"), + ("no_active_key", 400, "invalid_grant"), + ("unavailable", 503, "temporarily_unavailable"), + ], +) +async def test_token_exchange_relays_a_mint_refusal(failure, status, error): + client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"] + response = await _exchange_native(client_id, _Minter(failure), _Exchanger()) + assert response.status_code == status + assert json.loads(response.body)["error"] == error diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py new file mode 100644 index 00000000000..00440614120 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -0,0 +1,115 @@ +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal +from litellm.proxy._experimental.mcp_server.idp_token_exchange import identity_from_subject_token +from litellm.proxy._types import ProxyException +from litellm.proxy.auth.handle_jwt import JWTHandler + +IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" +REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} + + +def _authorized(user_id="u1", team_id="team-b"): + return { + "is_proxy_admin": False, + "team_object": None, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": IDP_JWT, + "team_id": team_id, + "user_id": user_id, + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": user_id}, + "agent_id": None, + } + + +class _Authorizer: + def __init__(self, result=None, raises=None): + self.calls = [] + self.result = result if result is not None else _authorized() + self.raises = raises + + async def __call__(self, subject_token, request_headers): + self.calls.append((subject_token, dict(request_headers))) + if self.raises is not None: + raise self.raises + return self.result + + +async def _identity(authorizer, subject_token=IDP_JWT, **overrides): + arguments = { + "request_headers": REQUEST_HEADERS, + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "is_jwt": JWTHandler.is_jwt, + "authorize": authorizer, + } + return await identity_from_subject_token(subject_token, **{**arguments, **overrides}) + + +@pytest.mark.asyncio +async def test_a_jwt_the_proxy_accepts_names_its_user_and_team(): + """The subject token goes to the proxy's own JWT auth with the caller's headers (that is + where the team header is read), and the identity it resolved is what gets minted.""" + authorizer = _Authorizer() + assert await _identity(authorizer) == SubjectIdentity(user_id="u1", team_id="team-b") + assert authorizer.calls == [(IDP_JWT, REQUEST_HEADERS)] + + +@pytest.mark.asyncio +async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): + assert await _identity(_Authorizer(_authorized(team_id=None))) == SubjectIdentity(user_id="u1", team_id=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, subject_token, error, mentions", + [ + ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), + ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), + ({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"), + ({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"), + ], +) +async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification( + overrides, subject_token, error, mentions +): + authorizer = _Authorizer() + refusal = await _identity(authorizer, subject_token=subject_token, **overrides) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == error + assert mentions in refusal.description + assert authorizer.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raised, mentions", + [ + (HTTPException(status_code=403, detail="User not allowed to access this route"), "not allowed"), + (ProxyException(message="Token expired", type="auth_error", param="token", code=401), "Token expired"), + (Exception("Validation fails: signature verification failed"), "signature verification failed"), + (Exception("Invalid JWT Submitted"), "Invalid JWT"), + ], +) +async def test_a_jwt_the_proxy_rejects_is_an_invalid_subject_token(raised, mentions): + refusal = await _identity(_Authorizer(raises=raised)) + assert isinstance(refusal, SubjectTokenRefusal) + assert refusal.error == "invalid_request" + assert refusal.description.startswith("subject_token was rejected: ") + assert mentions in refusal.description + + +@pytest.mark.asyncio +async def test_a_jwt_that_resolves_no_user_cannot_be_exchanged(): + refusal = await _identity(_Authorizer(_authorized(user_id=None))) + assert refusal == SubjectTokenRefusal( + error="invalid_request", description="subject_token names no user the gateway knows" + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26ae28a57d2..e4c03399b30 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8388,3 +8388,21 @@ async def test_access_group_model_fallback_uses_the_injected_database(channel: s llm_router=None, prisma_client=client, ) is True reader.assert_awaited_once_with(where={"access_group_id": "group-a"}) + + +def test_jwt_team_role_reaches_the_gateway_token_endpoint_by_default(): + """The RFC 8693 token exchange authorizes the IdP JWT against ``POST /token`` itself, and JWT + auth only binds a team from a multi-team claim when that team may call the route, so the + default team allowlist has to cover the gateway's token endpoint or the exchange would mint + teamless credentials for every ``team_ids_jwt_field`` deployment.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/token", litellm_proxy_roles=LiteLLM_JWTAuth() + ) + assert not allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route="/token", + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=[]), + ) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 806c55d51ce..630f69bf0f1 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -627,6 +627,7 @@ def test_virtual_key_llm_api_routes_denies_spend_logs_v2(): "/mcp/tools/call", "/mcp-rest/tools/call", "/mcp/tools/list", + "/token", ], ) def test_mcp_inference_routes_classified_as_llm_api(route): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 17ec8367324..b896f6901cb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24505,10 +24505,16 @@ export interface components { redirect_uri?: string; /** Refresh Token */ refresh_token?: string | null; + /** Requested Token Type */ + requested_token_type?: string | null; /** Resource */ resource?: string | null; /** Scope */ scope?: string | null; + /** Subject Token */ + subject_token?: string | null; + /** Subject Token Type */ + subject_token_type?: string | null; }; /** Body_token_endpoint_token_post */ Body_token_endpoint_token_post: { @@ -24526,10 +24532,16 @@ export interface components { redirect_uri?: string; /** Refresh Token */ refresh_token?: string | null; + /** Requested Token Type */ + requested_token_type?: string | null; /** Resource */ resource?: string | null; /** Scope */ scope?: string | null; + /** Subject Token */ + subject_token?: string | null; + /** Subject Token Type */ + subject_token_type?: string | null; }; /** Body_upload_logo_upload_logo_post */ Body_upload_logo_upload_logo_post: { From 0bacf26b1e147f709ea8e5c638ef81f6b7288599 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:32:51 -0700 Subject: [PATCH 050/393] chore(proxy): keep the lazy OpenAPI snapshot as the CI Python renders it --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 6893c06d2d7..0ff155eee8e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 7cb01cf47f67e91137b4457e5c107104b10055c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:53:10 -0700 Subject: [PATCH 051/393] fix(proxy): exchange role default, logged rejections, gated grant listing The exchange refused the very user JWT auth upserts (its row has no user_role) as "no longer active". The credential now carries the role the proxy already enforces for that user on every request, internal_user when the row has none, the same rule _get_user_role applies on the data plane. A rejected subject_token no longer echoes JWT auth's wording on the public /token endpoint: the response is a fixed invalid_request and the reason goes to the proxy log, since that wording can name the JWKS URL or relay the IdP's reply. The exchange grant is listed on /register, /.well-known/litellm-cli-auth, and the aggregate authorization-server metadata only when JWT auth is on, backed by a database, and licensed, so a client never selects a grant the gateway would then refuse. --- .../mcp_server/discoverable_endpoints.py | 17 ++-- .../mcp_server/gateway_dcr_flow.py | 21 +++- .../mcp_server/idp_token_exchange.py | 95 ++++++++++++++----- .../mcp_server/proxy_api_credentials.py | 12 ++- litellm/proxy/auth/auth_checks.py | 18 ++-- .../mcp_server/test_discoverable_endpoints.py | 19 ++++ .../mcp_server/test_gateway_dcr_flow.py | 27 +++++- .../mcp_server/test_idp_token_exchange.py | 74 +++++++++++---- .../mcp_server/test_proxy_api_credentials.py | 20 +++- 9 files changed, 225 insertions(+), 78 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ebd5f43bf02..873c3d1baec 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -46,7 +46,6 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( - TOKEN_EXCHANGE_GRANT_TYPE, VendorCredentialState, aggregate_authorize, aggregate_token, @@ -60,9 +59,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( register_aggregate_client, relative_request_url, revoke_refresh_token, + supported_grant_types, ) from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( exchange_idp_subject_token, + token_exchange_available, ) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, @@ -2142,7 +2143,9 @@ async def introspect_endpoint(token: str = Form(...)) -> Response: async def native_client_auth_discovery(request: Request) -> JSONResponse: """The versioned contract a native client (``lite login --pkce``, or a CLI in any other language) reads to sign a user in through the browser and obtain a proxy credential.""" - return JSONResponse(native_client_auth_contract(request), headers=TOKEN_NO_CACHE_HEADERS) + return JSONResponse( + native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS + ) # Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request @@ -2630,7 +2633,7 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict: } -def _build_aggregate_authorization_server_response(request: Request) -> dict: +def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict: """RFC 8414 metadata for the gateway as the aggregate authorization server. The issuer is ``{base}/mcp`` and must stay equal to the value the @@ -2649,7 +2652,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], - "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), + "grant_types_supported": supported_grant_types(token_exchange_available), "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], } @@ -2687,7 +2690,7 @@ async def oauth_authorization_server_aggregate(request: Request): per-server row win here instead would serve an issuer of {base} against a resource that advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - return _build_aggregate_authorization_server_response(request) + return _build_aggregate_authorization_server_response(request, token_exchange_available()) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} @@ -2913,7 +2916,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None): # advertises that), so this does not affect it. A request without redirect_uris is not # a DCR request, so the legacy single-server-or-dummy fallback is kept for it. if data.get("redirect_uris"): - return await register_aggregate_client(request=request, request_body=data) + return await register_aggregate_client( + request=request, request_body=data, token_exchange_available=token_exchange_available() + ) resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 9bdde3c5edc..ba24e861d6e 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -188,6 +188,17 @@ class MintProxyCredential(Protocol): TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange" + + +def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]: + """The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed + only where the JWT auth that proves a subject token is on, backed by a database, and + licensed, so a client never selects a grant the gateway would then refuse.""" + if token_exchange_available: + return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE) + return ("authorization_code", "refresh_token") + + """RFC 8693: a native client that already holds a token from the customer's identity provider trades it for the proxy-API credential without a browser round trip.""" @@ -359,7 +370,9 @@ def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None: return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY) -async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response: +async def register_aggregate_client( + request: Request, request_body: Mapping[str, object], token_exchange_available: bool +) -> Response: """RFC 7591 dynamic registration against the gateway itself, statelessly. Only ``redirect_uris`` is authoritative; every client is registered as a public @@ -423,7 +436,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str, "client_id_issued_at": int(now.timestamp()), "redirect_uris": list(raw_uris), "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE], + "grant_types": list(supported_grant_types(token_exchange_available)), "response_types": ["code"], }, ) @@ -621,7 +634,7 @@ class NativeClientAuthContract(TypedDict): revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]] -def native_client_auth_contract(request: Request) -> NativeClientAuthContract: +def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract: """The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a native client (in any language) needs to run the sign-in without reading LiteLLM source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter @@ -636,7 +649,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract: "revocation_endpoint": f"{base_url}/revoke", "resource": base_url, "response_types_supported": ("code",), - "grant_types_supported": ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE), + "grant_types_supported": supported_grant_types(token_exchange_available), "code_challenge_methods_supported": ("S256",), "token_endpoint_auth_methods_supported": ("none",), "revocation_endpoint_auth_methods_supported": ("none",), diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 69b81713810..cdefaf76d49 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -6,15 +6,69 @@ and the user and team sync it performs), so no browser round trip is needed.""" from __future__ import annotations from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from typing import Final, Protocol from fastapi import HTTPException, Request +from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._types import JWTAuthBuilderResult, ProxyException from litellm.proxy.auth.handle_jwt import JWTAuthManager EXCHANGE_ROUTE: Final = "/token" +REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" + + +@dataclass(frozen=True, slots=True) +class TokenExchangePrerequisites: + """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT + bearer. Discovery and registration advertise the exchange grant only when every one of + them holds, and an exchange attempt is refused naming the first one that does not.""" + + jwt_auth_enabled: bool + has_database: bool + licensed: bool + + @property + def available(self) -> bool: + return self.jwt_auth_enabled and self.has_database and self.licensed + + def refusal(self) -> SubjectTokenRefusal | None: + if not self.jwt_auth_enabled: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", + ) + if not self.has_database: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway has no database, so it cannot exchange IdP tokens", + ) + if not self.licensed: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="JWT auth is an enterprise only feature; no license is set", + ) + return None + + +def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call + general_settings, + premium_user, + prisma_client, + ) + + return TokenExchangePrerequisites( + jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, + has_database=prisma_client is not None, + licensed=premium_user is True, + ) + + +def token_exchange_available() -> bool: + return read_token_exchange_prerequisites().available class AuthorizeSubjectToken(Protocol): @@ -31,7 +85,6 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call general_settings, jwt_handler, - premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -55,9 +108,7 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su return await identity_from_subject_token( subject_token, request_headers=request.headers, - jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, - has_database=prisma_client is not None, - licensed=premium_user is True, + prerequisites=read_token_exchange_prerequisites(), is_jwt=jwt_handler.is_jwt, authorize=authorize, ) @@ -66,40 +117,34 @@ async def exchange_idp_subject_token(subject_token: str, request: Request) -> Su async def identity_from_subject_token( subject_token: str, request_headers: Mapping[str, str], - jwt_auth_enabled: bool, - has_database: bool, - licensed: bool, + prerequisites: TokenExchangePrerequisites, is_jwt: Callable[[str], bool], authorize: AuthorizeSubjectToken, ) -> SubjectIdentity | SubjectTokenRefusal: """Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which - RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token.""" - if not jwt_auth_enabled: - return SubjectTokenRefusal( - error="unsupported_grant_type", - description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens", - ) - if not has_database: - return SubjectTokenRefusal( - error="unsupported_grant_type", - description="this gateway has no database, so it cannot exchange IdP tokens", - ) + RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token. The + reason stays in the proxy log: this endpoint is public and JWT auth's own wording can + name the JWKS URL it fetched or quote the IdP's response.""" + unmet: Final = prerequisites.refusal() + if unmet is not None: + return unmet if not is_jwt(subject_token): return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT") - if not licensed: - return SubjectTokenRefusal( - error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set" - ) try: result: Final = await authorize(subject_token, request_headers) except HTTPException as denied: - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.detail}") + return _rejected_by_jwt_auth(denied.detail) except ProxyException as denied: - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied.message}") + return _rejected_by_jwt_auth(denied.message) except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures - return SubjectTokenRefusal(error="invalid_request", description=f"subject_token was rejected: {denied}") + return _rejected_by_jwt_auth(denied) user_id: Final = result["user_id"] if user_id is None: return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows") return SubjectIdentity(user_id=user_id, team_id=result["team_id"]) + + +def _rejected_by_jwt_auth(reason: object) -> SubjectTokenRefusal: + verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason) + return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index 27d0ebbd5e6..a34119edf10 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -16,7 +16,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( ReloadUserFailure, ) from litellm.proxy._types import LiteLLM_UserTable -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, effective_user_role from litellm.proxy.management_endpoints.ui_sso import ( CliSsoTeamDetail, fetch_cli_sso_team_details, @@ -51,12 +51,12 @@ async def mint_proxy_credential( posting the consent form without one. Memberships whose team rows are gone count as no team at all, the way ``lite login`` treats them, so they can never lock a user out. The user row handed to the minter carries no team list, exactly like ``lite login``'s, so - the minter's own first-team fallback stays inert.""" + the minter's own first-team fallback stays inert. The credential carries the role the + proxy already enforces for the user on every request, so a row with no role (JWT auth's + upsert writes none) mints as an internal user instead of being refused.""" user: Final = await load_active_user_by_id(user_id) if isinstance(user, str): return user - if user.user_role is None: - return "no_active_key" if team_id is not None and team_id not in user.teams: return "not_a_member" details: Final = await _team_details(user.teams) if user.teams else () @@ -68,7 +68,9 @@ async def mint_proxy_credential( if selected is None: return "not_a_member" key: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=LiteLLM_UserTable(user_id=user.user_id, user_role=user.user_role, models=user.models), + user_info=LiteLLM_UserTable( + user_id=user.user_id, user_role=effective_user_role(user.user_role).value, models=user.models + ), team_id=team_id, team_alias=selected.team_alias, team_models=selected.team_models, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ba68dc8a17f..d90553db72b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1207,21 +1207,19 @@ async def common_checks( return True +def effective_user_role(user_role: str | None) -> LitellmUserRoles: + try: + return LitellmUserRoles(user_role) + except ValueError: + return LitellmUserRoles.INTERNAL_USER + + def _get_user_role( user_obj: LiteLLM_UserTable | None, ) -> LitellmUserRoles | None: if user_obj is None: return None - - _user: Final = user_obj - - _user_role: Final = _user.user_role - try: - role: Final = LitellmUserRoles(_user_role) - except ValueError: - return LitellmUserRoles.INTERNAL_USER - - return role + return effective_user_role(user_obj.user_role) def _is_api_route_allowed( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index aa45b2f6793..e4edb5fb4dd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11048,6 +11048,25 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo assert stranger.json()["error"] == "invalid_client" +@pytest.mark.parametrize("exchange_servable", [True, False]) +def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(monkeypatch, exchange_servable): + """Every document a native client reads before it picks a grant (the versioned contract, the + aggregate authorization-server metadata, and the registration response) lists the RFC 8693 + exchange exactly when the running proxy can serve it: JWT auth on, a database, and a license.""" + client, _session_cookie, _minted = _native_client_app(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": exchange_servable}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] + expected = ["authorization_code", "refresh_token", *exchange_grant] + + assert client.get("/.well-known/litellm-cli-auth").json()["grant_types_supported"] == expected + assert client.get("/.well-known/oauth-authorization-server/mcp").json()["grant_types_supported"] == expected + registered = client.post("/register", json={"redirect_uris": ["http://127.0.0.1:51234/callback"]}) + assert registered.status_code == 201 + assert registered.json()["grant_types"] == expected + + def test_native_client_authorize_without_the_proxy_resource_keeps_the_mcp_flow(monkeypatch): """A registered client asking for the MCP resource (or no resource) never sees the consent page, so existing MCP clients are untouched by the native-client arm.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1be6ebb6e22..c42f8763e74 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -95,9 +95,11 @@ def _request(path="/authorize", query="", cookies=None, method="GET"): ) -async def _register(redirect_uris) -> dict: +async def _register(redirect_uris, token_exchange_available=True) -> dict: response = await register_aggregate_client( - request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": redirect_uris}, + token_exchange_available=token_exchange_available, ) return json.loads(response.body) @@ -119,11 +121,18 @@ async def test_register_mints_stateless_public_client(): assert record.redirect_uris == (REDIRECT_URI,) +@pytest.mark.asyncio +async def test_register_omits_the_exchange_grant_where_the_gateway_cannot_serve_it(): + body = await _register([REDIRECT_URI], token_exchange_available=False) + assert body["grant_types"] == ["authorization_code", "refresh_token"] + + @pytest.mark.asyncio @pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS]) async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={ "client_name": "Visual Studio Code", "client_uri": "https://code.visualstudio.com", @@ -149,6 +158,7 @@ async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[s async def test_register_rejects_five_valid_callbacks() -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]}, ) assert response.status_code == 400 @@ -162,6 +172,7 @@ async def test_register_rejects_five_valid_callbacks() -> None: async def test_register_four_callbacks_preserves_encoded_size_guard() -> None: response: Final = await register_aggregate_client( request=_request(path="/register", method="POST"), + token_exchange_available=True, request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]}, ) assert response.status_code == 400 @@ -214,6 +225,7 @@ async def test_register_rejects_userinfo_spoofed_origin(): response = await register_aggregate_client( request=_request(path="/register", method="POST"), request_body={"redirect_uris": ["https://claude.ai@attacker.example/callback"]}, + token_exchange_available=True, ) assert response.status_code == 400 assert json.loads(response.body)["error"] == "invalid_redirect_uri" @@ -234,7 +246,9 @@ async def test_register_rejects_userinfo_spoofed_origin(): ) async def test_register_rejects_bad_redirect_uris(redirect_uris): response = await register_aggregate_client( - request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris} + request=_request(path="/register", method="POST"), + request_body={"redirect_uris": redirect_uris}, + token_exchange_available=True, ) assert response.status_code == 400 assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata") @@ -1954,7 +1968,7 @@ async def test_revoke_refuses_unknown_clients_and_a_missing_master_key(): def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): - assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth")))) == { + assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), True))) == { "contract_version": 1, "issuer": "https://llm.example.com", "authorization_endpoint": "https://llm.example.com/authorize", @@ -1974,6 +1988,11 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): } +def test_native_client_auth_contract_omits_the_exchange_grant_where_the_gateway_cannot_serve_it(): + contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), False) + assert list(contract["grant_types_supported"]) == ["authorization_code", "refresh_token"] + + @pytest.mark.parametrize( "resource, expected", [ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py index 00440614120..d1b049dddd5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -1,13 +1,22 @@ +import logging + import pytest from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal -from litellm.proxy._experimental.mcp_server.idp_token_exchange import identity_from_subject_token +from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( + REJECTED_SUBJECT_TOKEN, + TokenExchangePrerequisites, + identity_from_subject_token, + token_exchange_available, +) from litellm.proxy._types import ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} +EVERY_GATE_HOLDS = {"jwt_auth_enabled": True, "has_database": True, "licensed": True} +JWKS_URL = "https://idp.example.com/.well-known/jwks.json" def _authorized(user_id="u1", team_id="team-b"): @@ -42,16 +51,14 @@ class _Authorizer: return self.result -async def _identity(authorizer, subject_token=IDP_JWT, **overrides): - arguments = { - "request_headers": REQUEST_HEADERS, - "jwt_auth_enabled": True, - "has_database": True, - "licensed": True, - "is_jwt": JWTHandler.is_jwt, - "authorize": authorizer, - } - return await identity_from_subject_token(subject_token, **{**arguments, **overrides}) +async def _identity(authorizer, subject_token=IDP_JWT, **unmet): + return await identity_from_subject_token( + subject_token, + request_headers=REQUEST_HEADERS, + prerequisites=TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}), + is_jwt=JWTHandler.is_jwt, + authorize=authorizer, + ) @pytest.mark.asyncio @@ -70,7 +77,7 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): @pytest.mark.asyncio @pytest.mark.parametrize( - "overrides, subject_token, error, mentions", + "unmet, subject_token, error, mentions", [ ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), @@ -79,32 +86,59 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): ], ) async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification( - overrides, subject_token, error, mentions + unmet, subject_token, error, mentions ): authorizer = _Authorizer() - refusal = await _identity(authorizer, subject_token=subject_token, **overrides) + refusal = await _identity(authorizer, subject_token=subject_token, **unmet) assert isinstance(refusal, SubjectTokenRefusal) assert refusal.error == error assert mentions in refusal.description assert authorizer.calls == [] +@pytest.mark.parametrize("unmet", [{}, {"jwt_auth_enabled": False}, {"has_database": False}, {"licensed": False}]) +def test_the_grant_is_available_exactly_when_every_gate_holds(unmet): + prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}) + assert prerequisites.available is (unmet == {}) + assert (prerequisites.refusal() is None) is prerequisites.available + + +@pytest.mark.parametrize( + "general_settings, prisma_client, premium_user, expected", + [ + ({"enable_jwt_auth": True}, object(), True, True), + ({}, object(), True, False), + ({"enable_jwt_auth": True}, None, True, False), + ({"enable_jwt_auth": True}, object(), False, False), + ], +) +def test_availability_is_read_from_the_running_proxy( + monkeypatch, general_settings, prisma_client, premium_user, expected +): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + assert token_exchange_available() is expected + + @pytest.mark.asyncio @pytest.mark.parametrize( - "raised, mentions", + "raised, reason", [ (HTTPException(status_code=403, detail="User not allowed to access this route"), "not allowed"), (ProxyException(message="Token expired", type="auth_error", param="token", code=401), "Token expired"), (Exception("Validation fails: signature verification failed"), "signature verification failed"), (Exception("Invalid JWT Submitted"), "Invalid JWT"), + (Exception(f"Failed to fetch keys from {JWKS_URL}: 502 Bad Gateway from the IdP"), JWKS_URL), ], ) -async def test_a_jwt_the_proxy_rejects_is_an_invalid_subject_token(raised, mentions): +async def test_a_jwt_the_proxy_rejects_is_refused_with_the_reason_kept_in_the_log(raised, reason, caplog): + """The endpoint is public, so the response never quotes JWT auth's wording (it can name + the JWKS URL or relay the IdP's reply); the operator reads the reason in the proxy log.""" + caplog.set_level(logging.WARNING, logger="LiteLLM Proxy") refusal = await _identity(_Authorizer(raises=raised)) - assert isinstance(refusal, SubjectTokenRefusal) - assert refusal.error == "invalid_request" - assert refusal.description.startswith("subject_token was rejected: ") - assert mentions in refusal.description + assert refusal == SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN) + assert reason in caplog.text @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 8bb8bdada7d..85650c6a05a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -8,6 +8,7 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.models.user import LiteLLM_UserTable from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential +from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail @@ -67,10 +68,21 @@ async def test_mint_passes_user_lookup_failures_through(failure, load_user, fetc @pytest.mark.asyncio -async def test_mint_refuses_a_user_without_a_role(load_user, fetch_teams): - load_user.return_value = _user(user_role=None) - assert await mint_proxy_credential("u1", None) == "no_active_key" - fetch_teams.assert_not_awaited() +@pytest.mark.parametrize( + "stored_role, minted_role", + [ + (None, LitellmUserRoles.INTERNAL_USER), + ("made_up_role", LitellmUserRoles.INTERNAL_USER), + ("proxy_admin", LitellmUserRoles.PROXY_ADMIN), + ], +) +async def test_mint_carries_the_role_the_proxy_enforces_for_the_user(load_user, fetch_teams, stored_role, minted_role): + """A user JWT auth upserted has no role in the database, and the proxy already treats + such a user as an internal user on every request, so the credential says the same.""" + load_user.return_value = _user(user_role=stored_role) + minted = await mint_proxy_credential("u1", "team-a") + assert isinstance(minted, MintedProxyCredential) + assert _decoded(minted).user_role == minted_role @pytest.mark.asyncio From 3a8ac47e99b1383a8f3231e2ef88494e669a069c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:12:15 -0700 Subject: [PATCH 052/393] fix(proxy): mint the exchange credential off the database user row, not the cache JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a team member never evicts the cached user row, so the mint read a row with no teams and refused the very first token exchange for a never-seen user as not a member. The loader now reads the row from the database and leaves the fresh row in the cache for the requests the credential makes next --- .../mcp_server/bridge_token_flow.py | 7 +++- .../mcp_server/test_discoverable_endpoints.py | 42 +++++++++++++++++-- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 2b13baa624b..4235471f2d9 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -273,7 +273,11 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.""" + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. The + row is read from the database, never the cache: JWT auth caches the user it creates before it adds + that user to the JWT's team and adding a member never evicts the cached row, so a credential minted + off the cache would refuse the very first exchange as not a member. The fresh row replaces the cached + one.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) @@ -296,6 +300,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, + check_db_only=True, ) except (ProxyException, HTTPException): return "no_active_key" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index e4edb5fb4dd..30b179f1a26 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7568,6 +7568,36 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_ assert await _reload_active_user_by_id("sso-user-7") == "faulted" +@pytest.mark.asyncio +async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_cache(proxy_globals): + """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a + member never evicts the cached row, so a credential minted off the cached row refused the very first + token exchange as not a member. The loader has to read the row from the database and leave the fresh + row in the cache for the requests the credential makes next.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="fresh-jwt-user", value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=[]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"]) + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = prisma + + loaded = await load_active_user_by_id("fresh-jwt-user") + + assert not isinstance(loaded, str) + assert loaded.teams == ["team-a"] + cached = await cache.async_get_cache(key="fresh-jwt-user", model_type=LiteLLM_UserTable) + assert cached is not None + assert cached.teams == ["team-a"] + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the @@ -11866,13 +11896,17 @@ async def test_oauth_refresh_revalidates_the_same_active_user_rule( from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id handler, _ = jwt_oauth_identity - handler.user_api_key_cache.set_cache( - "jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": state != "inactive"}) - ) + user_id: Final = f"jwt-owner-{state}" + row: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": state != "inactive"}) + proxy_server.prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row) if state == "missing_database": monkeypatch.setattr(proxy_server, "prisma_client", None) expected: Final = None if state == "active" else "no_active_key" if state == "inactive" else "unresolvable" - assert await _reload_active_user_by_id("jwt-owner") == expected + assert await _reload_active_user_by_id(user_id) == expected + if state != "missing_database": + cached: Final = handler.user_api_key_cache.get_cache(user_id, model_type=LiteLLM_UserTable) + assert cached is not None + assert cached.metadata == row.metadata @pytest.mark.asyncio From 167edf2769b28f76870b654014f99312843cc327 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:10:03 -0700 Subject: [PATCH 053/393] fix(proxy): read the database user row only in the credential mint The token exchange mint keeps reading the user row from the database, since JWT auth caches the user it creates before adding it to the JWT's team and a mint off that cached row refused the first exchange for a new user. Introspection and the refresh revalidation go back to the cache read, so a resource server calling /introspect per request pays no database read. --- .../mcp_server/bridge_token_flow.py | 18 ++++++--- .../mcp_server/proxy_api_credentials.py | 4 +- .../mcp_server/test_discoverable_endpoints.py | 39 +++++++++++++++++-- .../mcp_server/test_proxy_api_credentials.py | 30 +++++++++++++- 4 files changed, 78 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 4235471f2d9..f19cb87ae18 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -262,7 +262,12 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No return loaded if isinstance(loaded, str) else None -async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResolutionFailure": +UserRowSource = Literal["cache", "database"] + + +async def load_active_user_by_id( + user_id: str, source: UserRowSource = "cache" +) -> "LiteLLM_UserTable | _KeyResolutionFailure": """Load a live litellm user by id, returning the record when the user is active or a precise failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a @@ -273,11 +278,12 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. The - row is read from the database, never the cache: JWT auth caches the user it creates before it adds + chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. + ``source="database"`` reads the row from the database, never the cache, and leaves the fresh row in the + cache for the requests the credential makes next: JWT auth caches the user it creates before it adds that user to the JWT's team and adding a member never evicts the cached row, so a credential minted - off the cache would refuse the very first exchange as not a member. The fresh row replaces the cached - one.""" + off the cache would refuse the very first exchange as not a member. Every other caller keeps the cache + read, so introspection, which a resource server may call per request, stays off the database.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) @@ -300,7 +306,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, - check_db_only=True, + check_db_only=source == "database", ) except (ProxyException, HTTPException): return "no_active_key" diff --git a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py index a34119edf10..2f7fcaef645 100644 --- a/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py @@ -42,7 +42,7 @@ async def mint_proxy_credential( user_id: str, team_id: str | None ) -> MintedProxyCredential | ProxyCredentialMintFailure: """Mint the ``lite login`` credential for a consented grant. Membership is checked - live, so a team the user left between consent and redemption (or between refreshes) + live against the database row, so a team the user left between consent and redemption (or between refreshes) refuses the grant instead of minting a credential attributed to a team they are no longer on. The team is exactly the one the consent page sealed into the grant; nothing is picked on the user's behalf here, so a refresh can never move the credential, and a @@ -54,7 +54,7 @@ async def mint_proxy_credential( the minter's own first-team fallback stays inert. The credential carries the role the proxy already enforces for the user on every request, so a row with no role (JWT auth's upsert writes none) mints as an internal user instead of being refused.""" - user: Final = await load_active_user_by_id(user_id) + user: Final = await load_active_user_by_id(user_id, source="database") if isinstance(user, str): return user if team_id is not None and team_id not in user.teams: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 30b179f1a26..6965b3b4ebe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7572,8 +7572,8 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_cache(proxy_globals): """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member never evicts the cached row, so a credential minted off the cached row refused the very first - token exchange as not a member. The loader has to read the row from the database and leave the fresh - row in the cache for the requests the credential makes next.""" + token exchange as not a member. The database source has to read the row from the database and leave + the fresh row in the cache for the requests the credential makes next.""" from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -7589,7 +7589,7 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca proxy_globals.user_api_key_cache = cache proxy_globals.prisma_client = prisma - loaded = await load_active_user_by_id("fresh-jwt-user") + loaded = await load_active_user_by_id("fresh-jwt-user", source="database") assert not isinstance(loaded, str) assert loaded.teams == ["team-a"] @@ -7598,6 +7598,39 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca assert cached.teams == ["team-a"] +@pytest.mark.asyncio +async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_read(proxy_globals): + """Introspection and refresh revalidation run per call, so the loader's default source is the cache: a + cached row answers without a database read, and only a caller that asks for the database row pays for + one.""" + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _reload_active_user_by_id, + load_active_user_by_id, + ) + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="cached-jwt-user", + value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=["team-a"]), + model_type=LiteLLM_UserTable, + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=[]) + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = prisma + + loaded = await load_active_user_by_id("cached-jwt-user") + + assert not isinstance(loaded, str) + assert loaded.teams == ["team-a"] + assert await _reload_active_user_by_id("cached-jwt-user") is None + prisma.db.litellm_usertable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 85650c6a05a..04fbbe4a6ce 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -1,6 +1,6 @@ """Tests for minting the ``lite login`` credential from a consented native-client grant.""" -from unittest.mock import ANY, AsyncMock +from unittest.mock import ANY, AsyncMock, MagicMock import pytest @@ -10,6 +10,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail _LOAD_USER = "litellm.proxy._experimental.mcp_server.proxy_api_credentials.load_active_user_by_id" @@ -91,7 +92,7 @@ async def test_mint_refuses_a_teamless_grant_for_a_team_member(load_user, fetch_ is refused for a user with teams instead of minting an unscoped credential or drifting onto the first team, on redemption and on every refresh alike.""" assert await mint_proxy_credential("u1", None) == "team_required" - load_user.assert_awaited_once_with("u1") + load_user.assert_awaited_once_with("u1", source="database") fetch_teams.assert_awaited_once_with(ANY, ["team-a", "team-b"]) @@ -126,6 +127,31 @@ async def test_mint_honors_the_consented_team(load_user, fetch_teams): assert decoded.team_model_aliases == {"fast": "gpt-5.4-mini"} +@pytest.mark.asyncio +async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_row(fetch_teams, monkeypatch): + """JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member + never evicts the cached row, so a mint off the cached row refused the very first token exchange as not + a member. The mint has to read the database row, whatever the cache holds.""" + from litellm.proxy import proxy_server + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="stale-cache-user", value=_user(user_id="stale-cache-user", teams=[]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=_user(user_id="stale-cache-user", teams=["team-a"]) + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + minted = await mint_proxy_credential("stale-cache-user", "team-a") + + assert isinstance(minted, MintedProxyCredential) + assert minted.team_id == "team-a" + assert _decoded(minted).team_id == "team-a" + + @pytest.mark.asyncio async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams): assert await mint_proxy_credential("u1", "team-c") == "not_a_member" From 9885dc89621697e235fc65e0e86e147da87398c1 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:52:42 +0000 Subject: [PATCH 054/393] test(e2e): add azure, bedrock converse and vertex wires to the cost suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 54 +++- tests/e2e/cost_calculation/cost_matrix.py | 144 ++++++++- .../e2e/cost_calculation/scripted_provider.py | 284 +++++++++++++++++- .../cost_calculation/test_wire_formats_e2e.py | 64 ++++ tests/e2e/cost_map.json | 158 ++++++++++ tests/e2e/models.py | 1 + 6 files changed, 681 insertions(+), 24 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 345ca26f7e3..8c6db7c0010 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -12,6 +12,7 @@ Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). from __future__ import annotations import importlib.util +import json import sys from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -22,7 +23,7 @@ from typing import Final, Protocol, cast import pytest from cost_matrix import Case, FrontierModel -from e2e_config import COST_MAP_PROXY_URL +from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE from lifecycle import ResourceManager from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody from proxy_client import ProxyClient, build_proxy_client @@ -111,6 +112,41 @@ def client() -> CostCalcClient: return CostCalcClient(proxy=proxy) +_vertex_key_pem: str | None = None + + +def _vertex_service_account_json() -> str: + """A service-account credential JSON whose token_uri is the sidecar's + /_oauth/token route: the proxy's google-auth refresh then gets a scripted + access token without touching Google. One generated RSA key per process.""" + global _vertex_key_pem # mutable-ok: session-scoped key generation cached for reuse + if _vertex_key_pem is None: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + _vertex_key_pem = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": "cc-scripted-project", + "private_key_id": "scripted", + "private_key": _vertex_key_pem, + "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize", + "token_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/token", + } + ) + + def register_scenario_deployment( client: CostCalcClient, resources: ResourceManager, @@ -126,15 +162,21 @@ def register_scenario_deployment( handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) model_name: Final = f"{model.model_name}-{marker}" + extra_params: Final[dict[str, str]] = dict(model.litellm_params) + if model.wire == "vertex_generate": + extra_params["vertex_credentials"] = _vertex_service_account_json() model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, - litellm_params=LiteLLMParamsBody( - model=model.litellm_model, - api_key=model.api_key, - api_base=handle.api_base(), + litellm_params=LiteLLMParamsBody.model_validate( + { + "model": model.litellm_model, + "api_key": model.api_key, + "api_base": handle.api_base(), + **extra_params, + } ), - model_info=ModelInfoBody(), + model_info=ModelInfoBody(base_model=model.base_model), ) ) resources.defer(lambda: client.proxy.delete_model(model_id)) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 5f634712778..37495f37b0b 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -88,7 +88,14 @@ class FrontierModel: litellm_model: str wire: Wire map_key: str - override_model: str + override_model: str | None = None + override_map_key: str | None = None + # Registered as model_info.base_model; when set, the provider-reported + # model loses to it and every case bills at this deployment's own rates. + base_model: str | None = None + # Extra litellm_params merged into the /model/new registration (api_version, + # aws_* credentials, vertex_* auth). + litellm_params: Mapping[str, str] = MappingProxyType({}) @property def rates(self) -> CostMapEntry: @@ -96,11 +103,16 @@ class FrontierModel: @property def override_rates(self) -> CostMapEntry: + if self.base_model is not None or self.override_map_key is None: + return self.rates return _COST_MAP[self.override_map_key] @property - def override_map_key(self) -> str: - return _OVERRIDE_MAP_KEYS[self.override_model] + def provider_model(self) -> str: + """The bare provider-facing model name: litellm_model minus the provider + prefix and any routing segment (converse/, responses/).""" + tail: Final = self.litellm_model.split("/")[1:] + return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) @property def provider(self) -> str: @@ -166,6 +178,92 @@ _FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( ) +@dataclass(frozen=True, slots=True) +class _ExtendedSpec: + """A frontier entry whose override target, model_info.base_model or extra + litellm_params can't be derived from the map key alone.""" + + map_key: str + litellm_model: str + wire: Wire + override_model: str | None = None + override_map_key: str | None = None + base_model: str | None = None + litellm_params: Mapping[str, str] = MappingProxyType({}) + + +_AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"}) +_BEDROCK_PARAMS: Final[Mapping[str, str]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType( + { + "vertex_project": "cc-scripted-project", + "vertex_location": "us-central1", + } +) + +_EXTENDED_SPECS: Final[tuple[_ExtendedSpec, ...]] = ( + _ExtendedSpec( + map_key="azure/gpt-5.6", + litellm_model="azure/gpt-5.6", + wire="azure_chat", + override_model="gpt-5.4-mini", + override_map_key="azure/gpt-5.4-mini", + litellm_params=_AZURE_PARAMS, + ), + _ExtendedSpec( + # Deployment name is not a model; base_model pins billing so the + # response's model field loses, proving base_model wins. + map_key="azure/gpt-5.4-mini", + litellm_model="azure/cc-pinned-deployment", + wire="azure_chat", + override_model="gpt-5.6", + override_map_key="azure/gpt-5.6", + base_model="azure/gpt-5.4-mini", + litellm_params=_AZURE_PARAMS, + ), + _ExtendedSpec( + map_key="anthropic.claude-sonnet-5-v1:0", + litellm_model="bedrock/converse/anthropic.claude-sonnet-5-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="us.anthropic.claude-opus-5-v1:0", + litellm_model="bedrock/converse/us.anthropic.claude-opus-5-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="meta.llama4-maverick-17b-instruct-v1:0", + litellm_model="bedrock/converse/meta.llama4-maverick-17b-instruct-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="gemini-3.8-flash", + litellm_model="vertex_ai/gemini-3.8-flash", + wire="vertex_generate", + override_model="gemini-3.1-pro-preview", + override_map_key="gemini-3.1-pro-preview", + litellm_params=_VERTEX_PARAMS, + ), + _ExtendedSpec( + map_key="gemini-3.1-pro-preview", + litellm_model="vertex_ai/gemini-3.1-pro-preview", + wire="vertex_generate", + override_model="gemini-3.8-flash", + override_map_key="gemini-3.8-flash", + litellm_params=_VERTEX_PARAMS, + ), +) + + def _frontier() -> tuple[FrontierModel, ...]: return tuple( FrontierModel( @@ -174,8 +272,21 @@ def _frontier() -> tuple[FrontierModel, ...]: wire=wire, map_key=map_key, override_model=_OVERRIDE_MODELS[map_key], + override_map_key=_OVERRIDE_MAP_KEYS[_OVERRIDE_MODELS[map_key]], ) for map_key, litellm_model, wire in _FRONTIER_SPECS + ) + tuple( + FrontierModel( + model_name=f"cc-{spec.map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", + litellm_model=spec.litellm_model, + wire=spec.wire, + map_key=spec.map_key, + override_model=spec.override_model, + override_map_key=spec.override_map_key, + base_model=spec.base_model, + litellm_params=spec.litellm_params, + ) + for spec in _EXTENDED_SPECS ) @@ -219,6 +330,24 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), + "azure_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", + } + ), + "bedrock_converse": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "absent_usage", + "tool_call", "image_input", + } + ), + "vertex_generate": frozenset( + { + "cache_read", "reasoning", "audio", "web_search", "response_model", + "absent_usage", "tool_call", "image_input", "prompt_blocked", + } + ), }) CaseName: TypeAlias = Literal[ @@ -270,6 +399,7 @@ class Case: scenario_id=scenario_id, wire=model.wire, usage=self.usage, + model=model.provider_model, output=ScriptedOutput( text=text, response_model=model.override_model if self.response_model_override else None, @@ -296,7 +426,9 @@ _PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tok def _web_search_case(model: FrontierModel) -> Case: - counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + counts_exactly: Final = model.wire in ( + "openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate" + ) return Case( name="web_search", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), @@ -611,12 +743,12 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: wire's normalization: Anthropic folds cache read/write into prompt_tokens, everyone else reports the totals the wire emitted.""" u: Final = case.usage - if model.wire == "anthropic_messages": + if model.wire in ("anthropic_messages", "bedrock_converse"): return ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, u.output_tokens, ) - if model.wire == "gemini_generate": + if model.wire in ("gemini_generate", "vertex_generate"): return ( u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index e1a6c430307..00230fabeba 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -15,10 +15,15 @@ Layout on one port: - ``GET /health`` liveness - ``POST /_scenarios`` register a Scenario JSON, returns its id - ``DELETE /_scenarios/`` remove it +- ``POST /_oauth/token`` fake Google OAuth token endpoint for the + Vertex service-account credential's refresh call - ``POST ///`` provider wire; mount is one of - ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks`` and the - remainder is whatever path the provider client appends (``chat/completions``, - ``responses``, ``v1/messages``, ``models/:generateContent`` ...) + ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks``, ``azure``, + ``bedrock``, ``vertex`` and the remainder is whatever path the provider + client appends (``chat/completions``, ``responses``, ``v1/messages``, + ``models/:generateContent`` ...). Vertex appends ``:generateContent`` / + ``:streamGenerateContent`` to the mount segment itself, and Bedrock Converse + targets ``model//converse`` / ``converse-stream`` A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the @@ -28,15 +33,17 @@ final stream chunk carries usage or the provider reports none. from __future__ import annotations import json +import struct import sys import threading import time +import zlib from collections.abc import Mapping from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from types import MappingProxyType from typing import Final, Literal, TypeAlias -from urllib.parse import urlsplit +from urllib.parse import unquote, urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator @@ -47,6 +54,9 @@ Wire: TypeAlias = Literal[ "gemini_generate", "together_chat", "fireworks_chat", + "azure_chat", + "bedrock_converse", + "vertex_generate", ] WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( @@ -57,6 +67,9 @@ WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( "gemini_generate": "gemini", "together_chat": "together", "fireworks_chat": "fireworks", + "azure_chat": "azure", + "bedrock_converse": "bedrock", + "vertex_generate": "vertex", } ) @@ -69,6 +82,7 @@ _TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( { "openai_responses": frozenset({"incomplete", "unvalidated"}), "gemini_generate": frozenset({"prompt_blocked"}), + "vertex_generate": frozenset({"prompt_blocked"}), } ) @@ -131,6 +145,10 @@ class Scenario(BaseModel): wire: Wire usage: ScriptedUsage output: ScriptedOutput + # The bare provider-facing model name the renderer echoes when the request + # carries no model of its own (Vertex and Bedrock name the model in the URL + # path, not the body). + model: str stream_usage: StreamUsage = "final_chunk" service_tier: ServiceTier | None = None @@ -904,7 +922,208 @@ def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: ) -def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> RenderedResponse: +def _bedrock_usage(u: ScriptedUsage) -> Mapping[str, object]: + # Converse reports uncached input in inputTokens and rides cache reads and + # writes on top-level fields; totalTokens covers every input kind + output. + cache_writes: Final = u.cache_write_5m_tokens + u.cache_write_1h_tokens + return _jobj_opt( + ("inputTokens", u.fresh_input_tokens), + ("outputTokens", u.output_tokens), + ( + "totalTokens", + u.fresh_input_tokens + u.cache_read_tokens + cache_writes + u.output_tokens, + ), + ("cacheReadInputTokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ("cacheWriteInputTokens", cache_writes) if cache_writes else None, + ( + ( + "cacheDetails", + tuple( + _jobj(("inputTokens", count), ("ttl", ttl)) + for count, ttl in ( + (u.cache_write_5m_tokens, "5m"), + (u.cache_write_1h_tokens, "1h"), + ) + if count + ), + ) + if cache_writes + else None + ), + ) + + +def _bedrock_stop_reason(scenario: Scenario) -> str: + if scenario.output.tool_call is not None: + return "tool_use" + return "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + + +def _bedrock_content(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ( + "toolUse", + _jobj( + ("toolUseId", f"tooluse_{scenario.scenario_id}"), + ("name", tool_call.name), + ("input", json.loads(tool_call.arguments)), + ), + ), + ), + ) + return (_jobj(("text", scenario.output.text)),) + + +def _bedrock_body(scenario: Scenario) -> Mapping[str, object]: + return _jobj( + ( + "output", + _jobj( + ( + "message", + _jobj( + ("role", "assistant"), + ("content", _bedrock_content(scenario)), + ), + ), + ), + ), + ("stopReason", _bedrock_stop_reason(scenario)), + ("usage", _bedrock_usage(scenario.usage)), + ("metrics", _jobj(("latencyMs", 42))), + ) + + +def _aws_event_frame(event_type: str, payload: Mapping[str, object]) -> bytes: + """One application/vnd.amazon.eventstream frame: prelude + prelude CRC32 + + headers + JSON payload + message CRC32, matching botocore EventStreamBuffer.""" + try: + from botocore.eventstream import crc32 as _crc32 + except ImportError: + _crc32 = zlib.crc32 + + def _str_header(name: str, value: str) -> bytes: + name_b: Final = name.encode() + value_b: Final = value.encode() + return ( + struct.pack("!B", len(name_b)) + + name_b + + struct.pack("!B", 7) + + struct.pack("!H", len(value_b)) + + value_b + ) + + payload_bytes: Final = json.dumps(payload, default=dict, separators=(",", ":")).encode() + headers_bytes: Final = ( + _str_header(":event-type", event_type) + + _str_header(":content-type", "application/json") + + _str_header(":message-type", "event") + ) + total_length: Final = 12 + len(headers_bytes) + len(payload_bytes) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers_bytes)) + prelude_crc: Final = struct.pack("!I", _crc32(prelude) & 0xFFFFFFFF) + message: Final = prelude + prelude_crc + headers_bytes + payload_bytes + return message + struct.pack("!I", _crc32(message, 0) & 0xFFFFFFFF) + + +def _bedrock_eventstream(scenario: Scenario) -> bytes: + tool_call: Final = scenario.output.tool_call + block_start: Final[tuple[bytes, ...]] = ( + ( + _aws_event_frame( + "contentBlockStart", + _jobj( + ( + "start", + _jobj( + ( + "toolUse", + _jobj( + ("toolUseId", f"tooluse_{scenario.scenario_id}"), + ("name", tool_call.name), + ), + ), + ), + ), + ("contentBlockIndex", 0), + ), + ), + ) + if tool_call is not None + else () + ) + deltas: Final[tuple[bytes, ...]] = ( + tuple( + _aws_event_frame( + "contentBlockDelta", + _jobj( + ("delta", _jobj(("toolUse", _jobj(("input", arguments_slice))))), + ("contentBlockIndex", 0), + ), + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ) + if tool_call is not None + else ( + _aws_event_frame( + "contentBlockDelta", + _jobj( + ("delta", _jobj(("text", scenario.output.text))), + ("contentBlockIndex", 0), + ), + ), + ) + ) + return b"".join( + ( + _aws_event_frame("messageStart", _jobj(("role", "assistant"))), + *block_start, + *deltas, + _aws_event_frame("contentBlockStop", _jobj(("contentBlockIndex", 0))), + _aws_event_frame("messageStop", _jobj(("stopReason", _bedrock_stop_reason(scenario)))), + *( + ( + _aws_event_frame( + "metadata", + _jobj( + ("usage", _bedrock_usage(scenario.usage)), + ("metrics", _jobj(("latencyMs", 42))), + ), + ), + ) + if scenario.stream_usage == "final_chunk" + else () + ), + ) + ) + + +def _render( + scenario: Scenario, *, stream: bool, requested_model: str, path_tail: str +) -> RenderedResponse: + # Azure bridges gpt-5.4+ chat requests carrying function tools onto the + # Responses API, which lands on the same mount at openai/responses. + if scenario.wire == "azure_chat" and path_tail.endswith("openai/responses"): + if stream: + return RenderedResponse( + 200, "text/event-stream", _responses_sse(scenario, requested_model) + ) + return RenderedResponse( + 200, "application/json", _json_bytes(_responses_body(scenario, requested_model)) + ) + if scenario.wire == "bedrock_converse": + if stream: + return RenderedResponse( + 200, "application/vnd.amazon.eventstream", _bedrock_eventstream(scenario) + ) + return RenderedResponse(200, "application/json", _json_bytes(_bedrock_body(scenario))) + if scenario.wire == "vertex_generate": + if stream: + return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model))) if scenario.wire == "anthropic_messages": if stream: return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model)) @@ -917,7 +1136,8 @@ def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> Render if stream: return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model)) return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model))) - # openai_chat, together_chat, fireworks_chat share the OpenAI chat shape. + # openai_chat, together_chat, fireworks_chat and azure_chat share the + # OpenAI chat shape. if stream: return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model)) return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model))) @@ -954,17 +1174,28 @@ def _request_body(body: bytes) -> Mapping[str, object]: return MappingProxyType({}) -def _request_wants_stream(path_tail: str, body: bytes) -> bool: - if ":streamGenerateContent" in path_tail: +def _request_wants_stream(mount_endpoint: str | None, path_tail: str, body: bytes) -> bool: + if mount_endpoint == "streamGenerateContent" or ":streamGenerateContent" in path_tail: + return True + if path_tail.endswith("converse-stream"): return True if not body: return False return _request_body(body).get("stream") is True -def _request_model(body: bytes) -> str: +def _request_model(body: bytes, path_tail: str, scenario: Scenario) -> str: model: Final = _request_body(body).get("model") - return model if isinstance(model, str) else "unknown" + if isinstance(model, str): + return model + # Bedrock Converse names the model in the path: model//converse[-stream]. + if path_tail.startswith("model/"): + path_model: Final = path_tail.split("/", 2)[1] if path_tail.count("/") >= 2 else "" + if path_model: + return unquote(path_model) + # Vertex names it in the URL too, but the mount segment swallowed it when + # the api_base carried a path; fall back to the scenario's declared model. + return scenario.model def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: @@ -972,6 +1203,22 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte segments: Final = tuple(segment for segment in path.split("/") if segment) if method == "GET" and segments == ("health",): return RenderedResponse(200, "application/json", _json_bytes(_jobj(("status", "ok")))) + if segments and segments[0] == "_oauth": + if method == "POST" and segments == ("_oauth", "token"): + return RenderedResponse( + 200, + "application/json", + _json_bytes( + _jobj( + ("access_token", "scripted-token"), + ("token_type", "Bearer"), + ("expires_in", 3600), + ) + ), + ) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", "unknown control route"))) + ) if segments and segments[0] == "_scenarios": if method == "POST" and len(segments) == 1: try: @@ -998,7 +1245,15 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte return RenderedResponse( 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) ) - scenario_id, mount = segments[0], segments[1] + scenario_id: Final = segments[0] + # Vertex builds {api_base}:{endpoint}, so the mount segment can carry a + # :generateContent / :streamGenerateContent suffix. + mount_segment: Final = segments[1] + mount, mount_endpoint = ( + mount_segment.split(":", 1) + if ":" in mount_segment + else (mount_segment, None) + ) found: Final = store.get(scenario_id) if found is None: return RenderedResponse( @@ -1013,7 +1268,12 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte ), ) tail: Final = "/".join(segments[2:]) - return _render(found, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + return _render( + found, + stream=_request_wants_stream(mount_endpoint, tail, body), + requested_model=_request_model(body, tail, found), + path_tail=tail, + ) class _ScriptedHandler(BaseHTTPRequestHandler): diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index 4da7b31a6ef..a36bb1a8662 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -94,6 +94,40 @@ _WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ "fireworks_ai/kimi-k3", ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), ), + "azure_chat": ( + "azure/gpt-5.6", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "bedrock_converse": ( + "anthropic.claude-sonnet-5-v1:0", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + ), + ), + "vertex_generate": ( + "gemini-3.8-flash", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), }) _SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) @@ -141,6 +175,36 @@ _SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( response_model_override=True, ), ), + ( + "vertex_prompt_blocked", + "vertex_generate", + Case( + name="prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "vertex_prompt_blocked_stream", + "vertex_generate", + Case( + name="stream_prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "azure_served_model_override", + "azure_chat", + Case( + name="response_model_override", + usage=_SHAPE_USAGE, + response_model_override=True, + ), + ), ) diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index 68d840870c9..4fba337b701 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -304,6 +304,149 @@ "supports_reasoning": true, "supports_web_search": true }, + "anthropic.claude-sonnet-5-v1:0": { + "cache_creation_input_token_cost": 0.00051, + "cache_creation_input_token_cost_above_1hr": 0.00068, + "cache_read_input_token_cost": 1.7e-05, + "input_cost_per_token": 0.00017, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00034, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00048, + "cache_creation_input_token_cost_above_1hr": 0.00064, + "cache_read_input_token_cost": 1.6e-05, + "input_cost_per_audio_token": 0.00096, + "input_cost_per_token": 0.00016, + "input_cost_per_token_above_200k_tokens": 0.00128, + "input_cost_per_token_flex": 0.00024, + "input_cost_per_token_priority": 0.000272, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00112, + "output_cost_per_reasoning_token": 0.0008, + "output_cost_per_token": 0.00032, + "output_cost_per_token_above_200k_tokens": 0.00144, + "output_cost_per_token_flex": 0.0004, + "output_cost_per_token_priority": 0.000432, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "azure/gpt-5.6": { + "cache_creation_input_token_cost": 0.00044999999999999996, + "cache_creation_input_token_cost_above_1hr": 0.0006000000000000001, + "cache_read_input_token_cost": 1.5e-05, + "input_cost_per_audio_token": 0.0009000000000000001, + "input_cost_per_token": 0.00015000000000000001, + "input_cost_per_token_above_200k_tokens": 0.0012000000000000001, + "input_cost_per_token_flex": 0.000225, + "input_cost_per_token_priority": 0.000255, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0010500000000000002, + "output_cost_per_reasoning_token": 0.00075, + "output_cost_per_token": 0.00030000000000000003, + "output_cost_per_token_above_200k_tokens": 0.00135, + "output_cost_per_token_flex": 0.000375, + "output_cost_per_token_priority": 0.00040499999999999996, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.1e-05, + "input_cost_per_audio_token": 0.00126, + "input_cost_per_token": 0.00021, + "input_cost_per_token_above_200k_tokens": 0.00168, + "input_cost_per_token_flex": 0.000315, + "input_cost_per_token_priority": 0.000357, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00147, + "output_cost_per_reasoning_token": 0.0010500000000000002, + "output_cost_per_token": 0.00042, + "output_cost_per_token_above_200k_tokens": 0.0018900000000000001, + "output_cost_per_token_flex": 0.000525, + "output_cost_per_token_priority": 0.000567, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.8-flash": { + "cache_read_input_token_cost": 2e-05, + "input_cost_per_audio_token": 0.0012, + "input_cost_per_token": 0.0002, + "input_cost_per_token_above_200k_tokens": 0.0016, + "input_cost_per_token_flex": 0.0003, + "input_cost_per_token_priority": 0.00034, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0014000000000000002, + "output_cost_per_reasoning_token": 0.001, + "output_cost_per_token": 0.0004, + "output_cost_per_token_above_200k_tokens": 0.0018000000000000001, + "output_cost_per_token_flex": 0.0005, + "output_cost_per_token_priority": 0.00054, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "meta.llama4-maverick-17b-instruct-v1:0": { + "input_cost_per_token": 0.00019, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00038, + "supports_function_calling": true + }, "together_ai/moonshotai/Kimi-K3": { "cache_creation_input_token_cost": 0.00030000000000000003, "cache_creation_input_token_cost_above_1hr": 0.0004, @@ -363,5 +506,20 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_web_search": true + }, + "us.anthropic.claude-opus-5-v1:0": { + "cache_creation_input_token_cost": 0.0005400000000000001, + "cache_creation_input_token_cost_above_1hr": 0.00072, + "cache_read_input_token_cost": 1.8e-05, + "input_cost_per_token": 0.00018, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00036000000000000004, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true } } diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..d96478de1c4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1000,6 +1000,7 @@ class ModelInfoBody(BaseModel): access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None + base_model: str | None = None class ModelNewBody(BaseModel): From 2466975d290576de9e89a1d5d69c9ca9a6aab1ab Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:57:23 +0000 Subject: [PATCH 055/393] test(e2e): clean cost map decimals and simplify scripted wire helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 52 ++-- .../e2e/cost_calculation/scripted_provider.py | 39 ++- tests/e2e/cost_map.json | 270 +++++++++--------- 3 files changed, 177 insertions(+), 184 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 8c6db7c0010..3f3e9fd9243 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -11,6 +11,7 @@ Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). from __future__ import annotations +import functools import importlib.util import json import sys @@ -21,6 +22,8 @@ from types import ModuleType from typing import Final, Protocol, cast import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa from cost_matrix import Case, FrontierModel from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE @@ -112,33 +115,25 @@ def client() -> CostCalcClient: return CostCalcClient(proxy=proxy) -_vertex_key_pem: str | None = None +@functools.cache +def _vertex_private_key_pem() -> str: + return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() def _vertex_service_account_json() -> str: """A service-account credential JSON whose token_uri is the sidecar's /_oauth/token route: the proxy's google-auth refresh then gets a scripted - access token without touching Google. One generated RSA key per process.""" - global _vertex_key_pem # mutable-ok: session-scoped key generation cached for reuse - if _vertex_key_pem is None: - from cryptography.hazmat.primitives import serialization - from cryptography.hazmat.primitives.asymmetric import rsa - - _vertex_key_pem = ( - rsa.generate_private_key(public_exponent=65537, key_size=2048) - .private_bytes( - serialization.Encoding.PEM, - serialization.PrivateFormat.PKCS8, - serialization.NoEncryption(), - ) - .decode() - ) + access token without touching Google.""" return json.dumps( { "type": "service_account", "project_id": "cc-scripted-project", "private_key_id": "scripted", - "private_key": _vertex_key_pem, + "private_key": _vertex_private_key_pem(), "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com", "client_id": "0", "auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize", @@ -162,20 +157,21 @@ def register_scenario_deployment( handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) model_name: Final = f"{model.model_name}-{marker}" - extra_params: Final[dict[str, str]] = dict(model.litellm_params) - if model.wire == "vertex_generate": - extra_params["vertex_credentials"] = _vertex_service_account_json() + params: Final = { + "model": model.litellm_model, + "api_key": model.api_key, + "api_base": handle.api_base(), + **model.litellm_params, + **( + {"vertex_credentials": _vertex_service_account_json()} + if model.wire == "vertex_generate" + else {} + ), + } model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, - litellm_params=LiteLLMParamsBody.model_validate( - { - "model": model.litellm_model, - "api_key": model.api_key, - "api_base": handle.api_base(), - **extra_params, - } - ), + litellm_params=LiteLLMParamsBody.model_validate(params), model_info=ModelInfoBody(base_model=model.base_model), ) ) diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 00230fabeba..982132ed8df 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -997,36 +997,33 @@ def _bedrock_body(scenario: Scenario) -> Mapping[str, object]: ) +def _aws_str_header(name: str, value: str) -> bytes: + """One eventstream header: 1-byte name len + name + type-7 marker + value.""" + name_b: Final = name.encode() + value_b: Final = value.encode() + return ( + struct.pack("!B", len(name_b)) + + name_b + + struct.pack("!B", 7) + + struct.pack("!H", len(value_b)) + + value_b + ) + + def _aws_event_frame(event_type: str, payload: Mapping[str, object]) -> bytes: """One application/vnd.amazon.eventstream frame: prelude + prelude CRC32 + headers + JSON payload + message CRC32, matching botocore EventStreamBuffer.""" - try: - from botocore.eventstream import crc32 as _crc32 - except ImportError: - _crc32 = zlib.crc32 - - def _str_header(name: str, value: str) -> bytes: - name_b: Final = name.encode() - value_b: Final = value.encode() - return ( - struct.pack("!B", len(name_b)) - + name_b - + struct.pack("!B", 7) - + struct.pack("!H", len(value_b)) - + value_b - ) - payload_bytes: Final = json.dumps(payload, default=dict, separators=(",", ":")).encode() headers_bytes: Final = ( - _str_header(":event-type", event_type) - + _str_header(":content-type", "application/json") - + _str_header(":message-type", "event") + _aws_str_header(":event-type", event_type) + + _aws_str_header(":content-type", "application/json") + + _aws_str_header(":message-type", "event") ) total_length: Final = 12 + len(headers_bytes) + len(payload_bytes) + 4 prelude: Final = struct.pack("!II", total_length, len(headers_bytes)) - prelude_crc: Final = struct.pack("!I", _crc32(prelude) & 0xFFFFFFFF) + prelude_crc: Final = struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) message: Final = prelude + prelude_crc + headers_bytes + payload_bytes - return message + struct.pack("!I", _crc32(message, 0) & 0xFFFFFFFF) + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) def _bedrock_eventstream(scenario: Scenario) -> bytes: diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index 4fba337b701..85cd5ade3d5 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -1,4 +1,79 @@ { + "anthropic.claude-sonnet-5-v1:0": { + "cache_creation_input_token_cost": 0.00051, + "cache_creation_input_token_cost_above_1hr": 0.00068, + "cache_read_input_token_cost": 1.7e-05, + "input_cost_per_token": 0.00017, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00034, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00048, + "cache_creation_input_token_cost_above_1hr": 0.00064, + "cache_read_input_token_cost": 1.6e-05, + "input_cost_per_audio_token": 0.00096, + "input_cost_per_token": 0.00016, + "input_cost_per_token_above_200k_tokens": 0.00128, + "input_cost_per_token_flex": 0.00024, + "input_cost_per_token_priority": 0.000272, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00112, + "output_cost_per_reasoning_token": 0.0008, + "output_cost_per_token": 0.00032, + "output_cost_per_token_above_200k_tokens": 0.00144, + "output_cost_per_token_flex": 0.0004, + "output_cost_per_token_priority": 0.000432, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "azure/gpt-5.6": { + "cache_creation_input_token_cost": 0.00045, + "cache_creation_input_token_cost_above_1hr": 0.0006, + "cache_read_input_token_cost": 1.5e-05, + "input_cost_per_audio_token": 0.0009, + "input_cost_per_token": 0.00015, + "input_cost_per_token_above_200k_tokens": 0.0012, + "input_cost_per_token_flex": 0.000225, + "input_cost_per_token_priority": 0.000255, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00105, + "output_cost_per_reasoning_token": 0.00075, + "output_cost_per_token": 0.0003, + "output_cost_per_token_above_200k_tokens": 0.00135, + "output_cost_per_token_flex": 0.000375, + "output_cost_per_token_priority": 0.000405, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, "claude-haiku-4-5": { "cache_creation_input_token_cost": 0.00021, "cache_creation_input_token_cost_above_1hr": 0.00028000000000000003, @@ -134,6 +209,64 @@ "supports_reasoning": true, "supports_web_search": true }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.1e-05, + "input_cost_per_audio_token": 0.00126, + "input_cost_per_token": 0.00021, + "input_cost_per_token_above_200k_tokens": 0.00168, + "input_cost_per_token_flex": 0.000315, + "input_cost_per_token_priority": 0.000357, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00147, + "output_cost_per_reasoning_token": 0.00105, + "output_cost_per_token": 0.00042, + "output_cost_per_token_above_200k_tokens": 0.00189, + "output_cost_per_token_flex": 0.000525, + "output_cost_per_token_priority": 0.000567, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.8-flash": { + "cache_read_input_token_cost": 2e-05, + "input_cost_per_audio_token": 0.0012, + "input_cost_per_token": 0.0002, + "input_cost_per_token_above_200k_tokens": 0.0016, + "input_cost_per_token_flex": 0.0003, + "input_cost_per_token_priority": 0.00034, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0014, + "output_cost_per_reasoning_token": 0.001, + "output_cost_per_token": 0.0004, + "output_cost_per_token_above_200k_tokens": 0.0018, + "output_cost_per_token_flex": 0.0005, + "output_cost_per_token_priority": 0.00054, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 9e-06, "input_cost_per_audio_token": 0.00054, @@ -304,139 +437,6 @@ "supports_reasoning": true, "supports_web_search": true }, - "anthropic.claude-sonnet-5-v1:0": { - "cache_creation_input_token_cost": 0.00051, - "cache_creation_input_token_cost_above_1hr": 0.00068, - "cache_read_input_token_cost": 1.7e-05, - "input_cost_per_token": 0.00017, - "litellm_provider": "bedrock_converse", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 0.00034, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true - }, - "azure/gpt-5.4-mini": { - "cache_creation_input_token_cost": 0.00048, - "cache_creation_input_token_cost_above_1hr": 0.00064, - "cache_read_input_token_cost": 1.6e-05, - "input_cost_per_audio_token": 0.00096, - "input_cost_per_token": 0.00016, - "input_cost_per_token_above_200k_tokens": 0.00128, - "input_cost_per_token_flex": 0.00024, - "input_cost_per_token_priority": 0.000272, - "litellm_provider": "azure", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.00112, - "output_cost_per_reasoning_token": 0.0008, - "output_cost_per_token": 0.00032, - "output_cost_per_token_above_200k_tokens": 0.00144, - "output_cost_per_token_flex": 0.0004, - "output_cost_per_token_priority": 0.000432, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true - }, - "azure/gpt-5.6": { - "cache_creation_input_token_cost": 0.00044999999999999996, - "cache_creation_input_token_cost_above_1hr": 0.0006000000000000001, - "cache_read_input_token_cost": 1.5e-05, - "input_cost_per_audio_token": 0.0009000000000000001, - "input_cost_per_token": 0.00015000000000000001, - "input_cost_per_token_above_200k_tokens": 0.0012000000000000001, - "input_cost_per_token_flex": 0.000225, - "input_cost_per_token_priority": 0.000255, - "litellm_provider": "azure", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.0010500000000000002, - "output_cost_per_reasoning_token": 0.00075, - "output_cost_per_token": 0.00030000000000000003, - "output_cost_per_token_above_200k_tokens": 0.00135, - "output_cost_per_token_flex": 0.000375, - "output_cost_per_token_priority": 0.00040499999999999996, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true - }, - "gemini-3.1-pro-preview": { - "cache_read_input_token_cost": 2.1e-05, - "input_cost_per_audio_token": 0.00126, - "input_cost_per_token": 0.00021, - "input_cost_per_token_above_200k_tokens": 0.00168, - "input_cost_per_token_flex": 0.000315, - "input_cost_per_token_priority": 0.000357, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.00147, - "output_cost_per_reasoning_token": 0.0010500000000000002, - "output_cost_per_token": 0.00042, - "output_cost_per_token_above_200k_tokens": 0.0018900000000000001, - "output_cost_per_token_flex": 0.000525, - "output_cost_per_token_priority": 0.000567, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true, - "web_search_billing_unit": "per_query" - }, - "gemini-3.8-flash": { - "cache_read_input_token_cost": 2e-05, - "input_cost_per_audio_token": 0.0012, - "input_cost_per_token": 0.0002, - "input_cost_per_token_above_200k_tokens": 0.0016, - "input_cost_per_token_flex": 0.0003, - "input_cost_per_token_priority": 0.00034, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_audio_token": 0.0014000000000000002, - "output_cost_per_reasoning_token": 0.001, - "output_cost_per_token": 0.0004, - "output_cost_per_token_above_200k_tokens": 0.0018000000000000001, - "output_cost_per_token_flex": 0.0005, - "output_cost_per_token_priority": 0.00054, - "search_context_cost_per_query": { - "search_context_size_high": 0.03, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.02 - }, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_web_search": true, - "web_search_billing_unit": "per_query" - }, "meta.llama4-maverick-17b-instruct-v1:0": { "input_cost_per_token": 0.00019, "litellm_provider": "bedrock_converse", @@ -508,7 +508,7 @@ "supports_web_search": true }, "us.anthropic.claude-opus-5-v1:0": { - "cache_creation_input_token_cost": 0.0005400000000000001, + "cache_creation_input_token_cost": 0.00054, "cache_creation_input_token_cost_above_1hr": 0.00072, "cache_read_input_token_cost": 1.8e-05, "input_cost_per_token": 0.00018, @@ -517,7 +517,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 0.00036000000000000004, + "output_cost_per_token": 0.00036, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true From 813d96f26ea6780bacb4b5ad562f1aecc3cb5069 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:18:26 +0000 Subject: [PATCH 056/393] fix(e2e): resolve remaining merge markers in e2e_config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/e2e_config.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index b34eadd8744..e19cfaa684f 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,7 +145,6 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" -<<<<<<< HEAD # The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL # pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a # scripted-provider sidecar; deselected unless the opt-in env var is set. @@ -162,9 +161,6 @@ SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get( SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get( "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL ).rstrip("/") -||||||| 930ec9643a -======= -CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) From ce722ab1b30d4b1331504364eea7adb362887c38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:31:28 -0700 Subject: [PATCH 057/393] fix(proxy): evict the member's cached user row on team member add JWT auth caches the user row before it adds the user to the JWT's team, and admission checks the credential's team against the cached row on whichever worker takes the next request. On a two-worker gateway the credential minted for a newly joined team answered 403 "not in your team memberships" until the management-object TTL ran out, because /team/member_add only evicted the membership spend sentinel. The add now evicts the added members' cached user rows and broadcasts the eviction to the other workers, the way /team/member_delete already did The mint test now also covers a user SCIM deactivated after the cache last saw them active: the database read refuses the mint while the cached row still says active --- .../mcp_server/bridge_token_flow.py | 9 ++- .../management_endpoints/team_endpoints.py | 5 ++ .../mcp_server/test_proxy_api_credentials.py | 22 ++++++ .../test_team_endpoints.py | 71 +++++++++++++++++++ 4 files changed, 102 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index f19cb87ae18..37a893973e3 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -279,11 +279,10 @@ async def load_active_user_by_id( catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look identical, the original error surviving only as ``__context__``), so the outage check walks the cause chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. - ``source="database"`` reads the row from the database, never the cache, and leaves the fresh row in the - cache for the requests the credential makes next: JWT auth caches the user it creates before it adds - that user to the JWT's team and adding a member never evicts the cached row, so a credential minted - off the cache would refuse the very first exchange as not a member. Every other caller keeps the cache - read, so introspection, which a resource server may call per request, stays off the database.""" + ``source="database"`` reads the row from the database, never the cache, so the credential mint refuses + a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh + row in the cache for the requests the credential makes next. Every other caller keeps the cache read, + so introspection, which a resource server may call per request, stays off the database.""" from litellm.proxy._types import ( ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e719d6d761a..84e3e9f87e6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3154,6 +3154,7 @@ async def team_member_add( ``` """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, premium_user, @@ -3248,6 +3249,10 @@ async def team_member_add( litellm_proxy_admin_name=litellm_proxy_admin_name, ) + await evict_and_broadcast( + cache_keys=tuple(sorted(user.user_id for user in updated_users)), + user_api_key_cache=user_api_key_cache, + ) await _evict_created_membership_caches( user_ids=(tm.user_id for tm in updated_team_memberships), team_id=data.team_id, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index 04fbbe4a6ce..ed3e5f48516 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -152,6 +152,28 @@ async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_r assert _decoded(minted).team_id == "team-a" +@pytest.mark.asyncio +async def test_mint_refuses_a_user_scim_deactivated_after_the_cache_last_saw_them_active(fetch_teams, monkeypatch): + """SCIM deactivation writes the user row without evicting the cached copy, so a mint off the cache would + keep issuing credentials for the management-object TTL. The mint reads the database row, so the + deactivated user is refused on the first refresh after the deactivation.""" + from litellm.proxy import proxy_server + + cache = UserApiKeyCache() + await cache.async_set_cache( + key="deactivated-user", value=_user(user_id="deactivated-user", teams=["team-a"]), model_type=LiteLLM_UserTable + ) + prisma = MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=_user(user_id="deactivated-user", teams=["team-a"], metadata={"scim_active": False}) + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + assert await mint_proxy_credential("deactivated-user", "team-a") == "no_active_key" + fetch_teams.assert_not_awaited() + + @pytest.mark.asyncio async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams): assert await mint_proxy_credential("u1", "team-c") == "not_a_member" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ebbedc6541e..748cdbbc175 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13090,6 +13090,77 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"] +@pytest.mark.asyncio +async def test_team_member_add_evicts_the_new_members_cached_user_row_on_every_worker(monkeypatch): + """Auth admits a team-bound credential off the teams list of the cached user row. The add wrote the + new team to the database row only, so a worker still holding the old row refused the member's + credential with 403 until the management-object TTL expired. The add now evicts the row here and + broadcasts the eviction to the other workers, the way /team/member_delete already does.""" + from litellm.proxy._types import TeamMemberAddRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_member_add + + team_id = "team-b" + user_id = "dev-1" + cache = UserApiKeyCache() + await cache.async_set_cache( + key=user_id, value=LiteLLM_UserTable(user_id=user_id, teams=["team-a"]), model_type=LiteLLM_UserTable + ) + broadcast = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id") + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", broadcast + ) + + updated_team = MagicMock() + updated_team.model_dump.return_value = { + "team_id": team_id, + "members_with_roles": [{"user_id": user_id, "role": "user"}], + } + + async def fake_add_team_members_to_team(**kwargs): + return updated_team, [LiteLLM_UserTable(user_id=user_id, teams=["team-a", team_id])], [] + + with ( + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]), + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids", + new_callable=AsyncMock, + return_value=frozenset({user_id}), + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + side_effect=fake_add_team_members_to_team, + ), + patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers + "litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs", + new_callable=AsyncMock, + ), + ): + await team_member_add( + data=TeamMemberAddRequest(team_id=team_id, member=Member(user_id=user_id, role="user")), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + assert await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) is None + broadcast.assert_awaited_once_with(cache_key=user_id) + + def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): """A large member list must not echo every id back in the error body.""" from litellm.proxy.management_endpoints.team_endpoints import ( From 769b47457ee2aa35ac20e6b2815bc7800ef95bf2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:22:27 -0700 Subject: [PATCH 058/393] fix(proxy): keep the token exchange off gateways that map JWTs to virtual keys --- .../mcp_server/idp_token_exchange.py | 26 ++++++++-- .../mcp_server/test_discoverable_endpoints.py | 26 ++++++++-- .../mcp_server/test_idp_token_exchange.py | 50 +++++++++++++++---- 3 files changed, 85 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index cdefaf76d49..7a453e85cce 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -14,7 +14,7 @@ from fastapi import HTTPException, Request from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._types import JWTAuthBuilderResult, ProxyException -from litellm.proxy.auth.handle_jwt import JWTAuthManager +from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler EXCHANGE_ROUTE: Final = "/token" REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth" @@ -23,16 +23,21 @@ REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT @dataclass(frozen=True, slots=True) class TokenExchangePrerequisites: """The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT - bearer. Discovery and registration advertise the exchange grant only when every one of - them holds, and an exchange attempt is refused naming the first one that does not.""" + bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps + tokens authenticates a JWT as its mapped key, with that key's models and budget, or + refuses an unmapped one, and the exchange proves the token through ``auth_builder`` + alone, so it would mint the user's own credential past that policy. Discovery and + registration advertise the exchange grant only when every gate holds, and an exchange + attempt is refused naming the first one that does not.""" jwt_auth_enabled: bool has_database: bool licensed: bool + maps_jwts_to_virtual_keys: bool @property def available(self) -> bool: - return self.jwt_auth_enabled and self.has_database and self.licensed + return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys def refusal(self) -> SubjectTokenRefusal | None: if not self.jwt_auth_enabled: @@ -50,12 +55,18 @@ class TokenExchangePrerequisites: error="unsupported_grant_type", description="JWT auth is an enterprise only feature; no license is set", ) + if self.maps_jwts_to_virtual_keys: + return SubjectTokenRefusal( + error="unsupported_grant_type", + description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve", + ) return None def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call general_settings, + jwt_handler, premium_user, prisma_client, ) @@ -64,9 +75,16 @@ def read_token_exchange_prerequisites() -> TokenExchangePrerequisites: jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True, has_database=prisma_client is not None, licensed=premium_user is True, + maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler), ) +def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool: + if not hasattr(jwt_handler, "litellm_jwtauth"): + return False + return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured() + + def token_exchange_available() -> bool: return read_token_exchange_prerequisites().available diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 6965b3b4ebe..d7666f5e694 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11111,13 +11111,31 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo assert stranger.json()["error"] == "invalid_client" -@pytest.mark.parametrize("exchange_servable", [True, False]) -def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(monkeypatch, exchange_servable): +@pytest.mark.parametrize( + "jwt_auth_enabled, virtual_key_claim_field, exchange_servable", + [(True, None, True), (False, None, False), (True, "client_id", False)], + ids=["jwt auth on", "jwt auth off", "jwts mapped to virtual keys"], +) +def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it( + monkeypatch, jwt_auth_enabled, virtual_key_claim_field, exchange_servable +): """Every document a native client reads before it picks a grant (the versioned contract, the aggregate authorization-server metadata, and the registration response) lists the RFC 8693 - exchange exactly when the running proxy can serve it: JWT auth on, a database, and a license.""" + exchange exactly when the running proxy can serve it: JWT auth on, a database, a license, and + no JWT-to-virtual-key mapping, since the exchange would mint past the mapped key's policy.""" + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + client, _session_cookie, _minted = _native_client_app(monkeypatch) - monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": exchange_servable}) + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(virtual_key_claim_field=virtual_key_claim_field), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", handler) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled}) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py index d1b049dddd5..e12c8823f99 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py @@ -3,6 +3,7 @@ import logging import pytest from fastapi import HTTPException +from litellm.caching.caching import DualCache from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, @@ -10,12 +11,17 @@ from litellm.proxy._experimental.mcp_server.idp_token_exchange import ( identity_from_subject_token, token_exchange_available, ) -from litellm.proxy._types import ProxyException +from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException from litellm.proxy.auth.handle_jwt import JWTHandler IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature" REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"} -EVERY_GATE_HOLDS = {"jwt_auth_enabled": True, "has_database": True, "licensed": True} +EVERY_GATE_HOLDS = { + "jwt_auth_enabled": True, + "has_database": True, + "licensed": True, + "maps_jwts_to_virtual_keys": False, +} JWKS_URL = "https://idp.example.com/.well-known/jwks.json" @@ -82,6 +88,7 @@ async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity(): ({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"), ({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"), ({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"), + ({"maps_jwts_to_virtual_keys": True}, IDP_JWT, "unsupported_grant_type", "virtual keys"), ({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"), ], ) @@ -96,28 +103,53 @@ async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verificatio assert authorizer.calls == [] -@pytest.mark.parametrize("unmet", [{}, {"jwt_auth_enabled": False}, {"has_database": False}, {"licensed": False}]) +@pytest.mark.parametrize( + "unmet", + [ + {}, + {"jwt_auth_enabled": False}, + {"has_database": False}, + {"licensed": False}, + {"maps_jwts_to_virtual_keys": True}, + ], +) def test_the_grant_is_available_exactly_when_every_gate_holds(unmet): prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}) assert prerequisites.available is (unmet == {}) assert (prerequisites.refusal() is None) is prerequisites.available +MAPPED_ISSUER = JWTIssuerConfig( + issuer="https://idp.example.test", audience="litellm-gateway", virtual_key_claim_field="client_id" +) + + +def _running_jwt_handler(litellm_jwtauth): + handler = JWTHandler() + if litellm_jwtauth is not None: + handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth) + return handler + + @pytest.mark.parametrize( - "general_settings, prisma_client, premium_user, expected", + "general_settings, prisma_client, premium_user, litellm_jwtauth, expected", [ - ({"enable_jwt_auth": True}, object(), True, True), - ({}, object(), True, False), - ({"enable_jwt_auth": True}, None, True, False), - ({"enable_jwt_auth": True}, object(), False, False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), True), + ({"enable_jwt_auth": True}, object(), True, None, True), + ({}, object(), True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(), False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(virtual_key_claim_field="client_id"), False), + ({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(issuers=[MAPPED_ISSUER]), False), ], ) def test_availability_is_read_from_the_running_proxy( - monkeypatch, general_settings, prisma_client, premium_user, expected + monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected ): monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user) + monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth)) assert token_exchange_available() is expected From bdfff602fb0325f88768f7c4411cce921ab28fcb Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:47:55 +0000 Subject: [PATCH 059/393] test(e2e): drive the cost matrix from cases.json and expected.json goldens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/cost_calculation/cases.json | 231 ++ tests/e2e/cost_calculation/conftest.py | 10 +- tests/e2e/cost_calculation/cost_matrix.py | 788 ++----- tests/e2e/cost_calculation/expected.json | 2004 +++++++++++++++++ .../e2e/cost_calculation/generate_expected.py | 189 ++ .../e2e/cost_calculation/test_matrix_data.py | 64 + .../test_token_pricing_e2e.py | 62 +- .../cost_calculation/test_wire_formats_e2e.py | 368 --- 9 files changed, 2761 insertions(+), 957 deletions(-) create mode 100644 tests/e2e/cost_calculation/cases.json create mode 100644 tests/e2e/cost_calculation/expected.json create mode 100644 tests/e2e/cost_calculation/generate_expected.py create mode 100644 tests/e2e/cost_calculation/test_matrix_data.py delete mode 100644 tests/e2e/cost_calculation/test_wire_formats_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 49cfc29aa17..707d35b4aa6 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json new file mode 100644 index 00000000000..e898557ea35 --- /dev/null +++ b/tests/e2e/cost_calculation/cases.json @@ -0,0 +1,231 @@ +{ + "deployments": [ + { + "map_key": "azure/gpt-5.4-mini", + "litellm_model": "azure/cc-pinned-deployment", + "base_model": "azure/gpt-5.4-mini" + } + ], + "cases": [ + { + "name": "basic", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40} + }, + { + "name": "cache_read", + "usage": {"fresh_input_tokens": 100, "cache_read_tokens": 50, "output_tokens": 30}, + "requires_rates": ["cache_read_input_token_cost"], + "requires_caps": ["cache_read"] + }, + { + "name": "cache_write_5m", + "usage": {"fresh_input_tokens": 90, "cache_write_5m_tokens": 60, "output_tokens": 30}, + "requires_rates": ["cache_creation_input_token_cost"], + "requires_caps": ["cache_write_5m"] + }, + { + "name": "cache_write_1h", + "usage": {"fresh_input_tokens": 90, "cache_write_5m_tokens": 20, "cache_write_1h_tokens": 40, "output_tokens": 30}, + "requires_rates": ["cache_creation_input_token_cost_above_1hr", "cache_creation_input_token_cost"], + "requires_caps": ["cache_write_1h"] + }, + { + "name": "reasoning", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "reasoning_tokens": 70}, + "requires_rates": ["output_cost_per_reasoning_token"], + "requires_caps": ["reasoning"] + }, + { + "name": "audio", + "usage": {"fresh_input_tokens": 100, "audio_input_tokens": 25, "output_tokens": 30, "audio_output_tokens": 15}, + "requires_rates": ["input_cost_per_audio_token", "output_cost_per_audio_token"], + "requires_caps": ["audio"] + }, + { + "name": "tiered", + "usage": {"fresh_input_tokens": 200001, "output_tokens": 30}, + "requires_rates": ["input_cost_per_token_above_200k_tokens", "output_cost_per_token_above_200k_tokens"] + }, + { + "name": "service_tier_flex", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "service_tier": "flex", + "requires_rates": ["input_cost_per_token_flex", "output_cost_per_token_flex"] + }, + { + "name": "service_tier_priority", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "service_tier": "priority", + "requires_rates": ["input_cost_per_token_priority", "output_cost_per_token_priority"] + }, + { + "name": "web_search", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3}, + "requires_rates": ["search_context_cost_per_query"], + "requires_caps": ["web_search"] + }, + { + "name": "stream", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true + }, + { + "name": "stream_no_usage", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "exact_spend": false, + "requires_caps": ["absent_usage"] + }, + { + "name": "response_model_override", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "response_model_override": true, + "requires_caps": ["response_model"] + }, + { + "name": "stream_response_model_override", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "response_model_override": true, + "requires_caps": ["response_model"] + }, + { + "name": "tool_call", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "tool_call": true, + "requires_caps": ["tool_call"] + }, + { + "name": "stream_tool_call", + "usage": {"fresh_input_tokens": 80, "output_tokens": 25}, + "stream": true, + "tool_call": true, + "requires_caps": ["tool_call"] + }, + { + "name": "stream_no_usage_tool_call", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "tool_call": true, + "exact_spend": false, + "requires_caps": ["absent_usage", "tool_call"] + }, + { + "name": "stream_no_usage_image_input", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "image_input": true, + "exact_spend": false, + "requires_caps": ["absent_usage", "image_input"] + }, + { + "name": "stream_incomplete", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "terminal": "incomplete", + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_no_usage_incomplete", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "terminal": "incomplete", + "exact_spend": false, + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_unvalidated", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "terminal": "unvalidated", + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_no_usage_unvalidated", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "terminal": "unvalidated", + "exact_spend": false, + "requires_caps": ["responses_terminal"] + }, + { + "name": "prompt_blocked", + "usage": {"fresh_input_tokens": 1000, "output_tokens": 0}, + "terminal": "prompt_blocked", + "response_model_override": true, + "requires_caps": ["prompt_blocked"] + }, + { + "name": "stream_prompt_blocked", + "usage": {"fresh_input_tokens": 1000, "output_tokens": 0}, + "stream": true, + "terminal": "prompt_blocked", + "response_model_override": true, + "requires_caps": ["prompt_blocked"] + }, + { + "name": "all_components_chat", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25, + "reasoning_tokens": 15, + "audio_input_tokens": 5, + "audio_output_tokens": 3 + }, + "wires": ["openai_chat", "azure_chat", "together_chat"] + }, + { + "name": "all_components_fireworks", + "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, + "wires": ["fireworks_chat"] + }, + { + "name": "all_components_anthropic", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25 + }, + "wires": ["anthropic_messages", "bedrock_converse"] + }, + { + "name": "all_components_anthropic_stream", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25 + }, + "stream": true, + "wires": ["anthropic_messages"] + }, + { + "name": "all_components_gemini", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "output_tokens": 25, + "reasoning_tokens": 15, + "audio_input_tokens": 5, + "audio_output_tokens": 3 + }, + "wires": ["gemini_generate", "vertex_generate"] + }, + { + "name": "all_components_responses", + "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, + "wires": ["openai_responses"] + } + ] +} diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 3f3e9fd9243..3de9786854e 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -1,10 +1,12 @@ """Cost-calculation suite fixtures. Runs against a dedicated proxy whose whole model cost map is the test-owned -``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL), so every deployment -bills at rates the test asserts literal arithmetic on. Provider calls are -answered by the scripted-provider sidecar (``scripted_provider.py``), registered -per scenario over its control API. +``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL); every map entry is a +deployment under test, the request shapes live in ``cases.json``, and the +asserted goldens live in ``expected.json`` (regenerate proposals with +``generate_expected.py``). Provider calls are answered by the +scripted-provider sidecar (``scripted_provider.py``), registered per scenario +over its control API. Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). """ diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 37495f37b0b..b03c851d208 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -1,18 +1,16 @@ -"""The cost-calculation matrix: frontier model set, the pricing-component cases -each model runs, and the expected-cost arithmetic. +"""The cost-calculation matrix: the model set derived from the test cost map, +the request/response cases from ``cases.json``, and the loaders both use. -Rates come from ``tests/e2e/cost_map.json``, which the proxy under test loads as -its ENTIRE model cost map (LITELLM_MODEL_COST_MAP_URL), so an entry's rates are -exactly what the proxy bills and nothing in the suite depends on the bundled -map. Each model's rates are a distinct multiple of a shared base set, so a -component billed at the wrong model's rate (or the wrong case's rate) can never -coincidentally match. - -Case applicability is pricing-field-gated AND wire-gated: a case runs for a -model only when the entry carries the rate the case exercises and the wire can -report the token kind that rate prices. When the wire cannot report a kind -(e.g. Anthropic has no reasoning-token field, Responses reports no cache -creation), the case is absent from the matrix rather than silently zero. +Three data files drive the suite; nothing in Python lists models or cases: +- ``tests/e2e/cost_map.json`` is the proxy's ENTIRE model cost map + (LITELLM_MODEL_COST_MAP_URL); every entry becomes a deployment under test. +- ``tests/e2e/cost_calculation/cases.json`` is the case list; each case runs + for a model when the entry carries the rates it exercises (``requires_rates``) + and the wire can report the token kinds involved (``requires_caps`` / + ``wires``). +- ``tests/e2e/cost_calculation/expected.json`` holds the reviewed goldens; the + tests assert them verbatim and never compute a price themselves. The rate + arithmetic that proposes goldens lives in ``generate_expected.py``, not here. """ from __future__ import annotations @@ -26,13 +24,15 @@ from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from types import MappingProxyType -from typing import Final, Literal, TypeAlias +from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" +CASES_PATH: Final = Path(__file__).resolve().parent / "cases.json" +EXPECTED_PATH: Final = Path(__file__).resolve().parent / "expected.json" class SearchContextCostPerQuery(BaseModel): @@ -77,119 +77,90 @@ _COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( TIER_THRESHOLD_TOKENS: Final = 200_000 -@dataclass(frozen=True, slots=True) -class FrontierModel: - """One deployment under test: the model_name the suite registers, the - provider-prefixed litellm model string, the wire the scripted upstream - speaks, its cost-map key, and the sibling map model the response_model - override case reports.""" +class DeploymentSpec(BaseModel): + """A deployment-level fact from cases.json: when a map key needs a + registered deployment name that is not its provider model (or a + model_info.base_model pin), the matrix uses these instead of the defaults.""" + + model_config = ConfigDict(frozen=True) - model_name: str - litellm_model: str - wire: Wire map_key: str - override_model: str | None = None - override_map_key: str | None = None - # Registered as model_info.base_model; when set, the provider-reported - # model loses to it and every case bills at this deployment's own rates. + litellm_model: str | None = None base_model: str | None = None - # Extra litellm_params merged into the /model/new registration (api_version, - # aws_* credentials, vertex_* auth). - litellm_params: Mapping[str, str] = MappingProxyType({}) - - @property - def rates(self) -> CostMapEntry: - return _COST_MAP[self.map_key] - - @property - def override_rates(self) -> CostMapEntry: - if self.base_model is not None or self.override_map_key is None: - return self.rates - return _COST_MAP[self.override_map_key] - - @property - def provider_model(self) -> str: - """The bare provider-facing model name: litellm_model minus the provider - prefix and any routing segment (converse/, responses/).""" - tail: Final = self.litellm_model.split("/")[1:] - return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) - - @property - def provider(self) -> str: - return self.rates.litellm_provider - - @property - def api_key(self) -> str: - # The scripted upstream ignores auth; a fixed bogus key proves the suite - # spends zero real provider calls. - return "sk-scripted-provider" -# Response-model override targets: emit a sibling's bare provider-facing name so -# the biller's provider-prefixed lookup lands on that sibling's map key. -_OVERRIDE_MODELS: Final[Mapping[str, str]] = MappingProxyType({ - "gpt-5.6": "gpt-5.4-mini", - "gpt-5.5-pro": "gpt-5.3-codex", - "gpt-5.3-codex": "gpt-5.5-pro", - "gpt-5.4-mini": "gpt-5.6", - "claude-opus-5": "claude-sonnet-5", - "claude-sonnet-5": "claude-opus-5", - "claude-haiku-4-5": "claude-sonnet-5", - "gemini/gemini-3.8-flash": "gemini-3.1-pro-preview", - "gemini/gemini-3.1-pro-preview": "gemini-3.8-flash", - "together_ai/moonshotai/Kimi-K3": "zai-org/GLM-5.3", - "together_ai/zai-org/GLM-5.3": "moonshotai/Kimi-K3", - "fireworks_ai/kimi-k3": "qwen3p8-max", - "fireworks_ai/qwen3p8-max": "kimi-k3", - "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", -}) +class Case(BaseModel): + """One request/response shape from cases.json; gated onto a model by + ``requires_rates`` (entry must carry each rate field), ``requires_caps`` + (the wire must report the token kind) and ``wires`` (shape is wire-specific).""" -_OVERRIDE_MAP_KEYS: Final[Mapping[str, str]] = MappingProxyType({ - "gpt-5.4-mini": "gpt-5.4-mini", - "gpt-5.6": "gpt-5.6", - "gpt-5.3-codex": "gpt-5.3-codex", - "gpt-5.5-pro": "gpt-5.5-pro", - "claude-sonnet-5": "claude-sonnet-5", - "claude-opus-5": "claude-opus-5", - "gemini-3.1-pro-preview": "gemini/gemini-3.1-pro-preview", - "gemini-3.8-flash": "gemini/gemini-3.8-flash", - "zai-org/GLM-5.3": "together_ai/zai-org/GLM-5.3", - "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", - "qwen3p8-max": "fireworks_ai/qwen3p8-max", - "kimi-k3": "fireworks_ai/kimi-k3", -}) + model_config = ConfigDict(frozen=True) + + name: str + usage: ScriptedUsage + stream: bool = False + stream_usage: Literal["final_chunk", "absent"] = "final_chunk" + service_tier: Literal["flex", "priority"] | None = None + response_model_override: bool = False + exact_spend: bool = True + tool_call: bool = False + image_input: bool = False + terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" + requires_rates: tuple[str, ...] = () + requires_caps: tuple[str, ...] = () + wires: tuple[Wire, ...] | None = None + + def applies_to(self, model: FrontierModel) -> bool: + if self.wires is not None and model.wire not in self.wires: + return False + caps: Final = _WIRE_CAPS[model.wire] + if not frozenset(self.requires_caps) <= caps: + return False + return all( + getattr(model.rates, field, None) is not None for field in self.requires_rates + ) + + def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: + return Scenario( + scenario_id=scenario_id, + wire=model.wire, + usage=self.usage, + model=model.provider_model, + output=ScriptedOutput( + text=text, + response_model=model.override_model if self.response_model_override else None, + tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) + if self.tool_call + else None, + terminal=self.terminal, + ), + stream_usage=self.stream_usage, + service_tier=self.service_tier, + ) -_FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( - ("gpt-5.6", "openai/gpt-5.6", "openai_chat"), - ("gpt-5.5-pro", "openai/gpt-5.5-pro", "openai_responses"), - ("gpt-5.3-codex", "openai/gpt-5.3-codex", "openai_responses"), - ("gpt-5.4-mini", "openai/gpt-5.4-mini", "openai_chat"), - ("claude-opus-5", "anthropic/claude-opus-5", "anthropic_messages"), - ("claude-sonnet-5", "anthropic/claude-sonnet-5", "anthropic_messages"), - ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "anthropic_messages"), - ("gemini/gemini-3.8-flash", "gemini/gemini-3.8-flash", "gemini_generate"), - ("gemini/gemini-3.1-pro-preview", "gemini/gemini-3.1-pro-preview", "gemini_generate"), - ("together_ai/moonshotai/Kimi-K3", "together_ai/moonshotai/Kimi-K3", "together_chat"), - ("together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3", "together_chat"), - ("fireworks_ai/kimi-k3", "fireworks_ai/kimi-k3", "fireworks_chat"), - ("fireworks_ai/qwen3p8-max", "fireworks_ai/qwen3p8-max", "fireworks_chat"), - ("fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/deepseek-v4p1-flash", "fireworks_chat"), +class _CasesFile(BaseModel): + model_config = ConfigDict(frozen=True) + + deployments: tuple[DeploymentSpec, ...] = () + cases: tuple[Case, ...] = () + + +_CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) +CASES: Final[tuple[Case, ...]] = _CASES_FILE.cases +_DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( + {spec.map_key: spec for spec in _CASES_FILE.deployments} ) @dataclass(frozen=True, slots=True) -class _ExtendedSpec: - """A frontier entry whose override target, model_info.base_model or extra - litellm_params can't be derived from the map key alone.""" +class _ProviderWiring: + """How a (litellm_provider, mode) pair maps to a sidecar wire, the provider + prefix on the registered litellm model string, and extra litellm_params.""" - map_key: str - litellm_model: str wire: Wire - override_model: str | None = None - override_map_key: str | None = None - base_model: str | None = None - litellm_params: Mapping[str, str] = MappingProxyType({}) + model_prefix: str | None + litellm_params: Mapping[str, str] _AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"}) @@ -207,87 +178,134 @@ _VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType( } ) -_EXTENDED_SPECS: Final[tuple[_ExtendedSpec, ...]] = ( - _ExtendedSpec( - map_key="azure/gpt-5.6", - litellm_model="azure/gpt-5.6", - wire="azure_chat", - override_model="gpt-5.4-mini", - override_map_key="azure/gpt-5.4-mini", - litellm_params=_AZURE_PARAMS, - ), - _ExtendedSpec( - # Deployment name is not a model; base_model pins billing so the - # response's model field loses, proving base_model wins. - map_key="azure/gpt-5.4-mini", - litellm_model="azure/cc-pinned-deployment", - wire="azure_chat", - override_model="gpt-5.6", - override_map_key="azure/gpt-5.6", - base_model="azure/gpt-5.4-mini", - litellm_params=_AZURE_PARAMS, - ), - _ExtendedSpec( - map_key="anthropic.claude-sonnet-5-v1:0", - litellm_model="bedrock/converse/anthropic.claude-sonnet-5-v1:0", - wire="bedrock_converse", - litellm_params=_BEDROCK_PARAMS, - ), - _ExtendedSpec( - map_key="us.anthropic.claude-opus-5-v1:0", - litellm_model="bedrock/converse/us.anthropic.claude-opus-5-v1:0", - wire="bedrock_converse", - litellm_params=_BEDROCK_PARAMS, - ), - _ExtendedSpec( - map_key="meta.llama4-maverick-17b-instruct-v1:0", - litellm_model="bedrock/converse/meta.llama4-maverick-17b-instruct-v1:0", - wire="bedrock_converse", - litellm_params=_BEDROCK_PARAMS, - ), - _ExtendedSpec( - map_key="gemini-3.8-flash", - litellm_model="vertex_ai/gemini-3.8-flash", - wire="vertex_generate", - override_model="gemini-3.1-pro-preview", - override_map_key="gemini-3.1-pro-preview", - litellm_params=_VERTEX_PARAMS, - ), - _ExtendedSpec( - map_key="gemini-3.1-pro-preview", - litellm_model="vertex_ai/gemini-3.1-pro-preview", - wire="vertex_generate", - override_model="gemini-3.8-flash", - override_map_key="gemini-3.8-flash", - litellm_params=_VERTEX_PARAMS, - ), +_PROVIDER_WIRING: Final[Mapping[tuple[str, str], _ProviderWiring]] = MappingProxyType( + { + ("openai", "chat"): _ProviderWiring("openai_chat", "openai", MappingProxyType({})), + ("openai", "responses"): _ProviderWiring( + "openai_responses", "openai", MappingProxyType({}) + ), + ("anthropic", "chat"): _ProviderWiring( + "anthropic_messages", "anthropic", MappingProxyType({}) + ), + ("gemini", "chat"): _ProviderWiring("gemini_generate", None, MappingProxyType({})), + ("together_ai", "chat"): _ProviderWiring("together_chat", None, MappingProxyType({})), + ("fireworks_ai", "chat"): _ProviderWiring("fireworks_chat", None, MappingProxyType({})), + ("azure", "chat"): _ProviderWiring("azure_chat", None, _AZURE_PARAMS), + ("bedrock_converse", "chat"): _ProviderWiring( + "bedrock_converse", "bedrock/converse", _BEDROCK_PARAMS + ), + ("vertex_ai-language-models", "chat"): _ProviderWiring( + "vertex_generate", "vertex_ai", _VERTEX_PARAMS + ), + } ) +@dataclass(frozen=True, slots=True) +class FrontierModel: + """One deployment under test, derived from a cost-map entry: the model_name + the suite registers, the provider-prefixed litellm model string, the wire + the scripted upstream speaks, and the sibling map model the response_model + override case reports.""" + + model_name: str + litellm_model: str + wire: Wire + map_key: str + override_model: str | None = None + override_map_key: str | None = None + # Registered as model_info.base_model; when set, the provider-reported + # model loses to it and every case bills at this deployment's own rates. + base_model: str | None = None + litellm_params: Mapping[str, str] = MappingProxyType({}) + + @property + def rates(self) -> CostMapEntry: + return _COST_MAP[self.map_key] + + @property + def override_rates(self) -> CostMapEntry: + if self.base_model is not None or self.override_map_key is None: + return self.rates + return _COST_MAP[self.override_map_key] + + @property + def provider_model(self) -> str: + """The bare provider-facing model name: litellm_model minus the provider + prefix and any routing segment (converse/, responses/).""" + return _provider_model(self.litellm_model) + + @property + def provider(self) -> str: + return self.rates.litellm_provider + + @property + def api_key(self) -> str: + # The scripted upstream ignores auth; a fixed bogus key proves the suite + # spends zero real provider calls. + return "sk-scripted-provider" + + +def _provider_model(litellm_model: str) -> str: + tail: Final = litellm_model.split("/")[1:] + return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) + + +def _litellm_model_for(map_key: str, wiring: _ProviderWiring) -> str: + if wiring.model_prefix is None: + return map_key + if map_key.startswith(f"{wiring.model_prefix}/"): + return map_key + return f"{wiring.model_prefix}/{map_key}" + + def _frontier() -> tuple[FrontierModel, ...]: - return tuple( - FrontierModel( - model_name=f"cc-{map_key.replace('/', '-').lower()}", - litellm_model=litellm_model, - wire=wire, - map_key=map_key, - override_model=_OVERRIDE_MODELS[map_key], - override_map_key=_OVERRIDE_MAP_KEYS[_OVERRIDE_MODELS[map_key]], - ) - for map_key, litellm_model, wire in _FRONTIER_SPECS - ) + tuple( - FrontierModel( - model_name=f"cc-{spec.map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", - litellm_model=spec.litellm_model, - wire=spec.wire, - map_key=spec.map_key, - override_model=spec.override_model, - override_map_key=spec.override_map_key, - base_model=spec.base_model, - litellm_params=spec.litellm_params, - ) - for spec in _EXTENDED_SPECS + groups: Final[Mapping[tuple[str, str], tuple[str, ...]]] = MappingProxyType( + { + pair: tuple(sorted(k for k, e in _COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) + for pair in {(e.litellm_provider, e.mode) for e in _COST_MAP.values()} + } ) + models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple + for map_key in sorted(_COST_MAP): + entry: Final = _COST_MAP[map_key] + pair: Final = (entry.litellm_provider, entry.mode) + wiring: Final = _PROVIDER_WIRING.get(pair) + if wiring is None: + raise ValueError( + f"cost_map entry {map_key} has no wiring for " + f"(litellm_provider={pair[0]}, mode={pair[1]}); add a " + f"_ProviderWiring row in cost_matrix.py" + ) + siblings: Final = groups[pair] + override_key: Final = ( + siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None + ) + override_litellm: Final = ( + _litellm_model_for(override_key, wiring) if override_key is not None else None + ) + deployment: Final = _DEPLOYMENTS.get(map_key) + models.append( + FrontierModel( + model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", + litellm_model=( + deployment.litellm_model + if deployment is not None and deployment.litellm_model is not None + else _litellm_model_for(map_key, wiring) + ), + wire=wiring.wire, + map_key=map_key, + override_model=( + _provider_model(override_litellm) + if override_litellm is not None + else None + ), + override_map_key=override_key, + base_model=deployment.base_model if deployment is not None else None, + litellm_params=wiring.litellm_params, + ) + ) + return tuple(models) FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() @@ -350,71 +368,6 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ ), }) -CaseName: TypeAlias = Literal[ - "basic", - "cache_read", - "cache_write_5m", - "cache_write_1h", - "reasoning", - "audio", - "tiered", - "service_tier_flex", - "service_tier_priority", - "web_search", - "stream", - "stream_no_usage", - "response_model_override", - "stream_response_model_override", - "tool_call", - "stream_no_usage_tool_call", - "stream_no_usage_image_input", - "stream_no_usage_incomplete", - "stream_unvalidated", - "stream_no_usage_unvalidated", - "prompt_blocked", - "stream_prompt_blocked", -] - - -@dataclass(frozen=True, slots=True) -class Case: - name: CaseName - usage: ScriptedUsage - stream: bool = False - stream_usage: Literal["final_chunk", "absent"] = "final_chunk" - service_tier: Literal["flex", "priority"] | None = None - # For web_search the wire's reported call count is not always what gets - # billed: chat-completions surfaces only expose url_citation annotations, so - # the biller floors to one call; responses/messages/gemini report a real - # count. - billed_web_search_calls: int = 0 - response_model_override: bool = False - exact_spend: bool = True - tool_call: bool = False - image_input: bool = False - terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" - - def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: - return Scenario( - scenario_id=scenario_id, - wire=model.wire, - usage=self.usage, - model=model.provider_model, - output=ScriptedOutput( - text=text, - response_model=model.override_model if self.response_model_override else None, - tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) - if self.tool_call - else None, - terminal=self.terminal, - ), - stream_usage=self.stream_usage, - service_tier=self.service_tier, - ) - - -_BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) - TOOL_CALL_ARGUMENTS: Final = json.dumps({ "city": "Berlin", "days": 7, @@ -422,284 +375,9 @@ TOOL_CALL_ARGUMENTS: Final = json.dumps({ "notes": "filler " * 30, }) -_PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tokens=0) - - -def _web_search_case(model: FrontierModel) -> Case: - counts_exactly: Final = model.wire in ( - "openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate" - ) - return Case( - name="web_search", - usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), - billed_web_search_calls=3 if counts_exactly else 1, - ) - def cases_for(model: FrontierModel) -> tuple[Case, ...]: - rates: Final = model.rates - caps: Final = _WIRE_CAPS[model.wire] - candidates: Final[tuple[Case | None, ...]] = ( - Case(name="basic", usage=_BASIC_USAGE), - ( - Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) - if rates.cache_read_input_token_cost is not None and "cache_read" in caps - else None - ), - ( - Case( - name="cache_write_5m", - usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), - ) - if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps - else None - ), - ( - Case( - name="cache_write_1h", - usage=ScriptedUsage( - fresh_input_tokens=90, - cache_write_5m_tokens=20, - cache_write_1h_tokens=40, - output_tokens=30, - ), - ) - if ( - rates.cache_creation_input_token_cost_above_1hr is not None - and rates.cache_creation_input_token_cost is not None - and "cache_write_1h" in caps - ) - else None - ), - ( - Case( - name="reasoning", - usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), - ) - if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps - else None - ), - ( - Case( - name="audio", - usage=ScriptedUsage( - fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 - ), - ) - if ( - rates.input_cost_per_audio_token is not None - and rates.output_cost_per_audio_token is not None - and "audio" in caps - ) - else None - ), - ( - Case( - name="tiered", - usage=ScriptedUsage( - fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 - ), - ) - if ( - rates.input_cost_per_token_above_200k_tokens is not None - and rates.output_cost_per_token_above_200k_tokens is not None - ) - else None - ), - ( - Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") - if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None - else None - ), - ( - Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") - if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None - else None - ), - _web_search_case(model) if rates.search_context_cost_per_query is not None and "web_search" in caps else None, - Case(name="stream", usage=_BASIC_USAGE, stream=True), - ( - Case( - name="stream_no_usage", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - exact_spend=False, - ) - if "absent_usage" in caps - else None - ), - ( - Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True) - if "response_model" in caps - else None - ), - ( - Case( - name="stream_response_model_override", - usage=_BASIC_USAGE, - stream=True, - response_model_override=True, - ) - if "response_model" in caps - else None - ), - ( - Case(name="tool_call", usage=_BASIC_USAGE, tool_call=True) - if "tool_call" in caps - else None - ), - ( - Case( - name="stream_no_usage_tool_call", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - tool_call=True, - exact_spend=False, - ) - if "absent_usage" in caps and "tool_call" in caps - else None - ), - ( - Case( - name="stream_no_usage_image_input", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - image_input=True, - exact_spend=False, - ) - if "absent_usage" in caps and "image_input" in caps - else None - ), - ( - Case( - name="stream_no_usage_incomplete", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - terminal="incomplete", - exact_spend=False, - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="stream_unvalidated", - usage=_BASIC_USAGE, - stream=True, - terminal="unvalidated", - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="stream_no_usage_unvalidated", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - terminal="unvalidated", - exact_spend=False, - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="prompt_blocked", - usage=_PROMPT_BLOCKED_USAGE, - terminal="prompt_blocked", - response_model_override=True, - ) - if "prompt_blocked" in caps - else None - ), - ( - Case( - name="stream_prompt_blocked", - usage=_PROMPT_BLOCKED_USAGE, - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ) - if "prompt_blocked" in caps - else None - ), - ) - return tuple(case for case in candidates if case is not None) - - -@dataclass(frozen=True, slots=True) -class ExpectedCost: - """The expected bill split the way the spend row's cost_breakdown reports - it: the gross input component (cache reads/writes folded in), the output - component, and the tool-usage component.""" - - input_cost: float - output_cost: float - tool_cost: float - - @property - def total(self) -> float: - return self.input_cost + self.output_cost + self.tool_cost - - -def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: - """Literal arithmetic on the test-map rates over the scripted token counts. - - Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; - output = text*out + reasoning*reasoning + audio_out*audio_out; plus the - billed web-search calls at the medium search-context rate. Above-threshold - swaps every input/output rate to its ``_above_200k_tokens`` variant when - total prompt tokens exceed the threshold; a service tier swaps input/output - to the tier's variants, falling back to the base rate when a variant is - unset -- mirroring _get_token_base_cost in litellm's cost calculator. - """ - rates: Final = model.override_rates if case.response_model_override else model.rates - u: Final = case.usage - prompt_tokens: Final = ( - u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens - + u.cache_write_1h_tokens + u.audio_input_tokens - ) - tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS - in_rate: Final = ( - (rates.input_cost_per_token_above_200k_tokens if tiered else None) - or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.input_cost_per_token - or 0.0 - ) - out_rate: Final = ( - (rates.output_cost_per_token_above_200k_tokens if tiered else None) - or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.output_cost_per_token - or 0.0 - ) - input_cost: Final = ( - u.fresh_input_tokens * in_rate - + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) - + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) - ) - output_cost: Final = ( - u.output_tokens * out_rate - + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) - + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) - ) - search: Final = rates.search_context_cost_per_query - tool_cost: Final = case.billed_web_search_calls * ( - search.search_context_size_medium if search and search.search_context_size_medium else 0.0 - ) - return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) - - -def expected_cost(model: FrontierModel, case: Case) -> float: - return expected_breakdown(model, case).total + return tuple(case for case in CASES if case.applies_to(model)) def recount_cost( @@ -738,31 +416,23 @@ def image_input_data_url() -> str: IMAGE_INPUT_DATA_URL: Final = image_input_data_url() -def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: - """(prompt_tokens, completion_tokens) the spend row should carry, per the - wire's normalization: Anthropic folds cache read/write into prompt_tokens, - everyone else reports the totals the wire emitted.""" - u: Final = case.usage - if model.wire in ("anthropic_messages", "bedrock_converse"): - return ( - u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, - u.output_tokens, - ) - if model.wire in ("gemini_generate", "vertex_generate"): - return ( - u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, - u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, - ) - if model.wire == "openai_responses": - return ( - u.fresh_input_tokens + u.cache_read_tokens, - u.output_tokens + u.reasoning_tokens, - ) - return ( - u.fresh_input_tokens - + u.cache_read_tokens - + u.cache_write_5m_tokens - + u.cache_write_1h_tokens - + u.audio_input_tokens, - u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, - ) +class _ExpectedCell(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + input_cost: float + output_cost: float + prompt_tokens: int + completion_tokens: int + + +_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, _ExpectedCell]) +EXPECTED: Final[Mapping[str, _ExpectedCell]] = MappingProxyType( + _EXPECTED_ADAPTER.validate_python(json.loads(EXPECTED_PATH.read_text())) + if EXPECTED_PATH.exists() + else {} +) + + +def expected_key(model: FrontierModel, case: Case) -> str: + return f"{model.map_key}|{case.name}" diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json new file mode 100644 index 00000000000..7a92fb2476f --- /dev/null +++ b/tests/e2e/cost_calculation/expected.json @@ -0,0 +1,2004 @@ +{ + "anthropic.claude-sonnet-5-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.03128, + "output_cost": 0.0085, + "prompt_tokens": 150, + "spend": 0.03978 + }, + "anthropic.claude-sonnet-5-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "anthropic.claude-sonnet-5-v1:0|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01785, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.028050000000000002 + }, + "anthropic.claude-sonnet-5-v1:0|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.052700000000000004, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.06290000000000001 + }, + "anthropic.claude-sonnet-5-v1:0|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0459, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.056100000000000004 + }, + "anthropic.claude-sonnet-5-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "anthropic.claude-sonnet-5-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.013600000000000001, + "output_cost": 0.0085, + "prompt_tokens": 80, + "spend": 0.0221 + }, + "anthropic.claude-sonnet-5-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "azure/gpt-5.4-mini|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.03424, + "output_cost": 0.02336, + "prompt_tokens": 155, + "spend": 0.0576 + }, + "azure/gpt-5.4-mini|audio": { + "completion_tokens": 45, + "input_cost": 0.04, + "output_cost": 0.0264, + "prompt_tokens": 125, + "spend": 0.0664 + }, + "azure/gpt-5.4-mini|basic": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0168, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0264 + }, + "azure/gpt-5.4-mini|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.049600000000000005, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0592 + }, + "azure/gpt-5.4-mini|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0432, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0528 + }, + "azure/gpt-5.4-mini|reasoning": { + "completion_tokens": 100, + "input_cost": 0.016, + "output_cost": 0.0656, + "prompt_tokens": 100, + "spend": 0.0816 + }, + "azure/gpt-5.4-mini|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0288, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.0448 + }, + "azure/gpt-5.4-mini|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.03264, + "output_cost": 0.01728, + "prompt_tokens": 120, + "spend": 0.049920000000000006 + }, + "azure/gpt-5.4-mini|stream": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0128, + "output_cost": 0.008, + "prompt_tokens": 80, + "spend": 0.0208 + }, + "azure/gpt-5.4-mini|tiered": { + "completion_tokens": 30, + "input_cost": 256.00128, + "output_cost": 0.0432, + "prompt_tokens": 200001, + "spend": 256.04448 + }, + "azure/gpt-5.4-mini|tool_call": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|web_search": { + "completion_tokens": 30, + "input_cost": 0.016, + "output_cost": 0.009600000000000001, + "prompt_tokens": 100, + "spend": 0.0456 + }, + "azure/gpt-5.6|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.0321, + "output_cost": 0.0219, + "prompt_tokens": 155, + "spend": 0.05399999999999999 + }, + "azure/gpt-5.6|audio": { + "completion_tokens": 45, + "input_cost": 0.0375, + "output_cost": 0.02475, + "prompt_tokens": 125, + "spend": 0.06225 + }, + "azure/gpt-5.6|basic": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01575, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.02475 + }, + "azure/gpt-5.6|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0465, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.0555 + }, + "azure/gpt-5.6|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.040499999999999994, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.049499999999999995 + }, + "azure/gpt-5.6|reasoning": { + "completion_tokens": 100, + "input_cost": 0.015, + "output_cost": 0.0615, + "prompt_tokens": 100, + "spend": 0.0765 + }, + "azure/gpt-5.6|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.6|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.027, + "output_cost": 0.015, + "prompt_tokens": 120, + "spend": 0.041999999999999996 + }, + "azure/gpt-5.6|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.030600000000000002, + "output_cost": 0.0162, + "prompt_tokens": 120, + "spend": 0.0468 + }, + "azure/gpt-5.6|stream": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.6|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.011999999999999999, + "output_cost": 0.0075, + "prompt_tokens": 80, + "spend": 0.019499999999999997 + }, + "azure/gpt-5.6|tiered": { + "completion_tokens": 30, + "input_cost": 240.00119999999998, + "output_cost": 0.0405, + "prompt_tokens": 200001, + "spend": 240.0417 + }, + "azure/gpt-5.6|tool_call": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|web_search": { + "completion_tokens": 30, + "input_cost": 0.015, + "output_cost": 0.009, + "prompt_tokens": 100, + "spend": 0.044 + }, + "claude-haiku-4-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.012880000000000003, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 150, + "spend": 0.016380000000000002 + }, + "claude-haiku-4-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.012880000000000003, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 150, + "spend": 0.016380000000000002 + }, + "claude-haiku-4-5|basic": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.007350000000000001, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.011550000000000001 + }, + "claude-haiku-4-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.021700000000000004, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.025900000000000006 + }, + "claude-haiku-4-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0189, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.023100000000000002 + }, + "claude-haiku-4-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-haiku-4-5|stream": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-haiku-4-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.005600000000000001, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 80, + "spend": 0.0091 + }, + "claude-haiku-4-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.007000000000000001, + "output_cost": 0.004200000000000001, + "prompt_tokens": 100, + "spend": 0.0712 + }, + "claude-opus-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.0092, + "output_cost": 0.0025, + "prompt_tokens": 150, + "spend": 0.0117 + }, + "claude-opus-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.0092, + "output_cost": 0.0025, + "prompt_tokens": 150, + "spend": 0.0117 + }, + "claude-opus-5|basic": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.00525, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.00825 + }, + "claude-opus-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0155, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.0185 + }, + "claude-opus-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.013500000000000002, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.0165 + }, + "claude-opus-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-opus-5|stream": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-opus-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.004, + "output_cost": 0.0025, + "prompt_tokens": 80, + "spend": 0.006500000000000001 + }, + "claude-opus-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.005, + "output_cost": 0.003, + "prompt_tokens": 100, + "spend": 0.068 + }, + "claude-sonnet-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.011040000000000001, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 150, + "spend": 0.014040000000000002 + }, + "claude-sonnet-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.011040000000000001, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 150, + "spend": 0.014040000000000002 + }, + "claude-sonnet-5|basic": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.006300000000000001, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0099 + }, + "claude-sonnet-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.018600000000000002, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0222 + }, + "claude-sonnet-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.016200000000000003, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0198 + }, + "claude-sonnet-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-sonnet-5|stream": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-sonnet-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 80, + "spend": 0.007800000000000001 + }, + "claude-sonnet-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.006000000000000001, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 100, + "spend": 0.0696 + }, + "fireworks_ai/deepseek-v4p1-flash|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.011760000000000001, + "output_cost": 0.007000000000000001, + "prompt_tokens": 120, + "spend": 0.018760000000000002 + }, + "fireworks_ai/deepseek-v4p1-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.030500000000000003, + "output_cost": 0.019950000000000002, + "prompt_tokens": 125, + "spend": 0.05045000000000001 + }, + "fireworks_ai/deepseek-v4p1-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.014700000000000001, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.023100000000000002 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0368, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.045200000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0324, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.0408 + }, + "fireworks_ai/deepseek-v4p1-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.014000000000000002, + "output_cost": 0.0469, + "prompt_tokens": 100, + "spend": 0.060899999999999996 + }, + "fireworks_ai/deepseek-v4p1-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/deepseek-v4p1-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/deepseek-v4p1-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.011200000000000002, + "output_cost": 0.007000000000000001, + "prompt_tokens": 80, + "spend": 0.0182 + }, + "fireworks_ai/deepseek-v4p1-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.014000000000000002, + "output_cost": 0.008400000000000001, + "prompt_tokens": 100, + "spend": 0.04240000000000001 + }, + "fireworks_ai/kimi-k3|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.01008, + "output_cost": 0.006000000000000001, + "prompt_tokens": 120, + "spend": 0.01608 + }, + "fireworks_ai/kimi-k3|audio": { + "completion_tokens": 45, + "input_cost": 0.028500000000000004, + "output_cost": 0.01875, + "prompt_tokens": 125, + "spend": 0.04725 + }, + "fireworks_ai/kimi-k3|basic": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.012600000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0198 + }, + "fireworks_ai/kimi-k3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.035, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0422 + }, + "fireworks_ai/kimi-k3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.030600000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0378 + }, + "fireworks_ai/kimi-k3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.012000000000000002, + "output_cost": 0.0457, + "prompt_tokens": 100, + "spend": 0.0577 + }, + "fireworks_ai/kimi-k3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/kimi-k3|stream": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/kimi-k3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.009600000000000001, + "output_cost": 0.006000000000000001, + "prompt_tokens": 80, + "spend": 0.015600000000000003 + }, + "fireworks_ai/kimi-k3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|web_search": { + "completion_tokens": 30, + "input_cost": 0.012000000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 100, + "spend": 0.0392 + }, + "fireworks_ai/qwen3p8-max|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.010920000000000001, + "output_cost": 0.006500000000000001, + "prompt_tokens": 120, + "spend": 0.01742 + }, + "fireworks_ai/qwen3p8-max|audio": { + "completion_tokens": 45, + "input_cost": 0.029500000000000002, + "output_cost": 0.01935, + "prompt_tokens": 125, + "spend": 0.048850000000000005 + }, + "fireworks_ai/qwen3p8-max|basic": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01365, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.021450000000000004 + }, + "fireworks_ai/qwen3p8-max|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0359, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.0437 + }, + "fireworks_ai/qwen3p8-max|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0315, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.0393 + }, + "fireworks_ai/qwen3p8-max|reasoning": { + "completion_tokens": 100, + "input_cost": 0.013000000000000001, + "output_cost": 0.0463, + "prompt_tokens": 100, + "spend": 0.059300000000000005 + }, + "fireworks_ai/qwen3p8-max|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/qwen3p8-max|stream": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/qwen3p8-max|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.010400000000000001, + "output_cost": 0.006500000000000001, + "prompt_tokens": 80, + "spend": 0.016900000000000002 + }, + "fireworks_ai/qwen3p8-max|tool_call": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|web_search": { + "completion_tokens": 30, + "input_cost": 0.013000000000000001, + "output_cost": 0.007800000000000001, + "prompt_tokens": 100, + "spend": 0.0408 + }, + "gemini-3.1-pro-preview|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.023940000000000003, + "output_cost": 0.030660000000000003, + "prompt_tokens": 125, + "spend": 0.05460000000000001 + }, + "gemini-3.1-pro-preview|audio": { + "completion_tokens": 45, + "input_cost": 0.052500000000000005, + "output_cost": 0.03465, + "prompt_tokens": 125, + "spend": 0.08715 + }, + "gemini-3.1-pro-preview|basic": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|cache_read": { + "completion_tokens": 30, + "input_cost": 0.02205, + "output_cost": 0.0126, + "prompt_tokens": 150, + "spend": 0.03465 + }, + "gemini-3.1-pro-preview|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.2, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.2 + }, + "gemini-3.1-pro-preview|reasoning": { + "completion_tokens": 100, + "input_cost": 0.021, + "output_cost": 0.0861, + "prompt_tokens": 100, + "spend": 0.1071 + }, + "gemini-3.1-pro-preview|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.1-pro-preview|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0378, + "output_cost": 0.020999999999999998, + "prompt_tokens": 120, + "spend": 0.0588 + }, + "gemini-3.1-pro-preview|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.04284, + "output_cost": 0.02268, + "prompt_tokens": 120, + "spend": 0.06552 + }, + "gemini-3.1-pro-preview|stream": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.2, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.2 + }, + "gemini-3.1-pro-preview|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.1-pro-preview|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.016800000000000002, + "output_cost": 0.0105, + "prompt_tokens": 80, + "spend": 0.027300000000000005 + }, + "gemini-3.1-pro-preview|tiered": { + "completion_tokens": 30, + "input_cost": 336.00168, + "output_cost": 0.0567, + "prompt_tokens": 200001, + "spend": 336.05838 + }, + "gemini-3.1-pro-preview|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|web_search": { + "completion_tokens": 30, + "input_cost": 0.021, + "output_cost": 0.0126, + "prompt_tokens": 100, + "spend": 0.0936 + }, + "gemini-3.8-flash|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.022799999999999997, + "output_cost": 0.0292, + "prompt_tokens": 125, + "spend": 0.052 + }, + "gemini-3.8-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.05, + "output_cost": 0.033, + "prompt_tokens": 125, + "spend": 0.083 + }, + "gemini-3.8-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.021, + "output_cost": 0.012, + "prompt_tokens": 150, + "spend": 0.033 + }, + "gemini-3.8-flash|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.21000000000000002, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.21000000000000002 + }, + "gemini-3.8-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.02, + "output_cost": 0.082, + "prompt_tokens": 100, + "spend": 0.10200000000000001 + }, + "gemini-3.8-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.8-flash|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.036, + "output_cost": 0.02, + "prompt_tokens": 120, + "spend": 0.055999999999999994 + }, + "gemini-3.8-flash|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.0408, + "output_cost": 0.0216, + "prompt_tokens": 120, + "spend": 0.062400000000000004 + }, + "gemini-3.8-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.21000000000000002, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.21000000000000002 + }, + "gemini-3.8-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.8-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.016, + "output_cost": 0.01, + "prompt_tokens": 80, + "spend": 0.026000000000000002 + }, + "gemini-3.8-flash|tiered": { + "completion_tokens": 30, + "input_cost": 320.0016, + "output_cost": 0.054, + "prompt_tokens": 200001, + "spend": 320.05559999999997 + }, + "gemini-3.8-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.02, + "output_cost": 0.012, + "prompt_tokens": 100, + "spend": 0.092 + }, + "gemini/gemini-3.1-pro-preview|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.010260000000000002, + "output_cost": 0.01314, + "prompt_tokens": 125, + "spend": 0.023400000000000004 + }, + "gemini/gemini-3.1-pro-preview|audio": { + "completion_tokens": 45, + "input_cost": 0.0225, + "output_cost": 0.014849999999999999, + "prompt_tokens": 125, + "spend": 0.037349999999999994 + }, + "gemini/gemini-3.1-pro-preview|basic": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|cache_read": { + "completion_tokens": 30, + "input_cost": 0.009450000000000002, + "output_cost": 0.0054, + "prompt_tokens": 150, + "spend": 0.014850000000000002 + }, + "gemini/gemini-3.1-pro-preview|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.08, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.08 + }, + "gemini/gemini-3.1-pro-preview|reasoning": { + "completion_tokens": 100, + "input_cost": 0.009000000000000001, + "output_cost": 0.0369, + "prompt_tokens": 100, + "spend": 0.0459 + }, + "gemini/gemini-3.1-pro-preview|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.1-pro-preview|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0162, + "output_cost": 0.009000000000000001, + "prompt_tokens": 120, + "spend": 0.0252 + }, + "gemini/gemini-3.1-pro-preview|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.01836, + "output_cost": 0.00972, + "prompt_tokens": 120, + "spend": 0.02808 + }, + "gemini/gemini-3.1-pro-preview|stream": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.08, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.08 + }, + "gemini/gemini-3.1-pro-preview|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.1-pro-preview|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.007200000000000001, + "output_cost": 0.0045000000000000005, + "prompt_tokens": 80, + "spend": 0.011700000000000002 + }, + "gemini/gemini-3.1-pro-preview|tiered": { + "completion_tokens": 30, + "input_cost": 144.00072, + "output_cost": 0.024300000000000002, + "prompt_tokens": 200001, + "spend": 144.02502 + }, + "gemini/gemini-3.1-pro-preview|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|web_search": { + "completion_tokens": 30, + "input_cost": 0.009000000000000001, + "output_cost": 0.0054, + "prompt_tokens": 100, + "spend": 0.0744 + }, + "gemini/gemini-3.8-flash|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.00912, + "output_cost": 0.01168, + "prompt_tokens": 125, + "spend": 0.0208 + }, + "gemini/gemini-3.8-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.02, + "output_cost": 0.0132, + "prompt_tokens": 125, + "spend": 0.0332 + }, + "gemini/gemini-3.8-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0084, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 150, + "spend": 0.0132 + }, + "gemini/gemini-3.8-flash|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.09000000000000001, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.09000000000000001 + }, + "gemini/gemini-3.8-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.008, + "output_cost": 0.0328, + "prompt_tokens": 100, + "spend": 0.0408 + }, + "gemini/gemini-3.8-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.8-flash|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0144, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.0224 + }, + "gemini/gemini-3.8-flash|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.01632, + "output_cost": 0.00864, + "prompt_tokens": 120, + "spend": 0.024960000000000003 + }, + "gemini/gemini-3.8-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.09000000000000001, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.09000000000000001 + }, + "gemini/gemini-3.8-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.8-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0064, + "output_cost": 0.004, + "prompt_tokens": 80, + "spend": 0.0104 + }, + "gemini/gemini-3.8-flash|tiered": { + "completion_tokens": 30, + "input_cost": 128.00064, + "output_cost": 0.0216, + "prompt_tokens": 200001, + "spend": 128.02224 + }, + "gemini/gemini-3.8-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.008, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 100, + "spend": 0.0728 + }, + "gpt-5.3-codex|all_components_responses": { + "completion_tokens": 40, + "input_cost": 0.00252, + "output_cost": 0.0037500000000000007, + "prompt_tokens": 120, + "spend": 0.006270000000000001 + }, + "gpt-5.3-codex|basic": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0031500000000000005, + "output_cost": 0.0018000000000000002, + "prompt_tokens": 150, + "spend": 0.00495 + }, + "gpt-5.3-codex|reasoning": { + "completion_tokens": 100, + "input_cost": 0.0030000000000000005, + "output_cost": 0.0123, + "prompt_tokens": 100, + "spend": 0.015300000000000001 + }, + "gpt-5.3-codex|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.3-codex|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0054, + "output_cost": 0.003, + "prompt_tokens": 120, + "spend": 0.008400000000000001 + }, + "gpt-5.3-codex|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00612, + "output_cost": 0.00324, + "prompt_tokens": 120, + "spend": 0.00936 + }, + "gpt-5.3-codex|stream": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|stream_incomplete": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.3-codex|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0015000000000000002, + "prompt_tokens": 80, + "spend": 0.0039000000000000007 + }, + "gpt-5.3-codex|stream_unvalidated": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|tiered": { + "completion_tokens": 30, + "input_cost": 48.000240000000005, + "output_cost": 0.0081, + "prompt_tokens": 200001, + "spend": 48.008340000000004 + }, + "gpt-5.3-codex|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|web_search": { + "completion_tokens": 30, + "input_cost": 0.0030000000000000005, + "output_cost": 0.0018000000000000002, + "prompt_tokens": 100, + "spend": 0.0648 + }, + "gpt-5.4-mini|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.00856, + "output_cost": 0.00584, + "prompt_tokens": 155, + "spend": 0.0144 + }, + "gpt-5.4-mini|audio": { + "completion_tokens": 45, + "input_cost": 0.01, + "output_cost": 0.0066, + "prompt_tokens": 125, + "spend": 0.0166 + }, + "gpt-5.4-mini|basic": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0042, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0066 + }, + "gpt-5.4-mini|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.012400000000000001, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0148 + }, + "gpt-5.4-mini|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0108, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0132 + }, + "gpt-5.4-mini|reasoning": { + "completion_tokens": 100, + "input_cost": 0.004, + "output_cost": 0.0164, + "prompt_tokens": 100, + "spend": 0.0204 + }, + "gpt-5.4-mini|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.4-mini|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0072, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.0112 + }, + "gpt-5.4-mini|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00816, + "output_cost": 0.00432, + "prompt_tokens": 120, + "spend": 0.012480000000000002 + }, + "gpt-5.4-mini|stream": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.4-mini|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0032, + "output_cost": 0.002, + "prompt_tokens": 80, + "spend": 0.0052 + }, + "gpt-5.4-mini|tiered": { + "completion_tokens": 30, + "input_cost": 64.00032, + "output_cost": 0.0108, + "prompt_tokens": 200001, + "spend": 64.01112 + }, + "gpt-5.4-mini|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|web_search": { + "completion_tokens": 30, + "input_cost": 0.004, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 100, + "spend": 0.0264 + }, + "gpt-5.5-pro|all_components_responses": { + "completion_tokens": 40, + "input_cost": 0.00168, + "output_cost": 0.0025, + "prompt_tokens": 120, + "spend": 0.00418 + }, + "gpt-5.5-pro|basic": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0021, + "output_cost": 0.0012000000000000001, + "prompt_tokens": 150, + "spend": 0.0033 + }, + "gpt-5.5-pro|reasoning": { + "completion_tokens": 100, + "input_cost": 0.002, + "output_cost": 0.0082, + "prompt_tokens": 100, + "spend": 0.0102 + }, + "gpt-5.5-pro|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.5-pro|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0036, + "output_cost": 0.002, + "prompt_tokens": 120, + "spend": 0.0056 + }, + "gpt-5.5-pro|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00408, + "output_cost": 0.00216, + "prompt_tokens": 120, + "spend": 0.006240000000000001 + }, + "gpt-5.5-pro|stream": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|stream_incomplete": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.5-pro|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0016, + "output_cost": 0.001, + "prompt_tokens": 80, + "spend": 0.0026 + }, + "gpt-5.5-pro|stream_unvalidated": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|tiered": { + "completion_tokens": 30, + "input_cost": 32.00016, + "output_cost": 0.0054, + "prompt_tokens": 200001, + "spend": 32.00556 + }, + "gpt-5.5-pro|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|web_search": { + "completion_tokens": 30, + "input_cost": 0.002, + "output_cost": 0.0012000000000000001, + "prompt_tokens": 100, + "spend": 0.06319999999999999 + }, + "gpt-5.6|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.00214, + "output_cost": 0.00146, + "prompt_tokens": 155, + "spend": 0.0036 + }, + "gpt-5.6|audio": { + "completion_tokens": 45, + "input_cost": 0.0025, + "output_cost": 0.00165, + "prompt_tokens": 125, + "spend": 0.00415 + }, + "gpt-5.6|basic": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|cache_read": { + "completion_tokens": 30, + "input_cost": 0.00105, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.00165 + }, + "gpt-5.6|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0031000000000000003, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.0037 + }, + "gpt-5.6|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0027, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.0033 + }, + "gpt-5.6|reasoning": { + "completion_tokens": 100, + "input_cost": 0.001, + "output_cost": 0.0041, + "prompt_tokens": 100, + "spend": 0.0051 + }, + "gpt-5.6|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.6|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0018, + "output_cost": 0.001, + "prompt_tokens": 120, + "spend": 0.0028 + }, + "gpt-5.6|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00204, + "output_cost": 0.00108, + "prompt_tokens": 120, + "spend": 0.0031200000000000004 + }, + "gpt-5.6|stream": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.6|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0008, + "output_cost": 0.0005, + "prompt_tokens": 80, + "spend": 0.0013 + }, + "gpt-5.6|tiered": { + "completion_tokens": 30, + "input_cost": 16.00008, + "output_cost": 0.0027, + "prompt_tokens": 200001, + "spend": 16.00278 + }, + "gpt-5.6|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|web_search": { + "completion_tokens": 30, + "input_cost": 0.001, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 100, + "spend": 0.0216 + }, + "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.020900000000000002, + "output_cost": 0.0095, + "prompt_tokens": 150, + "spend": 0.030400000000000003 + }, + "meta.llama4-maverick-17b-instruct-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "meta.llama4-maverick-17b-instruct-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "meta.llama4-maverick-17b-instruct-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.015200000000000002, + "output_cost": 0.0095, + "prompt_tokens": 80, + "spend": 0.0247 + }, + "meta.llama4-maverick-17b-instruct-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "together_ai/moonshotai/Kimi-K3|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.0214, + "output_cost": 0.0146, + "prompt_tokens": 155, + "spend": 0.036 + }, + "together_ai/moonshotai/Kimi-K3|audio": { + "completion_tokens": 45, + "input_cost": 0.025, + "output_cost": 0.0165, + "prompt_tokens": 125, + "spend": 0.0415 + }, + "together_ai/moonshotai/Kimi-K3|basic": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0105, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.0165 + }, + "together_ai/moonshotai/Kimi-K3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.031, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.037 + }, + "together_ai/moonshotai/Kimi-K3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.027000000000000003, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.033 + }, + "together_ai/moonshotai/Kimi-K3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.01, + "output_cost": 0.041, + "prompt_tokens": 100, + "spend": 0.051000000000000004 + }, + "together_ai/moonshotai/Kimi-K3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/moonshotai/Kimi-K3|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.018000000000000002, + "output_cost": 0.01, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "together_ai/moonshotai/Kimi-K3|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.0108, + "prompt_tokens": 120, + "spend": 0.031200000000000002 + }, + "together_ai/moonshotai/Kimi-K3|stream": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/moonshotai/Kimi-K3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.008, + "output_cost": 0.005, + "prompt_tokens": 80, + "spend": 0.013000000000000001 + }, + "together_ai/moonshotai/Kimi-K3|tiered": { + "completion_tokens": 30, + "input_cost": 160.0008, + "output_cost": 0.027000000000000003, + "prompt_tokens": 200001, + "spend": 160.02779999999998 + }, + "together_ai/moonshotai/Kimi-K3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|web_search": { + "completion_tokens": 30, + "input_cost": 0.01, + "output_cost": 0.006, + "prompt_tokens": 100, + "spend": 0.036000000000000004 + }, + "together_ai/zai-org/GLM-5.3|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.023540000000000002, + "output_cost": 0.01606, + "prompt_tokens": 155, + "spend": 0.0396 + }, + "together_ai/zai-org/GLM-5.3|audio": { + "completion_tokens": 45, + "input_cost": 0.027500000000000004, + "output_cost": 0.01815, + "prompt_tokens": 125, + "spend": 0.04565 + }, + "together_ai/zai-org/GLM-5.3|basic": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.011550000000000001, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.01815 + }, + "together_ai/zai-org/GLM-5.3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.034100000000000005, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.04070000000000001 + }, + "together_ai/zai-org/GLM-5.3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.029699999999999997, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.0363 + }, + "together_ai/zai-org/GLM-5.3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.011000000000000001, + "output_cost": 0.0451, + "prompt_tokens": 100, + "spend": 0.056100000000000004 + }, + "together_ai/zai-org/GLM-5.3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/zai-org/GLM-5.3|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.019799999999999998, + "output_cost": 0.011000000000000001, + "prompt_tokens": 120, + "spend": 0.0308 + }, + "together_ai/zai-org/GLM-5.3|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.022439999999999998, + "output_cost": 0.01188, + "prompt_tokens": 120, + "spend": 0.034319999999999996 + }, + "together_ai/zai-org/GLM-5.3|stream": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/zai-org/GLM-5.3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0088, + "output_cost": 0.0055000000000000005, + "prompt_tokens": 80, + "spend": 0.0143 + }, + "together_ai/zai-org/GLM-5.3|tiered": { + "completion_tokens": 30, + "input_cost": 176.00088, + "output_cost": 0.0297, + "prompt_tokens": 200001, + "spend": 176.03058 + }, + "together_ai/zai-org/GLM-5.3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|web_search": { + "completion_tokens": 30, + "input_cost": 0.011000000000000001, + "output_cost": 0.0066, + "prompt_tokens": 100, + "spend": 0.0376 + }, + "us.anthropic.claude-opus-5-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.033120000000000004, + "output_cost": 0.009000000000000001, + "prompt_tokens": 150, + "spend": 0.042120000000000005 + }, + "us.anthropic.claude-opus-5-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|cache_read": { + "completion_tokens": 30, + "input_cost": 0.018900000000000004, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.029700000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0558, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.0666 + }, + "us.anthropic.claude-opus-5-v1:0|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.048600000000000004, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.05940000000000001 + }, + "us.anthropic.claude-opus-5-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.014400000000000001, + "output_cost": 0.009000000000000001, + "prompt_tokens": 80, + "spend": 0.023400000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + } +} diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py new file mode 100644 index 00000000000..de979f272fe --- /dev/null +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -0,0 +1,189 @@ +"""Golden generator for the cost suite. Run: + + uv run python tests/e2e/cost_calculation/generate_expected.py + +Loads the derived matrix (models x applicable cases), computes the golden for +each exact-spend cell from the rate arithmetic, and writes ``expected.json`` +with sorted keys. Default behaviour adds missing cells and drops stale cells +but never overwrites an existing cell's values (a reviewed golden is +authoritative); ``--rewrite`` recomputes everything. Prints added/removed/kept +counts. +""" + +from __future__ import annotations + +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from cost_matrix import ( # noqa: E402 # path bootstrap before package-local imports + EXPECTED_PATH, + FRONTIER_MODELS, + TIER_THRESHOLD_TOKENS, + Case, + CostMapEntry, + FrontierModel, + cases_for, + expected_key, +) + +# Wires whose response surface reports a real web-search call count; the +# chat-completions wires only expose url_citation annotations, so their billed +# count floors to one. +_EXACT_WEB_SEARCH_WIRES: Final = frozenset( + {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"} +) + + +def billed_web_search_calls(model: FrontierModel, case: Case) -> int: + if case.usage.web_search_calls == 0: + return 0 + return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1 + + +@dataclass(frozen=True, slots=True) +class ExpectedCost: + """The expected bill split the way the spend row's cost_breakdown reports + it: the gross input component (cache reads/writes folded in), the output + component, and the tool-usage component.""" + + input_cost: float + output_cost: float + tool_cost: float + + @property + def total(self) -> float: + return self.input_cost + self.output_cost + self.tool_cost + + +def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: + """Literal arithmetic on the test-map rates over the scripted token counts. + + Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; + output = text*out + reasoning*reasoning + audio_out*audio_out; plus the + billed web-search calls at the medium search-context rate. Above-threshold + swaps every input/output rate to its ``_above_200k_tokens`` variant when + total prompt tokens exceed the threshold; a service tier swaps input/output + to the tier's variants, falling back to the base rate when a variant is + unset -- mirroring _get_token_base_cost in litellm's cost calculator. + """ + rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates + u: Final = case.usage + prompt_tokens: Final = ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + u.audio_input_tokens + ) + tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate: Final = ( + (rates.input_cost_per_token_above_200k_tokens if tiered else None) + or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.input_cost_per_token + or 0.0 + ) + out_rate: Final = ( + (rates.output_cost_per_token_above_200k_tokens if tiered else None) + or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.output_cost_per_token + or 0.0 + ) + # The biller charges cache writes at the input rate when the entry carries + # no cache_creation rate (cost_calculator.py:2452), and at the 5m write + # rate when the 1h variant is unset; cache reads bill only at their own + # rate (zero when the entry lacks one). + write_5m_rate: Final = rates.cache_creation_input_token_cost or in_rate + input_cost: Final = ( + u.fresh_input_tokens * in_rate + + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + + u.cache_write_5m_tokens * write_5m_rate + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or write_5m_rate) + + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + ) + output_cost: Final = ( + u.output_tokens * out_rate + + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + ) + search: Final = rates.search_context_cost_per_query + tool_cost: Final = billed_web_search_calls(model, case) * ( + search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + ) + return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) + + +def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: + """(prompt_tokens, completion_tokens) the spend row should carry, per the + wire's normalization: Anthropic folds cache read/write into prompt_tokens, + everyone else reports the totals the wire emitted.""" + u: Final = case.usage + if model.wire in ("anthropic_messages", "bedrock_converse"): + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, + u.output_tokens, + ) + if model.wire in ("gemini_generate", "vertex_generate"): + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + if model.wire == "openai_responses": + return ( + u.fresh_input_tokens + u.cache_read_tokens, + u.output_tokens + u.reasoning_tokens, + ) + return ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + + +def _proposed() -> dict[str, dict[str, object]]: + return { + expected_key(model, case): ( + lambda breakdown, tokens: { + "spend": breakdown.total, + "input_cost": breakdown.input_cost, + "output_cost": breakdown.output_cost, + "prompt_tokens": tokens[0], + "completion_tokens": tokens[1], + } + )(expected_breakdown(model, case), expected_token_columns(model, case)) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + + +def main() -> None: + rewrite: Final = "--rewrite" in sys.argv[1:] + proposed: Final = _proposed() + existing: Final = ( + json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} + ) + merged: Final = { + key: (proposed[key] if rewrite or key not in existing else existing[key]) + for key in sorted(proposed) + } + added: Final = sum(1 for key in proposed if key not in existing) + removed: Final = sum(1 for key in existing if key not in proposed) + kept: Final = sum(1 for key in proposed if key in existing and not rewrite) + rewritten: Final = sum(1 for key in proposed if key in existing and rewrite) + EXPECTED_PATH.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") + print( + f"expected.json: {added} added, {removed} removed, {kept} kept, " + f"{rewritten} rewritten ({len(merged)} cells)" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py new file mode 100644 index 00000000000..fdbb6ddd293 --- /dev/null +++ b/tests/e2e/cost_calculation/test_matrix_data.py @@ -0,0 +1,64 @@ +"""Freshness checks for the cost suite's data files; markerless, so it runs on +any pytest invocation of the folder without the stack. expected.json is the +oracle: these tests check its key set against the derived matrix, never its +values (the generator proposes, the file decides).""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from cost_matrix import ( + _CASES_FILE, + _COST_MAP, + CASES, + EXPECTED, + FRONTIER_MODELS, + CostMapEntry, + cases_for, + expected_key, +) + + +def test_expected_keys_match_derived_exact_cells() -> None: + derived: Final = { + expected_key(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + golden: Final = set(EXPECTED) + if derived != golden: + missing: Final = sorted(derived - golden) + stale: Final = sorted(golden - derived) + pytest.fail( + "expected.json is out of sync with the derived matrix; run " + "uv run python tests/e2e/cost_calculation/generate_expected.py " + f"(missing: {missing}; stale: {stale})" + ) + + +def test_deployments_reference_existing_map_keys() -> None: + unknown: Final = sorted( + spec.map_key for spec in _CASES_FILE.deployments if spec.map_key not in _COST_MAP + ) + assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" + + +def test_requires_rates_are_cost_map_fields() -> None: + fields: Final = set(CostMapEntry.model_fields) + unknown: Final = sorted( + {field for case in CASES for field in case.requires_rates} - fields + ) + assert not unknown, f"requires_rates names that are not CostMapEntry fields: {unknown}" + + +def test_no_two_entries_share_input_rate() -> None: + rates: Final = [ + entry.input_cost_per_token for entry in _COST_MAP.values() + ] + assert len(rates) == len(set(rates)), ( + "two cost_map entries share input_cost_per_token; the suite relies on " + "distinct rates so a wrong-model bill can never coincidentally match" + ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 0b4f3e1fd37..7cd128ad6fb 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -1,7 +1,7 @@ -"""Token-pricing e2e: every (frontier model, pricing-component case) cell runs a -scripted-usage call through a deployment registered on the cost-map proxy, and -the spend row plus response-cost header must equal literal arithmetic on the -test map's rates. +"""Token-pricing e2e: every (map entry, case) cell derived from cost_map.json x +cases.json runs a scripted-usage call through a deployment registered on the +cost-map proxy, and the spend row plus response-cost header must equal the +reviewed golden in expected.json verbatim -- no rate arithmetic lives here. Nothing here touches a real provider or the bundled cost map: the proxy's upstream is the scripted-provider sidecar and its entire cost map is @@ -15,13 +15,13 @@ from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( + EXPECTED, FRONTIER_MODELS, IMAGE_INPUT_DATA_URL, Case, FrontierModel, cases_for, - expected_cost, - expected_token_columns, + expected_key, recount_cost, ) from e2e_config import unique_marker @@ -110,17 +110,6 @@ class TestTokenPricing: ) assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected: Final = expected_cost(model, case) - if case.exact_spend and not case.stream: - # Streamed responses commit headers before the bill is computed, so - # the x-litellm-response-cost header is asserted only on non-stream - # calls. - assert response.response_cost is not None and cost_rows.approx_equal( - response.response_cost, expected - ), ( - f"x-litellm-response-cost {response.response_cost} != expected {expected}" - ) - row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, @@ -149,16 +138,39 @@ class TestTokenPricing: cost_rows.assert_total_is_sum_of_components(row) return - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( - f"{model.map_key}/{case.name}: spend {row.spend} != expected {expected} " + golden: Final = EXPECTED[expected_key(model, case)] + + if not case.stream: + # Streamed responses commit headers before the bill is computed, so + # the x-litellm-response-cost header is asserted only on non-stream + # calls. + assert response.response_cost is not None and cost_rows.approx_equal( + response.response_cost, golden.spend + ), ( + f"x-litellm-response-cost {response.response_cost} != golden {golden.spend}" + ) + + assert row.spend is not None and cost_rows.approx_equal(row.spend, golden.spend), ( + f"{model.map_key}/{case.name}: spend {row.spend} != golden {golden.spend} " f"(breakdown {row.breakdown.model_dump()})" ) - - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + breakdown: Final = row.breakdown + assert breakdown.input_cost is not None and cost_rows.approx_equal( + breakdown.input_cost, golden.input_cost + ), ( + f"{model.map_key}/{case.name}: gross input_cost {breakdown.input_cost} " + f"!= golden {golden.input_cost}; cached/written tokens billed at the input rate" ) - assert row.completion_tokens == completion_tokens, ( - f"completion_tokens {row.completion_tokens} != {completion_tokens}" + assert breakdown.output_cost is not None and cost_rows.approx_equal( + breakdown.output_cost, golden.output_cost + ), ( + f"{model.map_key}/{case.name}: output_cost {breakdown.output_cost} " + f"!= golden {golden.output_cost}" + ) + assert row.prompt_tokens == golden.prompt_tokens, ( + f"prompt_tokens {row.prompt_tokens} != {golden.prompt_tokens}" + ) + assert row.completion_tokens == golden.completion_tokens, ( + f"completion_tokens {row.completion_tokens} != {golden.completion_tokens}" ) cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py deleted file mode 100644 index a36bb1a8662..00000000000 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ /dev/null @@ -1,368 +0,0 @@ -"""Wire-format e2e: one scripted upstream per provider wire, answering with a -usage payload where every token kind the wire can report is nonzero. The spend -row's gross input cost must equal fresh tokens at the input rate plus each cache -and audio component at its own rate -- proving the wire's usage shape landed the -cached tokens inside the total (OpenAI/Gemini) or as separate fields -(Anthropic), and that the biller subtracted them before billing fresh tokens. - -Also covers the Responses API wire (an openai/gpt-5.5-pro deployment bridged by -the proxy to POST /responses) and a streamed Anthropic-messages case. -""" - -from __future__ import annotations - -import pytest -from collections.abc import Mapping -from types import MappingProxyType -from typing import Final - -from conftest import CostCalcClient, cost_rows, register_scenario_deployment -from cost_matrix import ( - FRONTIER_MODELS, - Case, - FrontierModel, - expected_breakdown, - expected_token_columns, -) -from e2e_config import unique_marker -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions, ChatTool, ChatToolFunction -from scripted_provider import ScriptedUsage - -pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark - -_MODELS: Final[Mapping[str, FrontierModel]] = MappingProxyType( - {model.map_key: model for model in FRONTIER_MODELS} -) - -# One scripted usage per wire, every reportable token kind nonzero. -_WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ - "openai_chat": ( - "gpt-5.6", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "openai_responses": ( - "gpt-5.5-pro", - ScriptedUsage( - fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25, reasoning_tokens=15 - ), - ), - "anthropic_messages": ( - "claude-sonnet-5", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - ), - ), - "gemini_generate": ( - "gemini/gemini-3.8-flash", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "together_chat": ( - "together_ai/moonshotai/Kimi-K3", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "fireworks_chat": ( - "fireworks_ai/kimi-k3", - ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), - ), - "azure_chat": ( - "azure/gpt-5.6", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), - "bedrock_converse": ( - "anthropic.claude-sonnet-5-v1:0", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - ), - ), - "vertex_generate": ( - "gemini-3.8-flash", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - output_tokens=25, - reasoning_tokens=15, - audio_input_tokens=5, - audio_output_tokens=3, - ), - ), -}) - -_SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) - -# Renderer-level shapes the pricing matrix gates per cap, pinned here once per -# wire so the sidecar emits prove they survive the proxy end to end. -_SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( - *( - ( - f"tool_call_{'stream' if stream else 'sync'}", - wire, - Case(name="tool_call", usage=_SHAPE_USAGE, stream=stream, tool_call=True), - ) - for wire in _WIRE_USAGE - for stream in (False, True) - ), - ( - "responses_incomplete", - "openai_responses", - Case(name="stream_no_usage_incomplete", usage=_SHAPE_USAGE, stream=True, terminal="incomplete"), - ), - ( - "responses_unvalidated", - "openai_responses", - Case(name="stream_unvalidated", usage=_SHAPE_USAGE, stream=True, terminal="unvalidated"), - ), - ( - "gemini_prompt_blocked", - "gemini_generate", - Case( - name="prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "gemini_prompt_blocked_stream", - "gemini_generate", - Case( - name="stream_prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "vertex_prompt_blocked", - "vertex_generate", - Case( - name="prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "vertex_prompt_blocked_stream", - "vertex_generate", - Case( - name="stream_prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "azure_served_model_override", - "azure_chat", - Case( - name="response_model_override", - usage=_SHAPE_USAGE, - response_model_override=True, - ), - ), -) - - -def _shape_id(entry: tuple[str, str, Case]) -> str: - return entry[0] - - -class TestWireFormats: - @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_wire_usage_shape_bills_each_component( - self, - client: CostCalcClient, - resources: ResourceManager, - scoped_key: str, - wire: str, - ) -> None: - map_key, usage = _WIRE_USAGE[wire] - model: Final = _MODELS[map_key] - case: Final = Case(name="basic", usage=usage) - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted wire call"),), - ), - ) - assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, f"{wire}: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"{wire}: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - breakdown: Final = row.breakdown - assert breakdown.input_cost is not None and cost_rows.approx_equal( - breakdown.input_cost, expected.input_cost - ), ( - f"{wire}: gross input_cost {breakdown.input_cost} != expected {expected.input_cost}; " - "cached/written tokens billed at the input rate" - ) - assert breakdown.output_cost is not None and cost_rows.approx_equal( - breakdown.output_cost, expected.output_cost - ), f"{wire}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" - - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"{wire}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" - ) - assert row.completion_tokens == completion_tokens, ( - f"{wire}: completion_tokens {row.completion_tokens} != {completion_tokens}" - ) - cost_rows.assert_total_is_sum_of_components(row) - - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_anthropic_streamed_usage_bills_each_component( - self, client: CostCalcClient, resources: ResourceManager, scoped_key: str - ) -> None: - map_key, usage = _WIRE_USAGE["anthropic_messages"] - model: Final = _MODELS[map_key] - case: Final = Case(name="stream", usage=usage, stream=True) - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted anthropic stream"),), - stream=True, - stream_options=ChatStreamOptions(include_usage=True), - ), - stream=True, - ) - assert response.ok, f"anthropic stream: proxy returned {response.status_code}: {response.body[:400]}" - assert response.stream_done, "anthropic stream did not reach its terminal event" - assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, "anthropic stream: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"anthropic stream: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - cost_rows.assert_total_is_sum_of_components(row) - - @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=_shape_id) - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_response_shape_bills_reported_usage( - self, - client: CostCalcClient, - resources: ResourceManager, - scoped_key: str, - shape_wire_case: tuple[str, str, Case], - ) -> None: - shape, wire, case = shape_wire_case - map_key, _usage = _WIRE_USAGE[wire] - model: Final = _MODELS[map_key] - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted {shape}"),), - stream=case.stream, - stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, - tools=( - ( - ChatTool( - function=ChatToolFunction( - name="get_weather", - parameters={"type": "object", "properties": {"city": {"type": "string"}}}, - ) - ), - ) - if case.tool_call - else None - ), - ), - stream=case.stream, - ) - assert response.ok, f"{shape}: proxy returned {response.status_code}: {response.body[:400]}" - if case.stream: - assert response.stream_done, f"{shape}: stream did not reach its terminal event" - assert response.stream_error is None, f"{shape}: stream error: {response.stream_error}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, f"{shape}: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"{shape}: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"{shape}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" - ) - assert row.completion_tokens == completion_tokens, ( - f"{shape}: completion_tokens {row.completion_tokens} != {completion_tokens}" - ) - cost_rows.assert_total_is_sum_of_components(row) From 522a7f569283b9a3bc0fed2a66f1c7545f30b12b Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:51:43 +0000 Subject: [PATCH 060/393] test(e2e): gate all_components cases by rates and tidy cost matrix names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 26 ++++++++ tests/e2e/cost_calculation/cost_matrix.py | 31 +++++---- tests/e2e/cost_calculation/expected.json | 7 --- .../e2e/cost_calculation/generate_expected.py | 63 ++++++++++--------- .../e2e/cost_calculation/test_matrix_data.py | 11 ++-- 5 files changed, 79 insertions(+), 59 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index e898557ea35..3dc4fc4d99c 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -180,11 +180,20 @@ "audio_input_tokens": 5, "audio_output_tokens": 3 }, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "output_cost_per_reasoning_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token" + ], "wires": ["openai_chat", "azure_chat", "together_chat"] }, { "name": "all_components_fireworks", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, + "requires_rates": ["cache_read_input_token_cost"], "wires": ["fireworks_chat"] }, { @@ -196,6 +205,11 @@ "cache_write_1h_tokens": 10, "output_tokens": 25 }, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr" + ], "wires": ["anthropic_messages", "bedrock_converse"] }, { @@ -208,6 +222,11 @@ "output_tokens": 25 }, "stream": true, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr" + ], "wires": ["anthropic_messages"] }, { @@ -220,11 +239,18 @@ "audio_input_tokens": 5, "audio_output_tokens": 3 }, + "requires_rates": [ + "cache_read_input_token_cost", + "output_cost_per_reasoning_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token" + ], "wires": ["gemini_generate", "vertex_generate"] }, { "name": "all_components_responses", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, + "requires_rates": ["cache_read_input_token_cost", "output_cost_per_reasoning_token"], "wires": ["openai_responses"] } ] diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index b03c851d208..a8f60b79ae7 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -27,7 +27,6 @@ from types import MappingProxyType from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter - from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" @@ -69,9 +68,9 @@ class CostMapEntry(BaseModel): web_search_billing_unit: str | None = None -_COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) -_COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( - _COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) +COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( + COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) ) TIER_THRESHOLD_TOKENS: Final = 200_000 @@ -146,10 +145,10 @@ class _CasesFile(BaseModel): cases: tuple[Case, ...] = () -_CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) -CASES: Final[tuple[Case, ...]] = _CASES_FILE.cases +CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) +CASES: Final[tuple[Case, ...]] = CASES_FILE.cases _DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( - {spec.map_key: spec for spec in _CASES_FILE.deployments} + {spec.map_key: spec for spec in CASES_FILE.deployments} ) @@ -221,13 +220,13 @@ class FrontierModel: @property def rates(self) -> CostMapEntry: - return _COST_MAP[self.map_key] + return COST_MAP[self.map_key] @property def override_rates(self) -> CostMapEntry: if self.base_model is not None or self.override_map_key is None: return self.rates - return _COST_MAP[self.override_map_key] + return COST_MAP[self.override_map_key] @property def provider_model(self) -> str: @@ -262,13 +261,13 @@ def _litellm_model_for(map_key: str, wiring: _ProviderWiring) -> str: def _frontier() -> tuple[FrontierModel, ...]: groups: Final[Mapping[tuple[str, str], tuple[str, ...]]] = MappingProxyType( { - pair: tuple(sorted(k for k, e in _COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) - for pair in {(e.litellm_provider, e.mode) for e in _COST_MAP.values()} + pair: tuple(sorted(k for k, e in COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) + for pair in {(e.litellm_provider, e.mode) for e in COST_MAP.values()} } ) models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple - for map_key in sorted(_COST_MAP): - entry: Final = _COST_MAP[map_key] + for map_key in sorted(COST_MAP): + entry: Final = COST_MAP[map_key] pair: Final = (entry.litellm_provider, entry.mode) wiring: Final = _PROVIDER_WIRING.get(pair) if wiring is None: @@ -416,7 +415,7 @@ def image_input_data_url() -> str: IMAGE_INPUT_DATA_URL: Final = image_input_data_url() -class _ExpectedCell(BaseModel): +class ExpectedCell(BaseModel): model_config = ConfigDict(frozen=True) spend: float @@ -426,8 +425,8 @@ class _ExpectedCell(BaseModel): completion_tokens: int -_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, _ExpectedCell]) -EXPECTED: Final[Mapping[str, _ExpectedCell]] = MappingProxyType( +_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, ExpectedCell]) +EXPECTED: Final[Mapping[str, ExpectedCell]] = MappingProxyType( _EXPECTED_ADAPTER.validate_python(json.loads(EXPECTED_PATH.read_text())) if EXPECTED_PATH.exists() else {} diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index 7a92fb2476f..3b18e9ed9f4 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -1686,13 +1686,6 @@ "prompt_tokens": 100, "spend": 0.0216 }, - "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { - "completion_tokens": 25, - "input_cost": 0.020900000000000002, - "output_cost": 0.0095, - "prompt_tokens": 150, - "spend": 0.030400000000000003 - }, "meta.llama4-maverick-17b-instruct-v1:0|basic": { "completion_tokens": 40, "input_cost": 0.0228, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index de979f272fe..c093ecbe0ea 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -14,8 +14,10 @@ from __future__ import annotations import json import sys +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -27,6 +29,7 @@ from cost_matrix import ( # noqa: E402 # path bootstrap before package-local i TIER_THRESHOLD_TOKENS, Case, CostMapEntry, + ExpectedCell, FrontierModel, cases_for, expected_key, @@ -93,16 +96,11 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: or rates.output_cost_per_token or 0.0 ) - # The biller charges cache writes at the input rate when the entry carries - # no cache_creation rate (cost_calculator.py:2452), and at the 5m write - # rate when the 1h variant is unset; cache reads bill only at their own - # rate (zero when the entry lacks one). - write_5m_rate: Final = rates.cache_creation_input_token_cost or in_rate input_cost: Final = ( u.fresh_input_tokens * in_rate + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * write_5m_rate - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or write_5m_rate) + + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) output_cost: Final = ( @@ -147,39 +145,46 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: ) -def _proposed() -> dict[str, dict[str, object]]: - return { - expected_key(model, case): ( - lambda breakdown, tokens: { - "spend": breakdown.total, - "input_cost": breakdown.input_cost, - "output_cost": breakdown.output_cost, - "prompt_tokens": tokens[0], - "completion_tokens": tokens[1], - } - )(expected_breakdown(model, case), expected_token_columns(model, case)) - for model in FRONTIER_MODELS - for case in cases_for(model) - if case.exact_spend - } +def _cell(model: FrontierModel, case: Case) -> ExpectedCell: + breakdown: Final = expected_breakdown(model, case) + prompt_tokens, completion_tokens = expected_token_columns(model, case) + return ExpectedCell( + spend=breakdown.total, + input_cost=breakdown.input_cost, + output_cost=breakdown.output_cost, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + + +def _proposed() -> Mapping[str, ExpectedCell]: + return MappingProxyType( + { + expected_key(model, case): _cell(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + ) def main() -> None: rewrite: Final = "--rewrite" in sys.argv[1:] proposed: Final = _proposed() + proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()} existing: Final = ( json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} ) merged: Final = { - key: (proposed[key] if rewrite or key not in existing else existing[key]) - for key in sorted(proposed) + key: (proposed_values[key] if rewrite or key not in existing else existing[key]) + for key in sorted(proposed_values) } - added: Final = sum(1 for key in proposed if key not in existing) - removed: Final = sum(1 for key in existing if key not in proposed) - kept: Final = sum(1 for key in proposed if key in existing and not rewrite) - rewritten: Final = sum(1 for key in proposed if key in existing and rewrite) + added: Final = sum(1 for key in proposed_values if key not in existing) + removed: Final = sum(1 for key in existing if key not in proposed_values) + kept: Final = sum(1 for key in proposed_values if key in existing and not rewrite) + rewritten: Final = sum(1 for key in proposed_values if key in existing and rewrite) EXPECTED_PATH.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") - print( + print( # noqa: T201 # CLI summary is the tool output f"expected.json: {added} added, {removed} removed, {kept} kept, " f"{rewritten} rewritten ({len(merged)} cells)" ) diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py index fdbb6ddd293..8340257939c 100644 --- a/tests/e2e/cost_calculation/test_matrix_data.py +++ b/tests/e2e/cost_calculation/test_matrix_data.py @@ -8,11 +8,10 @@ from __future__ import annotations from typing import Final import pytest - from cost_matrix import ( - _CASES_FILE, - _COST_MAP, CASES, + CASES_FILE, + COST_MAP, EXPECTED, FRONTIER_MODELS, CostMapEntry, @@ -41,7 +40,7 @@ def test_expected_keys_match_derived_exact_cells() -> None: def test_deployments_reference_existing_map_keys() -> None: unknown: Final = sorted( - spec.map_key for spec in _CASES_FILE.deployments if spec.map_key not in _COST_MAP + spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP ) assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" @@ -55,9 +54,7 @@ def test_requires_rates_are_cost_map_fields() -> None: def test_no_two_entries_share_input_rate() -> None: - rates: Final = [ - entry.input_cost_per_token for entry in _COST_MAP.values() - ] + rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) assert len(rates) == len(set(rates)), ( "two cost_map entries share input_cost_per_token; the suite relies on " "distinct rates so a wrong-model bill can never coincidentally match" From e1c9ae5ae45e3b66041a25a6ae6ca9c7633944b7 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:56:39 +0000 Subject: [PATCH 061/393] test(e2e): drop needless sys.path bootstrap from golden generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/generate_expected.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index c093ecbe0ea..e243e477839 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -16,14 +16,10 @@ import json import sys from collections.abc import Mapping from dataclasses import dataclass -from pathlib import Path from types import MappingProxyType from typing import Final -sys.path.insert(0, str(Path(__file__).resolve().parent)) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) - -from cost_matrix import ( # noqa: E402 # path bootstrap before package-local imports +from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, TIER_THRESHOLD_TOKENS, From 3e11c986766ed7a32ead704e5284fbfeaf889c6b Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:02:41 +0000 Subject: [PATCH 062/393] test(e2e): satisfy pyright in cost matrix derivation and golden generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 14 +++++++------- tests/e2e/cost_calculation/generate_expected.py | 16 +++++++++++++--- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index a8f60b79ae7..35344a099be 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -267,23 +267,23 @@ def _frontier() -> tuple[FrontierModel, ...]: ) models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple for map_key in sorted(COST_MAP): - entry: Final = COST_MAP[map_key] - pair: Final = (entry.litellm_provider, entry.mode) - wiring: Final = _PROVIDER_WIRING.get(pair) + entry = COST_MAP[map_key] + pair = (entry.litellm_provider, entry.mode) + wiring = _PROVIDER_WIRING.get(pair) if wiring is None: raise ValueError( f"cost_map entry {map_key} has no wiring for " f"(litellm_provider={pair[0]}, mode={pair[1]}); add a " f"_ProviderWiring row in cost_matrix.py" ) - siblings: Final = groups[pair] - override_key: Final = ( + siblings = groups[pair] + override_key = ( siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None ) - override_litellm: Final = ( + override_litellm = ( _litellm_model_for(override_key, wiring) if override_key is not None else None ) - deployment: Final = _DEPLOYMENTS.get(map_key) + deployment = _DEPLOYMENTS.get(map_key) models.append( FrontierModel( model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index e243e477839..f5514092c88 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -19,6 +19,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, @@ -168,11 +170,19 @@ def main() -> None: rewrite: Final = "--rewrite" in sys.argv[1:] proposed: Final = _proposed() proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()} - existing: Final = ( - json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} + existing: Final[Mapping[str, ExpectedCell]] = ( + TypeAdapter(dict[str, ExpectedCell]).validate_python( + json.loads(EXPECTED_PATH.read_text()) + ) + if EXPECTED_PATH.exists() + else {} ) merged: Final = { - key: (proposed_values[key] if rewrite or key not in existing else existing[key]) + key: ( + proposed_values[key] + if rewrite or key not in existing + else existing[key].model_dump() + ) for key in sorted(proposed_values) } added: Final = sum(1 for key in proposed_values if key not in existing) From fc0cce553a631e912e7892893bde188e9716b415 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:13:36 +0000 Subject: [PATCH 063/393] test(e2e): derive cache rates from first principles and ungate all_components cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 17 +---------------- tests/e2e/cost_calculation/expected.json | 7 +++++++ tests/e2e/cost_calculation/generate_expected.py | 17 ++++++++++++++--- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index 3dc4fc4d99c..e01ac97e9ff 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -181,9 +181,6 @@ "audio_output_tokens": 3 }, "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", "output_cost_per_reasoning_token", "input_cost_per_audio_token", "output_cost_per_audio_token" @@ -193,7 +190,6 @@ { "name": "all_components_fireworks", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, - "requires_rates": ["cache_read_input_token_cost"], "wires": ["fireworks_chat"] }, { @@ -205,11 +201,6 @@ "cache_write_1h_tokens": 10, "output_tokens": 25 }, - "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr" - ], "wires": ["anthropic_messages", "bedrock_converse"] }, { @@ -222,11 +213,6 @@ "output_tokens": 25 }, "stream": true, - "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr" - ], "wires": ["anthropic_messages"] }, { @@ -240,7 +226,6 @@ "audio_output_tokens": 3 }, "requires_rates": [ - "cache_read_input_token_cost", "output_cost_per_reasoning_token", "input_cost_per_audio_token", "output_cost_per_audio_token" @@ -250,7 +235,7 @@ { "name": "all_components_responses", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, - "requires_rates": ["cache_read_input_token_cost", "output_cost_per_reasoning_token"], + "requires_rates": ["output_cost_per_reasoning_token"], "wires": ["openai_responses"] } ] diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index 3b18e9ed9f4..caea2c3c764 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -1686,6 +1686,13 @@ "prompt_tokens": 100, "spend": 0.0216 }, + "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.0285, + "output_cost": 0.0095, + "prompt_tokens": 150, + "spend": 0.038 + }, "meta.llama4-maverick-17b-instruct-v1:0|basic": { "completion_tokens": 40, "input_cost": 0.0228, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index f5514092c88..a6eabcc7286 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -94,11 +94,22 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: or rates.output_cost_per_token or 0.0 ) + write_rate: Final = ( + rates.cache_creation_input_token_cost + if rates.cache_creation_input_token_cost is not None + else in_rate + ) input_cost: Final = ( u.fresh_input_tokens * in_rate - + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + + u.cache_read_tokens + * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate) + + u.cache_write_5m_tokens * write_rate + + u.cache_write_1h_tokens + * ( + rates.cache_creation_input_token_cost_above_1hr + if rates.cache_creation_input_token_cost_above_1hr is not None + else write_rate + ) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) output_cost: Final = ( From 072b32baf2097e5421672956ce34809106f592aa Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:18:02 +0000 Subject: [PATCH 064/393] test(e2e): derive goldens from first-principles rate selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 10 ++- tests/e2e/cost_calculation/expected.json | 18 ++-- .../e2e/cost_calculation/generate_expected.py | 86 +++++++++---------- 3 files changed, 61 insertions(+), 53 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index e01ac97e9ff..49eebc85231 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -62,7 +62,15 @@ "name": "web_search", "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3}, "requires_rates": ["search_context_cost_per_query"], - "requires_caps": ["web_search"] + "requires_caps": ["web_search"], + "wires": ["openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"] + }, + { + "name": "web_search_single", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 1}, + "requires_rates": ["search_context_cost_per_query"], + "requires_caps": ["web_search"], + "wires": ["openai_chat", "together_chat", "fireworks_chat", "azure_chat"] }, { "name": "stream", diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index caea2c3c764..984b670a82c 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -160,7 +160,7 @@ "prompt_tokens": 120, "spend": 0.032 }, - "azure/gpt-5.4-mini|web_search": { + "azure/gpt-5.4-mini|web_search_single": { "completion_tokens": 30, "input_cost": 0.016, "output_cost": 0.009600000000000001, @@ -272,7 +272,7 @@ "prompt_tokens": 120, "spend": 0.03 }, - "azure/gpt-5.6|web_search": { + "azure/gpt-5.6|web_search_single": { "completion_tokens": 30, "input_cost": 0.015, "output_cost": 0.009, @@ -615,7 +615,7 @@ "prompt_tokens": 120, "spend": 0.028000000000000004 }, - "fireworks_ai/deepseek-v4p1-flash|web_search": { + "fireworks_ai/deepseek-v4p1-flash|web_search_single": { "completion_tokens": 30, "input_cost": 0.014000000000000002, "output_cost": 0.008400000000000001, @@ -706,7 +706,7 @@ "prompt_tokens": 120, "spend": 0.024 }, - "fireworks_ai/kimi-k3|web_search": { + "fireworks_ai/kimi-k3|web_search_single": { "completion_tokens": 30, "input_cost": 0.012000000000000002, "output_cost": 0.007200000000000001, @@ -797,7 +797,7 @@ "prompt_tokens": 120, "spend": 0.026000000000000002 }, - "fireworks_ai/qwen3p8-max|web_search": { + "fireworks_ai/qwen3p8-max|web_search_single": { "completion_tokens": 30, "input_cost": 0.013000000000000001, "output_cost": 0.007800000000000001, @@ -1462,7 +1462,7 @@ "prompt_tokens": 120, "spend": 0.008 }, - "gpt-5.4-mini|web_search": { + "gpt-5.4-mini|web_search_single": { "completion_tokens": 30, "input_cost": 0.004, "output_cost": 0.0024000000000000002, @@ -1679,7 +1679,7 @@ "prompt_tokens": 120, "spend": 0.002 }, - "gpt-5.6|web_search": { + "gpt-5.6|web_search_single": { "completion_tokens": 30, "input_cost": 0.001, "output_cost": 0.0006000000000000001, @@ -1826,7 +1826,7 @@ "prompt_tokens": 120, "spend": 0.02 }, - "together_ai/moonshotai/Kimi-K3|web_search": { + "together_ai/moonshotai/Kimi-K3|web_search_single": { "completion_tokens": 30, "input_cost": 0.01, "output_cost": 0.006, @@ -1938,7 +1938,7 @@ "prompt_tokens": 120, "spend": 0.022 }, - "together_ai/zai-org/GLM-5.3|web_search": { + "together_ai/zai-org/GLM-5.3|web_search_single": { "completion_tokens": 30, "input_cost": 0.011000000000000001, "output_cost": 0.0066, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index a6eabcc7286..64abdb14c99 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -19,8 +19,6 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final -from pydantic import TypeAdapter - from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, @@ -32,19 +30,11 @@ from cost_matrix import ( cases_for, expected_key, ) - -# Wires whose response surface reports a real web-search call count; the -# chat-completions wires only expose url_citation annotations, so their billed -# count floors to one. -_EXACT_WEB_SEARCH_WIRES: Final = frozenset( - {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"} -) +from pydantic import TypeAdapter -def billed_web_search_calls(model: FrontierModel, case: Case) -> int: - if case.usage.web_search_calls == 0: - return 0 - return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1 +def _first_present(*rates: float | None) -> float | None: + return next((rate for rate in rates if rate is not None), None) @dataclass(frozen=True, slots=True) @@ -67,11 +57,14 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; output = text*out + reasoning*reasoning + audio_out*audio_out; plus the - billed web-search calls at the medium search-context rate. Above-threshold - swaps every input/output rate to its ``_above_200k_tokens`` variant when - total prompt tokens exceed the threshold; a service tier swaps input/output - to the tier's variants, falling back to the base rate when a variant is - unset -- mirroring _get_token_base_cost in litellm's cost calculator. + billed web-search calls at the medium search-context rate. Every billed + token is a token the provider charged for: a component whose entry has no + dedicated rate bills at the ordinary input or output rate, and a present + rate (including an explicit 0.0) is authoritative. When the total prompt + tokens exceed the threshold, input/output rates come from the + ``_above_200k_tokens`` variants; a service tier takes its ``_priority`` or + ``_flex`` variant when the entry carries one, and otherwise bills at the + base rate. """ rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates u: Final = case.usage @@ -81,46 +74,53 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: ) tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS in_rate: Final = ( - (rates.input_cost_per_token_above_200k_tokens if tiered else None) - or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.input_cost_per_token + _first_present( + rates.input_cost_per_token_above_200k_tokens if tiered else None, + rates.input_cost_per_token_priority if case.service_tier == "priority" else None, + rates.input_cost_per_token_flex if case.service_tier == "flex" else None, + rates.input_cost_per_token, + ) or 0.0 ) out_rate: Final = ( - (rates.output_cost_per_token_above_200k_tokens if tiered else None) - or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.output_cost_per_token + _first_present( + rates.output_cost_per_token_above_200k_tokens if tiered else None, + rates.output_cost_per_token_priority if case.service_tier == "priority" else None, + rates.output_cost_per_token_flex if case.service_tier == "flex" else None, + rates.output_cost_per_token, + ) or 0.0 ) - write_rate: Final = ( - rates.cache_creation_input_token_cost - if rates.cache_creation_input_token_cost is not None - else in_rate + read_rate: Final = _first_present(rates.cache_read_input_token_cost, in_rate) or 0.0 + write_rate: Final = _first_present(rates.cache_creation_input_token_cost, in_rate) or 0.0 + write_1h_rate: Final = ( + _first_present(rates.cache_creation_input_token_cost_above_1hr, write_rate) or 0.0 ) + audio_in_rate: Final = _first_present(rates.input_cost_per_audio_token, in_rate) or 0.0 + reasoning_rate: Final = _first_present(rates.output_cost_per_reasoning_token, out_rate) or 0.0 + audio_out_rate: Final = _first_present(rates.output_cost_per_audio_token, out_rate) or 0.0 input_cost: Final = ( u.fresh_input_tokens * in_rate - + u.cache_read_tokens - * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate) + + u.cache_read_tokens * read_rate + u.cache_write_5m_tokens * write_rate - + u.cache_write_1h_tokens - * ( - rates.cache_creation_input_token_cost_above_1hr - if rates.cache_creation_input_token_cost_above_1hr is not None - else write_rate - ) - + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + + u.cache_write_1h_tokens * write_1h_rate + + u.audio_input_tokens * audio_in_rate ) output_cost: Final = ( u.output_tokens * out_rate - + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) - + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + + u.reasoning_tokens * reasoning_rate + + u.audio_output_tokens * audio_out_rate ) search: Final = rates.search_context_cost_per_query - tool_cost: Final = billed_web_search_calls(model, case) * ( - search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + medium_rate: Final = ( + search.search_context_size_medium if search is not None else None ) + if u.web_search_calls and medium_rate is None: + raise ValueError( + f"{model.map_key}: case {case.name} bills {u.web_search_calls} web-search " + "calls but the entry has no search_context_cost_per_query medium rate" + ) + tool_cost: Final = u.web_search_calls * (medium_rate if medium_rate is not None else 0.0) return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) From 1de633ac36644c5774cf629a793b6716a98b7580 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:22:01 +0000 Subject: [PATCH 065/393] test(e2e): move matrix data freshness checks to collection time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/cost_calculation/cost_matrix.py | 48 +++++++++++++++ .../e2e/cost_calculation/test_matrix_data.py | 61 ------------------- .../test_token_pricing_e2e.py | 4 ++ 4 files changed, 53 insertions(+), 62 deletions(-) delete mode 100644 tests/e2e/cost_calculation/test_matrix_data.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 707d35b4aa6..f89b3203622 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`; `cost_matrix.matrix_data_errors()` runs at collection time so a stale key set fails the suite's collection loudly), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 35344a099be..68f3186809d 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -435,3 +435,51 @@ EXPECTED: Final[Mapping[str, ExpectedCell]] = MappingProxyType( def expected_key(model: FrontierModel, case: Case) -> str: return f"{model.map_key}|{case.name}" + + +def matrix_data_errors() -> tuple[str, ...]: + """Freshness findings for the data files, as human-readable strings. + + Called at collection time by the e2e suite; also usable from + generate_expected.py's context without importing pytest. + """ + derived: Final = { + expected_key(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + golden: Final = set(EXPECTED) + unknown_deployments: Final = sorted( + spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP + ) + unknown_rates: Final = sorted( + {field for case in CASES for field in case.requires_rates} - set(CostMapEntry.model_fields) + ) + input_rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) + findings: Final = ( + ( + "expected.json is out of sync with the derived matrix; run " + "uv run python tests/e2e/cost_calculation/generate_expected.py " + f"(missing: {sorted(derived - golden)}; stale: {sorted(golden - derived)})" + ) + if derived != golden + else None, + ( + f"deployments entries name map keys absent from cost_map.json: {unknown_deployments}" + if unknown_deployments + else None + ), + ( + f"requires_rates names that are not CostMapEntry fields: {unknown_rates}" + if unknown_rates + else None + ), + ( + "two cost_map entries share input_cost_per_token; the suite relies on " + "distinct rates so a wrong-model bill can never coincidentally match" + if len(input_rates) != len(set(input_rates)) + else None + ), + ) + return tuple(finding for finding in findings if finding is not None) diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py deleted file mode 100644 index 8340257939c..00000000000 --- a/tests/e2e/cost_calculation/test_matrix_data.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Freshness checks for the cost suite's data files; markerless, so it runs on -any pytest invocation of the folder without the stack. expected.json is the -oracle: these tests check its key set against the derived matrix, never its -values (the generator proposes, the file decides).""" - -from __future__ import annotations - -from typing import Final - -import pytest -from cost_matrix import ( - CASES, - CASES_FILE, - COST_MAP, - EXPECTED, - FRONTIER_MODELS, - CostMapEntry, - cases_for, - expected_key, -) - - -def test_expected_keys_match_derived_exact_cells() -> None: - derived: Final = { - expected_key(model, case) - for model in FRONTIER_MODELS - for case in cases_for(model) - if case.exact_spend - } - golden: Final = set(EXPECTED) - if derived != golden: - missing: Final = sorted(derived - golden) - stale: Final = sorted(golden - derived) - pytest.fail( - "expected.json is out of sync with the derived matrix; run " - "uv run python tests/e2e/cost_calculation/generate_expected.py " - f"(missing: {missing}; stale: {stale})" - ) - - -def test_deployments_reference_existing_map_keys() -> None: - unknown: Final = sorted( - spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP - ) - assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" - - -def test_requires_rates_are_cost_map_fields() -> None: - fields: Final = set(CostMapEntry.model_fields) - unknown: Final = sorted( - {field for case in CASES for field in case.requires_rates} - fields - ) - assert not unknown, f"requires_rates names that are not CostMapEntry fields: {unknown}" - - -def test_no_two_entries_share_input_rate() -> None: - rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) - assert len(rates) == len(set(rates)), ( - "two cost_map entries share input_cost_per_token; the suite relies on " - "distinct rates so a wrong-model bill can never coincidentally match" - ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 7cd128ad6fb..346a55aa22d 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -22,6 +22,7 @@ from cost_matrix import ( FrontierModel, cases_for, expected_key, + matrix_data_errors, recount_cost, ) from e2e_config import unique_marker @@ -39,6 +40,9 @@ from models import ( pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark +if _data_errors := matrix_data_errors(): + raise ValueError("\n".join(_data_errors)) + _MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple( (model, case) for model in FRONTIER_MODELS for case in cases_for(model) ) From ef1f306a7dc0777276166859b3fa4d2e6272cefb Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:53:52 +0000 Subject: [PATCH 066/393] test(e2e): emit gemini stream usage only on the final chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/scripted_provider.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 982132ed8df..90d95441e5c 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -750,10 +750,8 @@ def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, objec def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: emit_usage: Final = scenario.stream_usage == "final_chunk" - first: Final = ( - _jobj(*((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata")) - if scenario.stream_usage == "absent" - else _gemini_body(scenario, requested_model) + first: Final = _jobj( + *((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata") ) return _sse( ( From 441021fc96eb24680ec41f78f16000402deff8b9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 04:46:26 +0000 Subject: [PATCH 067/393] fix(responses): announce message item before text events in the chat completions bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 67 ++++++---- .../test_streaming_iterator_transformation.py | 118 ++++++++++++++++++ 2 files changed, 160 insertions(+), 25 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 1b9f39449cf..7af62e9bfef 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -102,6 +102,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.sent_response_created_event: bool = False self.sent_response_in_progress_event: bool = False self.sent_output_item_added_event: bool = False + self.sent_message_item_added_event: bool = False self.sent_content_part_added_event: bool = False self.sent_output_text_done_event: bool = False self.sent_output_content_part_done_event: bool = False @@ -592,6 +593,29 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): event.__dict__["sequence_number"] = self._sequence_number return event + def _queue_message_item_added_events(self) -> None: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{uuid.uuid4()}" + self.sent_message_item_added_event = True + self.sent_content_part_added_event = True + self._sequence_number += 1 + event: Final = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": self._cached_item_id, + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + } + ), + ) + event.__dict__["sequence_number"] = self._sequence_number + self._pending_response_events.append(event) + self._pending_response_events.append(self.create_content_part_added_event()) + def _merge_provider_specific_fields(self, src: dict) -> None: """Merge provider_specific_fields using last-value-wins for lists. @@ -832,6 +856,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def return_default_done_events( self, litellm_complete_object: ModelResponse ) -> BaseLiteLLMOpenAIResponseObject | None: + if self.sent_message_item_added_event is False: + final_content: Final = litellm_complete_object.choices[0].message.content or "" + if not final_content: + self.sent_output_text_done_event = True + self.sent_output_content_part_done_event = True + self.sent_output_item_done_event = True + return None + self._queue_message_item_added_events() + return self._pending_response_events.pop(0) if self.sent_output_text_done_event is False: self.sent_output_text_done_event = True return self.create_output_text_done_event(litellm_complete_object) @@ -898,6 +931,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _ensure_output_item_for_chunk(self, chunk: ModelResponseStream) -> None: # Change: Never return a value, just enqueue output item events if self.sent_output_item_added_event: + if ( + not self.sent_message_item_added_event + and chunk.choices + and self._get_delta_string_from_streaming_choices(chunk.choices) + ): + self._queue_message_item_added_events() return if not chunk.choices: return @@ -936,31 +975,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return # Default: message - self._cached_item_id = self._cached_item_id or f"msg_{uuid.uuid4()}" - event = OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=0, - item=BaseLiteLLMOpenAIResponseObject( - **{ - "id": self._cached_item_id, - "type": "message", - "role": "assistant", - "status": "in_progress", - "content": [], - } - ), - ) - event.__dict__["sequence_number"] = self._sequence_number - self._pending_response_events.append(event) - - # Emit content_part.added immediately after output_item.added for message - # items. The OpenAI Responses spec requires this event before any - # output_text.delta events so downstream parsers can initialize the - # text part structure. - if not self.sent_content_part_added_event: - self.sent_content_part_added_event = True - content_part_event: Final = self.create_content_part_added_event() - self._pending_response_events.append(content_part_event) + self._queue_message_item_added_events() return async def __anext__( @@ -1189,6 +1204,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 2: Handle text deltas delta_content: Final = self._get_delta_string_from_streaming_choices(chunk.choices) if delta_content: + if not self.sent_message_item_added_event: + self._queue_message_item_added_events() self._sequence_number += 1 text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 343fc873fa4..7a7500f666b 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -957,3 +957,121 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: ] assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] + + +def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", reasoning_content=reasoning), + finish_reason=finish_reason, + ) + ], + ) + + +async def _collect_events(iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool) -> list: + if sync_mode: + return list(iterator) + return [event async for event in iterator] + + +def _is_message_item(event) -> bool: + return getattr(getattr(event, "item", None), "type", None) == "message" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_only_stream_emits_no_message_item_events(sync_mode): + """ + A turn that only calls tools must not announce or close a message output item: + Vercel AI SDK clients reject text/item events that reference a message id they + never saw in response.output_item.added. + """ + iterator: Final = _build_iterator([_tool_call_chunk(), _chunk("", finish_reason="tool_calls")]) + + events: Final = await _collect_events(iterator, sync_mode) + + message_item_events = [ + event + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + and _is_message_item(event) + ] + assert message_item_events == [] + assert [ + event + for event in events + if str(getattr(event, "type", "")).startswith("response.output_text") + or getattr(event, "type", None) + in (ResponsesAPIStreamEvents.CONTENT_PART_ADDED, ResponsesAPIStreamEvents.CONTENT_PART_DONE) + ] == [] + assert any(getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED for event in events) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode): + """ + When reasoning is announced first, a later text delta still has to be preceded by + the message output_item.added/content_part.added, and every text-scoped event must + reference that announced message item id. + """ + iterator: Final = _build_iterator( + [ + _reasoning_chunk("let me think"), + _chunk("Hello"), + _chunk("!", finish_reason="stop"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + announced_message_ids: set[str] = set() + content_part_added_seen = False + saw_text_delta = False + for event in events: + event_type = getattr(event, "type", None) + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event): + announced_message_ids.add(event.item.id) + elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED: + content_part_added_seen = True + elif event_type in ( + ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + ResponsesAPIStreamEvents.CONTENT_PART_DONE, + ): + assert event.item_id in announced_message_ids + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA: + assert content_part_added_seen + saw_text_delta = True + elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event): + assert event.item.id in announced_message_ids + assert saw_text_delta + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode): + iterator: Final = _build_iterator([_chunk("Hel"), _chunk("lo", finish_reason="stop")]) + + events: Final = await _collect_events(iterator, sync_mode) + + message_item_adds = [ + event + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event) + ] + assert len(message_item_adds) == 1 + for event in events: + if getattr(event, "type", None) in ( + ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + ): + assert event.item_id == message_item_adds[0].item.id From d4e54a0f345aedcfca85f120441ff2a56f21d5f7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:05:54 +0000 Subject: [PATCH 068/393] fix(responses): keep sync text deltas and give the message item its own output index Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 26 ++++++++++--------- .../test_streaming_iterator_transformation.py | 11 ++++++-- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 7af62e9bfef..6b13c9d4297 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -112,6 +112,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.completed_response = None self.final_text: str = "" self._cached_item_id: str | None = None + self._message_output_index: int = 0 self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -564,7 +565,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=0, + output_index=self._message_output_index, item=BaseLiteLLMOpenAIResponseObject( **{ "id": self._cached_item_id, @@ -586,7 +587,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): event: Final = ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id=self._cached_item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}), ) @@ -598,10 +599,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{uuid.uuid4()}" self.sent_message_item_added_event = True self.sent_content_part_added_event = True + self._message_output_index = 1 if self._cached_reasoning_item_id is not None else 0 self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=0, + output_index=self._message_output_index, item=BaseLiteLLMOpenAIResponseObject( **{ "id": self._cached_item_id, @@ -735,7 +737,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return OutputTextDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, item_id=self._cached_item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, text=getattr(litellm_complete_object.choices[0].message, "content", "") or "", ) @@ -771,7 +773,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return ContentPartDoneEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, item_id=self._cached_item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, part=part, ) @@ -790,7 +792,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - output_index=0, + output_index=self._message_output_index, sequence_number=1, item=BaseLiteLLMOpenAIResponseObject( **{ @@ -951,6 +953,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_reasoning_item_id is None: self._cached_reasoning_item_id = f"rs_{uuid.uuid4()}" self._reasoning_item_id = self._cached_reasoning_item_id + self._next_tool_output_index = max(self._next_tool_output_index, 2) event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -1130,12 +1133,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.collected_chat_completion_chunks.append( self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk)) ) - # Emit any just-queued output_item event - if self._pending_response_events: - return self._pending_response_events.pop(0) response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk) if response_api_chunk: - return response_api_chunk + self._pending_response_events.append(response_api_chunk) + if self._pending_response_events: + return self._pending_response_events.pop(0) # Otherwise, loop to next chunk except StopIteration: return self.common_done_event_logic(sync_mode=True) @@ -1177,7 +1179,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): event = OutputTextAnnotationAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, item_id=item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, annotation_index=idx, annotation=annotation_dict, @@ -1210,7 +1212,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, - output_index=0, + output_index=self._message_output_index, content_index=0, delta=delta_content, ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 7a7500f666b..895f59632c7 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1034,12 +1034,15 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn events: Final = await _collect_events(iterator, sync_mode) announced_message_ids: set[str] = set() + announced_indexes_by_item_type: dict[str, int] = {} content_part_added_seen = False saw_text_delta = False for event in events: event_type = getattr(event, "type", None) - if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event): - announced_message_ids.add(event.item.id) + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED: + announced_indexes_by_item_type[event.item.type] = event.output_index + if _is_message_item(event): + announced_message_ids.add(event.item.id) elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED: content_part_added_seen = True elif event_type in ( @@ -1054,6 +1057,10 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event): assert event.item.id in announced_message_ids assert saw_text_delta + assert "".join( + event.delta for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + ) == "Hello!" + assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] @pytest.mark.parametrize("sync_mode", [True, False]) From 42c4c8163333328fa053f9ee8967ddf2d319c186 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:15:07 +0000 Subject: [PATCH 069/393] fix(responses): allocate the message output index from the shared item allocator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 5 ++-- .../test_streaming_iterator_transformation.py | 28 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 6b13c9d4297..beffd12a349 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -599,7 +599,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{uuid.uuid4()}" self.sent_message_item_added_event = True self.sent_content_part_added_event = True - self._message_output_index = 1 if self._cached_reasoning_item_id is not None else 0 + if self._cached_reasoning_item_id is not None: + self._message_output_index = self._next_tool_output_index + self._next_tool_output_index += 1 self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -953,7 +955,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_reasoning_item_id is None: self._cached_reasoning_item_id = f"rs_{uuid.uuid4()}" self._reasoning_item_id = self._cached_reasoning_item_id - self._next_tool_output_index = max(self._next_tool_output_index, 2) event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 895f59632c7..c727a5be4bc 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1063,6 +1063,34 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"] +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode): + iterator: Final = _build_iterator( + [ + _tool_call_chunk(), + _reasoning_chunk("thinking"), + _chunk("Hello"), + _chunk("!", finish_reason="stop"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + output_item_added_events: Final = [ + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + ] + message_item_adds: Final = [event for event in output_item_added_events if _is_message_item(event)] + function_call_adds: Final = [ + event for event in output_item_added_events if getattr(event.item, "type", None) == "function_call" + ] + + assert len(message_item_adds) == 1 + assert all(message_item_adds[0].output_index != event.output_index for event in function_call_adds) + + output_indexes_by_item_id: Final = {event.item.id: event.output_index for event in output_item_added_events} + assert len(output_indexes_by_item_id) == len(set(output_indexes_by_item_id.values())) + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode): From e9625ad069920a224be093cede8de0bb1f379c0a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:22:09 +0000 Subject: [PATCH 070/393] fix(responses): default the message output index when no reasoning item exists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_completion_transformation/streaming_iterator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index beffd12a349..71502e33d5c 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -602,6 +602,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_reasoning_item_id is not None: self._message_output_index = self._next_tool_output_index self._next_tool_output_index += 1 + else: + self._message_output_index = 0 self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, From 586124094667c30d9a81a810ec4faee1d9c68216 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:43:14 +0000 Subject: [PATCH 071/393] test(responses): type new streaming bridge test parameters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_streaming_iterator_transformation.py | 29 ++++++++----------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index c727a5be4bc..3581771bc63 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -20,7 +20,10 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo LiteLLMCompletionStreamingIterator, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, + ResponsesAPIStreamEvents, +) from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( Delta, @@ -975,24 +978,21 @@ def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelR ) -async def _collect_events(iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool) -> list: +async def _collect_events( + iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool +) -> list[BaseLiteLLMOpenAIResponseObject]: if sync_mode: return list(iterator) return [event async for event in iterator] -def _is_message_item(event) -> bool: +def _is_message_item(event: BaseLiteLLMOpenAIResponseObject) -> bool: return getattr(getattr(event, "item", None), "type", None) == "message" @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_tool_only_stream_emits_no_message_item_events(sync_mode): - """ - A turn that only calls tools must not announce or close a message output item: - Vercel AI SDK clients reject text/item events that reference a message id they - never saw in response.output_item.added. - """ +async def test_tool_only_stream_emits_no_message_item_events(sync_mode: bool): iterator: Final = _build_iterator([_tool_call_chunk(), _chunk("", finish_reason="tool_calls")]) events: Final = await _collect_events(iterator, sync_mode) @@ -1017,12 +1017,7 @@ async def test_tool_only_stream_emits_no_message_item_events(sync_mode): @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode): - """ - When reasoning is announced first, a later text delta still has to be preceded by - the message output_item.added/content_part.added, and every text-scoped event must - reference that announced message item id. - """ +async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode: bool): iterator: Final = _build_iterator( [ _reasoning_chunk("let me think"), @@ -1065,7 +1060,7 @@ async def test_reasoning_then_text_announces_message_item_before_text_events(syn @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode): +async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode: bool): iterator: Final = _build_iterator( [ _tool_call_chunk(), @@ -1093,7 +1088,7 @@ async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index( @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode): +async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode: bool): iterator: Final = _build_iterator([_chunk("Hel"), _chunk("lo", finish_reason="stop")]) events: Final = await _collect_events(iterator, sync_mode) From 8ce2887888648fbea603ae91deffdc6e794926e9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:56:28 +0000 Subject: [PATCH 072/393] fix(responses): close the message content part as output_text on reasoning turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_iterator.py | 30 +++++++------------ 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 71502e33d5c..0cda83d979d 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -22,7 +22,6 @@ from litellm.types.llms.openai import ( ContentPartAddedEvent, ContentPartDoneEvent, ContentPartDonePartOutputText, - ContentPartDonePartReasoningText, FunctionCallArgumentsDeltaEvent, FunctionCallArgumentsDoneEvent, OutputItemAddedEvent, @@ -751,28 +750,19 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id = f"msg_{uuid.uuid4()}" text: Final = getattr(litellm_complete_object.choices[0].message, "content", "") or "" - reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" annotations: Final = getattr(litellm_complete_object.choices[0].message, "annotations", None) - part: PART_UNION_TYPES | None = None - if reasoning_content: - part = ContentPartDonePartReasoningText( - type="reasoning_text", - reasoning=reasoning_content, - ) - - else: - response_annotations: Final = ( - LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( - annotations=annotations - ) - ) - part = ContentPartDonePartOutputText( - type="output_text", - text=text, - annotations=response_annotations, - logprobs=None, + response_annotations: Final = ( + LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( + annotations=annotations ) + ) + part: Final[PART_UNION_TYPES] = ContentPartDonePartOutputText( + type="output_text", + text=text, + annotations=response_annotations, + logprobs=None, + ) return ContentPartDoneEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, From e254377049ca6f7087693cff4b99b291c9ba6010 Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:07:57 +0100 Subject: [PATCH 073/393] fix(azure): drop tool_choice without tools DEVX-829 --- litellm/llms/azure/chat/gpt_transformation.py | 9 +- .../test_azure_chat_gpt_transformation.py | 123 ++++++++++++++++++ ...test_azure_chat_o_series_transformation.py | 3 +- 3 files changed, 133 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 6d17a1359bc..0debbe4f74d 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,10 +280,17 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) + request_params: Final = { + key: value + for key, value in optional_params.items() + if key != "tool_choice" + or optional_params.get("tools") + or optional_params.get("functions") + } return { "model": model, "messages": azure_messages, - **optional_params, + **request_params, **sanitized_tools_update(optional_params), } diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index e8b98c696e1..92a8124d8a9 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -333,3 +333,126 @@ class TestAzureToolSchemaCombinatorFlattening: ) assert "tools" not in request assert request["temperature"] == 0.2 + + +@pytest.mark.parametrize("tool_choice", ["none", "auto"]) +def test_azure_drops_tool_choice_without_tools_or_functions(tool_choice: str) -> None: + optional_params = {"tool_choice": tool_choice, "temperature": 0.2} + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert "tool_choice" not in request + assert request["temperature"] == 0.2 + assert optional_params["tool_choice"] == tool_choice + + +def test_azure_tools_empty_drops_tool_choice() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": [], "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["tools"] == [] + assert "tool_choice" not in request + + +def test_azure_functions_empty_drops_tool_choice() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"functions": [], "tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["functions"] == [] + assert "tool_choice" not in request + + +def test_azure_preserves_tool_choice_with_tools() -> None: + tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["tools"] == tools + assert request["tool_choice"] == "auto" + + +def test_azure_preserves_tool_choice_with_legacy_functions() -> None: + functions = [{"name": "get_weather", "parameters": {}}] + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"functions": functions, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["functions"] == functions + assert request["tool_choice"] == "auto" + + +def test_azure_preserves_function_call_without_tools() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"function_call": "none", "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["function_call"] == "none" + assert "tool_choice" not in request + + +def test_azure_gpt5_drops_tool_choice_without_tools() -> None: + request = AzureOpenAIGPT5Config().transform_request( + model="gpt5_series/gpt-5.6-sol", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["model"] == "gpt-5.6-sol" + assert "tool_choice" not in request + + +@pytest.mark.asyncio +async def test_azure_async_transform_drops_tool_choice_without_tools() -> None: + request = await AzureOpenAIConfig().async_transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert "tool_choice" not in request + + +@pytest.mark.asyncio +async def test_azure_gpt5_async_transform_drops_tool_choice_without_tools() -> None: + request = await AzureOpenAIGPT5Config().async_transform_request( + model="gpt5_series/gpt-5.6-sol", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["model"] == "gpt-5.6-sol" + assert "tool_choice" not in request diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py index 9db9ab971a0..57d60df3a11 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -14,7 +14,7 @@ async def test_azure_chat_o_series_transformation(): provider_config = AzureOpenAIO1Config() model = "o_series/web-interface-o1-mini" messages = [{"role": "user", "content": "Hello, how are you?"}] - optional_params = {} + optional_params = {"tool_choice": "none"} litellm_params = {} headers = {} @@ -23,6 +23,7 @@ async def test_azure_chat_o_series_transformation(): ) print(response) assert response["model"] == "web-interface-o1-mini" + assert "tool_choice" not in response def test_azure_o_series_transform_request_flattens_top_level_anyof(): From 48712f733a641f16b0b4fe60c221fd9f1c7076fa Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:30:06 +0100 Subject: [PATCH 074/393] style(azure): format request parameter filter DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0debbe4f74d..7cb50ee5348 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -283,9 +283,7 @@ class AzureOpenAIConfig(BaseConfig): request_params: Final = { key: value for key, value in optional_params.items() - if key != "tool_choice" - or optional_params.get("tools") - or optional_params.get("functions") + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") } return { "model": model, From bd222bd8d9f6d083b8058c5fef3e998b4f92b3af Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 08:18:36 +0100 Subject: [PATCH 075/393] fix(azure): avoid mutable request mapping DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 7cb50ee5348..424422612db 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,11 +280,13 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) - request_params: Final = { - key: value - for key, value in optional_params.items() - if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") - } + request_params: Final = MappingProxyType( + { + key: value + for key, value in optional_params.items() + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") + } + ) return { "model": model, "messages": azure_messages, From 2fa115db2bac68ac90b0b900dc9eecb728c3bd4d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 17 Sep 2026 12:53:15 -0400 Subject: [PATCH 076/393] feat(router): add maintained Fuse model and harness presets --- .../public_endpoints/public_endpoints.py | 9 + .../complexity_router/README.md | 45 +++++ .../complexity_router/fuse_presets.json | 100 +++++++++++ .../complexity_router/fuse_presets.py | 52 ++++++ .../complexity_router/llm_v2.py | 37 +++- pyproject.toml | 1 + .../public_endpoints/test_public_endpoints.py | 11 ++ .../router_strategy/test_fuse_presets.py | 43 +++++ .../router_strategy/test_llm_v2.py | 110 ++++++++++++ .../test_auto_router_model_naming.py | 49 ++++++ ...ecastClassifierConfig.integration.test.tsx | 163 +++++++++++++++++- .../add_model/ForecastClassifierConfig.tsx | 26 +-- .../add_model/FuseProfilePresets.tsx | 117 +++++++++++++ .../build_complexity_router_config.test.ts | 17 ++ .../forecast_classifier_config.test.ts | 65 ++++++- .../add_model/forecast_classifier_config.ts | 33 +++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 82 ++++++++- 17 files changed, 917 insertions(+), 43 deletions(-) create mode 100644 litellm/router_strategy/complexity_router/fuse_presets.json create mode 100644 litellm/router_strategy/complexity_router/fuse_presets.py create mode 100644 tests/test_litellm/router_strategy/test_fuse_presets.py create mode 100644 ui/litellm-dashboard/src/components/add_model/FuseProfilePresets.tsx diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 94a59828451..e395f56194f 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.utils import get_custom_url from litellm.repositories.table_repositories import ClaudeCodePluginRepository +from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -424,6 +425,14 @@ async def get_complexity_scorer_defaults() -> ComplexityScorerDefaults: ) +@router.get( + "/public/complexity_router/fuse_presets", + response_model=FusePresetCatalog, +) +async def get_public_fuse_presets() -> FusePresetCatalog: + return get_fuse_presets() + + @router.get( "/public/litellm_model_cost_map", tags=["public", "model management"], diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 6505746bca1..d9159cea426 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -179,6 +179,51 @@ Configure capability forecasting through YAML or the model-management API. The dashboard preserves its classifier and calibration on an untouched save; it does not provide a capability-card editor +### Fuse v2 profile presets + +Fuse v2 accepts maintained model and runtime descriptions instead of requiring +custom prose for both solvers and the harness. Select profiles explicitly for +all deployments behind your configured model groups and their actual settings. +Group names do not select profiles automatically + +```yaml +complexity_router_config: + classifier_type: llm_v2 + classifier_llm_config: + model: your-judge-group + tiers: + SIMPLE: your-efficient-group + REASONING: your-capable-group + llm_v2_config: + efficient_profile_preset: claude-sonnet-5-v1 + capable_profile_preset: claude-fable-5-1-v1 + harness_preset: claude-code-v1 + max_quality_gap: 0.05 +``` + +`GET /public/complexity_router/fuse_presets` returns the catalog version, model +profiles, and runtime descriptions, including source URLs. The bundled catalog +is loaded once per process without network requests. Sources are citations only + +Each of `efficient_profile`, `capable_profile`, and `harness` requires either +nonblank custom text or its corresponding preset reference. Custom text wins +when both are supplied, but an unknown or wrong-kind preset is still rejected. +Explicit blank text is invalid even with a valid preset. Custom text remains +limited to 4000 characters + +Saved configurations retain preset references and explicit text separately. +Preset text is resolved when building the classifier prompt, not copied into +stored custom fields. Existing all-custom configurations keep the same prompt. +Versioned preset IDs identify immutable content: revised wording receives a new +ID, and older referenced entries must remain available + +The runtime presets do not imply a repository, runnable tests, network access, +additional tools, or a step, time, or spending budget. mini-SWE-agent describes +an agent interface, not a SWE-bench task. Model descriptions summarize provider +positioning without solve rates or guaranteed rankings. Wording is an evaluation +input, not a calibrated quality claim. Existing Fuse licensing, policy, +calibration, and prompt version are unchanged + ### Heuristic v2 Set `classifier_type: heuristic_v2` to classify with the bundled calibrated diff --git a/litellm/router_strategy/complexity_router/fuse_presets.json b/litellm/router_strategy/complexity_router/fuse_presets.json new file mode 100644 index 00000000000..4006366dc25 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.json @@ -0,0 +1,100 @@ +{ + "version": "2026-09-17-v1", + "models": [ + { + "id": "gpt-6-astra-v1", + "label": "GPT-6 Astra", + "model": "gpt-6-astra", + "text": "OpenAI model for demanding end-to-end work, including reasoning, coding, research, and document tasks", + "sources": ["https://developers.openai.com/api/docs/models/gpt-6-astra"] + }, + { + "id": "gpt-5.6-sol-v1", + "label": "GPT-5.6 Sol", + "model": "gpt-5.6-sol", + "text": "OpenAI model for complex professional work, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-sol"] + }, + { + "id": "gpt-5.6-luna-v1", + "label": "GPT-5.6 Luna", + "model": "gpt-5.6-luna", + "text": "OpenAI model for high-volume workloads, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-luna"] + }, + { + "id": "gpt-5.6-terra-v1", + "label": "GPT-5.6 Terra", + "model": "gpt-5.6-terra", + "text": "OpenAI general-purpose model supporting reasoning, text and image input, and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-terra"] + }, + { + "id": "claude-haiku-4-5-v1", + "label": "Claude Haiku 4.5", + "model": "claude-haiku-4-5", + "text": "Anthropic latency-focused model supporting text and image input, tool use, and extended thinking", + "sources": ["https://platform.claude.com/docs/en/models/haiku-4-5/overview"] + }, + { + "id": "claude-sonnet-5-v1", + "label": "Claude Sonnet 5", + "model": "claude-sonnet-5", + "text": "Anthropic model balancing speed and capability, with adaptive thinking and tool use", + "sources": ["https://platform.claude.com/docs/en/models/sonnet-5/overview"] + }, + { + "id": "claude-opus-5-v1", + "label": "Claude Opus 5", + "model": "claude-opus-5", + "text": "Anthropic model for complex agentic coding and enterprise work, with adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/opus-5/overview"] + }, + { + "id": "claude-fable-5-v1", + "label": "Claude Fable 5", + "model": "claude-fable-5", + "text": "Anthropic model for demanding reasoning and long-running agent tasks, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5/introducing-claude-fable-5-and-claude-mythos-5"] + }, + { + "id": "claude-fable-5-1-v1", + "label": "Claude Fable 5.1", + "model": "claude-fable-5-1", + "text": "Anthropic model for demanding reasoning, long-running agentic coding, and multistep research, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5-1/overview"] + } + ], + "harnesses": [ + { + "id": "unspecified-v1", + "label": "Unspecified runtime", + "text": "Agent runtime is unspecified. Assess the task using the supplied context without assuming repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works", "https://mini-swe-agent.com/latest/faq/"] + }, + { + "id": "claude-code-v1", + "label": "Claude Code", + "text": "Claude Code supplies an agent loop with context management and configured tools. Available actions depend on the session's tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works"] + }, + { + "id": "codex-cli-v1", + "label": "Codex CLI", + "text": "Codex CLI supplies a terminal-based coding agent. File operations, command execution, and integrations depend on the session's tools, permissions, and sandbox. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://learn.chatgpt.com/docs/codex/cli", "https://learn.chatgpt.com/codex/permissions"] + }, + { + "id": "opencode-v1", + "label": "OpenCode", + "text": "OpenCode supplies a configurable agent runtime. Available actions depend on the selected agent, tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://opencode.ai/docs/agents/"] + }, + { + "id": "mini-swe-agent-v1", + "label": "mini-SWE-agent", + "text": "The standard mini-SWE-agent setup uses a bash-only action interface and separate command executions. Available commands and resources depend on its configured environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://mini-swe-agent.com/latest/faq/"] + } + ] +} diff --git a/litellm/router_strategy/complexity_router/fuse_presets.py b/litellm/router_strategy/complexity_router/fuse_presets.py new file mode 100644 index 00000000000..66a96ec5ad5 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.py @@ -0,0 +1,52 @@ +from functools import lru_cache +from importlib.resources import files +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, StringConstraints + +ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] + + +class FuseModelPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + model: str + + +class FuseHarnessPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + + +class FusePresetCatalog(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + version: str + models: tuple[FuseModelPreset, ...] + harnesses: tuple[FuseHarnessPreset, ...] + + +@lru_cache(maxsize=1) +def get_fuse_presets() -> FusePresetCatalog: + return FusePresetCatalog.model_validate_json( + files(__package__).joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + + +def resolve_fuse_profile(text: str | None, preset_id: str | None, kind: Literal["model", "harness"]) -> str | None: + if preset_id is None: + return text + catalog: Final = get_fuse_presets() + presets: Final = catalog.models if kind == "model" else catalog.harnesses + preset: Final = next((entry for entry in presets if entry.id == preset_id), None) + if preset is None: + return None + return text if text is not None else preset.text diff --git a/litellm/router_strategy/complexity_router/llm_v2.py b/litellm/router_strategy/complexity_router/llm_v2.py index 2f545a65aaa..18351237e65 100644 --- a/litellm/router_strategy/complexity_router/llm_v2.py +++ b/litellm/router_strategy/complexity_router/llm_v2.py @@ -7,15 +7,15 @@ from dataclasses import dataclass from sys import float_info from typing import Annotated, Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter, model_validator from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # legacy output validated below ) +from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)] -ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] class _SolverProfile(TypedDict): @@ -139,20 +139,41 @@ class LLMV2Config(BaseModel): efficient_tier: str = "SIMPLE" capable_tier: str = "REASONING" - efficient_profile: ProfileText - capable_profile: ProfileText - harness: ProfileText + efficient_profile: ProfileText | None = None + capable_profile: ProfileText | None = None + harness: ProfileText | None = None + efficient_profile_preset: str | None = None + capable_profile_preset: str | None = None + harness_preset: str | None = None max_quality_gap: float = Field(ge=0.0, le=1.0, description="Maximum estimated success loss allowed for efficient.") max_output_tokens: int = Field(default=1024, ge=1) response_format: Literal["json_schema", "json_object"] = "json_schema" calibration: LLMV2Calibration | None = None + @model_validator(mode="after") + def validate_profiles(self) -> LLMV2Config: + self._profile_texts() + return self + + def _profile_texts(self) -> tuple[str, str, str]: + efficient: Final = resolve_fuse_profile(self.efficient_profile, self.efficient_profile_preset, "model") + capable: Final = resolve_fuse_profile(self.capable_profile, self.capable_profile_preset, "model") + harness: Final = resolve_fuse_profile(self.harness, self.harness_preset, "harness") + if efficient is None: + raise ValueError("efficient_profile requires text or a known efficient_profile_preset") + if capable is None: + raise ValueError("capable_profile requires text or a known capable_profile_preset") + if harness is None: + raise ValueError("harness requires text or a known harness_preset") + return efficient, capable, harness + def system_prompt(self, efficient_model: str, capable_model: str) -> str: + efficient, capable, harness = self._profile_texts() profiles: Final[_SolverProfiles] = { "prompt_version": LLM_V2_PROMPT_VERSION, - "harness": self.harness, - "efficient": {"model": efficient_model, "profile": self.efficient_profile}, - "capable": {"model": capable_model, "profile": self.capable_profile}, + "harness": harness, + "efficient": {"model": efficient_model, "profile": efficient}, + "capable": {"model": capable_model, "profile": capable}, } schema: Final = ( "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) diff --git a/pyproject.toml b/pyproject.toml index 93ff55c4069..2d6d133cd0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -290,6 +290,7 @@ editable-profile = "dev" include = [ "litellm/proxy/_experimental/out/**", "litellm/router_strategy/complexity_router/artifacts/*.json", + "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/client/cli/commands/codex_base_instructions.md", ] exclude = [ diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 0d82ed778f5..b830eb588a5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -11,12 +11,23 @@ from fastapi.testclient import TestClient from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.public_endpoints import router +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) from litellm.types.utils import LlmProviders +def test_fuse_presets_route_serves_the_shared_catalog_without_authentication() -> None: + app: Final = FastAPI() + app.include_router(router) + client: Final = TestClient(app) + response: Final = client.get("/public/complexity_router/fuse_presets") + assert response.status_code == 200 + assert response.json() == get_fuse_presets().model_dump(mode="json") + assert client.get("/public/complexity_router/fuse_presets").json() == response.json() + + def test_get_supported_providers_returns_enum_values(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/test_litellm/router_strategy/test_fuse_presets.py b/tests/test_litellm/router_strategy/test_fuse_presets.py new file mode 100644 index 00000000000..0b8d936383b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_fuse_presets.py @@ -0,0 +1,43 @@ +import json +from importlib.resources import files +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets, resolve_fuse_profile + + +def test_catalog_is_loaded_once_and_preserves_bundled_content() -> None: + get_fuse_presets.cache_clear() + first: Final = get_fuse_presets() + second: Final = get_fuse_presets() + assert first is second + bundled: Final = json.loads( + files("litellm.router_strategy.complexity_router").joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + assert first.model_dump(mode="json") == bundled + entries: Final = (*first.models, *first.harnesses) + assert len({entry.id for entry in entries}) == len(entries) + assert len(first.models) == 9 + assert len(first.harnesses) == 5 + assert all(entry.sources and all(source.startswith("https://") for source in entry.sources) for entry in entries) + + +def test_every_catalog_entry_resolves_without_changing_custom_ownership() -> None: + catalog: Final = get_fuse_presets() + for entry in catalog.models: + assert resolve_fuse_profile(None, entry.id, "model") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "model") == "Custom text" + for entry in catalog.harnesses: + assert resolve_fuse_profile(None, entry.id, "harness") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "harness") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "model") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "harness") == "Custom text" + + +def test_cached_catalog_and_records_cannot_be_modified() -> None: + catalog: Final = get_fuse_presets() + for record, field in ((catalog, "version"), (catalog.models[0], "text"), (catalog.harnesses[0], "text")): + with pytest.raises(ValidationError, match="frozen"): + setattr(record, field, "Changed") diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index 5447c8b43ce..27d31cbe640 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -11,8 +11,10 @@ from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.router_strategy.complexity_router.llm_v2 import ( LLM_V2_PROMPT_VERSION, + LLM_V2_SYSTEM_PROMPT, LLMV2Calibration, LLMV2Config, LLMV2ProbabilityCalibration, @@ -174,6 +176,114 @@ def test_invalid_forecast_settings_are_rejected(overrides: dict[str, object]) -> LLMV2Config.model_validate({**base.model_dump(), **overrides}) +def _preset_config(**overrides: object) -> LLMV2Config: + catalog: Final = get_fuse_presets() + return LLMV2Config.model_validate( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[-1].id, + "max_quality_gap": 0.05, + **overrides, + } + ) + + +def test_preset_roundtrip_keeps_references_without_materializing_text() -> None: + config: Final = _preset_config() + serialized: Final = config.model_dump(exclude_none=True) + assert serialized["efficient_profile_preset"] == config.efficient_profile_preset + assert serialized["capable_profile_preset"] == config.capable_profile_preset + assert serialized["harness_preset"] == config.harness_preset + assert not {"efficient_profile", "capable_profile", "harness"}.intersection(serialized) + assert LLMV2Config.model_validate(config.model_dump()) == config + assert LLMV2Config.model_validate_json(config.model_dump_json()) == config + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_explicit_override_wins_and_survives_roundtrip(field: str) -> None: + config: Final = _preset_config(**{field: " Operator description "}) + roundtrip: Final = LLMV2Config.model_validate_json(config.model_dump_json()) + assert roundtrip.model_dump()[field] == "Operator description" + assert roundtrip.efficient_profile_preset == config.efficient_profile_preset + assert roundtrip.capable_profile_preset == config.capable_profile_preset + assert roundtrip.harness_preset == config.harness_preset + payload: Final = json.loads( + roundtrip.system_prompt("opaque-efficient", "opaque-capable").split("Configured solver profiles:\n")[1] + ) + if field == "harness": + assert payload["harness"] == "Operator description" + else: + assert payload[field.removesuffix("_profile")]["profile"] == "Operator description" + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("invalid", ("", " \n\t", "x" * 4001)) +def test_preset_does_not_bypass_supplied_text_bounds(field: str, invalid: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{field: invalid}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("override", (None, "Custom override")) +@pytest.mark.parametrize("invalid_id", ("missing-v1", "")) +def test_preset_unknown_reference_rejects_even_when_overridden( + field: str, override: str | None, invalid_id: str +) -> None: + with pytest.raises(ValidationError, match=f"{field}.*preset"): + _preset_config(**{field: override, f"{field}_preset": invalid_id}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_missing_text_and_reference_rejects(field: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": None}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_reference_rejects_the_wrong_catalog_kind(field: str) -> None: + catalog: Final = get_fuse_presets() + wrong_id: Final = catalog.models[0].id if field == "harness" else catalog.harnesses[0].id + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": wrong_id}) + + +@pytest.mark.parametrize("mode", ("json_schema", "json_object")) +def test_custom_profile_prompt_bytes_are_unchanged(mode: str) -> None: + base: Final = _config().llm_v2_config + assert base is not None + config: Final = LLMV2Config.model_validate({**base.model_dump(), "response_format": mode}) + old_payload: Final = { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": config.harness, + "efficient": {"model": "opaque-efficient", "profile": config.efficient_profile}, + "capable": {"model": "opaque-capable", "profile": config.capable_profile}, + } + schema: Final = ( + "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) if mode == "json_object" else "" + ) + assert config.system_prompt("opaque-efficient", "opaque-capable") == ( + LLM_V2_SYSTEM_PROMPT + "\n\nConfigured solver profiles:\n" + json.dumps(old_payload) + schema + ) + + +@pytest.mark.asyncio +async def test_preset_router_passes_catalog_text_and_opaque_group_names_to_judge() -> None: + catalog: Final = get_fuse_presets() + config: Final = _config(llm_v2_config=_preset_config().model_dump()) + router, client = _router(_verdict().model_dump_json(), config) + outcome: Final = await router.aclassify("Complete the supplied task") + assert outcome.tier == ComplexityTier.SIMPLE + prompt: Final = client.acompletion.call_args.kwargs["messages"][0]["content"] + payload: Final = json.loads(prompt.split("Configured solver profiles:\n")[1]) + assert payload == { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": catalog.harnesses[-1].text, + "efficient": {"model": "efficient", "profile": catalog.models[0].text}, + "capable": {"model": "capable", "profile": catalog.models[-1].text}, + } + + @pytest.mark.asyncio async def test_one_judge_fuses_whole_task_and_keeps_caller_text_out_of_system_prompt() -> None: router, client = _router(_verdict().model_dump_json()) diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 3dcb8d5af94..7d59a0590f2 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,7 +1,10 @@ from collections.abc import Mapping +from typing import Final import pytest +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets + from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, @@ -171,6 +174,52 @@ def test_validate_accepts_loadable_complexity_config(complexity_router_config): assert validate_complexity_router_config_write(complexity_router_config=complexity_router_config) is None +def _fuse_write_config(profiles: Mapping[str, object]) -> Mapping[str, object]: + return { + "classifier_type": "llm_v2", + "classifier_llm_config": {"model": "judge"}, + "tiers": {"SIMPLE": ["opaque-efficient"], "REASONING": ["opaque-capable"]}, + "llm_v2_config": {"max_quality_gap": 0.05, **profiles}, + } + + +def test_fuse_write_accepts_presets_and_custom_text_with_the_same_entitlement() -> None: + catalog: Final = get_fuse_presets() + presets: Final = _fuse_write_config( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[0].id, + } + ) + custom: Final = _fuse_write_config( + { + "efficient_profile": catalog.models[0].text, + "capable_profile": catalog.models[-1].text, + "harness": catalog.harnesses[0].text, + } + ) + assert validate_complexity_router_config_write(presets) is None + assert validate_complexity_router_config_write(custom) is None + assert claimed_capability(presets) is claimed_capability(custom) + assert claimed_capability(presets) is not None + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_fuse_write_rejects_unknown_preset_even_with_custom_text(field: str) -> None: + config: Final = _fuse_write_config( + { + "efficient_profile": "Custom efficient solver", + "capable_profile": "Custom capable solver", + "harness": "Custom runtime", + f"{field}_preset": "unknown-v1", + } + ) + violation: Final = validate_complexity_router_config_write(config) + assert violation is not None + assert f"{field}_preset" in violation + + def test_naming_check_ignores_the_config_entirely(): """The naming contract and the config's contents are separate questions with separate owners; a write may carry a config without naming a model, so neither can stand in for the other.""" diff --git a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx index 4a574ac736d..f9b0edf9508 100644 --- a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; -import { describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import userEvent from "@testing-library/user-event"; -import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs"; import ForecastClassifierConfig from "./ForecastClassifierConfig"; @@ -37,6 +37,53 @@ const fuseInitial: ComplexityRouterConfigValue = { }, }; const options = ["judge", "efficient", "capable"].map((model) => ({ value: model, label: model })); +const catalog = { + version: "catalog-v1", + models: [ + { + id: "efficient-v1", + label: "Efficient preset", + text: "Maintained efficient profile", + sources: ["https://example.com/efficient"], + model: "efficient-model", + }, + { + id: "capable-v1", + label: "Capable preset", + text: "Maintained capable profile", + sources: ["https://example.com/capable"], + model: "capable-model", + }, + ], + harnesses: [ + { + id: "runtime-v1", + label: "Runtime preset", + text: "Maintained runtime profile", + sources: ["https://example.com/runtime"], + }, + ], +}; +const presetConfig = { + efficient_profile_preset: catalog.models[0].id, + capable_profile_preset: catalog.models[1].id, + harness_preset: catalog.harnesses[0].id, + max_quality_gap: 0.05, +}; +const presetInitial = { ...fuseInitial, llm_v2_config: presetConfig }; + +beforeEach(() => { + testQueryClient.clear(); + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation(async () => Response.json(catalog)), + ); +}); + +afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); +}); function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) { const [value, setValue] = useState(initialValue); @@ -72,6 +119,118 @@ function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfi } describe("forecast classifier form", () => { + it("selects all three maintained presets, previews provenance, and saves only references", async () => { + const user = userEvent.setup(); + renderWithProviders(
); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(await screen.findByRole("option", { name: /^Efficient preset/ })); + await user.click(screen.getByRole("combobox", { name: "Capable solver profile preset" })); + await user.click(screen.getByRole("option", { name: /^Capable preset/ })); + await user.click(screen.getByRole("combobox", { name: "Harness and budget preset" })); + await user.click(screen.getByRole("option", { name: /^Runtime preset/ })); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(catalog.models[0].text); + expect(screen.getByLabelText("Efficient solver profile")).toHaveAttribute("readonly"); + expect(screen.getByLabelText("Capable solver profile")).toHaveValue(catalog.models[1].text); + expect(screen.getByLabelText("Harness and budget")).toHaveValue(catalog.harnesses[0].text); + expect(screen.getAllByText(`Catalog version: ${catalog.version}`)).toHaveLength(3); + expect(screen.getByText(`Model: ${catalog.models[0].model}`)).toBeInTheDocument(); + expect(screen.getAllByRole("link", { name: "Source 1" }).map((link) => link.getAttribute("href"))).toEqual([ + catalog.models[0].sources[0], + catalog.models[1].sources[0], + catalog.harnesses[0].sources[0], + ]); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + presetConfig, + ); + expect(fetch).toHaveBeenCalledTimes(1); + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ url: expect.stringMatching(/\/public\/complexity_router\/fuse_presets$/) }), + ); + }); + + it.each([undefined, null, "Explicit override"])( + "copies effective text to Custom and clears only that reference, override=%s", + async (override) => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + const effectiveText = override ?? catalog.models[0].text; + await waitFor(() => expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText)); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(screen.getByRole("option", { name: "Custom", exact: true })); + expect(screen.getByLabelText("Efficient solver profile")).not.toHaveAttribute("readonly"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Custom budget" } }); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + const { efficient_profile_preset: _preset, ...rest } = presetConfig; + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ + ...rest, + efficient_profile: "Custom budget", + }); + }, + ); + + it.each([ + { ...fuseInitial.llm_v2_config!, efficient_profile: catalog.models[0].text }, + { + ...presetConfig, + efficient_profile: "Explicit override", + capable_profile: "Capable override", + harness: "Harness override", + }, + ])("keeps existing custom ownership and references on an unchanged save: %j", async (settings) => { + renderWithProviders(); + await waitFor(() => expect(screen.queryByText(/Loading profile presets/)).not.toBeInTheDocument()); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("Custom"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(settings.efficient_profile); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + + it.each([true, false])( + "keeps edits and stored IDs while the pending catalog settles, success=%s", + async (success) => { + let resolveCatalog: (response: Response) => void = () => {}; + vi.mocked(fetch).mockReturnValue( + new Promise((resolve) => { + resolveCatalog = resolve; + }), + ); + const settings = { ...presetConfig, efficient_profile: "Original override" }; + renderWithProviders(); + expect(screen.getByText(/Loading profile presets/)).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Typed while loading" } }); + await act(async () => + resolveCatalog(success ? Response.json(catalog) : Response.json({ error: "unavailable" }, { status: 503 })), + ); + if (success) await screen.findAllByText(`Catalog version: ${catalog.version}`); + else expect(await screen.findByText(/Profile presets could not be loaded/)).toBeInTheDocument(); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue("Typed while loading"); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ ...settings, efficient_profile: "Typed while loading" }); + }, + ); + + it("keeps unknown saved IDs visible with unavailable previews rather than replacing them", async () => { + const settings = { ...presetConfig, efficient_profile_preset: "unavailable-v8" }; + renderWithProviders(); + await screen.findAllByText(`Catalog version: ${catalog.version}`); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("unavailable-v8"); + expect(screen.getByText("Preset preview unavailable. The saved reference is preserved")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + it("switches a populated standard router to Capability without saving hidden pools or their overrides", () => { renderWithProviders( ) : ( <> - {(["efficient_profile", "capable_profile", "harness"] as const).map((field) => { - const label = { - efficient_profile: "Efficient solver profile", - capable_profile: "Capable solver profile", - harness: "Harness and budget", - }[field]; - return ( -
- -