feat: enhance LiteLLM metadata extraction by adding cost retrieval from response headers

Author: viniciusdsmello

src/openlayer/lib/integrations/litellm_tracer.py

@@ -661,6 +661,7 @@ def extract_usage_from_chunk(chunk: Any) -> Dict[str, Optional[int]]:
 def extract_litellm_metadata(response: Any, model_name: str) -> Dict[str, Any]:
     """Extract LiteLLM-specific metadata from response."""
     metadata = {}
+    response_headers = {}
 
     try:
         # Extract hidden parameters
@@ -689,6 +690,7 @@ def extract_litellm_metadata(response: Any, model_name: str) -> Dict[str, Any]:
             if 'additional_headers' in hidden_params:
                 headers = hidden_params['additional_headers']
                 if headers:
+                    response_headers = headers
                     metadata['response_headers'] = headers
 
         # Extract system fingerprint if available
@@ -697,9 +699,48 @@ def extract_litellm_metadata(response: Any, model_name: str) -> Dict[str, Any]:
 
         # Extract response headers if available
         if hasattr(response, '_response_headers'):
-            metadata['response_headers'] = dict(response._response_headers)
+            response_headers = dict(response._response_headers)
+            metadata['response_headers'] = response_headers
+        
+        # Fallback: Extract cost from x-litellm-response-cost header if cost is missing or zero
+        if not metadata.get('cost') and response_headers:
+            cost_from_header = _extract_cost_from_headers(response_headers)
+            if cost_from_header is not None:
+                metadata['cost'] = cost_from_header
 
     except Exception as e:
         logger.debug("Error extracting LiteLLM metadata: %s", e)
 
     return metadata
+
+
+def _extract_cost_from_headers(headers: Dict[str, Any]) -> Optional[float]:
+    """Extract cost from LiteLLM response headers."""
+    try:
+        # Try to get cost from x-litellm-response-cost header
+        cost_str = headers.get('x-litellm-response-cost')
+        if cost_str is not None:
+            # Handle string values (headers are often strings)
+            if isinstance(cost_str, str):
+                cost = float(cost_str)
+            else:
+                cost = float(cost_str)
+            
+            if cost > 0:
+                return cost
+        
+        # Fallback to x-litellm-response-cost-original if primary is zero/missing
+        cost_original_str = headers.get('x-litellm-response-cost-original')
+        if cost_original_str is not None:
+            if isinstance(cost_original_str, str):
+                cost = float(cost_original_str)
+            else:
+                cost = float(cost_original_str)
+            
+            if cost > 0:
+                return cost
+                
+    except (ValueError, TypeError) as e:
+        logger.debug("Error parsing cost from headers: %s", e)
+    
+    return None