Add AI Trace observability panel to the common.ai provider#69880
Draft
kaxil wants to merge 1 commit into
Draft
Conversation
Adds an AI Trace AirflowPlugin: a FastAPI sub-app mounted on the API
server plus a React panel on the UI, so a GenAI agent run (AgentOperator
/ @task.agent / @task.llm) can be inspected in context -- conversation,
model, tokens, cost, latency, and the full observation tree with a
duration waterfall -- from the task instance and from a deployment-wide
list, without leaving Airflow.
The panel reads traces from two sources:
- A configured tracing backend (Langfuse today): the trace is resolved
from the task instance's own OTel span context and read back through
the backend's API. No new data is written anywhere.
- A backend-free ObjectStorage trace store: when [common.ai]
trace_store_path is set, agent spans are written as standard OTLP JSON
lines under {dag_id}/{run_id}/{task_id}/{map_index}/{try_number}.jsonl
and the panel reads them directly. This needs no Langfuse, no
collector, and no core tracing -- just a path -- and targets local dev
and quick testing. Files are standard OTLP JSON, so a collector's
otlpjsonfilereceiver can replay them into a real backend later.
Costs are estimated at read time via genai-prices (already a transitive
dependency of pydantic-ai-slim) from the model, provider, and token
counts on each span, so old files are priced retroactively and unknown
or self-hosted models simply show no cost. The OTLP JSON encoder is
vendored because the upstream opentelemetry-exporter-otlp-json-file
release is currently uninstallable from PyPI (its opentelemetry-proto-json
dependency was never published); it follows the OTLP JSON-Protobuf
encoding so files stay collector-replayable.
Security: every endpoint is RBAC-gated. The task-instance panel checks
TASK_INSTANCE access to the DAG; the bare-id lookups resolve the trace
(or observation) to its owning task instance and require access to that
DAG, so a user authorized for one DAG cannot read another's captured
content. Task-instance coordinates are validated before they reach an
ObjectStoragePath join (it does not normalize ".."), and trace ids are
validated as 32-hex before use in a reverse lookup or file scan.
Includes unit tests for the OTLP JSON encoder (wire-format round trip
against real SDK spans), the store reader and normalizer, read-time cost
estimation, and the endpoints in both modes (RBAC, path-traversal
rejection, trace-id validation, per-DAG authorization, and 404 paths).
Known limitations: the trace-store layout is a proof of concept, not a
stable on-disk contract; no retention is applied; object stores flush a
task's spans at process exit rather than live; cost is an estimate, not
billing data.
fb9f80c to
53e7a0e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds an AI Trace observability panel to the
common.aiprovider: anAirflowPluginthat mounts a FastAPI sub-app on the API server and a React panelon the UI, so a GenAI agent run (
AgentOperator/@task.agent/@task.llm)can be inspected — conversation, model, tokens, cost, latency, and the full
observation tree with a waterfall — from inside Airflow, in context on the task
instance and in a deployment-wide list.
It reads traces from two sources:
instance's trace via its own OTel span context and proxies the backend's read
API. No new data is written anywhere.
[common.ai] trace_store_pathis set, agent spans are written as standard OTLP JSON linesunder an
ObjectStoragePath, and the panel reads them back directly. Thisneeds no Langfuse, no collector, and no core tracing — just a path. It targets
local dev and quick testing.
Screenshots
The AI Trace panel in backend-free store mode — role-ordered conversation, the
observation tree with a duration waterfall, and per-generation cost estimated
from
genai-prices(no "Open in Langfuse" link, since there is no externalbackend in store mode):
Design rationale
worker's task span and inherit its
trace_id, so the panel resolves "this taskinstance's trace" straight from the stored traceparent — no separate push or id
override.
{dag_id}/{run_id}/{task_id}/{map_index}/{try_number}.jsonl, so store mode workswith core tracing off. Files are standard OTLP JSON, so a collector's
otlpjsonfilereceivercan replay them into a real backend later.genai-prices(already a transitivedependency of
pydantic-ai-slim), using the model, provider, and token countson each span. Unlike ingest-time pricing, this stays current with the library's
bundled price data and prices old files retroactively; unknown or self-hosted
models simply show no cost.
_otlp_json.py): the upstreamopentelemetry-exporter-otlp-json-filerelease is currently uninstallable fromPyPI (its
opentelemetry-proto-jsondependency was never published), so a harddependency would break resolution. The vendored encoder follows the OTLP
JSON-Protobuf encoding so the files remain collector-replayable; it can be
swapped for the upstream exporter once that ships.
Usage
With
trace_store_pathunset, the panel instead reads from the connection namedby
[common.ai] trace_backend_conn_id(a Langfuse connection). Open the panelfrom a task instance's AI Trace tab, or the deployment-wide AI Traces nav
entry.
Security
TASK_INSTANCEaccess to the DAG; the bare-id lookups resolve the trace (or observation) to its
owning task instance and require access to that DAG, so a user authorized for
one DAG cannot read another's captured content.
ObjectStoragePathjoin (it does not normalize
..), and trace ids are validated as 32-hex beforeuse in a reverse lookup or file scan.
Known limitations
not a stable on-disk contract. Marked as such in the config docs.
s3:///gs://, a task's spansland when the process exits, not live (local
file://is live).and a shown cost can drift slightly across
genai-pricesupgrades.no task instance (possible in Langfuse mode) returns 404 by direct id, since there
is no DAG to authorize against.