custom-langgraph-schema-agent (1)(Python)

Loading...

LangGraph custom schema agent notebook

This notebook shows you how to write a LangGraph AI agent compatible with Mosaic AI Agent Framework that accepts custom inputs and returns custom outputs.

To ensure compatibility, the agent must conform to Mosaic AI Agent Framework schema requirements, see (AWS | Azure).

Model-as-code notebook

Mosaic AI Agent Framework uses MLflow's Models-as-code development workflow, which requires two notebooks:

  • An agent notebook that defines the agent's logic (this notebook)
  • A driver notebook that logs, registers, and deploys the agent
    • You can find the driver notebook for this agent, custom-langgraph-schema-driver, on Databricks documentation (AWS | Azure)

For more information on Model-as-code, see MLflow's Models as code guide.

Requirements

This notebook requires a Unity Catalog enabled workspace.

NOTE: This notebook uses LangGraph, but Mosaic AI Agent Framework is compatible with other agent authoring frameworks, like LlamaIndex.

2

Define the chat model

Create a LangGraph chat model that supports LangGraph tool calling.

4

Parse LangGraph output

The following cell defines helper methods for converting LangGraph output messages into the recommended output schema for Mosaic AI agent framework. The wrap_output helper returns chat-completion compatible messages, with an additional custom_outputs field containing custom outputs.

6

Create the agent

Use the LangGraph create_react_agent function to build a simple graph. For more customization, you can create your own LangGraph agent by following LangGraph - Quick Start.

8

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.

Replace this placeholder input with an appropriate domain-specific example for your agent.

10

11

    Next steps

    You can rerun the cells above to iterate and test the agent.

    See the driver notebook, custom-langgraph-schema-driver (AWS | Azure), to learn how to log, register, and deploy this agent.