diff --git a/ddtrace/llmobs/_integrations/google_utils.py b/ddtrace/llmobs/_integrations/google_utils.py index 76c12daeac8..43542cbf67e 100644 --- a/ddtrace/llmobs/_integrations/google_utils.py +++ b/ddtrace/llmobs/_integrations/google_utils.py @@ -141,7 +141,9 @@ def extract_generation_metrics_google_genai(response) -> Dict[str, Any]: output_tokens = None cached_tokens = _get_attr(usage_metadata, "cached_content_token_count", None) - total_tokens = _get_attr(usage_metadata, "total_token_count", None) or input_tokens + output_tokens + total_tokens = _get_attr(usage_metadata, "total_token_count", None) + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens if input_tokens is not None: usage[INPUT_TOKENS_METRIC_KEY] = input_tokens diff --git a/tests/llmobs/test_google_utils.py b/tests/llmobs/test_google_utils.py new file mode 100644 index 00000000000..cca6749c3db --- /dev/null +++ b/tests/llmobs/test_google_utils.py @@ -0,0 +1,40 @@ +from types import SimpleNamespace + +from ddtrace.llmobs._constants import INPUT_TOKENS_METRIC_KEY +from ddtrace.llmobs._constants import OUTPUT_TOKENS_METRIC_KEY +from ddtrace.llmobs._constants import TOTAL_TOKENS_METRIC_KEY +from ddtrace.llmobs._integrations.google_utils import extract_generation_metrics_google_genai + + +def test_extract_generation_metrics_google_genai_handles_missing_token_counts(): + response = SimpleNamespace( + usage_metadata={ + "prompt_token_count": None, + "candidates_token_count": None, + "thoughts_token_count": None, + "total_token_count": None, + } + ) + + metrics = extract_generation_metrics_google_genai(response) + + assert metrics == {} + + +def test_extract_generation_metrics_google_genai_falls_back_to_sum_when_available(): + response = SimpleNamespace( + usage_metadata={ + "prompt_token_count": 6, + "candidates_token_count": 4, + "thoughts_token_count": None, + "total_token_count": None, + } + ) + + metrics = extract_generation_metrics_google_genai(response) + + assert metrics == { + INPUT_TOKENS_METRIC_KEY: 6, + OUTPUT_TOKENS_METRIC_KEY: 4, + TOTAL_TOKENS_METRIC_KEY: 10, + }