Build an AI Agent with LangGraph and Atlas Vector Search

On this page

Prerequisites
Set Up the Environment
Use Atlas as a Vector Database
Define Agent Tools
Build the Graph
Add Memory

You can integrate Atlas Vector Search with LangGraph to build AI agents. This tutorial demonstrates how to build a simple agent with LangGraph that answers questions about some sample data in Atlas. You can use the code in this tutorial as a starting point to build more complex AI agents.

Specifically, you perform the following actions:

Set up the environment.
Use Atlas as a vector database.
Define tools for the agent.
Build and run the graph.
Add memory to the agent.

Work with a runnable version of this tutorial as a Python notebook.

Prerequisites

To complete this tutorial, you must have the following:

An Atlas account with a cluster running MongoDB version 6.0.11, 7.0.2, or later (including RCs). Ensure that your IP address is included in your Atlas project's access list. To learn more, see Create a Cluster.
An OpenAI API Key. You must have an OpenAI account with credits available for API requests. To learn more about registering an OpenAI account, see the OpenAI API website.
An environment to run interactive Python notebooks such as Colab.

Set Up the Environment

Set up the environment for this tutorial. Create an interactive Python notebook by saving a file with the .ipynb extension. This notebook allows you to run Python code snippets individually, and you'll use it to run the code in this tutorial.

To set up your notebook environment:

Install and import dependencies.

Run the following command in your notebook:

pip install --quiet --upgrade langgraph langgraph-checkpoint-mongodb langchain langchain_mongodb langchain-openai pymongo

Set environment variables.

Run the following code to set the environment variables for this tutorial. Provide your OpenAI API Key and Atlas cluster's SRV connection string when prompted.

import os
os.environ["OPENAI_API_KEY"] = "<api-key>"
MONGODB_URI = "<connection-string>"

Note

Your connection string should use the following format:

mongodb+srv://<db_username>:<db_password>@<clusterName>.<hostname>.mongodb.net

Use Atlas as a Vector Database

You will use Atlas as the vector database to store and retrieve documents for the agent. To quickly start using Atlas as a vector database:

Load the sample data.

For this tutorial, you use one of our sample datasets as the data source, so you can start building the agent's workflow right away. If you haven't already, complete the steps to load sample data into your Atlas cluster.

Specifically, you will use the embedded_movies dataset, which contains documents about movies, including the vector embeddings of their plots.

Note

If you want to use your own data, see LangChain Get Started or How to Create Vector Embeddings to learn how to ingest vector embeddings into Atlas.

Instantiate the vector store.

In your notebook, paste the following code to configure Atlas as a vector database by using the LangChain integration. Specifically, the code instantiates a vector store object that you can use to interact with Atlas as a vector database. It specifies the following:

The sample_mflix.embedded_movies collection as the data source that contains the vector embeddings and text data.
OpenAI's text-embedding-ada-002 embedding model as the model used to convert text into embeddings during queries.
plot as the field in the collection that contains the text.
plot_embedding as the field in the collection that contains the embeddings.
dotProduct as the relevance score function to use for vector search.

from langchain_mongodb import MongoDBAtlasVectorSearch
from langchain_openai import OpenAIEmbeddings
from pymongo import MongoClient
# Connect to your Atlas cluster
client = MongoClient(MONGODB_URI)
collection = client["sample_mflix"]["embedded_movies"]
embedding_model = OpenAIEmbeddings(model="text-embedding-ada-002", disallowed_special=())
# Instantiate the vector store
vector_store = MongoDBAtlasVectorSearch(
   collection = collection,
   embedding = embedding_model,
   text_key = "plot",
   embedding_key = "plot_embedding",
   relevance_score_fn = "dotProduct"
)

Create the indexes.

Note

To create an Atlas Vector Search index, you must have Project Data Access Admin or higher access to the Atlas project.

To enable vector search and full-text search queries on your data in Atlas, create an Atlas Vector Search and Atlas Search index on the collection. You can create the indexes by using either the LangChain helper methods or the PyMongo Driver method:

Create the Atlas Vector Search index.

Run the following code to create a vector search index that indexes the plot_embedding field in the collection.

# Use helper method to create the vector search index
vector_store.create_vector_search_index(
   dimensions = 1536
)

Tip

create_vector_search_index API Reference

Create the Atlas Search index.

Run the following code in your notebook to create a search index that indexes the title field in the collection.

from langchain_mongodb.index import create_fulltext_search_index
from pymongo import MongoClient
# Connect to your cluster
client = MongoClient(ATLAS_CONNECTION_STRING)
# Use helper method to create the search index
create_fulltext_search_index(
   collection = client["sample_mflix"]["embedded_movies"],
   field = "title",
   index_name = "search_index"
)

Tip

create_fulltext_search_index API Reference

Create the Atlas Vector Search index.

Run the following code to create a vector search index that indexes the plot_embedding field in the collection.

from pymongo.operations import SearchIndexModel
# Create your vector search index model, then create the index
vector_index_model = SearchIndexModel(
   definition={
      "fields": [
         {
         "type": "vector",
         "path": "plot_embedding",
         "numDimensions": 1536,
         "similarity": "dotProduct"
         }
      ]
   },
   name="vector_index",
   type="vectorSearch"
)
collection.create_search_index(model=vector_index_model)

Create the Atlas Search index.

Run the following code to create a search index that indexes the title field in the collection.

1 # Create your search index model, then create the search index
2 search_index_model = SearchIndexModel(
3    definition={
4       "mappings": {
5             "dynamic": False,
6             "fields": {
7                "plot": {
8                   "type": "title"
9                }
10             }
11       }
12    },
13    name="search_index"
14 )
15 collection.create_search_index(model=search_index_model)

The indexes should take about one minute to build. While they build, the indexes are in an initial sync state. When they finish building, you can start querying the data in your collection.

Define Agent Tools

In this section, you define tools that the agent can use to perform specific tasks, and then bind these tools to the LLM. You define the following tools:

Vector search tool to retrieve movies that are semantically similar to the user query.
Full-text search tool to find a specific movie title and retrieve its plot.

Paste and run the following code in your notebook to define and test the tools:

Define a tool for vector search.

This tool uses the vector store object as a retriever. Under the hood, the retriever runs an Atlas Vector Search query to retrieve semantically similar documents. The tool then returns the titles and plots of the retrieved movie documents.

from langchain.agents import tool
# Define a vector search tool
@tool
def vector_search(user_query: str) -> str:
    """
    Retrieve information using vector search to answer a user query.
    """
    
    retriever = vector_store.as_retriever(
       search_type = "similarity",
       search_kwargs = { "k": 5 } # Retrieve top 5 most similar documents
    )
    results = retriever.invoke(user_query)
   
    # Concatenate the results into a string
    context = "\n\n".join([f"{doc.metadata['title']}: {doc.page_content}" for doc in results])
    return context
# Test the tool
test_results = vector_search.invoke("What are some movies that take place in the ocean?")
print(test_results)

20,000 Leagues Under the Sea: In the 19th century, an expert marine biologist is hired by the government to determine what's sinking ships all over the ocean. His daughter follows him. They are intercepted by a mysterious captain Nemo and his incredible submarine.
Deep Rising: A group of heavily armed hijackers board a luxury ocean liner in the South Pacific Ocean to loot it, only to do battle with a series of large-sized, tentacled, man-eating sea creatures who have taken over the ship first.
Lost River: A single mother is swept into a dark underworld, while her teenage son discovers a road that leads him to a secret underwater town.
Waterworld: In a future where the polar ice-caps have melted and Earth is almost entirely submerged, a mutated mariner fights starvation and outlaw "smokers," and reluctantly helps a woman and a young girl try to find dry land.
Poseidon: On New Year's Eve, the luxury ocean liner Poseidon capsizes after being swamped by a rogue wave. The survivors are left to fight for their lives as they attempt to escape the sinking ship.

Define a tool for full-text search.

This tool uses the full-text search retriever to retrieve movie documents that match the specified movie title. Then, the tool returns the plot of the specified movie.

from langchain_mongodb.retrievers.full_text_search import MongoDBAtlasFullTextSearchRetriever
# Define a full-text search tool
@tool
def full_text_search(user_query: str) -> str:
    """
    Retrieve movie plot content based on the provided title.
    """
    
    # Initialize the retriever
    retriever = MongoDBAtlasFullTextSearchRetriever(
       collection = collection,            # MongoDB Collection in Atlas
       search_field = "title",             # Name of the field to search
       search_index_name = "search_index", # Name of the search index
       top_k = 1,                          # Number of top results to return       
    ) 
    results = retriever.invoke(user_query)
   
    for doc in results:
      if doc:
          return doc.metadata["fullplot"]
      else:
          return "Movie not found"
  
# Test the tool     
full_text_search.invoke("What is the plot of Titanic?")

"The plot focuses on the romances of two couples upon the doomed ship's maiden voyage. Isabella Paradine (Catherine Zeta-Jones) is a wealthy woman mourning the loss of her aunt, who reignites a romance with former flame Wynn Park (Peter Gallagher). Meanwhile, a charming ne'er-do-well named Jamie Perse (Mike Doyle) steals a ticket for the ship, and falls for a sweet innocent Irish girl on board. But their romance is threatened by the villainous Simon Doonan (Tim Curry), who has discovered about the ticket and makes Jamie his unwilling accomplice, as well as having sinister plans for the girl."

Prepare the LLM.

The following code prepares the LLM for the agent by doing the following:

Specifies which LLM to use. By default, the ChatOpenAI class uses gpt-3.5-turbo.
Defines a LangChain prompt template to instruct the LLM on how to generate responses, including how to handle tool calls.
Binds the tools and prompt template to the LLM.

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
# Initialize the LLM
llm = ChatOpenAI()
# Create a chat prompt template for the agent, which includes a system prompt and a placeholder for `messages`
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "You are a helpful AI agent."
            " You are provided with tools to answer questions about movies."
            " Think step-by-step and use these tools to get the information required to answer the user query."
            " Do not re-run tools unless absolutely necessary."
            " If you are not able to get enough information using the tools, reply with I DON'T KNOW."
            " You have access to the following tools: {tool_names}."
        ),
        MessagesPlaceholder(variable_name="messages"),
    ]
)
 
tools = [
    vector_search,
    full_text_search
]
# Provide the tool names to the prompt
prompt = prompt.partial(tool_names=", ".join([tool.name for tool in tools]))
# Prepare the LLM by making the tools and prompt available to the model
bind_tools = llm.bind_tools(tools)
llm_with_tools = prompt | bind_tools

Test the tool calls.

You can run the following code snippets to test that the LLM makes the correct tool calls based on the query by checking the name of the tool that the LLM is calling:

# Here, we expect the LLM to use the 'vector_search' tool.
llm_with_tools.invoke(["What are some movies that take place in the ocean?"]).tool_calls

[{'name': 'vector_search',
 'args': {'user_query': 'movies that take place in the ocean'},
 'id': 'call_gBrzDrB35i3bafwWMt5YJQ3E',
 'type': 'tool_call'}]

# Here, we expect the LLM to use the 'full_text_search' tool.
llm_with_tools.invoke(["What's the plot of Titanic?"]).tool_calls

[{'name': 'full_text_search',
  'args': {'user_query': 'Titanic'},
  'id': 'call_rxrOG8DuHWzhVvaai7NHMNTU',
  'type': 'tool_call'}]

Note

You can define any tool that you need to perform a specific task. You can also define tools for other retrieval methods, such as hybrid search or parent-document retrieval.

Build the Graph

In this section, you build a graph to orchestrate the agentic workflow. The graph defines the sequence of steps that the agent takes to respond to a query.

This agent uses the following workflow:

The agent receives a user query.
In the agent node, the tool-bound LLM generates a response based on the query.
This response includes information about whether the agent should use a tool. If the agent determines a tool is needed, it adds the tool configuration to the graph state and proceeds to the tools node. Otherwise, the agent generates an answer directly without using a tool.
If the response indicates a tool is needed, the workflow continues to the tools node.
In this node, the agent reads the tool configuration from the graph state.
Back in the agent node, the LLM receives the retrieved context and generates a final response.

Diagram that shows the workflow of the LangGraph-MongoDB agent.

click to enlarge

Paste and run the following code in your notebook to build and run the graph:

Define the graph state.

The graph state maintains the state of the graph throughout the workflow. It can contain any shared data that needs to be tracked and modified across different nodes. In this example, the GraphState component uses a dictionary that tracks the agent's messages, which includes the user query, the LLM's responses, and the results of tool calls. However, you can customize your graph state to include any data relevant to your application.

from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
# Define the graph state
class GraphState(TypedDict):
    messages: Annotated[list, add_messages]
# Instantiate the graph
graph = StateGraph(GraphState)

Define the nodes.

For this agent, you define two custom nodes :

Add the agent node.

This node processes the messages in the current state, invokes the LLM with these messages, and updates the state with the LLM's response, which includes any tool calls.

from typing import Dict, List
# Define the agent node function
def agent(state: GraphState) -> Dict[str, List]:
    """
    Agent node
    Args:
        state (GraphState): Graph state
    Returns:
        Dict[str, List]: Updates to messages
    """
    # Get the messages from the graph `state`
    messages = state["messages"]
    # Invoke `llm_with_tools` with `messages`
    result = llm_with_tools.invoke(messages)
    # Write `result` to the `messages` attribute of the graph state
    return {"messages": [result]}
# Add "agent" node using the `add_node` function
graph.add_node("agent", agent)

Add the tools node.

This node processes tool calls, determines the appropriate tool to use based on the current state, and updates the message history with the results of the tool call.

from langchain_core.messages import ToolMessage
# Create a map of tool name to tool call
tools_by_name = {tool.name: tool for tool in tools}
# Define the tools node function
def tools_node(state: GraphState) -> Dict[str, List]:
    result = []
    # Get the list of tool calls from messages
    tool_calls = state["messages"][-1].tool_calls
    # Iterate through `tool_calls`
    for tool_call in tool_calls:
        # Get the tool from `tools_by_name` using the `name` attribute of the `tool_call`
        tool = tools_by_name[tool_call["name"]]
        # Invoke the `tool` using the `args` attribute of the `tool_call`
        observation = tool.invoke(tool_call["args"])
        # Append the result of executing the tool to the `result` list as a ToolMessage
        result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
    # Write `result` to the `messages` attribute of the graph state
    return {"messages": result}
# Add "tools" node using the `add_node` function
graph.add_node("tools", tools_node)

Define the edges.

Edges connect the nodes in the graph and define the flow of the agent. In this code, you define the following edges:

The following normal edges that route:
- Start node to agent node.
- Agent node to tools node.
A conditional edge that routes the tools node to the agent node if the state contains tool calls. Otherwise, routes to the end node.

from langgraph.graph import END
# Add an edge from the START node to the `agent` node
graph.add_edge(START, "agent")
# Add an edge from the `tools` node to the `agent` node
graph.add_edge("tools", "agent")
# Define a conditional edge
def route_tools(state: GraphState):
    """
    Uses a conditional_edge to route to the tools node if the last message
    has tool calls. Otherwise, route to the end.
    """
    # Get messages from graph state
    messages = state.get("messages", [])
    if len(messages) > 0:
        # Get the last AI message from messages
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    
    # Check if the last message has tool calls
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
    return END
# Add a conditional edge from the `agent` node to the `tools` node
graph.add_conditional_edges(
    "agent",
    route_tools,
    {"tools": "tools", END: END},
)

Compile the graph.

Run the following code to compile the graph:

app = graph.compile()

Optionally, you can also run the following code to visualize the graph:

from IPython.display import Image, display
try:
    display(Image(app.get_graph().draw_mermaid_png()))
except Exception:
    pass

Define the execution function.

Define the execution function to run the agent's workflow. The following execution function streams outputs from the graph as they progress through the nodes, so you can see the agent's outputs in real time:

# Stream outputs from the graph as they pass through its nodes
def execute_graph(user_input: str) -> None:
    # Add user input to the messages attribute of the graph state
    input = {"messages": [("user", user_input)]}
    # Pass input to the graph and stream the outputs
    for output in app.stream(input):
        for key, value in output.items():
            print(f"Node {key}:")
            print(value)
            
    print("\n---FINAL ANSWER---")
    print(value["messages"][-1].content)

Test the agent.

Run the following sample queries to test the agent. Your generated responses might vary.

execute_graph("What are some movies that take place in the ocean?")

Node agent:
{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_ubFQjnq0s4GKoRAuumpaLxSV', 'function': {'arguments': '{"user_query":"movies that take place in the ocean"}', 'name': 'vector_search'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 22, 'prompt_tokens': 241, 'total_tokens': 263, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-d2cd95a6-27b1-4f1e-a173-51535209c99d-0', tool_calls=[{'name': 'vector_search', 'args': {'user_query': 'movies that take place in the ocean'}, 'id': 'call_ubFQjnq0s4GKoRAuumpaLxSV', 'type': 'tool_call'}], usage_metadata={'input_tokens': 241, 'output_tokens': 22, 'total_tokens': 263, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
Node tools:
{'messages': [ToolMessage(content='20,000 Leagues Under the Sea: In the 19th century, an expert marine biologist is hired by the government to determine what\'s sinking ships all over the ocean. His daughter follows him. They are intercepted by a mysterious captain Nemo and his incredible submarine.\n\nDeep Rising: A group of heavily armed hijackers board a luxury ocean liner in the South Pacific Ocean to loot it, only to do battle with a series of large-sized, tentacled, man-eating sea creatures who have taken over the ship first.\n\nLost River: A single mother is swept into a dark underworld, while her teenage son discovers a road that leads him to a secret underwater town.\n\nWaterworld: In a future where the polar ice-caps have melted and Earth is almost entirely submerged, a mutated mariner fights starvation and outlaw "smokers," and reluctantly helps a woman and a young girl try to find dry land.\n\nDagon: A boating accident runs a young man and woman ashore in a decrepit Spanish fishing town which they discover is in the grips of an ancient sea god and its monstrous half human offspring.', id='11001e4d-fef1-4abe-8700-d720876b5dce', tool_call_id='call_ubFQjnq0s4GKoRAuumpaLxSV')]}
Node agent:
{'messages': [AIMessage(content='Some movies that take place in the ocean are:\n1. 20,000 Leagues Under the Sea\n2. Deep Rising\n3. Lost River\n4. Waterworld\n5. Dagon', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 43, 'prompt_tokens': 495, 'total_tokens': 538, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-7832ca28-bac1-4e80-a7b9-76cc85034ce7-0', usage_metadata={'input_tokens': 495, 'output_tokens': 43, 'total_tokens': 538, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
---FINAL ANSWER---
Some movies that take place in the ocean are:
1. 20,000 Leagues Under the Sea
2. Deep Rising
3. Lost River
4. Waterworld
5. Dagon

execute_graph("What's the plot of Titanic?")

Node agent:
{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_FovVlygLymvbxDzNeEfQGedG', 'function': {'arguments': '{"user_query":"Titanic"}', 'name': 'full_text_search'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 237, 'total_tokens': 255, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-6ee12cbd-0c4a-451f-b56b-83851359d0bb-0', tool_calls=[{'name': 'full_text_search', 'args': {'user_query': 'Titanic'}, 'id': 'call_FovVlygLymvbxDzNeEfQGedG', 'type': 'tool_call'}], usage_metadata={'input_tokens': 237, 'output_tokens': 18, 'total_tokens': 255, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
Node tools:
{'messages': [ToolMessage(content="The plot focuses on the romances of two couples upon the doomed ship's maiden voyage. Isabella Paradine (Catherine Zeta-Jones) is a wealthy woman mourning the loss of her aunt, who reignites a romance with former flame Wynn Park (Peter Gallagher). Meanwhile, a charming ne'er-do-well named Jamie Perse (Mike Doyle) steals a ticket for the ship, and falls for a sweet innocent Irish girl on board. But their romance is threatened by the villainous Simon Doonan (Tim Curry), who has discovered about the ticket and makes Jamie his unwilling accomplice, as well as having sinister plans for the girl.", id='2cd41281-d195-44af-9ae1-f3ff099194a9', tool_call_id='call_FovVlygLymvbxDzNeEfQGedG')]}
Node agent:
{'messages': [AIMessage(content='The plot of "Titanic" focuses on the romances of two couples on the doomed ship\'s maiden voyage. Isabella Paradine, a wealthy woman, reunites with her former flame, Wynn Park. Meanwhile, a charming ne\'er-do-well named Jamie Perse falls for an innocent Irish girl on board, but their romance is threatened by a villainous character named Simon Doonan.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 83, 'prompt_tokens': 395, 'total_tokens': 478, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-cdf65abe-7ce0-417a-8f5b-84989521f47e-0', usage_metadata={'input_tokens': 395, 'output_tokens': 83, 'total_tokens': 478, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
---FINAL ANSWER---
The plot of "Titanic" focuses on the romances of two couples on the doomed ship's maiden voyage. Isabella Paradine, a wealthy woman, reunites with her former flame, Wynn Park. Meanwhile, a charming ne'er-do-well named Jamie Perse falls for an innocent Irish girl on board, but their romance is threatened by a villainous character named Simon Doonan.

Add Memory

To improve the agent's performance, you can persist its state. Persistence allows the agent to store information about previous interactions, which the agent can use in future interactions to provide more contextually relevant responses.

Tip

Initialize the checkpointer.

To persist the state of the graph execution, you can use checkpointers to save the state to a specific thread, which can be accessed even after the graph execution ends.

The MongoDBSaver checkpointer allows you to use MongoDB as the backing database for persisting the checkpoint state. Run the following code in your notebook to initialize the checkpointer and use it in your graph:

from langgraph.checkpoint.mongodb import MongoDBSaver
# Initialize a MongoDB checkpointer
checkpointer = MongoDBSaver(client)
# Instantiate the graph with the checkpointer
app = graph.compile(checkpointer=checkpointer)

Update the execution function.

You must also update the execution function to reference the thread_id, which is the unique identifier for the thread that you want to persist. Run the following code in your notebook to update the execution function:

# Update the `execute_graph` function to include the `thread_id` argument
def execute_graph(thread_id: str, user_input: str) -> None:
    config = {"configurable": {"thread_id": thread_id}}
    input = {
        "messages": [
            (
                "user",
                user_input,
            )
        ]
    }
    for output in app.stream(input, config):
        for key, value in output.items():
            print(f"Node {key}:")
            print(value)
            
    print("\n---FINAL ANSWER---")
    print(value["messages"][-1].content)

Test the agent.

Run the following code snippets in your notebook to test the agent. Your generated responses might vary.

The first code snippet runs a query and saves the response to a thread with the specified thread_id of 1.
The second code snippet runs a query about the previous interaction, loading the state from the thread with the specified thread_id of 1. It generates a response that is aware of the previous interaction.

execute_graph("1", "What's the plot of Titanic?")

Node agent:
{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_ZgBYIdPqV720s3oN7TC61Sjn', 'function': {'arguments': '{"user_query":"Titanic"}', 'name': 'full_text_search'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 860, 'total_tokens': 878, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-91a84f0d-ddba-4753-8de6-6db1d059f238-0', tool_calls=[{'name': 'full_text_search', 'args': {'user_query': 'Titanic'}, 'id': 'call_ZgBYIdPqV720s3oN7TC61Sjn', 'type': 'tool_call'}], usage_metadata={'input_tokens': 860, 'output_tokens': 18, 'total_tokens': 878, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
Node tools:
{'messages': [ToolMessage(content="The plot focuses on the romances of two couples upon the doomed ship's maiden voyage. Isabella Paradine (Catherine Zeta-Jones) is a wealthy woman mourning the loss of her aunt, who reignites a romance with former flame Wynn Park (Peter Gallagher). Meanwhile, a charming ne'er-do-well named Jamie Perse (Mike Doyle) steals a ticket for the ship, and falls for a sweet innocent Irish girl on board. But their romance is threatened by the villainous Simon Doonan (Tim Curry), who has discovered about the ticket and makes Jamie his unwilling accomplice, as well as having sinister plans for the girl.", id='20507bc4-383f-4478-8ffc-9386e423509c', tool_call_id='call_ZgBYIdPqV720s3oN7TC61Sjn')]}
Node agent:
{'messages': [AIMessage(content='The plot of "Titanic" focuses on the romances of two couples aboard the doomed ship\'s maiden voyage. It tells the story of Isabella Paradine, who rekindles a romance with Wynn Park, and Jamie Perse, who falls in love with an Irish girl on board. Their romances are jeopardized by the villainous Simon Doonan\'s sinister plans.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 80, 'prompt_tokens': 1018, 'total_tokens': 1098, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-8b1916d2-b5b4-4d17-be04-589a701e17dc-0', usage_metadata={'input_tokens': 1018, 'output_tokens': 80, 'total_tokens': 1098, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
---FINAL ANSWER---
The plot of "Titanic" focuses on the romances of two couples aboard the doomed ship's maiden voyage. It tells the story of Isabella Paradine, who rekindles a romance with Wynn Park, and Jamie Perse, who falls in love with an Irish girl on board. Their romances are jeopardized by the villainous Simon Doonan's sinister plans.

execute_graph("1", "What movies are similar to the one I just asked about?")

Node agent:
{'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_7hzNqOU0hZBHrm7wihISMrEz', 'function': {'arguments': '{"user_query": "Movies similar to Titanic"}', 'name': 'vector_search'}, 'type': 'function'}, {'id': 'call_OHAkJsyjPGKcCpqye2M56Moy', 'function': {'arguments': '{"user_query": "Titanic"}', 'name': 'vector_search'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 50, 'prompt_tokens': 1394, 'total_tokens': 1444, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-e48b75c7-4493-4dcd-af2e-afb556882052-0', tool_calls=[{'name': 'vector_search', 'args': {'user_query': 'Movies similar to Titanic'}, 'id': 'call_7hzNqOU0hZBHrm7wihISMrEz', 'type': 'tool_call'}, {'name': 'vector_search', 'args': {'user_query': 'Titanic'}, 'id': 'call_OHAkJsyjPGKcCpqye2M56Moy', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1394, 'output_tokens': 50, 'total_tokens': 1444, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
Node tools:
{'messages': [ToolMessage(content="The Poseidon Adventure: A group of passengers struggle to survive and escape when their ocean liner completely capsizes at sea.\n\nPoseidon: On New Year's Eve, the luxury ocean liner Poseidon capsizes after being swamped by a rogue wave. The survivors are left to fight for their lives as they attempt to escape the sinking ship.\n\nLife of Pi: A young man who survives a disaster at sea is hurtled into an epic journey of adventure and discovery. While cast away, he forms an unexpected connection with another survivor: a fearsome Bengal tiger.\n\nTraffickers: A thriller about the passengers with different objectives on board a cruiser headed for China, being chased over and over again and unexpected happening of things.\n\nAfter the Storm: When a luxury yacht goes down in a violent storm the race is on to salvage the bounty at any cost, causing two couples to commit the ultimate betrayal.", id='f1b40d2d-eaf9-4dca-8f4d-0f69eb4b4f3d', tool_call_id='call_7hzNqOU0hZBHrm7wihISMrEz'), ToolMessage(content="Titanic: The story of the 1912 sinking of the largest luxury liner ever built, the tragedy that befell over two thousand of the rich and famous as well as of the poor and unknown passengers aboard the doomed ship.\n\nThe Poseidon Adventure: A group of passengers struggle to survive and escape when their ocean liner completely capsizes at sea.\n\nRaise the Titanic: To obtain a supply of a rare mineral, a ship raising operation is conducted for the only known source, the Titanic.\n\nPoseidon: On New Year's Eve, the luxury ocean liner Poseidon capsizes after being swamped by a rogue wave. The survivors are left to fight for their lives as they attempt to escape the sinking ship.\n\nAll Is Lost: After a collision with a shipping container at sea, a resourceful sailor finds himself, despite all efforts to the contrary, staring his mortality in the face.", id='3551a58d-44d7-4055-a997-97c09dc563ef', tool_call_id='call_OHAkJsyjPGKcCpqye2M56Moy')]}
Node agent:
{'messages': [AIMessage(content='Movies similar to "Titanic" include:\n1. The Poseidon Adventure\n2. Poseidon\n3. Life of Pi\n4. Traffickers\n5. After the Storm\n\nThese movies feature themes of survival, disasters at sea, and unexpected events similar to those in "Titanic."', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 62, 'prompt_tokens': 1827, 'total_tokens': 1889, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-8332baba-75d3-4d18-baf6-75b3cf68b552-0', usage_metadata={'input_tokens': 1827, 'output_tokens': 62, 'total_tokens': 1889, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}
---FINAL ANSWER---
Movies similar to "Titanic" include:
1. The Poseidon Adventure
2. Poseidon
3. Life of Pi
4. Traffickers
5. After the Storm
These movies feature themes of survival, disasters at sea, and unexpected events similar to those in "Titanic."

Back

LangGraph

LangChainGo

1	# Create your search index model, then create the search index
2	search_index_model = SearchIndexModel(
3	definition={
4	"mappings": {
5	"dynamic": False,
6	"fields": {
7	"plot": {
8	"type": "title"
9	}
10	}
11	}
12	},
13	name="search_index"
14	)
15	collection.create_search_index(model=search_index_model)