Skip to content

TokenUsageExtractor drops mapping-based usage metadata #1432

Description

@KXHXK

Description

TokenUsageExtractor.extract_from_response() silently drops all token counts when a provider exposes usage or usage_metadata as a mapping instead of an attribute-based object.

The shared streaming wrapper explicitly accepts chunks with usage_metadata, and the existing test fixture uses a dictionary shape, but _extract_from_usage_object() reads every field with getattr(). As a result, dictionary-backed usage produces a TokenUsage whose fields are all None, so the span receives no token usage attributes.

Minimal reproduction

from types import SimpleNamespace

from agentops.instrumentation.common.token_counting import TokenUsageExtractor

response = SimpleNamespace(
    usage_metadata={
        prompt_tokens: 10,
        completion_tokens: 5,
        total_tokens: 15,
    }
)

print(TokenUsageExtractor.extract_from_response(response))

Current output on main:

TokenUsage(prompt_tokens=None, completion_tokens=None, total_tokens=None, cached_prompt_tokens=None, cached_read_tokens=None, reasoning_tokens=None)

Expected behavior

Dictionary and attribute-based usage containers should be normalized consistently, including the cache and reasoning token fields already supported by TokenUsage.

Suggested fix

Use a small mapping-aware field accessor in _extract_from_usage_object() and add regression coverage for both usage and usage_metadata mappings. The existing object behavior should remain unchanged.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions