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 new file mode 100644 index 0000000000..dd772ffa50 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -0,0 +1,383 @@ +""" +OpenAI Passthrough Logging Handler + +Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. +""" + +from datetime import datetime +from typing import List, Optional, Union +from urllib.parse import urlparse + +import httpx + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, +) +from litellm.llms.openai.openai import OpenAIConfig +from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( + BasePassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + PassthroughStandardLoggingPayload, +) +from litellm.types.utils import LlmProviders +from litellm.utils import ModelResponse, TextCompletionResponse + + +class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): + """ + OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints. + """ + + @property + def llm_provider_name(self) -> LlmProviders: + return LlmProviders.OPENAI + + @staticmethod + def get_provider_config(model: str) -> OpenAIConfigType: + """Get OpenAI provider configuration for the given model.""" + return OpenAIConfig() + + @staticmethod + def is_openai_chat_completions_route(url_route: str) -> bool: + """Check if the URL route is an OpenAI chat completions 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/chat/completions" in parsed_url.path + ) + + @staticmethod + def _get_user_from_metadata( + passthrough_logging_payload: PassthroughStandardLoggingPayload, + ) -> Optional[str]: + """Extract user information from passthrough logging payload.""" + request_body = passthrough_logging_payload.get("request_body") + if request_body: + return request_body.get("user") + return None + + @staticmethod + def openai_passthrough_handler( + httpx_response: httpx.Response, + response_body: dict, + logging_obj: LiteLLMLoggingObj, + url_route: str, + result: str, + start_time: datetime, + end_time: datetime, + cache_hit: bool, + request_body: dict, + **kwargs, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Handle OpenAI passthrough logging with cost tracking for chat completions. + """ + # 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 + base_handler = OpenAIPassthroughLoggingHandler() + return base_handler.passthrough_chat_handler( + httpx_response=httpx_response, + response_body=response_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + + # Extract model from request or response + model = request_body.get("model", response_body.get("model", "")) + if not model: + verbose_proxy_logger.warning( + "No model found in request or response for OpenAI passthrough cost tracking" + ) + base_handler = OpenAIPassthroughLoggingHandler() + return base_handler.passthrough_chat_handler( + httpx_response=httpx_response, + response_body=response_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + + 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={}, + ) + + # Calculate cost using LiteLLM's cost calculator + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider="openai", + ) + + # Update kwargs with cost information + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = "openai" + + # Extract user information for tracking + passthrough_logging_payload: Optional[ + PassthroughStandardLoggingPayload + ] = kwargs.get("passthrough_logging_payload") + if passthrough_logging_payload: + user = OpenAIPassthroughLoggingHandler._get_user_from_metadata( + passthrough_logging_payload=passthrough_logging_payload, + ) + if user: + kwargs.setdefault("litellm_params", {}) + kwargs["litellm_params"].update( + {"proxy_server_request": {"body": {"user": user}}} + ) + + # 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", + ) + + # 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 + + verbose_proxy_logger.debug( + f"OpenAI passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}" + ) + + return { + "result": litellm_model_response, + "kwargs": kwargs, + } + + except Exception as e: + verbose_proxy_logger.error( + f"Error in OpenAI passthrough cost tracking: {str(e)}" + ) + # Fall back to base handler without cost tracking + base_handler = OpenAIPassthroughLoggingHandler() + return base_handler.passthrough_chat_handler( + httpx_response=httpx_response, + response_body=response_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + + def _build_complete_streaming_response( + self, + all_chunks: list, + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: + """ + Builds complete response from raw chunks for OpenAI streaming responses. + + - Converts str chunks to generic chunks + - Converts generic chunks to litellm chunks (OpenAI format) + - Builds complete response from litellm chunks + """ + try: + # OpenAI's response iterator to parse chunks + from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator + + openai_iterator = OpenAIChatCompletionResponseIterator( + streaming_response=None, + sync_stream=False, + ) + + all_openai_chunks = [] + for chunk_str in all_chunks: + try: + # Parse the string chunk using the base iterator's string parser + from litellm.llms.base_llm.base_model_iterator import ( + BaseModelResponseIterator, + ) + + # Convert string chunk to dict + stripped_json_chunk = ( + BaseModelResponseIterator._string_to_dict_parser( + str_line=chunk_str + ) + ) + + if stripped_json_chunk: + # Parse the chunk using OpenAI's chunk parser + transformed_chunk = openai_iterator.chunk_parser( + chunk=stripped_json_chunk + ) + if transformed_chunk is not None: + all_openai_chunks.append(transformed_chunk) + + except (StopIteration, StopAsyncIteration, Exception) as e: + verbose_proxy_logger.debug(f"Error parsing streaming chunk: {e}") + continue + + if not all_openai_chunks: + verbose_proxy_logger.warning( + "No valid chunks found in streaming response" + ) + return None + + # Build complete response from chunks + complete_streaming_response = litellm.stream_chunk_builder( + chunks=all_openai_chunks + ) + + return complete_streaming_response + + except Exception as e: + verbose_proxy_logger.error( + f"Error building complete streaming response: {str(e)}" + ) + return None + + @staticmethod + def _handle_logging_openai_collected_chunks( + litellm_logging_obj: LiteLLMLoggingObj, + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, + request_body: dict, + endpoint_type: EndpointType, + start_time: datetime, + all_chunks: List[str], + end_time: datetime, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Handle logging for collected OpenAI streaming chunks with cost tracking. + """ + try: + # Extract model from request body + model = request_body.get("model", "gpt-4o") + + # Build complete response from chunks using our streaming handler + handler = OpenAIPassthroughLoggingHandler() + complete_response = handler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + ) + + if complete_response is None: + verbose_proxy_logger.warning( + "Failed to build complete response from OpenAI streaming chunks" + ) + return { + "result": None, + "kwargs": {}, + } + + # Calculate cost using LiteLLM's cost calculator + response_cost = litellm.completion_cost( + completion_response=complete_response, + model=model, + custom_llm_provider="openai", + ) + + # Prepare kwargs for logging + kwargs = { + "response_cost": response_cost, + "model": model, + "custom_llm_provider": "openai", + } + + # Extract user information for tracking + passthrough_logging_payload: Optional[ + PassthroughStandardLoggingPayload + ] = litellm_logging_obj.model_call_details.get( + "passthrough_logging_payload" + ) + if passthrough_logging_payload: + user = OpenAIPassthroughLoggingHandler._get_user_from_metadata( + passthrough_logging_payload=passthrough_logging_payload, + ) + if user: + kwargs.setdefault("litellm_params", {}) + kwargs["litellm_params"].update( + {"proxy_server_request": {"body": {"user": user}}} + ) + + # Create standard logging object + get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj=complete_response, + start_time=start_time, + end_time=end_time, + logging_obj=litellm_logging_obj, + status="success", + ) + + # 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["response_cost"] = response_cost + + verbose_proxy_logger.debug( + f"OpenAI streaming passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}" + ) + + return { + "result": complete_response, + "kwargs": kwargs, + } + + except Exception as e: + verbose_proxy_logger.error( + f"Error in OpenAI streaming passthrough cost tracking: {str(e)}" + ) + return { + "result": None, + "kwargs": {}, + } diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index adedcaf781..fccc65b8bd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -314,6 +314,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): return EndpointType.VERTEX_AI elif parsed_url.hostname == "api.anthropic.com": return EndpointType.ANTHROPIC + elif ( + parsed_url.hostname == "api.openai.com" + or parsed_url.hostname == "openai.azure.com" + or (parsed_url.hostname and "openai.com" in parsed_url.hostname) + ): + return EndpointType.OPENAI return EndpointType.GENERIC @staticmethod @@ -415,10 +421,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): for field_name, field_value in form_data.items(): if isinstance(field_value, (StarletteUploadFile, UploadFile)): - files[field_name] = ( - await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file=field_value - ) + files[ + field_name + ] = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( + upload_file=field_value ) else: form_data_dict[field_name] = field_value @@ -497,9 +503,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "passthrough_logging_payload": passthrough_logging_payload, } - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details[ + "passthrough_logging_payload" + ] = passthrough_logging_payload return kwargs @@ -531,10 +537,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): subpath = subpath[1:] return base_target + subpath - + @staticmethod def _update_stream_param_based_on_request_body( - parsed_body: dict, + parsed_body: dict, stream: Optional[bool] = None, ) -> Optional[bool]: """ @@ -699,9 +705,11 @@ async def pass_through_request( # noqa: PLR0915 "headers": headers, }, ) - stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( - parsed_body=_parsed_body, - stream=stream, + stream = ( + HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( + parsed_body=_parsed_body, + stream=stream, + ) ) if stream: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 08b49bac38..2d5b0a686c 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -14,6 +14,9 @@ from litellm.types.utils import StandardPassThroughResponseObject from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) +from .llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) @@ -78,6 +81,7 @@ class PassThroughStreamingHandler: Supported endpoint types: - Anthropic - Vertex AI + - OpenAI """ all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines( raw_bytes @@ -119,6 +123,23 @@ class PassThroughStreamingHandler: vertex_passthrough_logging_handler_result["result"] ) kwargs = vertex_passthrough_logging_handler_result["kwargs"] + elif endpoint_type == EndpointType.OPENAI: + openai_passthrough_logging_handler_result = ( + OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + ) + ) + standard_logging_response_object = ( + openai_passthrough_logging_handler_result["result"] + ) + kwargs = openai_passthrough_logging_handler_result["kwargs"] if standard_logging_response_object is None: standard_logging_response_object = StandardPassThroughResponseObject( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index ce576d5ac7..5b54f94f8b 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -162,9 +162,32 @@ class PassThroughEndpointLogging: cohere_passthrough_logging_handler_result["result"] ) kwargs = cohere_passthrough_logging_handler_result["kwargs"] - return_dict["standard_logging_response_object"] = ( - standard_logging_response_object - ) + elif self.is_openai_route(url_route): + from .llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + openai_passthrough_logging_handler_result = ( + OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + ) + standard_logging_response_object = ( + openai_passthrough_logging_handler_result["result"] + ) + kwargs = openai_passthrough_logging_handler_result["kwargs"] + return_dict[ + "standard_logging_response_object" + ] = standard_logging_response_object return_dict["kwargs"] = kwargs return return_dict @@ -185,9 +208,9 @@ class PassThroughEndpointLogging: standard_logging_response_object: Optional[ PassThroughEndpointLoggingResultValues ] = None - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details[ + "passthrough_logging_payload" + ] = passthrough_logging_payload if self.is_assemblyai_route(url_route): if ( AssemblyAIPassthroughLoggingHandler._should_log_request( @@ -286,6 +309,16 @@ class PassThroughEndpointLogging: return True return False + def is_openai_route(self, url_route: str): + """Check if the URL route is an OpenAI API route.""" + if not url_route: + return False + parsed_url = urlparse(url_route) + return parsed_url.hostname and ( + "api.openai.com" in parsed_url.hostname + or "openai.azure.com" in parsed_url.hostname + ) + def _set_cost_per_request( self, logging_obj: LiteLLMLoggingObj, @@ -305,8 +338,8 @@ class PassThroughEndpointLogging: kwargs["response_cost"] = passthrough_logging_payload.get( "cost_per_request" ) - logging_obj.model_call_details["response_cost"] = ( - passthrough_logging_payload.get("cost_per_request") - ) + logging_obj.model_call_details[ + "response_cost" + ] = passthrough_logging_payload.get("cost_per_request") return kwargs diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 9584baf736..46743d7bb8 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1009,4 +1009,4 @@ def list_input_items( original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, - ) + ) \ No newline at end of file diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index ac59d28a50..20e1ee957c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -333,4 +333,4 @@ class ResponseAPILoggingUtils: prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, - ) + ) \ No newline at end of file diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 8a91dd14eb..39facdd8e6 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -5,6 +5,7 @@ from typing import Optional, TypedDict class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" ANTHROPIC = "anthropic" + OPENAI = "openai" GENERIC = "generic" diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 5cb8295b1a..c177a76cda 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -536,4 +536,4 @@ class BaseResponsesAPITest(ABC): # Validate final response structure validate_responses_api_response(final_response, final_chunk=True) assert final_response.output is not None - assert len(final_response.output) > 0 + assert len(final_response.output) > 0 \ No newline at end of file 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 new file mode 100644 index 0000000000..6d5e80910b --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -0,0 +1,451 @@ +import json +import os +import sys +from datetime import datetime +from typing import Any, Dict, List +from unittest.mock import MagicMock, patch, AsyncMock + +import pytest +import httpx + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, +) + + +class TestOpenAIPassthroughLoggingHandler: + """Test the OpenAI passthrough logging handler for cost tracking.""" + + def setup_method(self): + """Set up test fixtures""" + self.start_time = datetime.now() + self.end_time = datetime.now() + self.handler = OpenAIPassthroughLoggingHandler() + + # Mock OpenAI chat completions response + self.mock_openai_response = { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o-2024-08-06", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 20, + "completion_tokens": 15, + "total_tokens": 35 + } + } + + 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 = self.mock_openai_response + + 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_llm_provider_name(self): + """Test that the handler returns the correct provider name""" + assert self.handler.llm_provider_name == "openai" + + def test_get_provider_config(self): + """Test that the handler returns an OpenAI config""" + config = OpenAIPassthroughLoggingHandler.get_provider_config(model="gpt-4o") + assert config is not None + # Verify it's an OpenAI config by checking if it has the expected methods + assert hasattr(config, 'transform_response') + + def test_is_openai_chat_completions_route(self): + """Test OpenAI chat completions route detection""" + # Positive cases + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/chat/completions") == True + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://openai.azure.com/v1/chat/completions") == True + + # Negative cases + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/models") == False + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("http://localhost:4000/openai/v1/chat/completions") == False + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages") == False + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_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): + """Test successful cost tracking for OpenAI chat completions""" + # 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() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # 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://api.openai.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 + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000045 + assert result["kwargs"]["model"] == "gpt-4o" + assert result["kwargs"]["custom_llm_provider"] == "openai" + + # Verify cost calculation was called + mock_completion_cost.assert_called_once() + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == 0.000045 + assert mock_logging_obj.model_call_details["model"] == "gpt-4o" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + + @patch('litellm.completion_cost') + def test_openai_passthrough_handler_non_chat_completions(self, mock_completion_cost): + """Test that non-chat-completions routes fall back to base handler""" + # Arrange + mock_httpx_response = self._create_mock_httpx_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", + } + + # Act - Use a non-chat-completions route + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body={"id": "file-123", "object": "file"}, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/files", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"purpose": "fine-tune"}, + **kwargs + ) + + # Assert - Should fall back to base handler for non-chat-completions + assert result is not None + assert "result" in result + assert "kwargs" in result + # Cost calculation may be called by the base handler fallback + # The important thing is that our specific OpenAI handler logic didn't run + + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + def test_openai_passthrough_handler_with_user_tracking(self, mock_get_standard_logging, mock_completion_cost): + """Test cost tracking with user information""" + # Arrange + mock_completion_cost.return_value = 0.000123 + 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 user information + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/chat/completions", + request_body={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hello"}], + "user": "test_user_123" + }, + request_method="POST", + ) + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # 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://api.openai.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"}], "user": "test_user_123"}, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000123 + + # Verify user information is included in litellm_params + assert "litellm_params" in result["kwargs"] + assert "proxy_server_request" in result["kwargs"]["litellm_params"] + assert "body" in result["kwargs"]["litellm_params"]["proxy_server_request"] + assert result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] == "test_user_123" + + @patch('litellm.completion_cost') + def test_openai_passthrough_handler_cost_calculation_error(self, mock_completion_cost): + """Test error handling in cost calculation""" + # Arrange + mock_completion_cost.side_effect = Exception("Cost calculation failed") + + mock_httpx_response = self._create_mock_httpx_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", + } + + # 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://api.openai.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 - Should fall back to base handler when cost calculation fails + assert result is not None + assert "result" in result + assert "kwargs" in result + + def test_build_complete_streaming_response(self): + """Test the streaming response builder (placeholder implementation)""" + # This is a placeholder method that returns None for now + result = self.handler._build_complete_streaming_response( + all_chunks=["chunk1", "chunk2"], + litellm_logging_obj=self._create_mock_logging_obj(), + model="gpt-4o", + ) + + assert result is None # Placeholder implementation + + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + def test_different_models_cost_tracking(self, mock_get_standard_logging, mock_completion_cost): + """Test cost tracking for different OpenAI models""" + # Arrange + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + test_cases = [ + ("gpt-4o", 0.000045), + ("gpt-4o-mini", 0.000015), + ("gpt-3.5-turbo", 0.000002), + ] + + for model, expected_cost in test_cases: + mock_completion_cost.return_value = expected_cost + + mock_httpx_response = self._create_mock_httpx_response() + mock_httpx_response.json.return_value = { + **self.mock_openai_response, + "model": model + } + + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": model, + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body={**self.mock_openai_response, "model": model}, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": model, "messages": [{"role": "user", "content": "Hello"}]}, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == expected_cost + assert result["kwargs"]["model"] == model + assert result["kwargs"]["custom_llm_provider"] == "openai" + + def test_static_methods(self): + """Test that static methods work correctly""" + # Test static method calls + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/chat/completions") == True + assert OpenAIPassthroughLoggingHandler.get_provider_config("gpt-4o") is not None + + +class TestOpenAIPassthroughIntegration: + """Integration tests for OpenAI passthrough cost tracking""" + + def setup_method(self): + """Set up test fixtures""" + self.handler = PassThroughEndpointLogging() + + def test_is_openai_route_detection(self): + """Test OpenAI route detection in the main success handler""" + # Positive cases + assert self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") == True + assert self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") == True + assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True + + # Negative cases + assert self.handler.is_openai_route("http://localhost:4000/openai/v1/chat/completions") == False + assert self.handler.is_openai_route("https://api.anthropic.com/v1/messages") == False + assert self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False + assert self.handler.is_openai_route("") == False + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler') + @pytest.mark.asyncio + async def test_success_handler_calls_openai_handler(self, mock_openai_handler): + """Test that the success handler calls our OpenAI handler for OpenAI routes""" + # Arrange + mock_openai_handler.return_value = { + "result": {"id": "chatcmpl-123"}, + "kwargs": { + "response_cost": 0.000045, + "model": "gpt-4o", + "custom_llm_provider": "openai" + } + } + + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}' + + mock_logging_obj = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.async_success_handler = AsyncMock() + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/chat/completions", + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + request_method="POST", + ) + + # Act + result = await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + passthrough_logging_payload=passthrough_payload, + ) + + # Assert + mock_openai_handler.assert_called_once() + # The success handler returns None on success, which is expected + assert result is None + + @pytest.mark.asyncio + async def test_success_handler_falls_back_for_non_openai_routes(self): + """Test that non-OpenAI routes don't call our handler""" + # Arrange + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = '{"status": "success"}' + mock_httpx_response.headers = {"content-type": "application/json"} + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.anthropic.com/v1/messages", + request_body={"model": "claude-3-sonnet", "messages": [{"role": "user", "content": "Hello"}]}, + request_method="POST", + ) + + # Mock the _handle_logging method to capture calls + self.handler._handle_logging = AsyncMock() + + # Act + result = await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={"status": "success"}, + logging_obj=mock_logging_obj, + url_route="https://api.anthropic.com/v1/messages", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "claude-3-sonnet", "messages": [{"role": "user", "content": "Hello"}]}, + passthrough_logging_payload=passthrough_payload, + ) + + # Assert - Should call the base handler, not our OpenAI handler + self.handler._handle_logging.assert_called_once() + + +if __name__ == "__main__": + pytest.main([__file__])