This is an automated email from the ASF dual-hosted git repository.
pankajastro pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new af61379fc14 Include prompt cache token counts in AgentOperator's usage
XCom (#74277)
af61379fc14 is described below
commit af61379fc140b71378d2bff269de20748e6bdb2d
Author: Pankaj Singh <[email protected]>
AuthorDate: Tue Oct 6 11:44:28 2026 +0530
Include prompt cache token counts in AgentOperator's usage XCom (#74277)
* Include prompt cache token counts in AgentOperator's usage XCom
input_tokens already includes tokens billed at the cache-read or
cache-write price when cache_prompt is on, but the split was only
visible in the task log. Push cache_read_tokens/cache_write_tokens
in the usage XCom too, so downstream cost tracking can tell a
cache-priced run from one billed at the plain input-token rate.
Co-Authored-By: Claude <[email protected]>
* Document usage XCom cache fields in AgentOperator docs
The prompt-caching doc section showed cache token counts in the task
log and GenAI span but didn't mention the usage XCom, the surface
added for cost tracking. Also drop a test that could not fail without
a logic change, since format_usage_for_xcom builds a flat dict literal
with no conditional on the cache-token keys.
Co-Authored-By: Claude <[email protected]>
---------
Co-authored-by: Claude <[email protected]>
---
providers/common/ai/docs/operators/agent.rst | 3 ++-
.../ai/src/airflow/providers/common/ai/utils/logging.py | 2 ++
.../common/ai/tests/unit/common/ai/operators/test_agent.py | 4 ++++
.../common/ai/tests/unit/common/ai/utils/test_logging.py | 12 +++++++++++-
4 files changed, 19 insertions(+), 2 deletions(-)
diff --git a/providers/common/ai/docs/operators/agent.rst
b/providers/common/ai/docs/operators/agent.rst
index 80c87facfe5..4144c104bbc 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -354,7 +354,8 @@ When the provider reports cache activity, the task log
shows it under the run su
LLM prompt cache: cache_read_tokens=..., cache_write_tokens=...
``input_tokens`` includes both counts. With :doc:`../observability` turned on,
each
-request's GenAI span carries them too.
+request's GenAI span carries them too. The ``usage`` XCom carries them as
+``cache_read_tokens`` and ``cache_write_tokens``.
Parameters
----------
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
index 51edfbf8f05..b442c870f13 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
+++ b/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,
+ "cache_write_tokens": usage.cache_write_tokens,
# Decimal | None, stringified so XCom serialization stays lossless.
"cost": str(usage.cost) if usage.cost is not None else None,
}
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 2d3422f2479..318f5efe6fd 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -2278,6 +2278,8 @@ class TestAgentOperatorRunIdentity:
"output_tokens": 0,
"total_tokens": 0,
"tool_calls": 0,
+ "cache_read_tokens": 0,
+ "cache_write_tokens": 0,
"cost": None,
}
@@ -2385,6 +2387,8 @@ class TestAgentOperatorUsageBudget:
"output_tokens": 0,
"total_tokens": 10,
"tool_calls": 0,
+ "cache_read_tokens": 0,
+ "cache_write_tokens": 0,
"cost": "0.1",
}
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
index 828574de4e7..6733315f66b 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
@@ -281,7 +281,15 @@ class TestLogRunUsage:
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,
@@ -289,6 +297,8 @@ class TestFormatUsageForXcom:
"output_tokens": 5,
"total_tokens": 15,
"tool_calls": 1,
+ "cache_read_tokens": 4,
+ "cache_write_tokens": 2,
"cost": "0.25",
}
