diff --git a/agentops/instrumentation/agentic/google_adk/patch.py b/agentops/instrumentation/agentic/google_adk/patch.py index 172c41ab7..7c4aa0b5e 100644 --- a/agentops/instrumentation/agentic/google_adk/patch.py +++ b/agentops/instrumentation/agentic/google_adk/patch.py @@ -265,26 +265,47 @@ def _extract_llm_attributes(llm_request_dict: dict, llm_response: Any) -> dict: attributes[SpanAttributes.LLM_RESPONSE_MODEL] = response_dict["model"] # Usage metadata - if "usage_metadata" in response_dict: - usage = response_dict["usage_metadata"] - if "prompt_token_count" in usage: - attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = usage["prompt_token_count"] - if "candidates_token_count" in usage: - attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = usage["candidates_token_count"] - if "total_token_count" in usage: - attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = usage["total_token_count"] + usage = response_dict.get("usage_metadata") or response_dict.get("usage") + if usage and isinstance(usage, dict): + prompt_tokens = ( + usage.get("prompt_token_count") + or usage.get("input_token_count") + or usage.get("input_tokens") + or usage.get("prompt_tokens") + ) + completion_tokens = ( + usage.get("candidates_token_count") + or usage.get("output_token_count") + or usage.get("output_tokens") + or usage.get("completion_tokens") + ) + total_tokens = ( + usage.get("total_token_count") + or usage.get("total_tokens") + ) + + if prompt_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = prompt_tokens + if completion_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = completion_tokens + if total_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = total_tokens + elif prompt_tokens is not None and completion_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = prompt_tokens + completion_tokens # Additional token details if available - if "prompt_tokens_details" in usage: - for detail in usage["prompt_tokens_details"]: - if "modality" in detail and "token_count" in detail: + prompt_details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") + if prompt_details and isinstance(prompt_details, list): + for detail in prompt_details: + if isinstance(detail, dict) and "modality" in detail and "token_count" in detail: attributes[f"gen_ai.usage.prompt_tokens.{detail['modality'].lower()}"] = detail[ "token_count" ] - if "candidates_tokens_details" in usage: - for detail in usage["candidates_tokens_details"]: - if "modality" in detail and "token_count" in detail: + candidates_details = usage.get("candidates_tokens_details") or usage.get("output_tokens_details") + if candidates_details and isinstance(candidates_details, list): + for detail in candidates_details: + if isinstance(detail, dict) and "modality" in detail and "token_count" in detail: attributes[f"gen_ai.usage.completion_tokens.{detail['modality'].lower()}"] = detail[ "token_count" ] diff --git a/agentops/instrumentation/providers/google_genai/attributes/model.py b/agentops/instrumentation/providers/google_genai/attributes/model.py index 75b5eb20c..b584263db 100644 --- a/agentops/instrumentation/providers/google_genai/attributes/model.py +++ b/agentops/instrumentation/providers/google_genai/attributes/model.py @@ -139,14 +139,45 @@ def _set_response_attributes(attributes: AttributeMap, response: Any) -> None: if hasattr(response, "model"): attributes[SpanAttributes.LLM_RESPONSE_MODEL] = response.model - if hasattr(response, "usage_metadata"): - usage = response.usage_metadata - if hasattr(usage, "prompt_token_count"): - attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = usage.prompt_token_count - if hasattr(usage, "candidates_token_count"): - attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = usage.candidates_token_count - if hasattr(usage, "total_token_count"): - attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = usage.total_token_count + usage = getattr(response, "usage_metadata", None) or getattr(response, "usage", None) + if usage is not None: + if isinstance(usage, dict): + prompt_tokens = ( + usage.get("prompt_token_count") + or usage.get("input_token_count") + or usage.get("input_tokens") + or usage.get("prompt_tokens") + ) + completion_tokens = ( + usage.get("candidates_token_count") + or usage.get("output_token_count") + or usage.get("output_tokens") + or usage.get("completion_tokens") + ) + total_tokens = usage.get("total_token_count") or usage.get("total_tokens") + else: + prompt_tokens = ( + getattr(usage, "prompt_token_count", None) + or getattr(usage, "input_token_count", None) + or getattr(usage, "input_tokens", None) + or getattr(usage, "prompt_tokens", None) + ) + completion_tokens = ( + getattr(usage, "candidates_token_count", None) + or getattr(usage, "output_token_count", None) + or getattr(usage, "output_tokens", None) + or getattr(usage, "completion_tokens", None) + ) + total_tokens = getattr(usage, "total_token_count", None) or getattr(usage, "total_tokens", None) + + if prompt_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = prompt_tokens + if completion_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = completion_tokens + if total_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = total_tokens + elif prompt_tokens is not None and completion_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = prompt_tokens + completion_tokens try: if hasattr(response, "text"): diff --git a/tests/unit/instrumentation/test_google_token_usage.py b/tests/unit/instrumentation/test_google_token_usage.py new file mode 100644 index 000000000..281a9e9c5 --- /dev/null +++ b/tests/unit/instrumentation/test_google_token_usage.py @@ -0,0 +1,85 @@ +import json +from unittest.mock import MagicMock +import pytest + +from agentops.instrumentation.agentic.google_adk.patch import _extract_llm_attributes +from agentops.instrumentation.providers.google_genai.attributes.model import ( + _set_response_attributes, +) +from agentops.semconv import SpanAttributes + + +class TestGoogleTokenUsageExtraction: + """Test suite verifying token usage normalization across Google ADK and GenAI.""" + + def test_extract_llm_attributes_with_adk_token_names(self): + """Verify _extract_llm_attributes correctly normalizes input_tokens and output_tokens from ADK.""" + llm_request_dict = {"model": "gemini-2.5-flash-lite"} + llm_response = { + "model": "gemini-2.5-flash-lite", + "usage_metadata": { + "input_tokens": 142, + "output_tokens": 58, + }, + } + + attrs = _extract_llm_attributes(llm_request_dict, json.dumps(llm_response)) + + assert attrs[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 142 + assert attrs[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 58 + assert attrs[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 200 + + def test_extract_llm_attributes_with_standard_gemini_token_counts(self): + """Verify _extract_llm_attributes handles prompt_token_count and candidates_token_count.""" + llm_request_dict = {"model": "gemini-2.5-pro"} + llm_response = { + "model": "gemini-2.5-pro", + "usage_metadata": { + "prompt_token_count": 500, + "candidates_token_count": 120, + "total_token_count": 620, + }, + } + + attrs = _extract_llm_attributes(llm_request_dict, json.dumps(llm_response)) + + assert attrs[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 500 + assert attrs[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 120 + assert attrs[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 620 + + def test_set_response_attributes_dict(self): + """Verify _set_response_attributes with dictionary usage metadata.""" + attributes = {} + response = MagicMock() + response.model = "gemini-2.5-flash" + response.usage_metadata = { + "input_token_count": 350, + "output_token_count": 80, + } + del response.text + del response.candidates + + _set_response_attributes(attributes, response) + + assert attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 350 + assert attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 80 + assert attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 430 + + def test_set_response_attributes_object(self): + """Verify _set_response_attributes with object usage metadata.""" + attributes = {} + response = MagicMock() + response.model = "gemini-2.5-flash" + usage_obj = MagicMock(spec=["prompt_token_count", "candidates_token_count", "total_token_count"]) + usage_obj.prompt_token_count = 210 + usage_obj.candidates_token_count = 95 + usage_obj.total_token_count = 305 + response.usage_metadata = usage_obj + del response.text + del response.candidates + + _set_response_attributes(attributes, response) + + assert attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 210 + assert attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 95 + assert attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 305