Centralized representation and processing of GenAI token usage metadata.
| 36 | |
| 37 | @dataclasses.dataclass |
| 38 | class TokenUsage: |
| 39 | """Centralized representation and processing of GenAI token usage metadata.""" |
| 40 | |
| 41 | usage_metadata: types.GenerateContentResponseUsageMetadata | None |
| 42 | |
| 43 | @property |
| 44 | def input_token_count(self) -> int | None: |
| 45 | if self.usage_metadata is None: |
| 46 | return None |
| 47 | # OTel semconv for `gen_ai.client.token.usage` states that token counts should |
| 48 | # be categorized under `gen_ai.token.type` as either "input" or "output". |
| 49 | # We aggregate prompt and tool use tokens for "input". |
| 50 | prompt_tokens = self.usage_metadata.prompt_token_count |
| 51 | tool_tokens = self.usage_metadata.tool_use_prompt_token_count |
| 52 | if prompt_tokens is None and tool_tokens is None: |
| 53 | return None |
| 54 | return (prompt_tokens or 0) + (tool_tokens or 0) |
| 55 | |
| 56 | @property |
| 57 | def output_token_count(self) -> int | None: |
| 58 | if self.usage_metadata is None: |
| 59 | return None |
| 60 | # According to OpenTelemetry Semantic Conventions: |
| 61 | # https://github.com/open-telemetry/semantic-conventions/blob/v1.41.0/docs/registry/attributes/gen-ai.md |
| 62 | # gen_ai.usage.reasoning.output_tokens (thoughts_token_count) SHOULD be included in gen_ai.usage.output_tokens. |
| 63 | candidates_tokens = self.usage_metadata.candidates_token_count |
| 64 | thoughts_tokens = self.usage_metadata.thoughts_token_count |
| 65 | if candidates_tokens is None and thoughts_tokens is None: |
| 66 | return None |
| 67 | return (candidates_tokens or 0) + (thoughts_tokens or 0) |
| 68 | |
| 69 | def to_attributes(self) -> dict[str, AttributeValue]: |
| 70 | """Returns a dictionary of OpenTelemetry token usage attributes.""" |
| 71 | attrs: dict[str, AttributeValue] = {} |
| 72 | if self.input_token_count is not None: |
| 73 | attrs[GEN_AI_USAGE_INPUT_TOKENS] = self.input_token_count |
| 74 | if self.output_token_count is not None: |
| 75 | attrs[GEN_AI_USAGE_OUTPUT_TOKENS] = self.output_token_count |
| 76 | |
| 77 | if self.usage_metadata is not None: |
| 78 | cached_tokens = self.usage_metadata.cached_content_token_count |
| 79 | if cached_tokens is not None: |
| 80 | attrs[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] = cached_tokens |
| 81 | |
| 82 | thoughts_tokens = self.usage_metadata.thoughts_token_count |
| 83 | if thoughts_tokens is not None: |
| 84 | attrs[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] = thoughts_tokens |
| 85 | |
| 86 | system_instruction_tokens = getattr( |
| 87 | self.usage_metadata, 'system_instruction_tokens', None |
| 88 | ) |
| 89 | if system_instruction_tokens is not None: |
| 90 | attrs['gen_ai.usage.experimental.system_instruction_tokens'] = ( |
| 91 | system_instruction_tokens |
| 92 | ) |
| 93 | |
| 94 | return attrs |
no outgoing calls
no test coverage detected