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:

1. Author a multi-agent system using LangGraph.
2. Wrap the LangGraph agent with MLflow ChatAgent to ensure compatibility with Databricks features.
3. Manually test the multi-agent system's output.
4. 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

• Address all TODOs in this notebook.
• Create a Genie Space, see Databricks documentation (AWS | Azure).

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

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.

Deploy the agent

Next steps

After your agent is deployed, you can chat with it in AI playground to perform additional checks, share it with SMEs in your organization for feedback, or embed it in a production application. See Databricks documentation (AWS | Azure).