fix(instrumentation): support mapping-based usage metadata in TokenUsageExtractor

agentops/instrumentation/common/token_counting.py

@@ -4,7 +4,7 @@
 information from various response formats.
 """
 
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Mapping
 from dataclasses import dataclass
 
 from agentops.logging import logger
@@ -56,18 +56,32 @@ class TokenUsageExtractor:
 
     @staticmethod
     def extract_from_response(response: Any) -> TokenUsage:
-        """Extract token usage from a generic response object.
+        """Extract token usage from a generic response object or mapping.
 
         Handles various response formats from different providers.
         """
         usage = TokenUsage()
+        if response is None:
+            return usage
+
+        # Try mapping-based response
+        if isinstance(response, Mapping):
+            if "usage" in response:
+                usage = TokenUsageExtractor._extract_from_usage_object(response["usage"])
+            elif "usage_metadata" in response:
+                usage = TokenUsageExtractor._extract_from_usage_object(response["usage_metadata"])
+            elif "token_usage" in response:
+                usage = TokenUsageExtractor._extract_from_crewai_format(response["token_usage"])
+            elif any(k in response for k in ("prompt_tokens", "completion_tokens", "input_tokens", "output_tokens")):
+                usage = TokenUsageExtractor._extract_from_attributes(response)
+            return usage
 
         # Try direct usage attribute
         if hasattr(response, "usage"):
             usage_data = response.usage
             usage = TokenUsageExtractor._extract_from_usage_object(usage_data)
 
-        # Try usage_metadata (Anthropic style)
+        # Try usage_metadata (Anthropic / Gemini style)
         elif hasattr(response, "usage_metadata"):
             usage_data = response.usage_metadata
             usage = TokenUsageExtractor._extract_from_usage_object(usage_data)
@@ -77,34 +91,59 @@ def extract_from_response(response: Any) -> TokenUsage:
             usage = TokenUsageExtractor._extract_from_crewai_format(response.token_usage)
 
         # Try direct attributes on response
-        elif hasattr(response, "prompt_tokens") or hasattr(response, "completion_tokens"):
+        elif any(hasattr(response, attr) for attr in ("prompt_tokens", "completion_tokens", "input_tokens", "output_tokens")):
             usage = TokenUsageExtractor._extract_from_attributes(response)
 
         return usage
 
     @staticmethod
     def _extract_from_usage_object(usage_data: Any) -> TokenUsage:
-        """Extract from a usage object with standard attributes."""
+        """Extract from a usage object or mapping with standard attributes."""
         if not usage_data:
             return TokenUsage()
 
+        def _get_val(*keys: str) -> Optional[int]:
+            for key in keys:
+                val = None
+                if isinstance(usage_data, Mapping):
+                    val = usage_data.get(key)
+                elif hasattr(usage_data, key):
+                    val = getattr(usage_data, key)
+                if val is not None:
+                    try:
+                        return int(val)
+                    except (ValueError, TypeError):
+                        pass
+            return None
+
+        prompt_tokens = _get_val("prompt_tokens", "input_tokens")
+        completion_tokens = _get_val("completion_tokens", "output_tokens")
+        total_tokens = _get_val("total_tokens")
+        if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
+            total_tokens = prompt_tokens + completion_tokens
+
         return TokenUsage(
-            prompt_tokens=getattr(usage_data, "prompt_tokens", None),
-            completion_tokens=getattr(usage_data, "completion_tokens", None),
-            total_tokens=getattr(usage_data, "total_tokens", None),
-            cached_prompt_tokens=getattr(usage_data, "cached_prompt_tokens", None),
-            cached_read_tokens=getattr(usage_data, "cache_read_input_tokens", None),
-            reasoning_tokens=getattr(usage_data, "reasoning_tokens", None),
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            cached_prompt_tokens=_get_val("cached_prompt_tokens", "cache_creation_input_tokens"),
+            cached_read_tokens=_get_val("cache_read_input_tokens", "cached_read_tokens"),
+            reasoning_tokens=_get_val("reasoning_tokens"),
         )
 
     @staticmethod
-    def _extract_from_crewai_format(token_usage_str: str) -> TokenUsage:
-        """Extract from CrewAI's string format (e.g., 'prompt_tokens=100 completion_tokens=50')."""
+    def _extract_from_crewai_format(token_usage_input: Any) -> TokenUsage:
+        """Extract from CrewAI's format (dict or string format like 'prompt_tokens=100 completion_tokens=50')."""
         usage = TokenUsage()
+        if not token_usage_input:
+            return usage
+
+        if isinstance(token_usage_input, Mapping):
+            return TokenUsageExtractor._extract_from_usage_object(token_usage_input)
 
         try:
             metrics = {}
-            for item in str(token_usage_str).split():
+            for item in str(token_usage_input).split():
                 if "=" in item:
                     key, value = item.split("=")
                     try:
@@ -115,6 +154,8 @@ def _extract_from_crewai_format(token_usage_str: str) -> TokenUsage:
             usage.prompt_tokens = metrics.get("prompt_tokens")
             usage.completion_tokens = metrics.get("completion_tokens")
             usage.total_tokens = metrics.get("total_tokens")
+            if usage.total_tokens is None and usage.prompt_tokens is not None and usage.completion_tokens is not None:
+                usage.total_tokens = usage.prompt_tokens + usage.completion_tokens
             usage.cached_prompt_tokens = metrics.get("cached_prompt_tokens")
 
         except Exception as e:
@@ -124,11 +165,31 @@ def _extract_from_crewai_format(token_usage_str: str) -> TokenUsage:
 
     @staticmethod
     def _extract_from_attributes(response: Any) -> TokenUsage:
-        """Extract from direct attributes on the response."""
+        """Extract from direct attributes or mapping keys on the response."""
+        def _get_val(*keys: str) -> Optional[int]:
+            for key in keys:
+                val = None
+                if isinstance(response, Mapping):
+                    val = response.get(key)
+                elif hasattr(response, key):
+                    val = getattr(response, key)
+                if val is not None:
+                    try:
+                        return int(val)
+                    except (ValueError, TypeError):
+                        pass
+            return None
+
+        prompt_tokens = _get_val("prompt_tokens", "input_tokens")
+        completion_tokens = _get_val("completion_tokens", "output_tokens")
+        total_tokens = _get_val("total_tokens")
+        if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
+            total_tokens = prompt_tokens + completion_tokens
+
         return TokenUsage(
-            prompt_tokens=getattr(response, "prompt_tokens", None),
-            completion_tokens=getattr(response, "completion_tokens", None),
-            total_tokens=getattr(response, "total_tokens", None),
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
         )
 
 

tests/unit/instrumentation/common/test_token_counting.py

@@ -1,4 +1,5 @@
-from agentops.instrumentation.common.token_counting import TokenUsage
+from types import SimpleNamespace
+from agentops.instrumentation.common.token_counting import TokenUsage, TokenUsageExtractor
 from agentops.semconv import SpanAttributes
 
 
@@ -24,3 +25,96 @@ def test_includes_positive_values_only(self):
         assert SpanAttributes.LLM_USAGE_COMPLETION_TOKENS not in attrs
         assert SpanAttributes.LLM_USAGE_TOTAL_TOKENS in attrs
         assert attrs[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 5
+
+
+class TestTokenUsageExtractor:
+    def test_extract_from_object_usage(self):
+        usage_obj = SimpleNamespace(
+            prompt_tokens=10,
+            completion_tokens=20,
+            total_tokens=30,
+            cached_prompt_tokens=2,
+            cache_read_input_tokens=3,
+            reasoning_tokens=5,
+        )
+        response = SimpleNamespace(usage=usage_obj)
+        usage = TokenUsageExtractor.extract_from_response(response)
+
+        assert usage.prompt_tokens == 10
+        assert usage.completion_tokens == 20
+        assert usage.total_tokens == 30
+        assert usage.cached_prompt_tokens == 2
+        assert usage.cached_read_tokens == 3
+        assert usage.reasoning_tokens == 5
+
+    def test_extract_from_mapping_usage_metadata(self):
+        response = SimpleNamespace(
+            usage_metadata={
+                "prompt_tokens": 10,
+                "completion_tokens": 5,
+                "total_tokens": 15,
+                "cached_prompt_tokens": 4,
+                "cache_read_input_tokens": 6,
+                "reasoning_tokens": 2,
+            }
+        )
+        usage = TokenUsageExtractor.extract_from_response(response)
+
+        assert usage.prompt_tokens == 10
+        assert usage.completion_tokens == 5
+        assert usage.total_tokens == 15
+        assert usage.cached_prompt_tokens == 4
+        assert usage.cached_read_tokens == 6
+        assert usage.reasoning_tokens == 2
+
+    def test_extract_from_dict_response(self):
+        response = {
+            "usage": {
+                "prompt_tokens": 12,
+                "completion_tokens": 8,
+                "total_tokens": 20,
+            }
+        }
+        usage = TokenUsageExtractor.extract_from_response(response)
+        assert usage.prompt_tokens == 12
+        assert usage.completion_tokens == 8
+        assert usage.total_tokens == 20
+
+    def test_extract_otel_input_output_tokens(self):
+        response = SimpleNamespace(
+            usage_metadata={
+                "input_tokens": 42,
+                "output_tokens": 18,
+            }
+        )
+        usage = TokenUsageExtractor.extract_from_response(response)
+        assert usage.prompt_tokens == 42
+        assert usage.completion_tokens == 18
+        assert usage.total_tokens == 60
+
+    def test_extract_from_crewai_dict(self):
+        response = SimpleNamespace(
+            token_usage={
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+            }
+        )
+        usage = TokenUsageExtractor.extract_from_response(response)
+        assert usage.prompt_tokens == 100
+        assert usage.completion_tokens == 50
+        assert usage.total_tokens == 150
+
+    def test_extract_from_crewai_string(self):
+        response = SimpleNamespace(
+            token_usage="prompt_tokens=100 completion_tokens=50"
+        )
+        usage = TokenUsageExtractor.extract_from_response(response)
+        assert usage.prompt_tokens == 100
+        assert usage.completion_tokens == 50
+        assert usage.total_tokens == 150
+
+    def test_extract_none_response(self):
+        usage = TokenUsageExtractor.extract_from_response(None)
+        assert usage.prompt_tokens is None
+        assert usage.completion_tokens is None
+        assert usage.total_tokens is None