Databricks - AgentX

Agents built on Databricks - declaratively with Agent Bricks or code-first with the Mosaic AI Agent Framework - are auto-instrumented by MLflow 3 Tracing. AgentX plugs into that in three complementary ways; pick per environment, they compose freely:

Path	When	How
Push (live)	Notebooks, jobs, Model Serving endpoints that allow egress	MLflow’s native OTLP export pointed at AgentX’s OTel endpoint
Pull (batch)	Endpoints without egress config, backfills, Unity Catalog trace stores	`agentx-databricks sync` importer
Offline evals	Release gates, regression suites	AgentX dataset runs invoking the serving endpoint

Push: live OTLP export

MLflow 3 exports traces over OpenTelemetry natively. One helper sets the environment variables - call it before the first trace starts:

import os
from agentx.integrations.databricks import enable_mlflow_export

enable_mlflow_export(
    api_key=os.environ["AGENTX_API_KEY"],
    base_url="http://localhost:4700/api/v1",   # your AgentX engine
    service_name="my-databricks-agent",
)

# then trace as usual - @mlflow.trace, autolog, or the Agent Framework

dual=True (default) also keeps MLflow’s own export, so the Databricks MLflow UI and inference tables keep working (it sets MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT).
AgentX maps MLflow’s native span attributes (mlflow.spanInputs/spanOutputs/spanType) directly - inputs, outputs, tool calls, and the full span tree arrive without any flags. genai_semconv=True switches to OTel GenAI semantic conventions if you prefer them.
Deployed endpoints: enable_mlflow_export(dry_run=True, ...) returns the exact variables to paste into the Model Serving endpoint’s environment-variable configuration.

Pull: `agentx-databricks sync`

For serving endpoints where egress env vars are awkward, or to backfill history, the SDK installs a cron-friendly importer that reads finished traces from the MLflow tracking server (pip install "agentx-python[databricks]"):

export AGENTX_API_KEY=agtx_local_...
export AGENTX_API_BASE_URL=http://localhost:4700/api/v1
export DATABRICKS_HOST=... DATABRICKS_TOKEN=...   # or MLFLOW_TRACKING_URI

agentx-databricks sync --experiment-id 123456 --since 7d   # first backfill
agentx-databricks sync --experiment-id 123456              # cron this - incremental cursor

What each MLflow trace becomes:

MLflow	Becomes in AgentX
Trace	Root trace (deterministic `span_id` - re-syncing never duplicates)
Spans	Child spans with real timings - the trace dialog’s Timeline/Graph views
`TOOL` spans	`tool_calls` on the root - tool-failure checks and trajectory matching
`mlflow.trace.session` metadata	Session (`dbx_<id>`) - Sessions view, session judges

Flags mirror agentx-moveworks: --monitor opts imported traces into ingest-time checks (patterns, PII, tool failure), --judge-sessions judges each imported session afterwards (ifStale - never duplicates the engine’s own sweep), --dry-run prints payloads without ingesting, --agent-name pins every trace to one agent.

Offline evaluation of a served agent

The agent stays deployed; an AgentX dataset run drives it - see sample-scripts/sdk_eval_samples/databricks_agent_eval.py for the full version:

def run_agent(case):
    answer = invoke_databricks_agent(case.query)   # POST /serving-endpoints/<name>/invocations
    return {"output": answer}

run = client.evaluations.run(dataset_id=dataset.id, subject={...}).execute(run_agent).finalize()

Cases with expected_tools=[...] are trajectory-matched against the tools the deployed agent actually called, once traces are linked via either path above.

​Push: live OTLP export

​Pull: agentx-databricks sync

​Offline evaluation of a served agent

Push: live OTLP export

Pull: `agentx-databricks sync`

Offline evaluation of a served agent