Tracing a LangGraph Application with Agent Engine in Vertex AI

This notebook is adapted from Google's "Building and Deploying a LangGraph Application with Agent Engine in Vertex AI" | Original Author: Kristopher Overholt

Overview

Agent Engine is a managed service that helps you to build and deploy agent frameworks. LangGraph is a library for building stateful, multi-actor applications with LLMs, used to create agent and multi-agent workflows.

This notebook demonstrates how to build, deploy, and test a simple LangGraph application using Agent Engine in Vertex AI. You'll learn how to combine LangGraph's workflow orchestration with the scalability of Vertex AI, which enables you to build custom generative AI applications.

Note that the approach used in this notebook defines a custom application template in Agent Engine, which can be extended to LangChain or other orchestration frameworks. If just want to use Agent Engine to build agentic generative AI applications, refer to the documentation for developing with the LangChain template in Agent Engine.

This notebook covers the following steps:

• Define Tools: Create custom Python functions to act as tools your AI application can use.
• Define Router: Set up routing logic to control conversation flow and tool selection.
• Build a LangGraph Application: Structure your application using LangGraph, including the Gemini model and custom tools that you define.
• Local Testing: Test your LangGraph application locally to ensure functionality.
• Deploying to Vertex AI: Seamlessly deploy your LangGraph application to Agent Engine for scalable execution.
• Remote Testing: Interact with your deployed application through Vertex AI, testing its functionality in a production-like environment.
• Cleaning Up Resources: Delete your deployed application on Vertex AI to avoid incurring unnecessary charges.

By the end of this notebook, you'll have the skills and knowledge to build and deploy your own custom generative AI applications using LangGraph, Agent Engine, and Vertex AI.

Get started

Install Vertex AI SDK and other required packages

Restart runtime

To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.

The restart might take a minute or longer. After it's restarted, continue to the next step.

Authenticate your notebook environment (Colab only)

If you're running this notebook on Google Colab, run the cell below to authenticate your environment.

Set Google Cloud project information and initialize Vertex AI SDK

To get started using Vertex AI, you must have an existing Google Cloud project and enable the Vertex AI API.

Learn more about setting up a project and a development environment.

Building and deploying a LangGraph app on Agent Engine

In the following sections, we'll walk through the process of building and deploying a LangGraph application using Agent Engine in Vertex AI.

Import libraries

Import the necessary Python libraries. These libraries provide the tools we need to interact with LangGraph, Vertex AI, and other components of our application.

Define tools

You'll start by defining the a tool for your LangGraph application. You'll define a custom Python function that act as tools in our agentic application.

In this case, we'll define a simple tool that returns a product description based on the product that the user asks about. In reality, you can write functions to call APIs, query databases, or anything other tasks that you might want your agent to be able to use.

Define router

Then, you'll define a router to control the flow of the conversation, determining which tool to use based on user input or the state of the interaction. Here we'll use a simple router setup, and you can customize the behavior of your router to handle multiple tools, custom logic, or multi-agent workflows.

Set Arize AX variables: Space ID, Api Key.

You'll need to set Arize variables below to send traces to the Arize AX Platform. Sign up for free here.

Define LangGraph application

Now you'll bring everything together to define your LangGraph application as a custom template in Agent Engine.

This application will use the tool and router that you just defined. LangGraph provides a powerful way to structure these interactions and leverage the capabilities of LLMs.

Local testing

In this section, you'll test your LangGraph app locally before deploying it to ensure that it behaves as expected before deployment.

Deploy your LangGraph app

Now that you verified that your LangGraph application is working locally, it's time to deploy it to Agent Engine! This will make your application accessible remotely and allow you to integrate it into larger systems or provide it as a service.

Remote test

Now that your LangGraph app is running on Agent Engine, let's test it out by querying it in the remote environment:

Cleaning up

After you've finished experimenting, it's a good practice to clean up your cloud resources. You can delete the deployed Agent Engine instance to avoid any unexpected charges on your Google Cloud account.