[Feat] Add pass through image gen and image editing on OpenAI (#14292)

* add pass through image gen and image editing on OpenAI

* fix lint

litellm/litellm_core_utils/litellm_logging.py

@@ -1165,6 +1165,14 @@ class Logging(LiteLLMLoggingBaseClass):
         used for consistent cost calculation across response headers + logging integrations.
         """
 
+        # Check if response_cost is already calculated and stored in model_call_details
+        # This is used by passthrough endpoints that calculate costs manually
+        if (
+            hasattr(self, "model_call_details")
+            and self.model_call_details.get("response_cost") is not None
+        ):
+            return self.model_call_details["response_cost"]
+
         if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"):
             hidden_params = getattr(result, "_hidden_params", {})
             if (

litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py

@@ -29,7 +29,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
     EndpointType,
     PassthroughStandardLoggingPayload,
 )
-from litellm.types.utils import LlmProviders
+from litellm.types.utils import LlmProviders, PassthroughCallTypes
 from litellm.utils import ModelResponse, TextCompletionResponse
 
 
@@ -62,6 +62,36 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
             and "/v1/chat/completions" in parsed_url.path
         )
 
+    @staticmethod
+    def is_openai_image_generation_route(url_route: str) -> bool:
+        """Check if the URL route is an OpenAI image generation 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/images/generations" in parsed_url.path
+        )
+
+    @staticmethod
+    def is_openai_image_editing_route(url_route: str) -> bool:
+        """Check if the URL route is an OpenAI image editing 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/images/edits" in parsed_url.path
+        )
+
     @staticmethod
     def _get_user_from_metadata(
         passthrough_logging_payload: PassthroughStandardLoggingPayload,
@@ -73,7 +103,79 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
         return None
 
     @staticmethod
-    def openai_passthrough_handler(
+    def _calculate_image_generation_cost(
+        model: str,
+        response_body: dict,
+        request_body: dict,
+    ) -> float:
+        """Calculate cost for OpenAI image generation."""
+        try:
+            # Extract parameters from request
+            n = request_body.get("n", 1)
+            try:
+                n = int(n)
+            except Exception:
+                n = 1
+            size = request_body.get("size", "1024x1024")
+            quality = request_body.get("quality", None)
+
+            # Use LiteLLM's default image cost calculator
+            from litellm.cost_calculator import default_image_cost_calculator
+
+            cost = default_image_cost_calculator(
+                model=model,
+                custom_llm_provider="openai",
+                quality=quality,
+                n=n,
+                size=size,
+                optional_params=request_body,
+            )
+
+            return cost
+        except Exception as e:
+            verbose_proxy_logger.warning(
+                f"Error calculating image generation cost: {str(e)}"
+            )
+            return 0.0
+
+    @staticmethod
+    def _calculate_image_editing_cost(
+        model: str,
+        response_body: dict,
+        request_body: dict,
+    ) -> float:
+        """Calculate cost for OpenAI image editing."""
+        try:
+            # Extract parameters from request
+            n = request_body.get("n", 1)
+            # Image edit typically uses multipart/form-data (because of files), so all fields arrive as strings (e.g., n = "1").
+            try:
+                n = int(n)
+            except Exception:
+                n = 1
+            size = request_body.get("size", "1024x1024")
+
+            # Use LiteLLM's default image cost calculator
+            from litellm.cost_calculator import default_image_cost_calculator
+
+            cost = default_image_cost_calculator(
+                model=model,
+                custom_llm_provider="openai",
+                quality=None,  # Image editing doesn't have quality parameter
+                n=n,
+                size=size,
+                optional_params=request_body,
+            )
+
+            return cost
+        except Exception as e:
+            verbose_proxy_logger.warning(
+                f"Error calculating image editing cost: {str(e)}"
+            )
+            return 0.0
+
+    @staticmethod
+    def openai_passthrough_handler(  # noqa: PLR0915
         httpx_response: httpx.Response,
         response_body: dict,
         logging_obj: LiteLLMLoggingObj,
@@ -86,13 +188,21 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
         **kwargs,
     ) -> PassThroughEndpointLoggingTypedDict:
         """
-        Handle OpenAI passthrough logging with cost tracking for chat completions.
+        Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, and image editing.
         """
-        # Only handle chat completions endpoints
-        if not OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
-            url_route
-        ):
-            # For non-chat-completions endpoints, use the base handler without cost tracking
+        # Check if this is a supported endpoint for cost tracking
+        is_chat_completions = (
+            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route)
+        )
+        is_image_generation = (
+            OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route)
+        )
+        is_image_editing = (
+            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
+        )
+
+        if not (is_chat_completions or is_image_generation or is_image_editing):
+            # For unsupported endpoints, use the base handler without cost tracking
             base_handler = OpenAIPassthroughLoggingHandler()
             return base_handler.passthrough_chat_handler(
                 httpx_response=httpx_response,
@@ -128,31 +238,89 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
             )
 
         try:
-            # Transform the response to LiteLLM format for cost calculation
-            provider_config = OpenAIPassthroughLoggingHandler.get_provider_config(
-                model=model
-            )
-            litellm_model_response: ModelResponse = 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=request_body.get("response_format", {}).get("type")
-                == "json_object",
-                litellm_params={},
-            )
+            response_cost = 0.0
+            litellm_model_response = None
 
-            # Calculate cost using LiteLLM's cost calculator
-            response_cost = litellm.completion_cost(
-                completion_response=litellm_model_response,
-                model=model,
-                custom_llm_provider="openai",
-            )
+            if is_chat_completions:
+                # Handle chat completions with existing logic
+                provider_config = OpenAIPassthroughLoggingHandler.get_provider_config(
+                    model=model
+                )
+                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=request_body.get("response_format", {}).get("type")
+                    == "json_object",
+                    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",
+                )
+            elif is_image_generation:
+                # Handle image generation cost calculation
+                response_cost = (
+                    OpenAIPassthroughLoggingHandler._calculate_image_generation_cost(
+                        model=model,
+                        response_body=response_body,
+                        request_body=request_body,
+                    )
+                )
+                # Mark call type for downstream image-aware logic/metrics
+                try:
+                    logging_obj.call_type = (
+                        PassthroughCallTypes.passthrough_image_generation.value
+                    )
+                except Exception:
+                    pass
+                # Create a simple response object for logging
+                from litellm.types.utils import ImageResponse
+
+                litellm_model_response = ImageResponse(
+                    data=response_body.get("data", []),
+                    model=model,
+                )
+                # Set the calculated cost in _hidden_params to prevent recalculation
+                if not hasattr(litellm_model_response, "_hidden_params"):
+                    litellm_model_response._hidden_params = {}
+                litellm_model_response._hidden_params["response_cost"] = response_cost
+            elif is_image_editing:
+                # Handle image editing cost calculation
+                response_cost = (
+                    OpenAIPassthroughLoggingHandler._calculate_image_editing_cost(
+                        model=model,
+                        response_body=response_body,
+                        request_body=request_body,
+                    )
+                )
+                # Mark call type for downstream image-aware logic/metrics
+                try:
+                    logging_obj.call_type = (
+                        PassthroughCallTypes.passthrough_image_generation.value
+                    )
+                except Exception:
+                    pass
+                # Create a simple response object for logging
+                from litellm.types.utils import ImageResponse
+
+                litellm_model_response = ImageResponse(
+                    data=response_body.get("data", []),
+                    model=model,
+                )
+                # Set the calculated cost in _hidden_params to prevent recalculation
+                if not hasattr(litellm_model_response, "_hidden_params"):
+                    litellm_model_response._hidden_params = {}
+                litellm_model_response._hidden_params["response_cost"] = response_cost
 
             # Update kwargs with cost information
             kwargs["response_cost"] = response_cost
@@ -174,26 +342,34 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
                     )
 
             # Create standard logging object
-            get_standard_logging_object_payload(
-                kwargs=kwargs,
-                init_response_obj=litellm_model_response,
-                start_time=start_time,
-                end_time=end_time,
-                logging_obj=logging_obj,
-                status="success",
-            )
+            if litellm_model_response is not None:
+                get_standard_logging_object_payload(
+                    kwargs=kwargs,
+                    init_response_obj=litellm_model_response,
+                    start_time=start_time,
+                    end_time=end_time,
+                    logging_obj=logging_obj,
+                    status="success",
+                )
 
             # 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["response_cost"] = response_cost
 
+            endpoint_type = (
+                "chat_completions"
+                if is_chat_completions
+                else "image_generation"
+                if is_image_generation
+                else "image_editing"
+            )
             verbose_proxy_logger.debug(
-                f"OpenAI passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}"
+                f"OpenAI passthrough cost tracking - Endpoint: {endpoint_type}, Model: {model}, Cost: ${response_cost:.6f}"
             )
 
             return {
-                "result": litellm_model_response,
+                "result": litellm_model_response or response_body,
                 "kwargs": kwargs,
             }
 

tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py

@@ -105,6 +105,30 @@ class TestOpenAIPassthroughLoggingHandler:
         assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages") == False
         assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") == False
 
+    def test_is_openai_image_generation_route(self):
+        """Test OpenAI image generation route detection"""
+        # Positive cases
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/generations") == True
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://openai.azure.com/v1/images/generations") == True
+        
+        # Negative cases
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/chat/completions") == False
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/edits") == False
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("http://localhost:4000/openai/v1/images/generations") == False
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("") == False
+
+    def test_is_openai_image_editing_route(self):
+        """Test OpenAI image editing route detection"""
+        # Positive cases
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/edits") == True
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://openai.azure.com/v1/images/edits") == True
+        
+        # Negative cases
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/chat/completions") == False
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/generations") == False
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("http://localhost:4000/openai/v1/images/edits") == False
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_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):
@@ -349,6 +373,34 @@ class TestOpenAIPassthroughIntegration:
     def setup_method(self):
         """Set up test fixtures"""
         self.handler = PassThroughEndpointLogging()
+        self.start_time = datetime.now()
+        self.end_time = datetime.now()
+
+    def _create_mock_logging_obj(self) -> LiteLLMLoggingObj:
+        """Create a mock logging object"""
+        mock_logging_obj = MagicMock()
+        mock_logging_obj.model_call_details = {}
+        return mock_logging_obj
+
+    def _create_mock_httpx_response(self, response_data: dict = None) -> httpx.Response:
+        """Create a mock httpx response"""
+        if response_data is None:
+            response_data = {"id": "test", "choices": [{"message": {"content": "Hello"}}]}
+            
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.text = json.dumps(response_data)
+        mock_response.json.return_value = response_data
+        mock_response.headers = {"content-type": "application/json"}
+        return mock_response
+
+    def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload:
+        """Create a mock passthrough logging payload"""
+        return PassthroughStandardLoggingPayload(
+            url="https://api.openai.com/v1/chat/completions",
+            request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]},
+            request_method="POST",
+        )
 
     def test_is_openai_route_detection(self):
         """Test OpenAI route detection in the main success handler"""
@@ -446,6 +498,240 @@ class TestOpenAIPassthroughIntegration:
         # Assert - Should call the base handler, not our OpenAI handler
         self.handler._handle_logging.assert_called_once()
 
+    @patch('litellm.cost_calculator.default_image_cost_calculator')
+    def test_calculate_image_generation_cost(self, mock_image_cost_calculator):
+        """Test image generation cost calculation"""
+        # Arrange
+        mock_image_cost_calculator.return_value = 0.040
+        model = "dall-e-3"
+        response_body = {
+            "data": [
+                {
+                    "url": "https://example.com/image1.png",
+                    "revised_prompt": "A beautiful sunset over the ocean"
+                }
+            ]
+        }
+        request_body = {
+            "model": "dall-e-3",
+            "prompt": "A beautiful sunset over the ocean",
+            "n": 1,
+            "size": "1024x1024",
+            "quality": "standard"
+        }
+
+        # Act
+        cost = OpenAIPassthroughLoggingHandler._calculate_image_generation_cost(
+            model=model,
+            response_body=response_body,
+            request_body=request_body,
+        )
+
+        # Assert
+        assert cost == 0.040
+        mock_image_cost_calculator.assert_called_once_with(
+            model=model,
+            custom_llm_provider="openai",
+            quality="standard",
+            n=1,
+            size="1024x1024",
+            optional_params=request_body,
+        )
+
+    @patch('litellm.cost_calculator.default_image_cost_calculator')
+    def test_calculate_image_editing_cost(self, mock_image_cost_calculator):
+        """Test image editing cost calculation"""
+        # Arrange
+        mock_image_cost_calculator.return_value = 0.020
+        model = "dall-e-2"
+        response_body = {
+            "data": [
+                {
+                    "url": "https://example.com/edited_image.png",
+                    "revised_prompt": "A beautiful sunset over the ocean with added clouds"
+                }
+            ]
+        }
+        request_body = {
+            "model": "dall-e-2",
+            "prompt": "Add clouds to the sky",
+            "n": 1,
+            "size": "1024x1024"
+        }
+
+        # Act
+        cost = OpenAIPassthroughLoggingHandler._calculate_image_editing_cost(
+            model=model,
+            response_body=response_body,
+            request_body=request_body,
+        )
+
+        # Assert
+        assert cost == 0.020
+        mock_image_cost_calculator.assert_called_once_with(
+            model=model,
+            custom_llm_provider="openai",
+            quality=None,  # Image editing doesn't have quality parameter
+            n=1,
+            size="1024x1024",
+            optional_params=request_body,
+        )
+
+    def test_cost_calculation_preservation(self):
+        """Test that manually calculated costs are preserved and not overridden."""
+        # Create a logging object
+        logging_obj = LiteLLMLoggingObj(
+            model="dall-e-3",
+            messages=[{"role": "user", "content": "Generate an image"}],
+            stream=False,
+            call_type="pass_through_endpoint",
+            start_time=self.start_time,
+            litellm_call_id="test_123",
+            function_id="test_fn",
+        )
+        
+        # Set a manually calculated cost in model_call_details
+        test_cost = 0.040000
+        logging_obj.model_call_details["response_cost"] = test_cost
+        logging_obj.model_call_details["model"] = "dall-e-3"
+        logging_obj.model_call_details["custom_llm_provider"] = "openai"
+        
+        # Create an ImageResponse with cost in _hidden_params
+        from litellm.types.utils import ImageResponse
+        image_response = ImageResponse(
+            data=[{"url": "https://example.com/image.png"}],
+            model="dall-e-3",
+        )
+        image_response._hidden_params = {"response_cost": test_cost}
+        
+        # Test the _response_cost_calculator method
+        calculated_cost = logging_obj._response_cost_calculator(result=image_response)
+        
+        assert calculated_cost == test_cost, f"Expected {test_cost}, got {calculated_cost}"
+
+    @patch('litellm.cost_calculator.default_image_cost_calculator')
+    def test_openai_passthrough_handler_image_generation(self, mock_image_cost_calculator):
+        """Test successful cost tracking for OpenAI image generation"""
+        # Arrange
+        mock_image_cost_calculator.return_value = 0.040
+        
+        mock_image_response = {
+            "data": [
+                {
+                    "url": "https://example.com/image1.png",
+                    "revised_prompt": "A beautiful sunset over the ocean"
+                }
+            ]
+        }
+        
+        mock_httpx_response = self._create_mock_httpx_response(mock_image_response)
+        mock_logging_obj = self._create_mock_logging_obj()
+        passthrough_payload = self._create_passthrough_logging_payload()
+        
+        kwargs = {
+            "passthrough_logging_payload": passthrough_payload,
+            "model": "dall-e-3",
+        }
+
+        request_body = {
+            "model": "dall-e-3",
+            "prompt": "A beautiful sunset over the ocean",
+            "n": 1,
+            "size": "1024x1024",
+            "quality": "standard"
+        }
+
+        # Act
+        result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
+            httpx_response=mock_httpx_response,
+            response_body=mock_image_response,
+            logging_obj=mock_logging_obj,
+            url_route="https://api.openai.com/v1/images/generations",
+            result="",
+            start_time=self.start_time,
+            end_time=self.end_time,
+            cache_hit=False,
+            request_body=request_body,
+            **kwargs
+        )
+
+        # Assert
+        assert result is not None
+        assert "result" in result
+        assert "kwargs" in result
+        assert result["kwargs"]["response_cost"] == 0.040
+        assert result["kwargs"]["model"] == "dall-e-3"
+        assert result["kwargs"]["custom_llm_provider"] == "openai"
+        
+        # Verify cost calculation was called
+        mock_image_cost_calculator.assert_called_once()
+        
+        # Verify logging object was updated
+        assert mock_logging_obj.model_call_details["response_cost"] == 0.040
+        assert mock_logging_obj.model_call_details["model"] == "dall-e-3"
+        assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
+
+    @patch('litellm.cost_calculator.default_image_cost_calculator')
+    def test_openai_passthrough_handler_image_editing(self, mock_image_cost_calculator):
+        """Test successful cost tracking for OpenAI image editing"""
+        # Arrange
+        mock_image_cost_calculator.return_value = 0.020
+        
+        mock_image_response = {
+            "data": [
+                {
+                    "url": "https://example.com/edited_image.png",
+                    "revised_prompt": "A beautiful sunset over the ocean with added clouds"
+                }
+            ]
+        }
+        
+        mock_httpx_response = self._create_mock_httpx_response(mock_image_response)
+        mock_logging_obj = self._create_mock_logging_obj()
+        passthrough_payload = self._create_passthrough_logging_payload()
+        
+        kwargs = {
+            "passthrough_logging_payload": passthrough_payload,
+            "model": "dall-e-2",
+        }
+
+        request_body = {
+            "model": "dall-e-2",
+            "prompt": "Add clouds to the sky",
+            "n": 1,
+            "size": "1024x1024"
+        }
+
+        # Act
+        result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
+            httpx_response=mock_httpx_response,
+            response_body=mock_image_response,
+            logging_obj=mock_logging_obj,
+            url_route="https://api.openai.com/v1/images/edits",
+            result="",
+            start_time=self.start_time,
+            end_time=self.end_time,
+            cache_hit=False,
+            request_body=request_body,
+            **kwargs
+        )
+
+        # Assert
+        assert result is not None
+        assert "result" in result
+        assert "kwargs" in result
+        assert result["kwargs"]["response_cost"] == 0.020
+        assert result["kwargs"]["model"] == "dall-e-2"
+        assert result["kwargs"]["custom_llm_provider"] == "openai"
+        
+        # Verify cost calculation was called
+        mock_image_cost_calculator.assert_called_once()
+        
+        # Verify logging object was updated
+        assert mock_logging_obj.model_call_details["response_cost"] == 0.020
+        assert mock_logging_obj.model_call_details["model"] == "dall-e-2"
+        assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
+
 
 if __name__ == "__main__":
     pytest.main([__file__])