kaxil commented on code in PR #74355:
URL: https://github.com/apache/airflow/pull/74355#discussion_r4201901173
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -307,7 +311,10 @@ def execute(self, context: Context) -> Any:
log_run_summary(self.log, result)
output = result.output
+ self._push_xcom(context, "usage", format_usage_for_xcom(result.usage))
Review Comment:
This only runs once `run_agent_sync` returns, so a run that raises (say
`UsageLimitExceeded` from a `cost_limit`) pushes no `usage` at all.
`AgentOperator` passes its own `RunUsage` into `run_sync` and pushes the
partial usage from `_report_failed_run`, and its docstring promises that to
`all_done` tasks and failure callbacks. `run_agent_sync` already forwards
`**run_kwargs`, so could this pass `usage=RunUsage()` and push it in an
`except` before re-raising? Otherwise the cost-capped failure, the run where
cost matters most, is the one the Model tab can't show. If you'd rather keep it
success-only, the docs line below should say so.
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm.py:
##########
@@ -403,7 +404,46 @@ def
test_hand_built_context_skips_the_decision_push_with_a_warning(
output = op.execute(context)
assert Summary.model_validate(output).text == "t"
- assert "the decision record was not pushed to XCom" in caplog.text
+ assert "'decision' was not pushed to XCom" in caplog.text
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_resolved_model_name_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
+ """The model that actually answered is exposed on its own namespaced
XCom key."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value = self._result(make_mock_run_result,
Summary(text="t"), None)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=Summary)
+ context = MagicMock(spec=dict)
+
+ op.execute(context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
+ }
+ assert pushes[MODEL_NAME_XCOM_KEY] == "jev-1.13.0"
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_usage_pushed_to_xcom(self, mock_hook_cls, make_mock_run_result):
+ """Token usage/cost is exposed on the usage XCom key, the same shape
AgentOperator uses."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value = self._result(make_mock_run_result,
Summary(text="t"), None)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=Summary)
+ context = MagicMock(spec=dict)
+
+ op.execute(context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
+ }
+ assert pushes["usage"] == {
Review Comment:
The fixture's usage is all zeros with `cost=None`, so this only pins the key
and `requests: 1`. Passing a cost to `make_mock_run_result` and asserting the
stringified value (like `test_usage_cost_is_stringified_on_xcom` in
test_agent.py) would catch more. The setup is identical to
`test_resolved_model_name_pushed_to_xcom` above, so the two could also be one
test.
##########
providers/common/ai/docs/operators/llm.rst:
##########
@@ -284,3 +284,6 @@ After each LLM call, the operator logs a summary with model
name, token usage,
and request count at INFO level. At DEBUG level, the LLM output is also logged
(truncated to 500 characters). See :ref:`AgentOperator logging
<howto/operator:agent>`
for details on the log format.
+
+The same request/token/cost counts are also pushed to XCom under the ``usage``
Review Comment:
The Scope bullet in `observability.rst` still says the `run_id` / `usage`
XComs come only from `AgentOperator` and `@task.agent`. With this PR
`LLMOperator` and `@task.llm` push `usage` too, so that sentence needs updating
(`run_id` stays agent-only). The `LLMOperator` class docstring could also
mention the `usage` key the way `AgentOperator`'s does.
--
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]
