Mosaic AI Agent Framework: Author and deploy a multi-agent system with Genie
This notebook demonstrates how to build a multi-agent system using Mosaic AI Agent Framework and LangGraph, where Genie is one of the agents. In this notebook, you:
- Author a multi-agent system using LangGraph.
- Wrap the LangGraph agent with MLflow
ChatAgent
to ensure compatibility with Databricks features. - Manually test the multi-agent system's output.
- Log and deploy the multi-agent system.
This example is based on LangGraph documentation - Multi-agent supervisor example
Why use a Genie agent?
Multi-agent systems consist of multiple AI agents working together, each with specialized capabilities. As one of those agents, Genie allows users to interact with their structured data using natural language.
Unlike SQL functions which can only run pre-defined queries, Genie has the flexibility to create novel queries to answer user questions.
Prerequisites
Define the multi-agent system
Create a multi-agent system in LangGraph using a supervisor agent node directing the following agent nodes:
- GenieAgent: The Genie agent that queries and reasons over structured data.
- Tool-calling agent: An agent that calls Unity Catalog function tools.
In this example, the tool-calling agent uses the built-in Unity Catalog function system.ai.python_exec
to execute Python code.
For examples of other tools you can add to your agents, see Databricks documentation (AWS | Azure).
Wrap the LangGraph agent using the ChatAgent
interface
Databricks recommends using ChatAgent
to ensure compatibility with Databricks AI features and to simplify authoring multi-turn conversational agents using an open source standard.
The LangGraphChatAgent
class implements the ChatAgent
interface to wrap the LangGraph agent.
See MLflow's ChatAgent documentation.
Write agent code to file
Define the agent code in a single cell below. This lets you write the agent code to a local Python file, using the %%writefile
magic command, for subsequent logging and deployment.
Create a Personal Access Token (PAT) as a Databricks secret
In order to access the Genie Space and its underlying resources, we need to create a PAT
- This can either be your own PAT or that of a System Principal (AWS | Azure). You will have to rotate this token yourself upon expiry.
- Add secrets-based environment variables to a model serving endpoint (AWS | Azure).
- You can reference the table in the deploy docs for the right permissions level for each resource: (AWS | Azure).
- Provision with
CAN RUN
on the Genie Space - Provision with
CAN USE
on the SQL Warehouse powering the Genie Space - Provision with
SELECT
on underlying Unity Catalog Tables - Provision with
EXECUTE
on underyling Unity Catalog Functions
- Provision with
Test the agent
Interact with the agent to test its output. Since this notebook called mlflow.langchain.autolog()
you can view the trace for each step the agent takes.
TODO: Replace this placeholder input_example
with a domain-specific prompt for your agent.
Log the agent as an MLflow model
Log the agent as code from the agent.py
file. See MLflow - Models from Code.
Enable automatic authentication for Databricks resources
For the most common Databricks resource types, Databricks supports and recommends declaring resource dependencies for the agent upfront during logging. This enables automatic authentication passthrough when you deploy the agent. With automatic authentication passthrough, Databricks automatically provisions, rotates, and manages short-lived credentials to securely access these resource dependencies from within the agent endpoint.
To enable automatic authentication, specify the dependent Databricks resources when calling mlflow.pyfunc.log_model().
- TODO: If your Unity Catalog tool queries a vector search index or leverages external functions, you need to include the dependent vector search index and UC connection objects, respectively, as resources. See docs (AWS | Azure).
Pre-deployment agent validation
Before registering and deploying the agent, perform pre-deployment checks using the mlflow.models.predict() API. See Databricks documentation (AWS | Azure)."
Register the model to Unity Catalog
Update the catalog
, schema
, and model_name
below to register the MLflow model to Unity Catalog.