kaxil commented on code in PR #74355:
URL: https://github.com/apache/airflow/pull/74355#discussion_r4212193808
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -302,11 +321,45 @@ def execute(self, context: Context) -> Any:
agent: Agent[object, Any] = self.llm_hook.create_agent(
output_type=self.output_type, instructions=self.system_prompt,
**self.agent_params
)
- result = self.run_agent_sync(agent, self.prompt,
usage_limits=usage_limits)
+
+ # A per-attempt key (the task-instance id on Airflow 3, which is
regenerated on each
+ # retry; dag/run/task/map/try on Airflow 2) ties this run to its GenAI
trace, the same
+ # join key AgentOperator uses -- see the provider's observability
docs. ``ti`` is None
+ # only for a hand-built context with no task instance (e.g. a bare
dict in a unit
+ # test), in which case the run proceeds untracked, same as
``_push_xcom``'s fallback.
+ ti = self._resolve_task_instance(context)
+ run_kwargs: dict[str, Any] = {"usage_limits": usage_limits}
+ if ti is not None:
+ stamp_identity_on_agent_spans(agent,
build_run_identity_attributes(ti))
+ run_kwargs["run_id"] = make_task_instance_run_key(ti)
+
+ run_usage = RunUsage()
+ try:
+ result = self.run_agent_sync(agent, self.prompt, usage=run_usage,
**run_kwargs)
Review Comment:
`@task.llm` goes through `LLMOperator.execute`, so three tests in
`decorators/test_llm.py` now fail at this head.
`test_execute_calls_callable_and_returns_output` and
`test_execute_accepts_sequence_prompt` assert `run_sync` without `usage=`, and
`test_execute_merges_op_kwargs_into_callable` hits `AttributeError: Mock object
has no attribute 'instrument'` from `stamp_identity_on_agent_spans`. They need
the same `spec=["run_sync", "instrument"]` and `usage=ANY` updates as
`operators/test_llm.py`; CI hasn't run on this stacked PR, so nothing has
flagged it yet.
##########
providers/common/ai/docs/observability.rst:
##########
@@ -62,33 +62,47 @@ How it works
attempts share one trace and appear as repeated task-run spans on it,
distinguished by ``try number``. Only a manual clear or rerun regenerates the
context and starts a new trace.
-* **Run join key.** For an ``AgentOperator`` run, the task-instance id (unique
- per attempt, since Airflow regenerates it on each retry) is passed to
+* **Run join key.** For an ``AgentOperator`` or ``LLMOperator`` run (and their
+ ``@task.agent`` / ``@task.llm`` decorator equivalents), the task-instance id
+ (unique per attempt, since Airflow regenerates it on each retry) is passed to
pydantic-ai as the run's ``run_id``. It surfaces on the run's GenAI spans as
``gen_ai.agent.call.id``, and the operator also exposes it, alongside the
run's
token usage, on XCom under the ``run_id`` and ``usage`` keys. ``usage`` is
- this attempt's own usage -- not the cross-attempt cumulative total described
- under ``usage_limits`` in :ref:`howto/operator:agent` -- and it is pushed on
- a failed attempt too, so a downstream ``all_done`` task or failure callback
- can read what the last attempt spent. With ``durable=True``, steps an
- attempt replays from the cache are not part of it. XCom is cleared at the
start of every
- attempt, so only the most recent attempt's value survives, not each
+ this attempt's own usage -- for ``AgentOperator`` not the cross-attempt
cumulative
+ total described under ``usage_limits`` in :ref:`howto/operator:agent` -- and
it is
+ pushed on a failed attempt too, so a downstream ``all_done`` task or failure
+ callback can read what the last attempt spent. XCom is cleared at the start
of
+ every attempt, so only the most recent attempt's value survives, not each
historical attempt's. A downstream task can then reference the
- run (``ti.xcom_pull(task_ids="my_agent", key="run_id")``) and a trace
+ run (``ti.xcom_pull(task_ids="my_task", key="run_id")``) and a trace
backend can join a task's output to its agent trace without parsing logs.
- With ``enable_hitl_review`` the ``run_id`` and ``usage`` reflect the initial
- model run, not the human-feedback regenerations. A run that resumes after a
- tool approval (see :doc:`tool_approval`) continues as
- ``<task-instance id>-resumed``, which is the ``run_id`` the operator pushes;
- ``usage`` covers both parts.
+ ``AgentOperator`` additionally: with ``durable=True``, steps an attempt
replays
+ from the cache are not part of ``usage``; with ``enable_hitl_review`` the
+ ``run_id`` and ``usage`` reflect the initial model run, not the
human-feedback
+ regenerations, and a run that resumes after a tool approval (see
+ :doc:`tool_approval`) continues as ``<task-instance id>-resumed``, which is
the
+ ``run_id`` the operator pushes, with ``usage`` covering both parts.
+ ``LLMOperator`` has no durable caching, cross-attempt usage budget, or HITL
+ regeneration, so its ``run_id`` and ``usage`` always describe the one run the
+ attempt made.
Airflow 2 has no task-instance id, so there the key is
``<dag_id>/<run_id>/<task_id>/<map_index>/<try_number>``, and spans carry
the five
identity keys without ``airflow.task_instance.id``.
-* **Scope.** The ``run_id`` / ``usage`` XComs come only from ``AgentOperator``
and
- ``@task.agent``, and so do the ``airflow.*`` identity attributes, apart from
a Strands or
- ADK agent run inside ``agent_framework_tracing`` (see below). The other LLM
- operators still emit GenAI spans correlated to the task span by nesting, but
- without the identity attributes or the run join key.
+* **Scope.** The ``run_id`` XCom, the ``usage`` XCom, and the ``airflow.*``
identity
Review Comment:
The Correlation bullet above (around line 50) still says only
`AgentOperator` (and `@task.agent`) stamps the task-instance identity on GenAI
spans, which now contradicts this bullet and `llm.rst`. It needs `LLMOperator`
/ `@task.llm` added too.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -302,11 +321,45 @@ def execute(self, context: Context) -> Any:
agent: Agent[object, Any] = self.llm_hook.create_agent(
output_type=self.output_type, instructions=self.system_prompt,
**self.agent_params
)
- result = self.run_agent_sync(agent, self.prompt,
usage_limits=usage_limits)
+
+ # A per-attempt key (the task-instance id on Airflow 3, which is
regenerated on each
+ # retry; dag/run/task/map/try on Airflow 2) ties this run to its GenAI
trace, the same
+ # join key AgentOperator uses -- see the provider's observability
docs. ``ti`` is None
+ # only for a hand-built context with no task instance (e.g. a bare
dict in a unit
+ # test), in which case the run proceeds untracked, same as
``_push_xcom``'s fallback.
+ ti = self._resolve_task_instance(context)
+ run_kwargs: dict[str, Any] = {"usage_limits": usage_limits}
+ if ti is not None:
+ stamp_identity_on_agent_spans(agent,
build_run_identity_attributes(ti))
+ run_kwargs["run_id"] = make_task_instance_run_key(ti)
+
+ run_usage = RunUsage()
+ try:
+ result = self.run_agent_sync(agent, self.prompt, usage=run_usage,
**run_kwargs)
+ except BaseException:
Review Comment:
`AgentOperator._report_failed_run` also calls `log_run_usage(self.log, ...,
outcome="failed")` before it pushes. This path only pushes, so the task log of
a run stopped by `cost_limit` says nothing about what it spent, and with
`do_xcom_push=False` the spend isn't recorded anywhere. Could it log the
partial usage too?
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm.py:
##########
@@ -403,7 +443,87 @@ 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_and_usage_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
+ """The model that actually answered, the run id, and the run's
usage/cost are exposed on XCom."""
+ cost = Decimal("0.0042")
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ result = self._result(make_mock_run_result, Summary(text="t"), None,
cost=cost)
+ result.run_id = "ti-9"
+
+ def _run_sync(*args, **kwargs):
+ # Mirrors what a real pydantic-ai run does to the ``usage=``
object it was
+ # called with, so the usage pushed to XCom (read from that object,
not from
+ # the mock result's own ``.usage``) reflects the values configured
here.
+ kwargs["usage"].incr(
+ RunUsage(
+ requests=1,
+ input_tokens=120,
+ output_tokens=45,
+ tool_calls=2,
+ cache_read_tokens=5,
+ cache_write_tokens=3,
+ cost=cost,
+ )
+ )
+ return result
+
+ mock_agent.run_sync.side_effect = _run_sync
+ 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 = _make_context(ti_id="ti-9")
+
+ 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"
+ assert pushes["run_id"] == "ti-9"
+ assert pushes["usage"] == {
+ "requests": 1,
+ "input_tokens": 120,
+ "output_tokens": 45,
+ "total_tokens": 165,
+ "tool_calls": 2,
+ "cache_read_tokens": 5,
+ "cache_write_tokens": 3,
+ "cost": str(cost),
+ }
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_failure_pushes_partial_usage_to_xcom(self, mock_hook_cls):
Review Comment:
Nothing exercises the two inner `try/except` guards in the failure path: if
both were removed, a failing XCom backend would surface as its own error
instead of `UsageLimitExceeded` and the suite would still pass.
`test_failed_xcom_push_while_reporting_does_not_mask_the_original_exception` in
test_agent.py is the pattern to port.
`test_execute_raises_when_cost_cap_exceeded` above runs a real `FunctionModel`,
so asserting the `usage` push there would check this path against pydantic-ai's
own accounting rather than a hand-incremented `RunUsage`.
--
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]
