kaxil commented on code in PR #74769:
URL: https://github.com/apache/airflow/pull/74769#discussion_r4284775119
##########
providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py:
##########
@@ -152,3 +176,176 @@ async def call_airflow_tool(ctx: ToolContext[Any],
arguments: str) -> str:
on_invoke_tool=call_airflow_tool,
strict_json_schema=False,
)
+
+
+class AirflowTraceProcessor(TracingProcessor):
+ """
+ An OpenAI Agents SDK trace processor that records the SDK's traces as
OpenTelemetry spans.
+
+ The SDK has a tracing system of its own rather than OpenTelemetry.
Register this
+ processor with it and every trace becomes a tree of OpenTelemetry spans
under the span
+ that was current when the trace started, so a run inside a task nests
under the task's
+ span:
+
+ .. code-block:: python
+
+ from agents import set_trace_processors
+
+ from airflow.providers.common.ai.tools.openai_agents import
AirflowTraceProcessor
+ from airflow.providers.common.ai.tools.tracing import
agent_framework_tracing
+
+ set_trace_processors([AirflowTraceProcessor()])
+ with agent_framework_tracing():
+ answer = Runner.run_sync(agent, question).final_output
+
+ ``set_trace_processors`` replaces every processor registered with the SDK,
including
+ its default one, which sends every trace to OpenAI. To keep that export,
or a processor
+ another tracing library registered, add this one with
``add_trace_processor`` instead.
+ Run the agent inside
+ :func:`~airflow.providers.common.ai.tools.tracing.agent_framework_tracing`
so the spans
+ carry the task's identity. While an SDK span runs it is the current span,
so spans
+ other instrumentation starts inside a tool or model call, such as an HTTP
request,
+ nest under it.
+
+ Spans use the names of the OpenTelemetry GenAI conventions:
``invoke_agent`` for an
+ agent, ``chat`` for a model call with its model and token counts, and
``execute_tool``
+ for a tool call. The SDK's other spans, such as one per ``Runner.run`` and
one per turn
+ of the agent loop, are named ``openai_agents.<kind>``. Prompts,
completions, tool
+ arguments and results, and the SDK's error messages are recorded only when
+ ``[common.ai] otel_export_enabled`` and ``capture_content`` are both on,
and they pass
+ through Airflow's secret masker first. The SDK records that content only
while
+ ``RunConfig(trace_include_sensitive_data=True)``, its default; with it
off, a model
+ call through the Responses API also carries no model name, since the SDK
withholds the
+ whole response.
+
+ :param tracer_provider: The OpenTelemetry tracer provider to record spans
with. By
+ default, the global one, which core tracing sets up on the worker.
+ """
+
+ def __init__(self, tracer_provider: TracerProvider | None = None) -> None:
+ self._tracer = otel_trace.get_tracer(__name__,
tracer_provider=tracer_provider)
+ # The open OpenTelemetry span for each SDK trace and span, by the
SDK's ID, and the
+ # token that made each SDK span current. Each access is one dict
operation, which is
+ # safe across the threads the SDK may call from.
+ self._open: dict[str, OtelSpan] = {}
+ self._made_current: dict[str, Token[Context]] = {}
+
+ def on_trace_start(self, trace: Trace) -> None:
+ # Started as a child of whatever span is current, the task's span
inside a task.
+ self._open[trace.trace_id] = self._start(trace.name,
{"openai_agents.trace_id": trace.trace_id})
+
+ def on_trace_end(self, trace: Trace) -> None:
+ if (otel_span := self._open.pop(trace.trace_id, None)) is not None:
+ otel_span.end()
+
+ def on_span_start(self, span: Span[Any]) -> None:
+ parent = self._open.get(span.parent_id or span.trace_id)
+ context = otel_trace.set_span_in_context(parent) if parent is not None
else None
+ name, _ = _describe(span.span_data, capture=False)
+ otel_span = self._start(name, {}, context)
+ self._open[span.span_id] = otel_span
+ self._made_current[span.span_id] =
otel_context.attach(otel_trace.set_span_in_context(otel_span))
+
+ def on_span_end(self, span: Span[Any]) -> None:
+ if (token := self._made_current.pop(span.span_id, None)) is not None:
+ otel_context.detach(token)
+ if (otel_span := self._open.pop(span.span_id, None)) is None:
+ return
+ # Model, token counts and results are filled in while the span runs,
so describe it now.
+ capture = _content_capture_enabled()
+ name, attributes = _describe(span.span_data, capture=capture)
+ otel_span.update_name(name)
+ otel_span.set_attributes(attributes)
+ if span.error is not None:
+ # The message can be app-supplied text, such as the reason a tool
call was rejected.
+ message = _masked(span.error["message"]) if capture else None
+ otel_span.set_status(Status(StatusCode.ERROR, message))
+ otel_span.end()
+
+ def _start(self, name: str, attributes: dict[str, Any], context: Context |
None = None) -> OtelSpan:
+ # The task's identity, when inside agent_framework_tracing, whichever
provider records the span.
+ identity = _task_identity.get() or {}
+ return self._tracer.start_span(name, context=context,
attributes={**identity, **attributes})
+
+ def shutdown(self) -> None:
+ # The tracer provider owns export, so there is nothing of this
processor's to flush.
+ pass
+
+ def force_flush(self) -> None:
+ pass
+
+
+def _describe(data: SpanData, *, capture: bool) -> tuple[str, dict[str, Any]]:
+ """Return the OpenTelemetry span name and attributes for one SDK span, by
its kind."""
+ attributes: dict[str, Any] = {}
+ content: dict[str, Any] = {}
+ if isinstance(data, AgentSpanData):
+ name = f"invoke_agent {data.name}"
+ attributes = {"gen_ai.operation.name": "invoke_agent",
"gen_ai.agent.name": data.name}
+ elif isinstance(data, GenerationSpanData):
+ name = f"chat {data.model}" if data.model else "chat"
+ attributes = {"gen_ai.operation.name": "chat", **_usage(data.usage)}
+ if data.model:
+ attributes["gen_ai.request.model"] = data.model
+ content = {"gen_ai.input.messages": data.input,
"gen_ai.output.messages": data.output}
+ elif isinstance(data, ResponseSpanData):
+ model = data.response.model if data.response is not None else None
+ name = f"chat {model}" if model else "chat"
+ attributes = {"gen_ai.operation.name": "chat", **_usage(data.usage)}
+ if data.response is not None:
+ attributes |= {"gen_ai.response.id": data.response.id,
"gen_ai.response.model": model}
+ output = None
+ if data.response is not None:
+ output = [item.model_dump(mode="json") for item in
data.response.output]
+ content = {"gen_ai.input.messages": data.input,
"gen_ai.output.messages": output}
+ elif isinstance(data, FunctionSpanData):
+ name = f"execute_tool {data.name}"
+ attributes = {"gen_ai.operation.name": "execute_tool",
"gen_ai.tool.name": data.name}
+ content = {"gen_ai.tool.call.arguments": data.input,
"gen_ai.tool.call.result": data.output}
+ elif isinstance(data, TaskSpanData):
+ name = f"openai_agents.task {data.name}"
+ elif isinstance(data, TurnSpanData):
+ name = f"openai_agents.turn {data.turn}"
+ attributes = {"gen_ai.agent.name": data.agent_name,
"openai_agents.turn": data.turn}
+ elif isinstance(data, HandoffSpanData):
+ name = "openai_agents.handoff"
+ attributes = {
+ "openai_agents.handoff.from_agent": data.from_agent,
+ "openai_agents.handoff.to_agent": data.to_agent,
+ }
+ elif isinstance(data, GuardrailSpanData):
+ name = f"openai_agents.guardrail {data.name}"
+ attributes = {"openai_agents.guardrail.triggered": data.triggered}
+ else:
+ name = f"openai_agents.{data.type}"
+ if capture:
+ attributes |= {key: _masked(value) for key, value in content.items()
if value is not None}
+ # OpenTelemetry rejects None attribute values.
+ return name, {key: value for key, value in attributes.items() if value is
not None}
+
+
+def _masked(value: Any) -> str:
Review Comment:
Renamed to `_mask_content` and `_build_usage_attributes` in 6d5e7359798.
##########
providers/common/ai/docs/frameworks/openai_agents.rst:
##########
@@ -103,18 +103,47 @@ nested run shares the outer run's retry count for each
tool.
Tracing
-------
-The SDK has its own tracing, which is on by default and sends every run's
traces to
-OpenAI's Traces backend using ``OPENAI_API_KEY`` from the environment. It does
this
-whatever model the agent runs, including a self-hosted one, and the traces
hold the
-prompts and model output. Tool results are masked before the model sees them;
nothing
-else in a trace is.
-
-The example turns tracing off for its run with
``RunConfig(tracing_disabled=True)``. To
-turn it off for the whole process, call ``agents.set_tracing_disabled(True)``
or set
-``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. The SDK does
not emit
-OpenTelemetry spans, so
-:func:`~airflow.providers.common.ai.tools.tracing.agent_framework_tracing`
does not apply
-to it.
+The SDK has a tracing system of its own rather than OpenTelemetry. By default
it sends
+every run's traces to OpenAI's Traces backend using ``OPENAI_API_KEY`` from the
+environment, whatever model the agent runs, including a self-hosted one, and
the traces
+hold the prompts and model output. Tool results are masked before the model
sees them;
+nothing else in those traces is.
+
+:class:`~airflow.providers.common.ai.tools.openai_agents.AirflowTraceProcessor`
is a trace
+processor for the SDK that records each trace as OpenTelemetry spans under the
task's
+span instead. The example replaces the SDK's processors with it, so nothing
goes to
+OpenAI, and runs the agent inside
+:func:`~airflow.providers.common.ai.tools.tracing.agent_framework_tracing`, so
the spans
+carry the task's identity:
+
+.. code-block:: python
+
+ from agents import Runner, set_trace_processors
+
+ from airflow.providers.common.ai.tools.openai_agents import
AirflowTraceProcessor
+ from airflow.providers.common.ai.tools.tracing import
agent_framework_tracing
+
+ set_trace_processors([AirflowTraceProcessor()])
+ with agent_framework_tracing():
+ answer = Runner.run_sync(agent, question).final_output
+
+``set_trace_processors`` replaces every processor registered with the SDK. To
keep the
+export to OpenAI, or a processor another instrumentation registered, add this
one with
+``agents.add_trace_processor`` instead. The spans use the names of the
OpenTelemetry GenAI
+conventions (``invoke_agent``, ``chat`` with the model and token counts,
+``execute_tool``), and the SDK's other spans are named
``openai_agents.<kind>``. Prompts, completions, and tool arguments and results
are left
+out unless ``[common.ai] otel_export_enabled`` and ``capture_content`` are
both on, and pass
+through Airflow's secret masker when they are recorded; see
:doc:`../observability`.
+
+Third-party packages such as ``openinference-instrumentation-openai-agents``
and
Review Comment:
Agreed, it's now the one sentence you suggested (6d5e7359798).
##########
providers/common/ai/docs/frameworks/index.rst:
##########
@@ -214,9 +214,10 @@ adapters do.
Blocking hook calls, such as a SQL query, run one at a time in the task's
process, so an
agent that calls two database tools at once gets its answers one after the
other.
-Run a Strands or ADK agent inside
+Run a Strands, ADK or OpenAI Agents SDK agent inside
Review Comment:
Split into its own sentence in 6d5e7359798.
##########
providers/common/ai/docs/observability.rst:
##########
@@ -175,7 +175,8 @@ outputs (``gen_ai.input.messages`` /
``gen_ai.output.messages``), set:
exported to your tracing backend **without redaction**. Airflow's secret
masking
applies to logs and rendered template fields, not to OpenTelemetry span
attributes. The one exception is what a function tool returns or raises
while it
Review Comment:
Fixed in 6d5e7359798: the warning now lists the two exceptions as separate
items.
##########
providers/common/ai/tests/unit/common/ai/tools/test_openai_agents.py:
##########
@@ -264,3 +293,239 @@ def
test_the_toolsets_calls_are_counted_as_openai_agents(self):
_run_agent(AirflowTools(ts), "list_tables", {})
assert stats.incr.call_args.kwargs["tags"]["framework"] ==
"openai_agents"
+
+
[email protected]
+def recorded():
+ """Route the SDK's traces through an AirflowTraceProcessor into an
in-memory exporter."""
+ exporter = InMemorySpanExporter()
+ provider = TracerProvider()
+ provider.add_span_processor(SimpleSpanProcessor(exporter))
+ set_trace_processors([AirflowTraceProcessor(tracer_provider=provider)])
+ try:
+ yield exporter, provider
+ finally:
+ set_trace_processors([default_processor()])
+
+
+def _traced_lookup(arguments: str = '{"key": "a"}', result: ToolResult |
Exception = _OK) -> None:
+ """Run an agent that calls ``lookup`` once, with the SDK's tracing on."""
+ agent = Agent(
+ name="analyst", model=_ToolCallingModel("lookup", arguments),
tools=AirflowTools(_tool(result))
+ )
+ _run(agent, tracing_disabled=False)
+
+
+def _span(exporter: InMemorySpanExporter, name: str = "execute_tool lookup")
-> ReadableSpan:
+ (span,) = [span for span in exporter.get_finished_spans() if span.name ==
name]
+ return span
+
+
+class TestAirflowTraceProcessor:
+ def test_a_run_becomes_a_span_tree_under_the_current_span(self, recorded):
+ exporter, provider = recorded
+
+ with provider.get_tracer("test").start_as_current_span("task"):
+ _traced_lookup()
+
+ spans = exporter.get_finished_spans()
+ by_id = {span.context.span_id: span for span in spans}
+ ancestors = []
+ span: ReadableSpan | None = _span(exporter)
+ while span is not None and span.parent is not None:
+ span = by_id.get(span.parent.span_id)
+ ancestors.append(span.name if span is not None else None)
+ assert ancestors == [
+ "openai_agents.turn 1",
+ "invoke_agent analyst",
+ "openai_agents.task Agent workflow",
+ "Agent workflow",
+ "task",
+ ]
+ invoke, tool = _span(exporter, "invoke_agent analyst"), _span(exporter)
+ assert invoke.attributes["gen_ai.agent.name"] == "analyst"
+ assert tool.attributes["gen_ai.operation.name"] == "execute_tool"
+ assert tool.attributes["gen_ai.tool.name"] == "lookup"
+
+ def test_a_span_started_inside_a_tool_call_nests_under_it(self, recorded):
+ """Spans of other instrumentation, such as an HTTP client's, belong to
the call that made them."""
+ exporter, provider = recorded
+
+ async def function(arguments: dict[str, Any]) -> ToolResult:
+ with provider.get_tracer("http").start_as_current_span("GET"):
+ return ToolResult("ok")
+
+ tool = AirflowTool(
+ name="lookup", description="Look up.", parameters={"type":
"object"}, function=function
+ )
+ agent = Agent(name="analyst", model=_ToolCallingModel("lookup", "{}"),
tools=AirflowTools(tool))
+
+ _run(agent, tracing_disabled=False)
+
+ request, call = _span(exporter, "GET"), _span(exporter)
+ assert request.parent is not None
+ assert request.parent.span_id == call.context.span_id
+
+ @pytest.mark.parametrize(
+ ("export", "capture"),
+ [("False", "False"), ("True", "False"), ("False", "True")],
+ )
+ def test_leaves_content_out_unless_export_and_capture_are_both_on(self,
recorded, export, capture):
+ exporter, _ = recorded
+
+ with conf_vars(
+ {("common.ai", "otel_export_enabled"): export, ("common.ai",
"capture_content"): capture}
+ ):
+ _traced_lookup()
+
+ attributes = _span(exporter).attributes
+ assert "gen_ai.tool.call.arguments" not in attributes
+ assert "gen_ai.tool.call.result" not in attributes
+
+ @pytest.mark.enable_redact
+ @conf_vars(CAPTURE_CONTENT)
+ def test_captured_content_passes_through_the_secret_masker(self, recorded,
registered_secret):
+ exporter, _ = recorded
+
+ _traced_lookup(json.dumps({"key": registered_secret}),
ToolResult(f"seen {registered_secret}"))
+
+ attributes = _span(exporter).attributes
+ assert attributes["gen_ai.tool.call.arguments"] == json.dumps({"key":
"***"})
+ assert attributes["gen_ai.tool.call.result"] == "seen ***"
+
+ def test_spans_carry_the_task_whichever_provider_records_them(self,
recorded):
+ """The processor's own provider gets no identity processor from
agent_framework_tracing."""
+ exporter, _ = recorded
+
+ with (
+ patch(f"{TRACING}._live_tracer_provider", autospec=True,
return_value=None),
+ patch(f"{TRACING}.get_current_context", autospec=True,
return_value={"ti": TI}),
+ agent_framework_tracing(),
+ ):
+ _traced_lookup()
+
+ spans = exporter.get_finished_spans()
+ assert spans
+ assert {span.attributes["airflow.dag_id"] for span in spans} ==
{TI.dag_id}
+ assert {span.attributes["airflow.task_id"] for span in spans} ==
{TI.task_id}
+
+ @pytest.mark.parametrize(
+ ("settings", "description"),
+ [
+ pytest.param({}, None, id="content-off"),
+ pytest.param(CAPTURE_CONTENT, "Error running tool",
id="content-on"),
+ ],
+ )
+ def test_a_failed_tool_call_marks_its_span_as_an_error(self, recorded,
settings, description):
+ """The SDK's error message can be app-supplied text, so it follows the
content rule."""
+ exporter, _ = recorded
+
+ with conf_vars(settings), pytest.raises(ToolCallError):
+ _traced_lookup(result=PermissionError("role cannot read orders"))
+
+ status = _span(exporter).status
+ assert status.status_code == StatusCode.ERROR
+ assert status.description == description
+
+ def
test_a_response_span_is_named_for_its_model_once_the_response_arrives(self,
recorded):
+ exporter, _ = recorded
+
+ with sdk_trace("workflow"), response_span() as span:
+ span.span_data.response = Response.model_construct(id="resp_1",
model="gpt-5.6-luna", output=[])
+ span.span_data.usage = {"input_tokens": 12, "output_tokens": 3}
+
+ attributes = _span(exporter, "chat gpt-5.6-luna").attributes
+ assert attributes["gen_ai.usage.input_tokens"] == 12
+ assert attributes["gen_ai.usage.output_tokens"] == 3
+
+
+class TestDescribe:
Review Comment:
Added a parametrised test for the handoff, guardrail and fallback branches
in 6d5e7359798.
--
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]
