diff --git a/tests/test_llm_attributes_extractors.py b/tests/test_llm_attributes_extractors.py index 841c09d45..e97480b8e 100644 --- a/tests/test_llm_attributes_extractors.py +++ b/tests/test_llm_attributes_extractors.py @@ -14,14 +14,16 @@ import json from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, call from google.adk.models.llm_request import LlmRequest from google.adk.tools.function_tool import FunctionTool from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extractors import ( llm_gen_ai_request_functions, + llm_gen_ai_usage_output_tokens, ) +from veadk.tracing.telemetry.attributes.extractors.types import ExtractorResponse def test_request_functions_reuses_adk_request_declaration(monkeypatch): @@ -99,3 +101,35 @@ def test_request_functions_builds_missing_declaration_once(): ) assert parameters == {"type": "object"} get_declaration.assert_called_once_with() + + +def test_missing_output_token_count_is_not_written_to_span(): + params = SimpleNamespace( + llm_response=SimpleNamespace( + usage_metadata=SimpleNamespace(candidates_token_count=None) + ) + ) + response = llm_gen_ai_usage_output_tokens(params) + span = Mock() + + ExtractorResponse.update_span(span, "gen_ai.usage.output_tokens", response) + + span.set_attribute.assert_not_called() + + +def test_falsy_attribute_values_are_written_to_span(): + span = Mock() + + ExtractorResponse.update_span(span, "zero", ExtractorResponse(content=0)) + ExtractorResponse.update_span(span, "false", ExtractorResponse(content=False)) + + assert span.set_attribute.call_args_list == [call("zero", 0), call("false", False)] + + +def test_none_values_in_attribute_mappings_are_not_written_to_span(): + span = Mock() + response = ExtractorResponse(content=[{"present": 0, "missing": None}]) + + ExtractorResponse.update_span(span, "unused", response) + + span.set_attribute.assert_called_once_with("present", 0) diff --git a/veadk/tracing/telemetry/attributes/extractors/types.py b/veadk/tracing/telemetry/attributes/extractors/types.py index cb3d9b85c..a16bb5d64 100644 --- a/veadk/tracing/telemetry/attributes/extractors/types.py +++ b/veadk/tracing/telemetry/attributes/extractors/types.py @@ -86,13 +86,15 @@ def update_span( """ if response.type == "attribute": res = response.content + if res is None: + return if isinstance(res, list): for _res in res: if isinstance(_res, dict): for k, v in _res.items(): - span.set_attribute(k, v) + if v is not None: + span.set_attribute(k, v) else: - # set anyway span.set_attribute(attr_name, res) # type: ignore elif response.type == "event": if isinstance(response.content, dict):