Merge pull request #15102 from BerriAI/litellm_message_api_cost_tracking

(Feat) Add cost tracking for /v1/messages

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

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

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

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")