phanikumv commented on code in PR #74769:
URL: https://github.com/apache/airflow/pull/74769#discussion_r4282736583
##########
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:
```suggestion
attributes = {"gen_ai.operation.name": "chat",
"gen_ai.provider.name": "openai", **_usage(data.usage)}
```
--
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]
