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
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.
Define the chat model
Create a LangGraph chat model that supports LangGraph tool calling.
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.
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.
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.