Merge pull request #17258 from BerriAI/litellm_add_tags_in_ui

Fix metadata tags and model name display in UI for Azure passthrough + Add cost tracking for responses API

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,

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(

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
 
 

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"""