diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index d2ed7de92c1..0d564310c4a 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -464,17 +464,20 @@ def add_usage_details(usage1: UsageDetails | None, usage2: UsageDetails | None) combined = add_usage_details(usage1, usage2) # Result: {'input_token_count': 8, 'output_token_count': 16} """ - if usage1 is None: - return usage2 or UsageDetails() - if usage2 is None: - return usage1 + u1 = usage1 or UsageDetails() + u2 = usage2 or UsageDetails() result = UsageDetails() # Combine all keys from both dictionaries - all_keys = set(usage1.keys()) | set(usage2.keys()) + all_keys = set(u1.keys()) | set(u2.keys()) for key in all_keys: - if not isinstance((val1 := usage1.get(key, 0)), (int | None)) or not isinstance( - (val2 := usage2.get(key, 0)), (int | None) + val1 = u1.get(key, 0) + val2 = u2.get(key, 0) + if ( + isinstance(val1, bool) + or isinstance(val2, bool) + or not isinstance(val1, (int, type(None))) + or not isinstance(val2, (int, type(None))) ): logger.warning("Non `int` value found in usage details, skipping.") continue diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index df2f55efb63..a6797b30449 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -808,12 +808,43 @@ def test_usage_details_add_with_none_and_type_errors(): def test_usage_details_add_skips_non_int(): - u1 = cast(UsageDetails, {"input_token_count": 10, "other": "test"}) - u2 = cast(UsageDetails, {"input_token_count": 10, "another": "test"}) + u1 = cast(UsageDetails, {"input_token_count": 10, "other": "test", "flag": True}) + u2 = cast(UsageDetails, {"input_token_count": 10, "another": "test", "flag": 5}) u3 = add_usage_details(u1, u2) assert len(u3.keys()) == 1 assert "input_token_count" in u3 assert u3["input_token_count"] == 20 + assert "flag" not in u3 + + +def test_usage_details_add_with_none_returns_copy(): + u = UsageDetails(input_token_count=1) + v1 = add_usage_details(u, None) + v2 = add_usage_details(None, u) + assert v1 == u + assert v1 is not u + assert v2 == u + assert v2 is not u + + +def test_usage_details_add_with_none_filters_booleans_and_non_ints(): + payload = cast(UsageDetails, {"input_token_count": 5, "flag": True, "text": "skip_me"}) + v1 = add_usage_details(payload, None) + v2 = add_usage_details(None, payload) + assert v1 == {"input_token_count": 5} + assert v1 is not payload + assert "flag" not in v1 + assert "text" not in v1 + assert v2 == {"input_token_count": 5} + assert v2 is not payload + assert "flag" not in v2 + assert "text" not in v2 + + +def test_usage_details_add_both_none_returns_empty(): + result = add_usage_details(None, None) + assert result == {} + assert isinstance(result, dict) # region UserInputRequest and Response