Developing AI Agents on Databricks with Databricks Apps and MLFlow

Databricks announced new capabilities for developing and deploying AI agents using its platform, integrating with LangGraph and MLFlow to enable agent-based systems that extend LLMs with custom tools and reasoning. The approach addresses LLM knowledge limitations through tool integration and Unity Catalog functions, allowing organizations to build autonomous applications on Databricks.

As organizations seek to unlock the full potential of AI, they are increasingly adopting agent-based systems to enable more sophisticated and autonomous applications and workflows. A traditional LLM based chat application is generally limited to generating responses in reply to user prompts and queries. Agents can augment these foundational LLM apps by incorporating reasoning capabilities and providing access to customized tools. There are many different frameworks and platforms that provide agentic AI development capabilities. Frameworks such as LangGraph, LangChain, CrewAI, Strands and AutoGen provide developers with patterns and packages to build agents from scratch. Once developed, these agents can be hosted and deployed on a variety of different platforms including the major cloud providers and other vendors such as Databricks. This article will walk through the Agentic AI development lifecycle on Databricks from beginning to end, demonstrating how to build an agent using LangGraph, connect it to various tools, and deploy and integrate it with monitoring and tracing. LLMs suffer from a knowledge cutoff problem, where the models knowledge is limited to the data it was trained on along with any supplemental information provided as context in prompts. There are several techniques to overcome this problem, to essentially grant the model more knowledge that extends beyond its training data. One approach is Retrieval Augmented Generation RAG , which allows the LLM to retrieve and information for external ‘knowledge bases’ and incorporate that information in its responses. Another way is the use of agents, which extend the LLM with tools that can be custom developed to perform virtually any task, such as interacting with APIs or querying data sources. Once a tool has been developed, it must be exposed in a way that allows the agent to discover and invoke it. Within the Databricks ecosystem, there are two common approaches for making tools available to AI agents: This section will walk through how to develop and register tools as Unity Catalog functions. The following section will show how to build and expose tools through a custom MCP server. In Databricks, unity catalog functions can be written in SQL or Python. We will write our tool as a python unity catalog function. This is a simple tool that makes a request to the open-meteo api to get the upcoming weeks weather forecast for a specific region: The core databricks package use to create unity catalog functionsfrom unitycatalog.ai.core.databricks import DatabricksFunctionClient Set the catalog and schema to populate the function undercatalog name = "main"schema name = "default"client = DatabricksFunctionClient execution mode="serverless" Custom logic to get weather forecastdef get weekly weather forecast latitude: float, longitude: float - str: """ Returns the upcoming 7-day weather forecast for a given location using the Open-Meteo API. Args: latitude: Latitude of the location e.g. 47.6062 for Seattle . longitude: Longitude of the location e.g. -122.3321 for Seattle . Returns: A JSON string containing daily forecasts for the next 7 days, including max/min temperature Celsius and weather condition codes. """ import json import urllib.request url = f"https://api.open-meteo.com/v1/forecast" f"?latitude={latitude}&longitude={longitude}" f"&daily=weathercode,temperature 2m max,temperature 2m min" f"&forecast days=7&timezone=auto" with urllib.request.urlopen url as response: data = json.loads response.read .decode daily = data "daily" forecast = for i in range len daily "time" : forecast.append { "date": daily "time" i , "weather code": daily "weathercode" i , "temp max c": daily "temperature 2m max" i , "temp min c": daily "temperature 2m min" i , } return json.dumps {"location": {"latitude": latitude, "longitude": longitude}, "forecast": forecast} Create the unity catalog function under the specified catalog and schemafunction info = client.create python function func=get weekly weather forecast, catalog=catalog name, schema=schema name, replace=True, Used for local test executionsif name == " main ": function info = client.create python function func=get weekly weather forecast, catalog=catalog name, schema=schema name, replace=True, print f"Registered: {function info.full name}" result = client.execute function function name=f"{catalog name}.{schema name}.get weekly weather forecast", parameters={ "latitude": 47.6062, "longitude": -122.3321, }, print result.value Take note of the lengthy docstring in the get-weekly-weather-forecast function. An agent may have access to dozens upon dozens of different tools, and it is imperative that the agent select the most appropriate tool to respond to a user’s prompt. Metadata, including the function name, parameter definitions, type hints, and especially the docstring, will help the agent understand the tool’s purpose and capabilities. Informative docstrings greatly improve tool selection accuracy because they provide semantic descriptions of what the tool does, the inputs it expects, and the outputs it returns. A clear and descriptive docstring can be the difference between an agent selecting the correct tool versus choosing a less relevant one. You can verify that your function is populating in the correct catalog and schema via the Databricks portal: Databricks offers two modalities to deploy custom agents and MCP servers: The official Databricks documentation recommends using Databricks Apps over Model Serving https://docs.databricks.com/aws/en/generative-ai/agent-framework/migrate-agent-to-apps for your agent hosting needs. Databricks Apps provides a managed hosting and compute environment that allows developers to package and deploy custom applications in their Databricks workspace. This makes Databricks Apps a common deployment engine for AI agents, MCP servers, internal tools, and any other user-facing web application. At a minimum, a Databricks App typically contains the following core components: my-mcp-server-app/├── app.py Application code entrypoint├── server.py FastMCP code├── pyproject.toml Project metadata, dependencies and build settings└── app.yaml Specifies things like runtime settings and env vars In this section we will build a FastMCP server on Databricks Apps in accordance with the above structure: python app.py import osimport uvicornfrom fastapi import FastAPIfrom server import mcpapp = FastAPI title="MCP Tutorial Server" @app.get "/health" def health : return {"status": "ok"}app.mount "/", mcp.http app transport="sse" if name == " main ": port = int os.environ.get "DATABRICKS APP PORT", 8000 uvicorn.run app, host="0.0.0.0", port=port python server.py import jsonimport urllib.requestfrom fastmcp import FastMCPmcp = FastMCP "tutorial-server" @mcp.tool def get exchange rate base currency: str, target currency: str - dict: """ Fetches the live exchange rate between two currencies using the Open Exchange Rates API. Use this tool whenever a user asks about current currency conversion rates. LLMs cannot know real-time exchange rates — always use this tool for accurate results. Args: base currency: The currency to convert from e.g. 'USD', 'EUR', 'GBP' . target currency: The currency to convert to e.g. 'JPY', 'CAD', 'AUD' . Returns: A dictionary with the base currency, target currency, and current exchange rate. """ base = base currency.upper target = target currency.upper url = f"https://open.er-api.com/v6/latest/{base}" with urllib.request.urlopen url as response: data = json.loads response.read .decode if data.get "result" = "success": raise ValueError f"Failed to fetch exchange rates for {base}" if target not in data "rates" : raise ValueError f"Unknown currency: {target}" return { "base currency": base, "target currency": target, "rate": data "rates" target , "last updated": data "time last update utc" , }if name == " main ": mcp.run app.yaml command: 'uv', 'run', 'python', 'app.py', pyproject.toml project name = "mcp-server-app"version = "0.1.0"requires-python = " =3.10"dependencies = "fastmcp =2.0.0", "mcp =1.0.0", "fastapi", "uvicorn", tool.uv package = false a. Place those files in your Databricks workspace b. Create your Databricks App through the UI. Apps can also be created through the Databricks CLI or SDK c. Once your Databricks App is created, you can deploy your source code to it. Pass the path to the folder containing your FastMCP files: If everything in your source code is configured correctly, then your app should have started successfully. Please note that this article is not intended to give a detailed overview of what MCP is. At a high level, an MCP server exposes tools that can be discovered and invoked by agents. Many software vendors provide their own MCP servers that expose platform-specific capabilities, enabling agents to interact with their products through a standardized interface. But for this tutorial we are developing a basic MCP server that exposes one get-exchange-rate tool that fetches real-time exchange rate info. In the upcoming two sections, we will develop our LangGraph agent and connect that agent to our unity catalog tool that fetches weather info as well as our MCP server tool that fetches exchange rate info. We will now create the actual agent with LangGraph. For deployment to Databricks Apps, our LangGraph agent project will follow an identical structure to the MCP server project: my-langgraph-agent-app/├── app.py Application code entrypoint├── agent.py LangGraph code├── pyproject.toml Project metadata, dependencies and build settings└── app.yaml Specifies things like runtime settings and env vars LangGraph provides a framework for building stateful AI applications. While LLMs are inherently stateless, LangGraph introduces explicit state management that allows information to persist across multiple reasoning steps. This enables agents to plan, invoke tools, process results, and continue reasoning until a task is complete. In LangGraph, workflows are modeled as graphs consisting of nodes and edges: The code below defines a very simple ReAct Reason and Act agent. The agent starts by invoking a foundation model hosted through a Databricks Model Serving endpoint. In this tutorial we are using databricks-meta-llama-3–3–70b-instruct If model decides a tool is required to answer the user request, then execution is routed to a ToolNode. After the tool completes, control returns to the model so it can reason over the tool output and determine whether additional actions and tool calls are necessary. The process continues until the model produces a final answer without requesting another tool. This is the agentic loop in a nutshell. python agent.py import operatorimport osfrom typing import Annotated, TypedDictfrom langchain core.messages import SystemMessagefrom langchain databricks import ChatDatabricksfrom langgraph.graph import END, StateGraphfrom langgraph.prebuilt import ToolNodeclass AgentState TypedDict : messages: Annotated list, operator.add def build agent tools: list : llm = ChatDatabricks endpoint=os.environ.get "MODEL ENDPOINT", "databricks-meta-llama-3-3-70b-instruct" llm with tools = llm.bind tools tools system = SystemMessage content="You are a helpful assistant. Only call tools when the user's request clearly requires them. If no tool is relevant, answer directly from your knowledge." def call model state : response = llm with tools.invoke system + state "messages" return {"messages": response } def should continue state : if state "messages" -1 .tool calls: return "tools" return END graph = StateGraph AgentState graph.add node "model", call model graph.add node "tools", ToolNode tools graph.set entry point "model" graph.add conditional edges "model", should continue, {"tools": "tools", END: END} graph.add edge "tools", "model" return graph.compile The resulting graph consists of only two nodes: model and tools but illustrates the core concepts used by many production AI agents. Now that our agent is in place, we will deploy it to Databricks apps in the next section and test it’s functionality with access tools and without access to tools. We will create our Databricks app and deploy our LangGraph agent to it. The complete sample code for the agent can be found here: https://github.com/junnkim9393/databricks-agent-tutorial/tree/main/databricks-apps/langgraph-agent https://github.com/junnkim9393/databricks-agent-tutorial/tree/main/databricks-apps/langgraph-agent That code should be synced to folder in your Databricks workspace, and the path to that folder will be passed to the Databricks app: The agent is exposed through a FastAPI application defined in app.py https://github.com/junnkim9393/databricks-agent-tutorial/blob/main/databricks-apps/langgraph-agent/app.py , which acts as the entrypoint for the Databricks App. Once deployed, Databricks Apps hosts the application and provides a URL that can be used to access the agent through the browser. First, let’s prove that the LangGraph agent, without access to either the unity catalog function or custom MCP server , cannot provide information about upcoming weather or real time exchange rate info. We can easily turn on and off access to the tools and MCP server by commenting out the env var values here in app.yaml https://github.com/junnkim9393/databricks-agent-tutorial/blob/main/databricks-apps/langgraph-agent/app.yaml When asked about weather without access to the get-weekly-weather-forecast tool: When asked about exchange rates without access to the get-exchange-rate-info tool: After exposing our agent to the tools, it can now provide exact answers about weather: and exchange rates: As you can see, tools are a powerful way to augment LLMs with real-time information and external capabilities. To integrate a Databricks-hosted LangGraph agent with external tools, two key libraries are typically used: one for Unity Catalog functions and another for MCP Model Context Protocol servers. For Unity Catalog functions , the primary integration package is databricks-langchain. This library provides the UCFunctionToolkit https://github.com/junnkim9393/databricks-agent-tutorial/blob/main/databricks-apps/langgraph-agent/app.py L32 , which wraps Unity Catalog functions and exposes them as LangChain-compatible tools. This allows the agent to discover and invoke governed Databricks functions as standard agent tools within a LangGraph workflow. For MCP servers , the key package is langchain-mcp-adapters. This library provides the MultiServerMCPClient https://github.com/junnkim9393/databricks-agent-tutorial/blob/main/databricks-apps/langgraph-agent/app.py L40 , which connects to one or more MCP servers and automatically converts the tools exposed by those servers into LangChain tools. Once connected, the agent can access the MCP capabilities as tools within its reasoning loop. Both these patterns are implemented in app.py https://github.com/junnkim9393/databricks-agent-tutorial/blob/main/databricks-apps/langgraph-agent/app.py and can be modified as needed. NOTE: In this tutorial, both the LangGraph agent and the custom MCP server are deployed as Databricks Apps. Each Databricks App is associated with a service principal that governs its access permissions. To allow the LangGraph agent to connect to the MCP server, we must grant the LangGraph App’s service principal Can Use permission on the MCP App. Now that our agent is working end to end, we want to add some baseline tracing and observability to it. In Databricks we can leverage MLFlow for this purpose. MLFlow is an observability framework that logs agent calls. It can be used for agent debugging and evaluation. To get started with MLFlow: a . Add mlflow as a dependency to be installed in the Databricks app. b . Add the following lines to your app’s entrypoint. This is the bare minimum needed to setup mlflow tracing. python app.py import mlflowmlflow.langchain.autolog c. Create an MLFlow experiment. An experiment is a container for traces and logs. There are different ways to create an experiment. In this tutorial, we are using mlflow.set tracking uri “databricks" to point MLflow at the Databricks workspace, then reading the experiment name from the MLFLOW EXPERIMENT NAMEenvironment variable and passing it to mlflow.set experiement experiment name d . View your experiment and agent traces here: e . For the “ weather in la tomorrow” request, you can see the tools that were invoked along with other information: f . For the “ what’s the status of southwest flights” request, the LLM should be unable to handle that request and no tools should be invoked: I hope this blog has been helpful in showcasing the Agentic AI capabilities available on Databricks, and how they can be used to develop production ready AI agents. Developing AI Agents on Databricks with Databricks Apps and MLFlow https://pub.towardsai.net/developing-ai-agents-on-databricks-with-databricks-apps-and-mlflow-3b935d8496b7 was originally published in Towards AI https://pub.towardsai.net on Medium, where people are continuing the conversation by highlighting and responding to this story.