Skip to main content

Get started with AI agents

Build your first AI agent using Mosaic AI Agent Framework. In this tutorial, you will:

  • Author an agent using Agent Framework.
  • Add a tool to your agent.
  • Deploy your agent to a Databricks model serving endpoint.

For a conceptual introduction to agents and other gen AI apps, see What are gen AI apps?

Requirements​

Your workspace must have the following features enabled:

Example notebook​

This notebook contains all the code you need to author and deploy your first AI agent. Import the notebook to your Databricks workspace to run it.

Mosaic AI agent demo

Open notebook in new tab

Define the agent​

An AI agent consists of the following:

  • A large language model (LLM) that can reason and make decisions
  • Tools that the LLM can use to do more than just generate text, such as running Python code or fetching data

Run the following code in a Databricks notebook to define a simple tool-calling agent:

  1. Install the required Python packages:

    Python
    %pip install -U -qqqq mlflow databricks-openai databricks-agents
    dbutils.library.restartPython()
    • mlflow: Used for agent development and agent tracing.
    • databricks-openai: Used to connect to the Databricks-hosted LLM.
    • databricks-agent: Used to package and deploy the agent.

  2. Define the agent. This code snippet does the following:

    • Connects to the Databricks model serving endpoint using the OpenAI client.
    • Enables MLflow tracing using autolog(). This adds instrumentation so you can see what your agent does when you submit a query.
    • Adds the system.ai.python_exec tool to your agent. This built-in Unity Catalog function allows your agent to run Python code.
    • Defines a function that queries the LLM with a prompt and handles responses.
    Python
    import mlflow
    import json
    from databricks.sdk import WorkspaceClient
    from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient

    # Get an OpenAI client configured to connect to Databricks model serving endpoints
    # Use this client to query the LLM
    openai_client = WorkspaceClient().serving_endpoints.get_open_ai_client()

    # Enable automatic tracing for easier debugging
    mlflow.openai.autolog()

    # Load Databricks built-in tools (Python code interpreter)
    client = DatabricksFunctionClient()
    builtin_tools = UCFunctionToolkit(function_names=["system.ai.python_exec"], client=client).tools
    for tool in builtin_tools:
      del tool["function"]["strict"]


    def call_tool(tool_name, parameters):
      if tool_name == "system__ai__python_exec":
        return DatabricksFunctionClient().execute_function("system.ai.python_exec", parameters=parameters)
      raise ValueError(f"Unknown tool: {tool_name}")

    def run_agent(prompt):
      """
      Send a user prompt to the LLM and return a list of LLM response messages
      The LLM is allowed to call the code interpreter tool, if needed, to respond to the user
      """
      result_msgs = []
      response = openai_client.chat.completions.create(
        model="databricks-claude-3-7-sonnet",
        messages=[{"role": "user", "content": prompt}],
        tools=builtin_tools,
      )
      msg = response.choices[0].message
      result_msgs.append(msg.to_dict())

      # If the model executed a tool, call it
      if msg.tool_calls:
        call = msg.tool_calls[0]
        tool_result = call_tool(call.function.name, json.loads(call.function.arguments))
        result_msgs.append({"role": "tool", "content": tool_result.value, "name": call.function.name, "tool_call_id": call.id})
      return result_msgs

Test the agent​

Test the agent by querying it with a prompt that requires running Python code:

Python
answer = run_agent("What is the 100th fibonacci number?")
for message in answer:
  print(f'{message["role"]}: {message["content"]}')

In addition to the LLM's output, you will see detailed trace information directly in your notebook. These traces help you debug slow or failed agent calls. These traces were automatically added using mlflow.openai.autolog() .

Deploy the agent​

Now that you have an agent, you can package and deploy it to a Databricks serving endpoint. Start collecting feedback on a deployed agent by sharing it with others and chatting with it using a built-in chat UI.

Prepare agent code for deployment​

To prepare your agent code for deployment, wrap it using MLflow's ChatAgent interface. The ChatAgent interface is the recommended way to package agents for deployment on Databricks.

  1. To implement the ChatAgent interface, you must define a predict() function that sends the user's message to the agent, collects the agent's response, and returns it in the ChatAgentResponses format.

    Python
    import uuid
    from typing import Any, Optional

    from mlflow.pyfunc import ChatAgent
    from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext

    class QuickstartAgent(ChatAgent):
      def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
      ) -> ChatAgentResponse:
        # 1. Extract the last user prompt from the input messages
        prompt = messages[-1].content

        # 2. Call run_agent to get back a list of response messages
        raw_msgs = run_agent(prompt)

        # 3. Map each response message into a ChatAgentMessage and return
        # the response
        out = []
        for m in raw_msgs:
          out.append(ChatAgentMessage(
            id=uuid.uuid4().hex,
            **m
          ))

        return ChatAgentResponse(messages=out)

  2. Add the following code to your notebook to test your ChatAgent class:

    Python
    AGENT = QuickstartAgent()
    for response_message in AGENT.predict({"messages": [{"role": "user", "content": "What's the 100th fibonacci number?"}]}).messages:
      print(f"role: {response_message.role}, content: {response_message.content}")

  3. Combine all of your agent code into a single file so you can log and deploy it.

  • Consolidate all of your agent code into one notebook cell.
  • At the top of the cell, add the %%writefile quickstart_agent.py magic command to save your agent to a file.
  • At the bottom of the cell, call mlflow.models.set_model() with your agent object. This tells MLflow which agent object to use when serving predictions. This step effectively configures the entry point to our agent code.

Your notebook cell should look like the following:

Python
%%writefile quickstart_agent.py

import json
import uuid
from databricks.sdk import WorkspaceClient
from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient
from typing import Any, Optional

import mlflow
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext

# Add an mlflow.openai.autolog() call to capture traces in the serving endpoint

# Get an OpenAI client configured to talk to Databricks model serving endpoints
# We'll use this to query an LLM in our agent
openai_client = WorkspaceClient().serving_endpoints.get_open_ai_client()

# Load Databricks built-in tools (a stateless Python code interpreter tool)
client = DatabricksFunctionClient()
builtin_tools = UCFunctionToolkit(function_names=["system.ai.python_exec"], client=client).tools
for tool in builtin_tools:
  del tool["function"]["strict"]


def call_tool(tool_name, parameters):
  if tool_name == "system__ai__python_exec":
    return DatabricksFunctionClient().execute_function("system.ai.python_exec", parameters=parameters)
  raise ValueError(f"Unknown tool: {tool_name}")

def run_agent(prompt):
  """
  Send a user prompt to the LLM, and return a list of LLM response messages
  The LLM is allowed to call the code interpreter tool if needed, to respond to the user
  """
  result_msgs = []
  response = openai_client.chat.completions.create(
    model="databricks-claude-3-7-sonnet",
    messages=[{"role": "user", "content": prompt}],
    tools=builtin_tools,
  )
  msg = response.choices[0].message
  result_msgs.append(msg.to_dict())

  # If the model executed a tool, call it
  if msg.tool_calls:
    call = msg.tool_calls[0]
    tool_result = call_tool(call.function.name, json.loads(call.function.arguments))
    result_msgs.append({"role": "tool", "content": tool_result.value, "name": call.function.name, "tool_call_id": call.id})
  return result_msgs

class QuickstartAgent(ChatAgent):
  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    prompt = messages[-1].content
    raw_msgs = run_agent(prompt)
    out = []
    for m in raw_msgs:
      out.append(ChatAgentMessage(
        id=uuid.uuid4().hex,
        **m
      ))

    return ChatAgentResponse(messages=out)

AGENT = QuickstartAgent()
mlflow.models.set_model(AGENT)

Log the agent​

Log your agent and register it to Unity Catalog. This packages your agent and its dependencies into a single artifact for deployment.

Python
import mlflow
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint
from pkg_resources import get_distribution

# Change the catalog name ("main") and schema name ("default") to register the agent to a different location
registered_model_name = "main.default.quickstart_agent"

# Specify Databricks resources that the agent needs to access.
# This step lets Databricks automatically configure authentication
# so the agent can access these resources when it's deployed.
resources = [
  DatabricksServingEndpoint(endpoint_name="databricks-claude-3-7-sonnet"),
  DatabricksFunction(function_name="system.ai.python_exec"),
]

mlflow.set_registry_uri("databricks-uc")
logged_agent_info = mlflow.pyfunc.log_model(
  artifact_path="agent",
  python_model="quickstart_agent.py",
  extra_pip_requirements=[f"databricks-connect=={get_distribution('databricks-connect').version}"],
  resources=resources,
  registered_model_name=registered_model_name
)

Deploy the agent​

Deploy your registered agent to a serving endpoint:

Python
from databricks import agents

deployment_info = agents.deploy(
  model_name=registered_model_name, model_version=logged_agent_info.registered_model_version
)

After the agent endpoint starts, you can chat with it using AI Playground or share it with stakeholders for feedback.

Next steps​

Choose where to go next based on your goals:

Measure and improve your agent’s quality: See Agent Evaluation quickstart.

Build more advanced agents: Create an agent that performs RAG using unstructured data, handles multi-turn conversations, and uses Agent Evaluation to measure quality. See Tutorial: Build, evaluate, and deploy a retrieval agent.

Learn how to build agents using other frameworks: Learn how to build agents using popular libraries like LangGraph, pure Python, and OpenAI. See Use ChatAgent to author agents

Last updated on
Was this article helpful?