diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f07a61c544..4a08b05e86 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]: @@ -734,7 +735,6 @@ class ProxyBaseLLMRequestProcessing: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events """ - from litellm.types.utils import ModelResponse, ModelResponseStream verbose_proxy_logger.debug("inside generator") try: @@ -759,6 +759,10 @@ 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 + 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) except Exception as e: @@ -790,3 +794,141 @@ 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) + ) + + # 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(**usage_kwargs) + ) + + 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 diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 2e3057cdba..5fca0d1f90 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -27,6 +27,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 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")