kaxil commented on code in PR #74769:
URL: https://github.com/apache/airflow/pull/74769#discussion_r4284763492
##########
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)}
Review Comment:
Done in 6d5e7359798. Model-call spans now start as `CLIENT`. The Responses
path sets `gen_ai.provider.name: openai`, and the Chat Completions path sets
`server.address` from the client's base URL, without assuming a provider.
##########
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:
+ """
+ Render content as text with registered secrets masked.
+
+ Tool arguments, and many tool results, are JSON text. A secret holding a
quote or a
+ backslash appears escaped inside it and would not match, so JSON text is
parsed and
+ masked as values before it is rendered again.
+ """
+ if isinstance(value, str):
+ try:
+ value = json.loads(value)
Review Comment:
Good catch, fixed in 6d5e7359798. Only a JSON object or array is parsed now;
any other text, including a bare number, is masked as it is, which also keeps
`1.50` as written. There's a test with a numeric secret.
##########
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]
Review Comment:
Fixed in 6d5e7359798: the Responses output is serialized only when content
is captured, and a test checks it isn't otherwise.
##########
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:
Review Comment:
Fixed in 6d5e7359798: building the attributes is in `try/finally`, so the
span always ends. There's a test where it raises.
##########
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:
Review Comment:
Yes, fixed in 6d5e7359798. `on_trace_end` ends any span of that trace the
SDK never finished. Their context tokens are dropped rather than detached,
since the context they belong to is gone.
--
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]
