ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Litellm passthrough cost tracking chat completion (#14256)

Browse Source 
* feat: add structured output for sdk

* Add support for cost tracking for chat completion in passthrough

* remove not required changes
This commit is contained in:
Sameer Kankute 2025-09-04 22:27:48 +05:30 committed by GitHub
parent 1237be04a5
commit 5f79e8aac6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 921 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

383
litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py Normal file
View File

@ -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": {},
}

32
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View File

@ -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:

21
litellm/proxy/pass_through_endpoints/streaming_handler.py
View File

@ -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(

51
litellm/proxy/pass_through_endpoints/success_handler.py
View File

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

2
litellm/responses/main.py
View File

@ -1009,4 +1009,4 @@ def list_input_items(
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
)

2
litellm/responses/utils.py
View File

@ -333,4 +333,4 @@ class ResponseAPILoggingUtils:
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
)

1
litellm/types/passthrough_endpoints/pass_through_endpoints.py
View File

@ -5,6 +5,7 @@ from typing import Optional, TypedDict
class EndpointType(str, Enum):
VERTEX_AI = "vertex-ai"
ANTHROPIC = "anthropic"
OPENAI = "openai"
GENERIC = "generic"

2
tests/llm_responses_api_testing/base_responses_api.py
View File

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

451
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py Normal file
View File

@ -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__])