Skip to content

fix(google): normalize token usage extraction across ADK and GenAI instrumentation #1465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
himanshu-xyz-1 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(google): normalize token usage extraction across ADK and GenAI instrumentation #1465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 35 additions & 14 deletions agentops/instrumentation/agentic/google_adk/patch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
]
Expand Down
47 changes: 39 additions & 8 deletions agentops/instrumentation/providers/google_genai/attributes/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down
85 changes: 85 additions & 0 deletions tests/unit/instrumentation/test_google_token_usage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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