kaxil commented on code in PR #73275:
URL: https://github.com/apache/airflow/pull/73275#discussion_r4036670575
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -573,10 +594,29 @@ def _emit_message_history(self, context: Context, result:
Any) -> None:
transcript =
ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode()
context["task_instance"].xcom_push(key="message_history",
value=transcript)
+ def _emit_run_metadata(self, context: Context, result: Any) -> None:
+ """Expose the pydantic-ai run id and token usage on XCom for
downstream tasks."""
+ usage = result.usage
+ ti = context["task_instance"]
+ ti.xcom_push(key="run_id", value=result.run_id)
Review Comment:
These two go out unconditionally, so `do_xcom_push=False` no longer stops
this operator writing XCom (`RuntimeTaskInstance.xcom_push` doesn't consult it,
and nothing in the provider does either). Every other push in the class sits
behind `message_history` or `enable_hitl_review`, so this is the first
always-on one. Worth gating on `self.do_xcom_push`?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -470,7 +481,16 @@ def execute(self, context: Context) -> Any:
agent = self._build_agent()
+ ti = context["task_instance"]
+ self._run_identity_attrs = build_run_identity_attributes(ti)
+ stamp_identity_on_agent_spans(agent, self._run_identity_attrs)
+
run_kwargs: dict[str, Any] = {"usage_limits": usage_limits}
+ if ti.id is not None:
Review Comment:
Can `ti.id` actually be None here? It's a required non-nullable `UUID` on
the Task SDK datamodel
([_generated.py#L554](https://github.com/apache/airflow/blob/11a3892daa9dfd5e84fb037d49373b174a372fb8/task-sdk/src/airflow/sdk/api/datamodels/_generated.py#L554)),
so this branch looks unreachable and
`test_run_id_omitted_when_task_instance_has_no_id` is asserting on a state the
runtime can't produce. `build_run_identity_attributes` takes the opposite view
of the same question: it calls `str(ti.id)` with no guard, so if None were
reachable every GenAI span would carry the literal
`airflow.task_instance.id="None"`.
##########
providers/common/ai/src/airflow/providers/common/ai/observability.py:
##########
@@ -106,3 +111,71 @@ def genai_instrumentation_settings() ->
InstrumentationSettings | None:
include_binary_content=False,
tracer_provider=provider,
)
+
+
+def build_run_identity_attributes(ti: Any) -> dict[str, Any]:
+ """
+ Build the Airflow identity attributes to stamp on a run's GenAI spans.
+
+ Reuses core's task-span attribute keys (see ``_make_task_span``) so agent
+ spans filter identically to the task span they nest under, plus the
+ per-attempt task-instance id as the run join key carried on every span.
+ """
+ return {
+ "airflow.dag_id": ti.dag_id,
+ "airflow.task_id": ti.task_id,
+ "airflow.dag_run.run_id": ti.run_id,
+ "airflow.task_instance.try_number": ti.try_number,
+ "airflow.task_instance.map_index": ti.map_index if ti.map_index is not
None else -1,
+ "airflow.task_instance.id": str(ti.id),
+ }
+
+
+def stamp_identity_on_agent_spans(agent: Agent, attributes: dict[str, Any]) ->
None:
+ """
+ Stamp *attributes* on every GenAI span *agent* emits during its run.
+
+ pydantic-ai opens all of a run's agent/model/tool spans from
+ ``InstrumentationSettings.tracer``, so wrapping that one tracer reaches
them
+ all without touching the shared core ``TracerProvider``. No-op when the
+ agent is not instrumented with an ``InstrumentationSettings`` (tracing off,
+ or the caller supplied its own non-settings ``instrument`` value).
+ """
+ from pydantic_ai.models.instrumented import InstrumentationSettings
+
+ instrument = agent.instrument
+ if isinstance(instrument, InstrumentationSettings):
+ # _IdentityTracer implements the Tracer surface structurally (it cannot
+ # subclass Tracer without importing opentelemetry at module load).
+ instrument.tracer = cast("Tracer", _IdentityTracer(instrument.tracer,
attributes))
Review Comment:
This writes through to an object the caller may own. `create_agent` pops
`instrument` out of `agent_params`, so a DAG that builds one module-level
`InstrumentationSettings` and hands it to several agent tasks ends up with that
object carrying whichever task stamped last. The HITL path shows the other
half: `regenerate_with_feedback` rebuilds the agent every round and stamps
again, and against a caller-supplied settings object the `_IdentityTracer`s
nest one deeper each time (depth 4 after 4 stamps when I tried it). Copying
first (`settings = copy.copy(instrument)`, set `.tracer` on the copy, then
`agent.instrument = settings`) keeps it per-run. `dataclasses.replace` won't
work, `tracer` isn't accepted by `__init__`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -573,10 +594,29 @@ def _emit_message_history(self, context: Context, result:
Any) -> None:
transcript =
ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode()
context["task_instance"].xcom_push(key="message_history",
value=transcript)
+ def _emit_run_metadata(self, context: Context, result: Any) -> None:
+ """Expose the pydantic-ai run id and token usage on XCom for
downstream tasks."""
+ usage = result.usage
+ ti = context["task_instance"]
+ ti.xcom_push(key="run_id", value=result.run_id)
+ ti.xcom_push(
+ key="usage",
+ value={
+ "requests": usage.requests,
+ "input_tokens": usage.input_tokens,
+ "output_tokens": usage.output_tokens,
+ "total_tokens": usage.total_tokens,
+ "tool_calls": usage.tool_calls,
Review Comment:
The class docstring up at line 133 says a downstream task can reference the
run "and its cost", but `cost` isn't in this dict. `RunUsage.cost` is the field
the `pydantic-ai-slim>=2.23.0` floor is pinned for, and `log_run_summary`
already prints it. It's a `Decimal | None` so it needs stringifying before it
goes to XCom, but either add it here or drop the cost half of the docstring
promise.
--
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]
