diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7bc7702684..397858060d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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 ( 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 dd772ffa50..d230023a23 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 @@ -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, } 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 6d5e80910b..6f808c9759 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 @@ -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__])