kaxil commented on code in PR #74277:
URL: https://github.com/apache/airflow/pull/74277#discussion_r4190010150
##########
providers/common/ai/src/airflow/providers/common/ai/utils/logging.py:
##########
@@ -102,6 +102,8 @@ def format_usage_for_xcom(usage: RunUsage) -> dict[str,
Any]:
"output_tokens": usage.output_tokens,
"total_tokens": usage.total_tokens,
"tool_calls": usage.tool_calls,
+ "cache_read_tokens": usage.cache_read_tokens,
Review Comment:
Could you update `docs/operators/agent.rst` (around line 350) along with
this? It still says the cache split shows up in the task log and on the GenAI
span, and doesn't mention the `usage` XCom, which is the surface this PR adds
for cost tracking. One sentence after "``input_tokens`` includes both counts"
would do, e.g. "The ``usage`` XCom carries them as ``cache_read_tokens`` and
``cache_write_tokens``."
##########
providers/common/ai/tests/unit/common/ai/utils/test_logging.py:
##########
@@ -281,20 +281,37 @@ def test_cost_set_logs_cost_line_with_plain_decimal(self):
class TestFormatUsageForXcom:
def test_builds_expected_dict_shape(self):
- usage = RunUsage(requests=3, tool_calls=1, input_tokens=10,
output_tokens=5, cost=Decimal("0.25"))
+ usage = RunUsage(
+ requests=3,
+ tool_calls=1,
+ input_tokens=10,
+ output_tokens=5,
+ cache_read_tokens=4,
+ cache_write_tokens=2,
+ cost=Decimal("0.25"),
+ )
assert format_usage_for_xcom(usage) == {
"requests": 3,
"input_tokens": 10,
"output_tokens": 5,
"total_tokens": 15,
"tool_calls": 1,
+ "cache_read_tokens": 4,
+ "cache_write_tokens": 2,
"cost": "0.25",
}
def test_none_cost_stays_none_not_stringified(self):
assert format_usage_for_xcom(RunUsage(cost=None))["cost"] is None
+ def test_zero_cache_tokens_are_still_pushed(self):
Review Comment:
I don't think this one adds coverage. The dict literal has no conditionals,
so the keys can't go missing at 0, and the zero case is already pinned by the
exact-dict asserts in `test_run_id_and_usage_pushed_to_xcom` and the
failure-path test in `test_agent.py`. Fine to drop.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
