From c32f42098cf94c2cfa44826a8769c4364993b658 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 1 Oct 2025 19:57:16 +0530 Subject: [PATCH 1/4] Add cost tracking for /v1/messages --- litellm/proxy/common_request_processing.py | 83 +++++++++++- .../anthropic_passthrough_logging_handler.py | 5 +- .../test_anthropic_passthrough.py | 120 ++++++++++++++++++ 3 files changed, 205 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f07a61c544..e7bad1e19d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -734,7 +734,7 @@ class ProxyBaseLLMRequestProcessing: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events """ - from litellm.types.utils import ModelResponse, ModelResponseStream + from litellm.types.utils import ModelResponse, ModelResponseStream, Usage verbose_proxy_logger.debug("inside generator") try: @@ -759,6 +759,87 @@ class ProxyBaseLLMRequestProcessing: response_str = litellm.get_response_string(response_obj=chunk) str_so_far += response_str + # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider + # Handle both dict SSE events and pre-formatted string SSE lines + if getattr(litellm, "include_cost_in_streaming_usage", False) is True: + try: + def _inject_cost_into_usage_dict(obj: dict) -> Optional[dict]: + if ( + obj.get("type") == "message_delta" + and isinstance(obj.get("usage"), dict) + ): + _usage = obj["usage"] + prompt_tokens = int(_usage.get("input_tokens", 0) or 0) + completion_tokens = int(_usage.get("output_tokens", 0) or 0) + total_tokens = int( + _usage.get("total_tokens", prompt_tokens + completion_tokens) + or (prompt_tokens + completion_tokens) + ) + + _mr = ModelResponse( + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + ) + model_name = request_data.get("model", "") + try: + cost_val = litellm.completion_cost( + completion_response=_mr, + model=model_name, + ) + except Exception: + cost_val = None + if cost_val is not None: + obj.setdefault("usage", {})["cost"] = cost_val + return obj + return None + + def _inject_cost_into_sse_frame_str(frame_str: str) -> Optional[str]: + # frame_str may contain multiple lines like 'event: ...\ndata: {...}\n\n' + # We only modify the JSON in the 'data:' line + try: + # Split preserving lines + lines = frame_str.split("\n") + for idx, ln in enumerate(lines): + stripped_ln = ln.strip() + if stripped_ln.startswith("data:"): + json_part = stripped_ln.split("data:", 1)[1].strip() + if json_part and json_part != "[DONE]": + obj = json.loads(json_part) + maybe_modified = _inject_cost_into_usage_dict(obj) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) + return None + except Exception: + return None + + if isinstance(chunk, dict): + maybe_modified = _inject_cost_into_usage_dict(chunk) + if maybe_modified is not None: + chunk = maybe_modified + elif isinstance(chunk, (bytes, bytearray)): + # Decode to str, inject, and rebuild as bytes + try: + s = chunk.decode("utf-8", errors="ignore") + maybe_mod = _inject_cost_into_sse_frame_str(s) + if maybe_mod is not None: + chunk = (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + except Exception: + pass + elif isinstance(chunk, str): + # Try to parse SSE frame and inject cost into the data line + maybe_mod = _inject_cost_into_sse_frame_str(chunk) + if maybe_mod is not None: + # Ensure trailing frame separator + chunk = maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") + except Exception: + # Never break streaming on optional cost injection + pass + # Format chunk using helper function yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index b9858202bf..9b6e22b819 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -20,7 +20,7 @@ from litellm.types.utils import ModelResponse, TextCompletionResponse if TYPE_CHECKING: from ..success_handler import PassThroughEndpointLogging - from ..types import EndpointType + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType else: PassThroughEndpointLogging = Any EndpointType = Any @@ -228,6 +228,7 @@ class AnthropicPassthroughLoggingHandler: except (StopIteration, StopAsyncIteration): break complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, ) return complete_streaming_response diff --git a/tests/pass_through_tests/test_anthropic_passthrough.py b/tests/pass_through_tests/test_anthropic_passthrough.py index 002fb20e9e..6e819f9971 100644 --- a/tests/pass_through_tests/test_anthropic_passthrough.py +++ b/tests/pass_through_tests/test_anthropic_passthrough.py @@ -296,3 +296,123 @@ async def test_anthropic_streaming_with_headers(): assert log_entry["end_user"] == "test-user-1" assert log_entry["custom_llm_provider"] == "anthropic" + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_anthropic_messages_streaming_cost_injection(): + """ + Test that cost is injected into message_delta usage for Anthropic Messages API streaming + """ + print("Testing cost injection in Anthropic Messages API streaming response") + + headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + + payload = { + "model": "claude-3-7-sonnet-20250219", + "max_tokens": 10, + "stream": True, + "messages": [{"role": "user", "content": "Say 'Hi'"}], + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "http://0.0.0.0:4000/v1/messages", + json=payload, + headers=headers + ) as response: + assert response.status == 200 + + # Collect all SSE events + events = [] + async for line in response.content: + line_str = line.decode('utf-8').strip() + if line_str.startswith('data: '): + try: + data = json.loads(line_str[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + + # Find message_delta event with usage + message_delta_events = [ + event for event in events + if event.get('type') == 'message_delta' and 'usage' in event + ] + + assert len(message_delta_events) > 0, "No message_delta events with usage found" + + # Check that cost is included in usage + for event in message_delta_events: + usage = event.get('usage', {}) + assert 'cost' in usage, f"Cost not found in usage: {usage}" + assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" + assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" + + print(f"✅ Found message_delta with cost: {usage}") + + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_anthropic_messages_openai_model_streaming_cost_injection(): + """ + Test that cost is injected into message_delta usage for OpenAI model via Anthropic Messages API + """ + print("Testing cost injection in Anthropic Messages API with OpenAI model") + + headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + + payload = { + "model": "openai/gpt-4o", + "max_tokens": 10, + "stream": True, + "messages": [{"role": "user", "content": "Say 'Hi'"}], + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "http://0.0.0.0:4000/v1/messages", + json=payload, + headers=headers + ) as response: + assert response.status == 200 + + # Collect all SSE events + events = [] + async for line in response.content: + line_str = line.decode('utf-8').strip() + if line_str.startswith('data: '): + try: + data = json.loads(line_str[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + + # Find message_delta event with usage + message_delta_events = [ + event for event in events + if event.get('type') == 'message_delta' and 'usage' in event + ] + + assert len(message_delta_events) > 0, "No message_delta events with usage found" + + # Check that cost is included in usage + for event in message_delta_events: + usage = event.get('usage', {}) + assert 'cost' in usage, f"Cost not found in usage: {usage}" + assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" + assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" + + print(f"✅ Found message_delta with cost: {usage}") + + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") From 56e429e33d036548e75ce4ded5e9a782c97b9684 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 1 Oct 2025 23:38:53 +0530 Subject: [PATCH 2/4] refactor code for better handling cost --- litellm/proxy/common_request_processing.py | 198 +++++++++++++-------- 1 file changed, 119 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e7bad1e19d..091077b6e7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -46,6 +46,7 @@ if TYPE_CHECKING: else: ProxyConfig = Any from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: @@ -760,85 +761,8 @@ class ProxyBaseLLMRequestProcessing: str_so_far += response_str # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider - # Handle both dict SSE events and pre-formatted string SSE lines - if getattr(litellm, "include_cost_in_streaming_usage", False) is True: - try: - def _inject_cost_into_usage_dict(obj: dict) -> Optional[dict]: - if ( - obj.get("type") == "message_delta" - and isinstance(obj.get("usage"), dict) - ): - _usage = obj["usage"] - prompt_tokens = int(_usage.get("input_tokens", 0) or 0) - completion_tokens = int(_usage.get("output_tokens", 0) or 0) - total_tokens = int( - _usage.get("total_tokens", prompt_tokens + completion_tokens) - or (prompt_tokens + completion_tokens) - ) - - _mr = ModelResponse( - usage=Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=total_tokens, - ) - ) - model_name = request_data.get("model", "") - try: - cost_val = litellm.completion_cost( - completion_response=_mr, - model=model_name, - ) - except Exception: - cost_val = None - if cost_val is not None: - obj.setdefault("usage", {})["cost"] = cost_val - return obj - return None - - def _inject_cost_into_sse_frame_str(frame_str: str) -> Optional[str]: - # frame_str may contain multiple lines like 'event: ...\ndata: {...}\n\n' - # We only modify the JSON in the 'data:' line - try: - # Split preserving lines - lines = frame_str.split("\n") - for idx, ln in enumerate(lines): - stripped_ln = ln.strip() - if stripped_ln.startswith("data:"): - json_part = stripped_ln.split("data:", 1)[1].strip() - if json_part and json_part != "[DONE]": - obj = json.loads(json_part) - maybe_modified = _inject_cost_into_usage_dict(obj) - if maybe_modified is not None: - # Replace just this line with updated JSON using safe_dumps - lines[idx] = f"data: {safe_dumps(maybe_modified)}" - return "\n".join(lines) - return None - except Exception: - return None - - if isinstance(chunk, dict): - maybe_modified = _inject_cost_into_usage_dict(chunk) - if maybe_modified is not None: - chunk = maybe_modified - elif isinstance(chunk, (bytes, bytearray)): - # Decode to str, inject, and rebuild as bytes - try: - s = chunk.decode("utf-8", errors="ignore") - maybe_mod = _inject_cost_into_sse_frame_str(s) - if maybe_mod is not None: - chunk = (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") - except Exception: - pass - elif isinstance(chunk, str): - # Try to parse SSE frame and inject cost into the data line - maybe_mod = _inject_cost_into_sse_frame_str(chunk) - if maybe_mod is not None: - # Ensure trailing frame separator - chunk = maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") - except Exception: - # Never break streaming on optional cost injection - pass + model_name = request_data.get("model", "") + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) # Format chunk using helper function yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) @@ -871,3 +795,119 @@ class ProxyBaseLLMRequestProcessing: ) error_returned = json.dumps({"error": proxy_exception.to_dict()}) yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n" + + @staticmethod + def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: + """ + Process a streaming chunk and inject cost information if enabled. + + Args: + chunk: The streaming chunk (dict, str, bytes, or bytearray) + model_name: Model name for cost calculation + + Returns: + The processed chunk with cost information injected if applicable + """ + if not getattr(litellm, "include_cost_in_streaming_usage", False): + return chunk + + try: + if isinstance(chunk, dict): + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + if maybe_modified is not None: + return maybe_modified + elif isinstance(chunk, (bytes, bytearray)): + # Decode to str, inject, and rebuild as bytes + try: + s = chunk.decode("utf-8", errors="ignore") + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + if maybe_mod is not None: + return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + except Exception: + pass + elif isinstance(chunk, str): + # Try to parse SSE frame and inject cost into the data line + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + if maybe_mod is not None: + # Ensure trailing frame separator + return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") + except Exception: + # Never break streaming on optional cost injection + pass + + return chunk + + @staticmethod + def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> Optional[str]: + """ + Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. + + Args: + frame_str: SSE frame string that may contain multiple lines + model_name: Model name for cost calculation + + Returns: + Modified SSE frame string with cost injected, or None if no modification needed + """ + try: + # Split preserving lines + lines = frame_str.split("\n") + for idx, ln in enumerate(lines): + stripped_ln = ln.strip() + if stripped_ln.startswith("data:"): + json_part = stripped_ln.split("data:", 1)[1].strip() + if json_part and json_part != "[DONE]": + obj = json.loads(json_part) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) + return None + except Exception: + return None + + @staticmethod + def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]: + """ + Inject cost information into a usage dictionary for message_delta events. + + Args: + obj: Dictionary containing the SSE event data + model_name: Model name for cost calculation + + Returns: + Modified dictionary with cost injected, or None if no modification needed + """ + if ( + obj.get("type") == "message_delta" + and isinstance(obj.get("usage"), dict) + ): + _usage = obj["usage"] + prompt_tokens = int(_usage.get("input_tokens", 0) or 0) + completion_tokens = int(_usage.get("output_tokens", 0) or 0) + total_tokens = int( + _usage.get("total_tokens", prompt_tokens + completion_tokens) + or (prompt_tokens + completion_tokens) + ) + + _mr = ModelResponse( + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + ) + + try: + cost_val = litellm.completion_cost( + completion_response=_mr, + model=model_name, + ) + except Exception: + cost_val = None + + if cost_val is not None: + obj.setdefault("usage", {})["cost"] = cost_val + return obj + return None \ No newline at end of file From fb664b0f762fa01f261dca1df8743327ea9b1697 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 2 Oct 2025 16:13:03 +0530 Subject: [PATCH 3/4] add other fields in cost calculations --- litellm/proxy/common_request_processing.py | 32 ++++++++++++++++++---- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 091077b6e7..db9b895031 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -891,12 +891,34 @@ class ProxyBaseLLMRequestProcessing: or (prompt_tokens + completion_tokens) ) + # Extract additional usage fields + cache_creation_input_tokens = _usage.get("cache_creation_input_tokens") + cache_read_input_tokens = _usage.get("cache_read_input_tokens") + web_search_requests = _usage.get("web_search_requests") + completion_tokens_details = _usage.get("completion_tokens_details") + prompt_tokens_details = _usage.get("prompt_tokens_details") + + # Build usage kwargs with only non-None values + usage_kwargs = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + } + + # Add optional fields if they exist + if cache_creation_input_tokens is not None: + usage_kwargs["cache_creation_input_tokens"] = cache_creation_input_tokens + if cache_read_input_tokens is not None: + usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens + if web_search_requests is not None: + usage_kwargs["web_search_requests"] = web_search_requests + if completion_tokens_details is not None: + usage_kwargs["completion_tokens_details"] = completion_tokens_details + if prompt_tokens_details is not None: + usage_kwargs["prompt_tokens_details"] = prompt_tokens_details + _mr = ModelResponse( - usage=Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=total_tokens, - ) + usage=Usage(**usage_kwargs) ) try: From 507b0973b4ec8049dce71cc6d27f14a1c4491f15 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 3 Oct 2025 08:21:07 +0530 Subject: [PATCH 4/4] fix lint code --- litellm/proxy/common_request_processing.py | 1 - litellm/proxy/hooks/parallel_request_limiter_v3.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index db9b895031..4a08b05e86 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -735,7 +735,6 @@ class ProxyBaseLLMRequestProcessing: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events """ - from litellm.types.utils import ModelResponse, ModelResponseStream, Usage verbose_proxy_logger.debug("inside generator") try: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 0a49d7f675..9f9b49dcb6 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -25,6 +25,7 @@ from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject +from fastapi import HTTPException if TYPE_CHECKING: from opentelemetry.trace import Span as _Span