pankajkoti commented on code in PR #73275:
URL: https://github.com/apache/airflow/pull/73275#discussion_r4037288867
##########
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:
Good point, gated it. `_emit_run_metadata` now returns early when
`do_xcom_push` is False, so the `run_id` and `usage` pushes respect the flag
the same way the rest of the operator's XCom does. Done in bcc285d7af.
##########
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:
You're right, `TaskInstance.id` is a non-nullable `UUID` on the Task SDK
datamodel, so the guard was dead code and the test was asserting a state the
runtime cannot produce. Dropped the guard so `run_id` is always `str(ti.id)`,
which also makes it consistent with `build_run_identity_attributes`, and
removed that test. Done in bcc285d7af.
##########
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:
Added `cost` to the `usage` payload rather than trimming the docstring. The
floor is already at `pydantic-ai-slim>=2.23.0` (bumped for per-run cost limits
in #71403), and `RunUsage.cost` ships in that same 2.23.0 cost API, so
surfacing it needs no further bump. It is `Decimal | None`, so it is
stringified before it goes to XCom. Done in bcc285d7af.
##########
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:
Confirmed both the shared-settings leak and the HITL nesting. Now
`copy.copy` the settings, wrap the copy's tracer, and assign it to
`agent.instrument`, leaving the caller's object untouched, so each run gets its
own single wrapper. Verified `agent.instrument` is read at run time via
`_resolve_instrumentation_settings`, so the swapped copy is the one used. Done
in bcc285d7af.
--
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]
