Databricks Agent Framework in Practice: A Guide to ChatAgent Wrapping and Production Deployment

Why You Need Databricks Agent Framework

When building AI agents, developers commonly face a thorny problem: different frameworks operate in silos. OpenAI, LangGraph, and Anthropic each have their own invocation methods, inconsistent input formats, and even more divergent output formats. This fragmentation stems from the historical evolution paths of the major LLM providers. When OpenAI introduced Function Calling in June 2023, it prioritized backward compatibility with its Chat Completions API, choosing to embed a tool_calls array structure within message objects. Anthropic, starting from Claude 3, introduced a content-block-based tool use protocol that treats tool calls as a distinct content type—a choice rooted in a different design philosophy regarding message structure atomicity. LangChain, as an upper-layer framework, abstracted its own AIMessage and ToolCall object systems in order to support multiple underlying models simultaneously. These three systems are incompatible in terms of the JSON Schema format for tool descriptions, the incremental granularity of streaming outputs, and error handling mechanisms—forcing any system that attempts to integrate multiple LLM providers to maintain a complex adaptation layer.

It's worth noting that this protocol fragmentation is not accidental, but rather the combined result of commercial competition and diverging technical roadmaps. OpenAI uses the tool_calls field to describe tool calls, while Anthropic uses tool_use content blocks; in streaming output, OpenAI is based on deltas, while LangChain relies on an event stream. The deeper divergence lies in this: OpenAI's function calling design treats tools as part of the conversation context, whereas Anthropic's content block design treats tool calls as an independent content primitive parallel to text—this design choice affects how tool results are managed in the context of multi-turn conversations. From an engineering cost perspective, this underlying incompatibility forces developers to write dedicated adaptation layers for each framework. Take a medium-sized agent system that integrates three LLM providers simultaneously: the adaptation layer code often accounts for more than 30% of the entire system's codebase, and as provider API versions iterate, this code requires ongoing maintenance. In multi-agent systems, this cost scales exponentially with the number of agents—when you need to mix multiple frameworks in the same system, or build multi-agent collaboration systems, handling these heterogeneous inputs and outputs becomes a real and heavy engineering burden.

Databricks Agent Framework (i.e., the Mosaic AI Agent Framework) is designed precisely to solve this pain point. Its core idea is simple and powerful: wrap agents from various frameworks using a unified ChatAgent interface. Whether you create an agent with LangGraph or OpenAI, you simply wrap it in a ChatAgent to obtain a consistent input/output interface—with no changes required to the core agent logic.

The value of this wrapping mechanism goes beyond interface unification—it naturally supports the construction of multi-agent systems. You can create different agents using multiple LLMs such as OpenAI and Anthropic, and after uniform wrapping, the input/output behavior of the entire system remains consistent, greatly reducing integration complexity.

Implementation Details of ChatAgent Wrapping

Based on the official documentation examples, the wrapping process follows a clear pattern. First, import the required components from MLflow, including ChatAgentRequest, ChatAgentResponse, the streaming event classes, and the ChatAgent base class. Then define the LLM Endpoint, system prompt, bound tools (such as the UC Function Toolkit and Vector Search tools), and build the agent state—this part of the code remains completely unchanged.

The actual wrapping is reflected in four key methods within the ChatAgent class:

• responses_to_cc: Converts MLflow-style messages to LangChain-style messages for the agent to process
• langchain_to_responses: The reverse conversion, turning LangGraph responses back into MLflow format for output
• predict: Executes batch prediction
• predict_stream: Supports triggering the agent in streaming mode

These four methods together constitute the Adapter Pattern implementation of ChatAgent—responses_to_cc and langchain_to_responses are inverse operations of each other, essentially establishing a bidirectional mapping between MLflow's OpenAI-compatible message format and LangChain's BaseMessage system. The classic software engineering definition of the adapter pattern is: enabling incompatible classes to work together without modifying existing interfaces. ChatAgent applies this pattern to the core contradiction of the LLM ecosystem, allowing the same set of agent logic to be reused across multiple format systems. This design keeps the responsibility boundaries of the adaptation layer extremely clear: format conversion logic is concentrated in the two conversion methods, while inference logic remains in the original agent code, with the two not interfering with each other.

This set of interface code all comes from the official documentation and can be used directly as a ready-made helper—developers need not invest much effort in the wrapper itself. In actual testing, when a Python code request to "calculate 6 times 7" is fed into the agent, the model automatically recognizes and invokes the ai_python_executable tool, ultimately returning the result 42, with the entire reasoning process fully recorded in MLflow.

Logging and Registering Models with MLflow

MLflow was first open-sourced by Databricks in 2018, positioned as a machine learning lifecycle management platform. Its early core consisted of four components: Tracking (experiment logging), Projects (code packaging), Models (model standardization), and Registry (model registry). However, the inputs and outputs of traditional ML models are structured tensor or tabular data, whereas the core of LLM applications is unstructured conversation flows, tool call chains, and multi-step reasoning processes—which the original system could not directly accommodate.

MLflow 3 underwent an architecture-level refactoring to address this shift: introducing native support for conversation message formats, incorporating tool call sequences as first-class citizens in the tracking system, and adding an evaluation framework designed specifically for LLM outputs (including an automatic scoring mechanism based on LLM-as-Judge). At the tracking level, MLflow 3 introduces the concept of a "Span" to record the hierarchical structure of the agent's reasoning process: the top-level Span corresponds to a complete user request, while child Spans correspond to each LLM call or tool execution, with each Span carrying complete input/output, latency, and token consumption information. This hierarchical tracking structure allows developers to pinpoint exactly which tool call produced an incorrect intermediate result when debugging multi-step reasoning problems. Notably, the design of Span is inspired by the OpenTelemetry specification from the distributed tracing field—MLflow 3 essentially models the LLM inference process as a distributed request, where each tool call is equivalent to a microservice invocation, allowing DevOps engineers to understand and debug AI systems using familiar distributed tracing thinking. The ChatAgent protocol was established as the standard agent interface during this upgrade—which is the foundation that allows all the features in this article to connect smoothly. A useful tip: use the notebook's magic methods to write the agent code into an agent.py file, achieving decoupling between code and logging.

When logging a model, the key is to explicitly declare all resources required by the agent. During its runtime, an agent interacts with numerous external services—Vector Search indexes, UC Function tools, etc.—all of which must be explicitly declared at the logging stage, along with defining the agent name, code path, and all dependencies. The essence of resource declaration is to make the agent's runtime dependency relationships static: when the model is deployed to a production environment, the Databricks runtime will automatically configure service account permissions based on the declared resource list, ensuring consistent permissions for the agent across different environments, while also providing the basis for building the asset dependency graph for Unity Catalog's lineage tracking. This is analogous to the practice of explicitly declaring dependency images in a Dockerfile during containerized deployment—converting implicit runtime dependencies into auditable, versionable static declarations is an important guarantee for the maintainability of production-grade systems.

Once logging is complete, you can directly make predictions on the logged model via the mlflow.models.predict function—pass in the run id, agent name, and input data, and the system will download the artifacts and return the result. This is an effective means of quickly verifying whether the agent's behavior meets expectations.

Next comes registering to Unity Catalog. Unity Catalog (UC) was officially launched in 2021 as Databricks' unified data governance layer, solving the problems of scattered data permissions and fragmented lineage tracking in multi-workspace environments. Its core design philosophy is to bring data assets and AI assets into the same permission model and lineage tracking system, expanding governance objects from data tables to functions, models, vector indexes, and various other asset types.

In AI agent scenarios, UC manages model versions with a three-level namespace of catalog.schema.model_name, ensuring that model assets across teams and workspaces do not have naming conflicts; UC Functions serve as a tool registry, allowing data engineers to define reusable tool functions in SQL or Python, and these functions automatically gain fine-grained column-level permission control, with agents securely accessing them through the UC permission model; in addition, UC's built-in data lineage capabilities can track the datasets and vector indexes accessed during agent reasoning, meeting the compliance audit requirements of heavily regulated industries such as finance and healthcare. From a practical operations perspective, the three-level namespace also provides natural permission isolation boundaries for enterprise-level teams: different business lines can independently manage agent versions under their own schemas, precisely controlling cross-team model access permissions through UC's GRANT/REVOKE syntax. Set up the Databricks Unity Catalog registry, specify the catalog, schema, and model name, and call the MLflow registration interface. Each registration automatically increments the version number—this is exactly the core capability of version management in LLMOps. After you make modifications to an agent, version control enables a complete CI/CD process.

Built-in Evaluation: Measuring Agent Quality in One Line of Code

MLflow 3 introduces various "scorers" that can quantify the quality of agent output across dimensions such as relevance and safety, with about three scorers built into the general version that can be used directly.

The evaluation approach is similar to writing test cases: provide the input (such as "calculate the 15th Fibonacci number") and the expected response, then use MLflow to run the evaluation, passing in the evaluation dataset, the predict function, and the scorers (such as relevance_to_query). After running, you can see a detailed report—the question, response content, token consumption, execution time, status, and whether relevance and safety pass.

The reason this evaluation capability is so easy to implement is fundamentally the deep integration between the Agent Framework and MLflow. MLflow 3's evaluation framework adopts the LLM-as-Judge mechanism, which invokes an independent large language model as a judge to automatically score agent outputs based on preset scoring criteria. The core assumption of LLM-as-Judge is: for open-domain agent outputs, judging semantic quality (relevance, completeness, factual accuracy) inherently requires language understanding capabilities, which is precisely the strength of LLMs. Traditional evaluation metrics based on string matching or BLEU/ROUGE were originally designed for machine translation scenarios and rely on the vocabulary overlap with reference answers to measure quality. Such metrics cannot capture cases of "different wording but semantically equivalent," nor can they determine whether the agent correctly invoked the tool chain—and these two points are exactly the core concerns of agent quality evaluation. The main limitation of LLM-as-Judge lies in the consistency of the judgment itself (different judge models may produce biased scores for the same output), so in production practice it is usually used in conjunction with human sampling verification. This enables large-scale evaluation of agent quality without manual annotation. In practice, it is recommended that the evaluation dataset cover three categories—normal questions, edge cases, and adversarial inputs—to more comprehensively reflect the agent's actual performance in a production environment. When used together, the evaluation process requires almost no additional configuration.

Deploying as a Model Serving Endpoint

After model registration is complete, it can be directly deployed as a Model Serving Endpoint. Databricks Model Serving is built on a Kubernetes containerized architecture, with specialized optimizations for agent-type workloads—which differs fundamentally from traditional ML inference services: the latency of traditional ML inference (such as image classification or text vectorization) is typically at the millisecond level, requests are stateless, and high-throughput batch processing is easily achievable; whereas the latency of agent inference can range from several seconds to several minutes, because a single user request may trigger dozens of LLM calls and tool executions.

For this reason, Model Serving enables asynchronous request queues and longer timeout windows (600 seconds by default) for agent Endpoints; streaming output is implemented based on the Server-Sent Events (SSE) protocol—SSE is a standard feature of HTTP/1.1, based on a unidirectional persistent connection, where the server can continuously push data frames in text/event-stream format without closing the connection. Compared to WebSocket's bidirectional full-duplex connection, SSE is lower in implementation complexity and can traverse enterprise firewalls without special proxy configuration, making it well-suited for unidirectional push scenarios like "the agent continuously outputs intermediate reasoning steps, and the user only needs to receive them"; in addition, the standard REST API exposed by the Endpoint follows a superset of the OpenAI Chat Completions format, allowing it to be directly integrated into third-party applications that support the OpenAI interface (such as Slack bots and enterprise knowledge base systems) without an additional adaptation layer. At the elastic scaling level, Model Serving supports auto-scaling based on the number of concurrent requests, and offers a "Scale to Zero" option—completely releasing computing resources when there is no traffic, which is especially important for cost control during the development and testing phase.

When deploying using the Agent Framework, you only need to provide the model name and version number in Unity Catalog, and Databricks will run the model on its managed infrastructure and automatically generate a set of standard REST API endpoints. After deployment, you can chat directly with the agent in the Playground, for example asking "how to cancel an order" or "how does the membership plan work," to quickly verify whether the online behavior meets expectations. Note that the Databricks free tier has a limit on the number of Serving Endpoints running simultaneously.

ReAct Agents: A Construction Paradigm for Iterative Reasoning

The ReAct agent is currently one of the most mainstream ways to build agents. The ReAct (Reasoning + Acting) paradigm was formally proposed by Princeton University and Google Research in 2022 in the paper "ReAct: Synergizing Reasoning and Acting in Language Models." Academically, this paradigm addresses the respective shortcomings of two earlier LLM application paradigms: pure Chain-of-Thought (CoT) reasoning, while improving complex reasoning capabilities, locks the model's knowledge in the period before the training data cutoff date, making it unable to access real-time information and prone to hallucinations; pure action sequences, while able to access external information, lack interpretable intermediate reasoning steps, making them difficult to debug and optimize. The core finding is: after combining the two, the model's performance on knowledge-intensive tasks (such as multi-hop question answering and fact-checking) improves significantly, and the hallucination rate also decreases markedly. The original paper's experiments on benchmarks such as HotpotQA and FEVER showed that ReAct improved accuracy by an average of about 10-20 percentage points over pure CoT, while producing fewer factual errors than pure action sequence methods. Its core is the "reasoning-acting-observing" loop: the LLM first analyzes the user's question and decides which tool to invoke; after executing the tool, it observes whether the returned result meets the need; if not, it continues reasoning until it obtains sufficient information before giving a final reply.

Take "what's the weather like in Paris, and what time is it in Tokyo now" as an example. The agent will first call get_weather(Paris) to obtain the sunny weather information, then call get_time(Tokyo) to get the local time, and finally integrate the two pieces of information into a complete answer. The core advantages of this pattern are: iterative reasoning, self-correction, process transparency, and high flexibility, enabling it to reliably handle complex multi-step problems.

LangGraph's create_react_agent greatly simplifies implementation difficulty—at the engineering implementation level, LangGraph models the reasoning loop as a directed graph rather than a simple recursive call: nodes represent LLM calls or tool executions, edges represent state transition conditions, and the graph's termination condition is when the LLM output no longer contains tool call requests. The graph structure design brings key advantages: it naturally supports state persistence (checkpointing), allowing the agent to resume from any node after an interruption; it supports conditional edges, allowing certain reasoning steps to be skipped under specific conditions; and it supports parallel node execution, providing performance optimization space for scenarios requiring multiple tools to be called simultaneously. Compared to early ReAct implementations based on recursive function calls, the graph structure also brings an important debuggability advantage: the input/output state of each node can be independently inspected and modified, allowing developers to inject test data at specific reasoning steps and precisely verify whether the behavior of a single node is correct. Developers only need to provide the model, system prompt, and tool array to automatically obtain complete handling of the reasoning loop, tool call logic, and state transitions.

Prompt Engineering: An Underrated Key Element

Prompt quality directly affects the actual performance of ReAct agents. ReAct relies on the LLM's Instruction Following and Few-shot Learning capabilities—by providing a "Thought → Action → Observation" example format in the system prompt, it guides the model to naturally generate a structured alternating sequence of reasoning and action during generation.

An effective ReAct prompt typically contains three layers: role definition (clarifying the agent's professional identity and capability boundaries), behavior rules (specifying when and how to invoke tools), and few-shot examples (demonstrating a complete example of the "Thought → Action → Observation" loop). Therefore, each tool function should include clear comments explaining its purpose, so the LLM can accurately determine when to invoke which tool. The system prompt should include clear behavior rules, for example: "Follow the ReAct pattern—first think about what information is needed, then invoke the correct tool, and finally integrate the results before returning them to the user." The importance of prompt engineering lies in this: even with the same combination of model and tools, the difference in tool call accuracy between a carefully designed prompt and a hastily written one can be as high as tens of percentage points. In addition, the tool description itself is part of prompt engineering—the tool's function name, parameter names, and docstring are automatically serialized into the LLM's context, and vague tool descriptions can cause the model to be uncertain in tool selection, thereby triggering unnecessary multi-turn reasoning loops.

By integrating LangChain community tools, you can also quickly extend the agent's capabilities. For example, after adding the Wikipedia tool, when the agent simultaneously possesses three capabilities—weather query, time query, and Wikipedia search—it will invoke the required tools in sequence and integrate the output, with the entire process fully traceable and reviewable in MLflow.

Summary

The core value of Databricks Agent Framework lies in providing a complete pipeline from development to production:

ChatAgent unified wrapping → MLflow logging and evaluation → Unity Catalog version management → Model Serving Endpoint deployment

Throughout this process, developers can obtain a standardized interface, comprehensive observability, and one-click deployment capabilities without modifying the core agent logic. The design philosophy of this framework embodies the concrete practice of the "separation of concerns" principle in the LLMOps domain: business logic (the agent's reasoning strategy and tool selection) is clearly isolated from infrastructure logic (format adaptation, permission management, version control, elastic scaling), allowing the two dimensions to evolve independently without interfering with each other. For teams looking to scale AI agents in an enterprise-level environment, this framework significantly reduces engineering complexity and is a technical path worth prioritizing as a reference for building production-grade agent systems.