diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index e165f96b66..dab5fb1bfd 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -80,10 +80,25 @@ class _ProxyDBLogger(CustomLogger): if "litellm_params" not in request_data: request_data["litellm_params"] = {} + + existing_litellm_params = request_data.get("litellm_params", {}) + existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} + + # Preserve tags from existing metadata + if existing_litellm_metadata.get("tags"): + existing_metadata["tags"] = existing_litellm_metadata.get("tags") + request_data["litellm_params"]["proxy_server_request"] = ( - request_data.get("proxy_server_request") or {} + request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) request_data["litellm_params"]["metadata"] = existing_metadata + + # Preserve model name and custom_llm_provider + if "model" not in request_data: + request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "") + if "custom_llm_provider" not in request_data: + request_data["custom_llm_provider"] = existing_litellm_params.get("custom_llm_provider") or request_data.get("custom_llm_provider", "") + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, response_cost=0.0, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index d6ab121096..6745c559cd 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -91,6 +91,21 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): and "/v1/images/edits" in parsed_url.path ) + @staticmethod + def is_openai_responses_route(url_route: str) -> bool: + """Check if the URL route is an OpenAI responses API endpoint.""" + if not url_route: + return False + parsed_url = urlparse(url_route) + return bool( + parsed_url.hostname + and ( + "api.openai.com" in parsed_url.hostname + or "openai.azure.com" in parsed_url.hostname + ) + and ("/v1/responses" in parsed_url.path or "/responses" in parsed_url.path) + ) + def _get_user_from_metadata( self, passthrough_logging_payload: PassthroughStandardLoggingPayload, @@ -187,7 +202,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): **kwargs, ) -> PassThroughEndpointLoggingTypedDict: """ - Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, and image editing. + Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, image editing, and responses API. """ # Check if this is a supported endpoint for cost tracking is_chat_completions = ( @@ -199,8 +214,11 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): is_image_editing = ( OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) ) + is_responses = ( + OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) + ) - if not (is_chat_completions or is_image_generation or is_image_editing): + if not (is_chat_completions or is_image_generation or is_image_editing or is_responses): # For unsupported endpoints, return None to let the system fall back to generic behavior return { "result": None, @@ -232,9 +250,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): litellm_model_response: Optional[Union[ModelResponse, TextCompletionResponse, ImageResponse]] = None handler_instance = OpenAIPassthroughLoggingHandler() + custom_llm_provider = kwargs.get("custom_llm_provider", "openai") + if is_chat_completions: # Handle chat completions with existing logic provider_config = handler_instance.get_provider_config(model=model) + # Preserve existing litellm_params to maintain metadata tags + existing_litellm_params = kwargs.get("litellm_params", {}) or {} litellm_model_response = provider_config.transform_response( raw_response=httpx_response, model_response=litellm.ModelResponse(), @@ -247,14 +269,14 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): encoding=litellm.encoding, json_mode=request_body.get("response_format", {}).get("type") == "json_object", - litellm_params={}, + litellm_params=existing_litellm_params, ) # Calculate cost using LiteLLM's cost calculator response_cost = litellm.completion_cost( completion_response=litellm_model_response, model=model, - custom_llm_provider="openai", + custom_llm_provider=custom_llm_provider, ) elif is_image_generation: # Handle image generation cost calculation @@ -306,11 +328,36 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not hasattr(litellm_model_response, "_hidden_params"): litellm_model_response._hidden_params = {} litellm_model_response._hidden_params["response_cost"] = response_cost + elif is_responses: + # Handle responses API cost calculation + provider_config = handler_instance.get_provider_config(model=model) + existing_litellm_params = kwargs.get("litellm_params", {}) or {} + litellm_model_response = provider_config.transform_response( + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + model=model, + messages=request_body.get("messages", []), + logging_obj=logging_obj, + optional_params=request_body.get("optional_params", {}), + api_key="", + request_data=request_body, + encoding=litellm.encoding, + json_mode=False, + litellm_params=existing_litellm_params, + ) + + # Calculate cost using LiteLLM's cost calculator with responses call type + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider=custom_llm_provider, + call_type="responses", + ) # Update kwargs with cost information kwargs["response_cost"] = response_cost kwargs["model"] = model - kwargs["custom_llm_provider"] = "openai" + kwargs["custom_llm_provider"] = custom_llm_provider # Extract user information for tracking passthrough_logging_payload: Optional[ @@ -321,10 +368,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs.setdefault("litellm_params", {}) - kwargs["litellm_params"].update( - {"proxy_server_request": {"body": {"user": user}}} - ) + kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = user # Create standard logging object if litellm_model_response is not None: @@ -339,7 +383,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Update logging object with cost information logging_obj.model_call_details["model"] = model - logging_obj.model_call_details["custom_llm_provider"] = "openai" + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider logging_obj.model_call_details["response_cost"] = response_cost endpoint_type = ( @@ -481,18 +525,27 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): "kwargs": {}, } + custom_llm_provider = litellm_logging_obj.model_call_details.get( + "custom_llm_provider", "openai" + ) # Calculate cost using LiteLLM's cost calculator response_cost = litellm.completion_cost( completion_response=complete_response, model=model, - custom_llm_provider="openai", + custom_llm_provider=custom_llm_provider, ) + # Preserve existing litellm_params to maintain metadata tags + existing_litellm_params = litellm_logging_obj.model_call_details.get( + "litellm_params", {} + ) or {} + # Prepare kwargs for logging kwargs = { "response_cost": response_cost, "model": model, - "custom_llm_provider": "openai", + "custom_llm_provider": custom_llm_provider, + "litellm_params": existing_litellm_params.copy(), } # Extract user information for tracking @@ -506,10 +559,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs.setdefault("litellm_params", {}) - kwargs["litellm_params"].update( - {"proxy_server_request": {"body": {"user": user}}} - ) + kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = user # Create standard logging object get_standard_logging_object_payload( @@ -523,7 +573,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Update logging object with cost information litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details["custom_llm_provider"] = "openai" + litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider litellm_logging_obj.model_call_details["response_cost"] = response_cost verbose_proxy_logger.debug( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 780a3d6dcc..5b47a8af7a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -735,6 +735,12 @@ async def pass_through_request( # noqa: PLR0915 logging_obj=logging_obj, ) + # Store custom_llm_provider in kwargs and logging object if provided + if custom_llm_provider: + kwargs["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["litellm_params"] = kwargs.get("litellm_params", {}) + # done for supporting 'parallel_request_limiter.py' with pass-through endpoints logging_obj.update_environment_variables( model="unknown", @@ -923,6 +929,12 @@ async def pass_through_request( # noqa: PLR0915 if kwargs: for key, value in kwargs.items(): request_payload[key] = value + + if "model" not in request_payload and _parsed_body and isinstance(_parsed_body, dict): + request_payload["model"] = _parsed_body.get("model", "") + if "custom_llm_provider" not in request_payload and custom_llm_provider: + request_payload["custom_llm_provider"] = custom_llm_provider + await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -957,11 +969,21 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di """ If tags are in the request headers, add them to the metadata - Used for google and vertex JS SDKs + Used for google and vertex JS SDKs, and Azure passthrough + Checks both 'tags' and 'x-litellm-tags' headers """ + # Initialize tags list if it doesn't exist + if "tags" not in metadata: + metadata["tags"] = [] + + # Check for 'tags' header first _tags = request.headers.get("tags") if _tags: - metadata["tags"] = _tags.split(",") + metadata["tags"].extend([tag.strip() for tag in _tags.split(",")]) + + _tags = request.headers.get("x-litellm-tags") + if _tags: + metadata["tags"].extend([tag.strip() for tag in _tags.split(",")]) return metadata diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 789b16f951..becb34409b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -130,6 +130,19 @@ class TestOpenAIPassthroughLoggingHandler: assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("http://localhost:4000/openai/v1/images/edits") == False assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False + def test_is_openai_responses_route(self): + """Test OpenAI responses API route detection""" + # Positive cases + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/responses") == True + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://openai.azure.com/v1/responses") == True + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/responses") == True + + # Negative cases + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/chat/completions") == False + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/images/generations") == False + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("http://localhost:4000/openai/v1/responses") == False + assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("") == False + @patch('litellm.completion_cost') @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') def test_openai_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost): @@ -369,6 +382,178 @@ class TestOpenAIPassthroughLoggingHandler: handler = OpenAIPassthroughLoggingHandler() assert handler.get_provider_config("gpt-4o") is not None + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + def test_azure_passthrough_tags_metadata_model_provider(self, mock_get_standard_logging, mock_completion_cost): + """Test that tags, metadata, model, and custom_llm_provider are preserved for Azure passthrough in UI""" + # Arrange + mock_completion_cost.return_value = 0.000045 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + + # Create payload with metadata tags + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://openai.azure.com/v1/chat/completions", + request_body={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hello"}] + }, + request_method="POST", + ) + + # Set up kwargs with existing litellm_params containing metadata tags + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + "custom_llm_provider": "azure", # Azure passthrough + "litellm_params": { + "metadata": { + "tags": ["production", "azure-deployment"], + "user_id": "user_123" + }, + "proxy_server_request": { + "body": { + "user": "test_user" + } + } + } + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_openai_response, + logging_obj=mock_logging_obj, + url_route="https://openai.azure.com/v1/chat/completions", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + **kwargs + ) + + # Assert - Verify tags, model, and custom_llm_provider are preserved + assert result is not None + assert "kwargs" in result + + # Verify model and custom_llm_provider are set correctly + assert result["kwargs"]["model"] == "gpt-4o" + assert result["kwargs"]["custom_llm_provider"] == "azure" # Should preserve Azure, not default to "openai" + assert result["kwargs"]["response_cost"] == 0.000045 + + # Verify metadata tags are preserved in litellm_params + assert "litellm_params" in result["kwargs"] + assert "metadata" in result["kwargs"]["litellm_params"] + assert "tags" in result["kwargs"]["litellm_params"]["metadata"] + assert result["kwargs"]["litellm_params"]["metadata"]["tags"] == ["production", "azure-deployment"] + assert result["kwargs"]["litellm_params"]["metadata"]["user_id"] == "user_123" + + # Verify logging object has correct values for UI display + assert mock_logging_obj.model_call_details["model"] == "gpt-4o" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "azure" + assert mock_logging_obj.model_call_details["response_cost"] == 0.000045 + + # Verify cost calculation was called with correct custom_llm_provider + mock_completion_cost.assert_called_once() + call_args = mock_completion_cost.call_args + assert call_args[1]["custom_llm_provider"] == "azure" + + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.get_provider_config') + def test_responses_api_cost_tracking(self, mock_get_provider_config, mock_get_standard_logging, mock_completion_cost): + """Test cost tracking for responses API route""" + # Arrange + mock_completion_cost.return_value = 0.000050 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + # Mock the provider config's transform_response to return a valid ModelResponse + from litellm import ModelResponse + mock_model_response = ModelResponse( + id="resp_abc123", + model="gpt-4o-2024-08-06", + choices=[{ + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?" + } + }], + usage={ + "prompt_tokens": 20, + "completion_tokens": 15, + "total_tokens": 35 + } + ) + + mock_provider_config = MagicMock() + mock_provider_config.transform_response.return_value = mock_model_response + mock_get_provider_config.return_value = mock_provider_config + + # Mock responses API response + mock_responses_response = { + "id": "resp_abc123", + "object": "response", + "created": 1677652288, + "model": "gpt-4o-2024-08-06", + "output": [ + { + "type": "text", + "text": "Hello! How can I help you today?" + } + ], + "usage": { + "input_tokens": 20, + "output_tokens": 15 + } + } + + mock_httpx_response = self._create_mock_httpx_response(mock_responses_response) + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + "custom_llm_provider": "openai", + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=mock_responses_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/responses", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "input": "Tell me about AI"}, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000050 + assert result["kwargs"]["model"] == "gpt-4o" + assert result["kwargs"]["custom_llm_provider"] == "openai" + + # Verify cost calculation was called with responses call type + mock_completion_cost.assert_called_once() + call_args = mock_completion_cost.call_args + assert call_args[1]["call_type"] == "responses" + assert call_args[1]["model"] == "gpt-4o" + assert call_args[1]["custom_llm_provider"] == "openai" + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == 0.000050 + assert mock_logging_obj.model_call_details["model"] == "gpt-4o" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + class TestOpenAIPassthroughIntegration: """Integration tests for OpenAI passthrough cost tracking"""