Databricks GenAI Engineer: Building Agents with MLflow, Agent Framework and Genie

You author a Databricks agent by wrapping your agent logic in a class that subclasses mlflow.pyfunc.ResponsesAgent and implementing predict (and, for streaming, predict_stream). ResponsesAgent is framework-agnostic: a LangGraph graph, an OpenAI SDK loop or plain Python stays exactly as it is, and the wrapper only translates between the Responses-style request/response format and whatever your agent uses internally.

In return for that thin wrapper, MLflow infers the model signature for you, traces predict and predict_stream automatically, and Databricks surfaces (AI Playground, the Review App, inference tables, monitoring) understand the agent's input and output. ChatAgent is the older interface that served the same purpose; new agents use ResponsesAgent.

Streaming

An agent that implements only predict returns its answer when generation is finished, so a chat window shows nothing and then everything. To stream token by token, implement predict_stream and yield events: a series of response.output_text.delta events that share one item_id, followed by a response.output_item.done event carrying the complete item. The done event matters beyond the UI: it is what MLflow uses to aggregate the streamed output into the trace. A client-side "stream" flag or a faster endpoint cannot stream an agent that never yields partial output.

Per-request extras: custom_inputs and custom_outputs

Values that are not chat messages but change on every request (a session ID, a customer tier, a feature flag) travel in the request's custom_inputs field, which the agent reads inside predict; it can return extra data in custom_outputs. AI Playground and the Review App let testers type that JSON. Compare the alternatives: model_config is read once when the model loads (fixed per logged version), endpoint environment variables are shared by all requests, and stuffing values into a system message makes them text the LLM sees rather than typed fields your code can rely on.

Databricks recommends logging agents with Models from Code: put the agent in a Python file (for example agent.py) that ends with mlflow.models.set_model(agent), and pass that file path to log_model (python_model= for pyfunc, lc_model= for LangChain). MLflow stores the code, and when the model is loaded or served it restores the environment and re-executes the file.

This sidesteps pickling. Agents typically hold objects that cannot be serialised (a WorkspaceClient, a compiled LangGraph graph, open connections); logging an instantiated object fails on them, and workarounds such as cloudpickle tricks or serialising a graph to JSON are brittle. With Models from Code those objects are simply re-created when the file runs.

import mlflow
from mlflow.models.resources import DatabricksServingEndpoint

with mlflow.start_run():
    info = mlflow.pyfunc.log_model(
        name="agent",
        python_model="agent.py",          # a file path, not an object
        code_paths=["helpers"],           # local modules agent.py imports
        resources=[DatabricksServingEndpoint(endpoint_name="my-llm-endpoint")],
        pip_requirements=["mlflow", "databricks-langchain"],
    )

Only the code file is captured automatically. If agent.py imports a sibling package such as helpers, loading the logged model fails with ModuleNotFoundError until that folder is passed in code_paths.

Validate before you deploy

Calling the in-memory agent in a notebook proves only that the notebook works. mlflow.models.predict(model_uri, input_data=...) loads the logged artifact (optionally in a fresh environment built from its recorded requirements) and runs a sample request, which is exactly what serving will do. It catches missing modules, missing dependencies and load-time errors before registration and deployment.

A deployed agent reaches other Databricks resources in one of two ways: automatic authentication passthrough, where the agent acts as a system identity with short-lived credentials for the resources you declared, or on-behalf-of-user (OBO) authorization, where it acts as the person who sent the request. Both avoid long-lived secrets; they differ in whose permissions apply.

Automatic passthrough: declare every resource at log time

You list the resources the agent calls when you log it, for example DatabricksServingEndpoint, DatabricksVectorSearchIndex, DatabricksFunction (a Unity Catalog function tool), DatabricksGenieSpace, DatabricksTable and DatabricksSQLWarehouse. At deployment Databricks provisions short-lived credentials scoped to exactly those resources and rotates them. Two consequences follow:

• Undeclared means unreachable. A call to anything not on the list fails with a permission error, even if the LLM and retrieval calls in the same agent work. Granting rights to a person, or giving someone CAN MANAGE, does not change the credential's scope.
• The list belongs to the logged version. If you add a tool (say a UC function) after logging, redeploying the same version cannot help; log a new version that declares it and deploy that.

A Genie space is not self-contained: running its generated SQL needs the space's tables, the SQL warehouse and any UC functions it uses, so those must be declared alongside DatabricksGenieSpace.

On-behalf-of-user authorization

When results must respect each person's permissions (row filters, column masks, per-user audit), use OBO after a workspace admin has enabled it. The agent declares the REST API scopes it needs in a UserAuthPolicy and, inside predict, builds a WorkspaceClient with ModelServingUserCredentials (from databricks_ai_bridge). It must be created per request because the caller is known only then; a client built in __init__ has no user context and would otherwise reuse one identity for everyone.

The user's credential is restricted to the declared scopes, so a downstream call succeeds only if the scope is declared and that user can access the resource. If some users succeed and others fail through the same version, the scopes and code are fine; the failing users need access to the resource. Granting the service principal changes nothing for user-authorized calls.

Mixing both

One agent can do both: pass AuthPolicy(system_auth_policy=SystemAuthPolicy(resources=[...]), user_auth_policy=UserAuthPolicy(api_scopes=[...])) to log_model. Shared resources that users should not access directly (a documentation index, the LLM endpoint) go through system credentials; resources whose results must vary per person (a Genie space over row-filtered tables) go through user credentials.

You instrument an agent by enabling autologging for the frameworks it uses and decorating or wrapping the custom code that autologging cannot see. A trace is a tree of spans, each recording inputs, outputs, latency and exceptions for one step, and it is the primary tool for debugging wrong answers.

• Autologging (mlflow.langchain.autolog(), mlflow.openai.autolog() and others) instruments every call into a supported library with one line: a span per LLM call, retriever and tool. Use the integration that matches the library actually called.
• @mlflow.trace turns a custom Python function into a span, nested in the same trace as the autologged spans. A span_type such as TOOL labels what kind of step it is.
• with mlflow.start_span(...) creates a span for any block of code, such as a loop inside a larger function, where a decorator does not fit.

Decorating only the top-level entry function gives a single span, which hides where time and errors occur. Calling autolog again inside a custom function adds nothing, because autologging never instruments non-library code.

databricks.agents.deploy() takes a Unity Catalog model name and version and, in one call, sets up a scalable Model Serving endpoint with short-lived credentials for the declared resources, a Review App, real-time tracing to the agent's MLflow experiment, and AI Gateway inference tables (production monitoring hooks into the same traces). It does not create data assets such as Vector Search indexes or Genie spaces, and it does not fine-tune anything.

The Review App

The Review App is a web chat interface for stakeholders who do not use notebooks. They talk to the agent and give feedback, which is stored as assessments on the agent's traces in its MLflow experiment, and it also hosts labeling sessions where experts attach expected facts to records. AI Playground, the experiment UI and the endpoint's query page are tools for engineers, not structured feedback collection.

Rolling out a new version

To replace version 3 with version 4 on a live endpoint, call agents.deploy() again with the same UC model name and the new version. The existing endpoint is updated and traffic moves when the new version is ready, so the URL stays the same and there is no gap. The principle is that the endpoint is tied to the UC model name: keep the name and the version change happens in place, while anything that tears down or renames the endpoint interrupts clients. A UC alias such as champion is a pointer for people and pipelines; the endpoint serves the version it was deployed with until you deploy again.

Current-docs note: Databricks now also promotes building new agents with Agent Bricks and Databricks Apps; agents.deploy() to Model Serving remains documented and is the path the exam guide's "Agent Framework" refers to.

Evaluation answers "is this version good enough to ship?" and runs before release on a curated dataset; monitoring answers "is the live app still good?" and runs after release on sampled production traces. They use the same scorers, which is what makes their numbers comparable.

Reuse your evaluation scorers in monitoring

Built-in judges, Guidelines scorers and custom scorers that ran in mlflow.genai.evaluate can be registered and started on the production experiment unchanged. The principle: a quality number is only comparable across phases if the same scorer definition produces it and it is applied to real traffic. Re-implementing a scorer elsewhere quietly changes what it measures, an offline dataset never shows how live users behave, and end-user feedback is a complementary signal rather than the same measurement.

Sampling and filtering

Each registered scorer is started with its own ScorerSamplingConfig:

• sample_rate: the fraction of eligible traces scored. Use 1.0 for checks policy requires on every response (safety) and a small fraction (Databricks suggests roughly 0.05 to 0.2) for expensive LLM judges.
• filter_string: which traces are eligible, in mlflow.search_traces filter syntax (for example a tag such as environment = 'production'), so test traffic sharing the experiment is excluded.

Results appear as assessments on each evaluated trace in the experiment's Traces tab (initial processing can take a few minutes), not as new evaluation runs, inference-table columns or labeling tasks.

Comparing versions fairly

Monitoring scores whatever traffic each version receives. If version 6 got easier questions than version 5, its higher live score proves nothing. To decide between versions, run mlflow.genai.evaluate for both on the same labeled dataset with the scorer that measures the quality in question.

Some built-in judges compare an answer with expectations (expected facts, an expected response, per-row guidelines); others judge the trace on its own. Production traces carry no expectations, so only reference-free judges belong in monitoring, while ground-truth judges stay in offline evaluation over a labeled dataset.

Note the near-twins: groundedness compares the answer with the retrieved chunks (catches unsupported claims), relevance compares chunks or answer with the request, and sufficiency compares the context with expected facts. Using the answer as its own reference ("expected response = actual response") makes every answer pass and measures nothing.

Closing the loop: traces to evaluation datasets

Monitoring finds problems; evaluation datasets keep them fixed. MLflow evaluation datasets are Unity Catalog tables, created with mlflow.genai.datasets.create_dataset and governed and versioned like other UC assets. The loop is:

1. Find problem traces with mlflow.search_traces (filter by tag, assessment, time).
2. Add them with dataset.merge_records(traces). This copies each trace's request (and response) into the dataset but adds no expectations unless the trace already carries expectation assessments (for example from an earlier labeling session); ordinary production traces have none.
3. Have domain experts add expected facts, typically in a labeling session in the Review App.
4. Evaluate every future version against the dataset with mlflow.genai.evaluate.

Until step 3 is done, Correctness cannot produce usable assessments for the new records, while reference-free judges such as Safety work immediately. The principle: a regression set is only useful if it is governed, versioned and re-run against every candidate version. Changing how much traffic monitoring scores, or keeping failing cases outside a managed dataset, does not give you that repeatable check.

An agent should use a Genie space for natural-language questions over governed, structured tables and a retriever (a Mosaic AI Vector Search index) for unstructured documents. Genie translates a question into SQL, runs it on a SQL warehouse under Unity Catalog permissions and returns exact results; similarity search returns the top few matching chunks.

The principle: aggregation needs every qualifying row, and similarity search by design returns only a few. Any variation that keeps retrieval in charge of structured data still asks an LLM to do arithmetic over a partial sample, or answers only the questions someone pre-computed. Unity Catalog SQL functions as tools work for known questions; a hand-built text-to-SQL tool is workable but you own its prompts, execution and conversation handling. For unpredictable questions with the least custom code, a Genie space is the fit. (Databricks' current docs also call Genie spaces "Genie agents"; the API paths still say spaces.)

Adding Genie to a code-based multi-agent system

In a LangGraph supervisor built in code, GenieAgent from databricks_langchain wraps a Genie space as a worker node. You create it with the space ID and a description, and it calls the Genie Conversation API for you (no polling code to maintain).

• Descriptions drive routing. The supervisor chooses workers by reading their descriptions; if none is given, the space's own description is used. With neither, data questions look no different from document questions and get misrouted. Say which data and question types the worker answers.
• One GenieAgent per space. A second space (say HR headcount, owned and permissioned by another team) becomes a second worker with its own description. Merging tables into one broad space breaks ownership and focus; Databricks recommends narrowly focused spaces.

Agent Bricks' no-code Multi-Agent Supervisor, which can also include Genie spaces, is covered with chains and tools in Section 1.

You improve Genie's answers by curating the space itself: example SQL queries, SQL expressions for metrics, table and column descriptions with synonyms, and (as a last resort) text instructions. Genie generates SQL from the knowledge in the space, not from the prompt of the agent that calls it, so supervisor prompts, GenieAgent descriptions and post-processing in one agent cannot fix wrong SQL; they also leave every other team using the space with wrong answers.

Databricks' Genie best-practice guidance ranks the levers: complete example SQL for hard, multi-part questions; SQL expressions for reusable metric definitions; clear descriptions and synonyms; and plain-text instructions only when the structured options do not apply. Changing a shared UC column comment or adding a renaming view would work technically but alter shared metadata or add objects to maintain.

The Genie Conversation API is asynchronous: you start a conversation, poll the resulting message until it reaches a final status, then fetch the query result from its attachment.

1. POST /api/2.0/genie/spaces/{space_id}/start-conversation with the question. The response holds a conversation ID and message ID; no result yet, because Genie is still generating and running SQL.
2. GET .../conversations/{conversation_id}/messages/{message_id} repeatedly until the status is conclusive: COMPLETED, FAILED or CANCELLED.
3. On COMPLETED, GET .../messages/{message_id}/attachments/{attachment_id}/query-result to read the rows.
4. Follow-ups ("and only for EMEA?") go to the same conversation with POST .../conversations/{conversation_id}/messages, so Genie interprets them in context.

Fetching the result once straight after starting hits an unfinished message; re-sending the question starts new work each time.

Polling and conversation hygiene

• Poll every 1 to 5 seconds, back off exponentially up to one minute between polls for long queries, and stop after about 10 minutes with a timeout message. Short initial polls keep quick answers quick, backoff limits request volume and rate-limit errors, and the timeout bounds the user's wait.
• One conversation per user session. Starting a conversation for every message loses follow-up context; reusing one thread across days drags stale context along and can reduce accuracy (spaces also cap stored conversations). Start a new conversation when a new session begins.

Permissions: Genie runs SQL as the caller

Genie executes its generated SQL with the caller's identity on the space's SQL warehouse, so the caller (user or service principal) needs:

• CAN RUN (or higher) on the Genie space, and the Databricks SQL entitlement;
• CAN USE on the space's pro or serverless SQL warehouse;
• Unity Catalog SELECT on the tables plus USE CATALOG on the catalog and USE SCHEMA on the schema (SELECT on a table alone is not enough; an error naming the catalog and schema points at the missing USE privileges).

CAN MANAGE on the space or warehouse, audit-table access and serving-endpoint permissions are not part of answering questions.

Because the caller's identity is what row filters evaluate, an app that sends every employee's question with one service principal (OAuth machine-to-machine) shows everyone that principal's rows. For per-person results and audit trails, authenticate each request with the employee's own OAuth user token (or, inside an agent, on-behalf-of-user authorization). Prompt text asking Genie to "filter by user" or per-department principals cannot enforce per-person access.