From ed104c8f031a977e508d02dbdd99a4a202bbb345 Mon Sep 17 00:00:00 2001 From: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com> Date: Tue, 27 Oct 2026 00:37:02 +0530 Subject: [PATCH] fix(instrumentation): support mapping-based usage metadata and normalize OTel token fields in TokenUsageExtractor When providers expose usage or usage_metadata as a dictionary/mapping rather than an attribute-based object, TokenUsageExtractor._extract_from_usage_object previously attempted getattr() access on all fields and returned None. This update: - Adds mapping-aware field access to extract token counts from dicts or objects - Supports input_tokens and output_tokens alongside prompt_tokens and completion_tokens - Computes total_tokens when absent if prompt/input and completion/output tokens exist - Adds full regression coverage for mappings, dict responses, CrewAI formats, and OTel conventions --- .../instrumentation/common/token_counting.py | 97 +++++++++++++++---- .../common/test_token_counting.py | 96 +++++++++++++++++- 2 files changed, 174 insertions(+), 19 deletions(-) diff --git a/agentops/instrumentation/common/token_counting.py b/agentops/instrumentation/common/token_counting.py index 7d2f19f07..445f7eaa0 100644 --- a/agentops/instrumentation/common/token_counting.py +++ b/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, ) diff --git a/tests/unit/instrumentation/common/test_token_counting.py b/tests/unit/instrumentation/common/test_token_counting.py index 2f00561b9..a48e689c9 100644 --- a/tests/unit/instrumentation/common/test_token_counting.py +++ b/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