Code Summary: Implementing Short-Term Memory in AI Applications

Short-Term Memory Implementation

Import Modules for Short-Term Memory Management

This block imports the necessary modules for short-term memory: MongoDBSaver from LangGraph's checkpoint module manages conversation state persistence. MongoClient from PyMongo establishes the MongoDB database connection, and create_agent from LangChain constructs the conversational AI agent.

from langgraph.checkpoint.mongodb import MongoDBSaver
from pymongo import MongoClient
from langchain.agents import create_agent

Create a Checkpointer

This establishes a connection to MongoDB using MongoClient and instantiates MongoDBSaver to create a checkpointer. The checkpointer automatically creates two collections: checkpoints (storing conversation state including messages and context) and checkpoint_writes (logging individual agent actions). This enables the agent to persist and retrieve conversation state across interactions.

client = MongoClient(mongodb_uri)
checkpointer = MongoDBSaver(client)

Create an Agent that uses the Checkpointer

This creates an agent using create_agent() with arguments: model (the LLM), system_prompt (defining agent behavior), tools (empty list indicating no additional tools), and checkpointer (enabling state persistence). By including the checkpointer during agent creation, the agent automatically saves its conversation state after every interaction, allowing it to resume conversations and maintain context within a thread.

system_prompt = """You are a helpful AI assistant."""
agent = create_agent(model, system_prompt=system_prompt, tools=[], checkpointer=checkpointer)

Create a Basic Chat Function

This function facilitates agent interaction with parameters user_id, thread_id, and message. It creates a configuration object {"configurable": {"user_id": user_id, "thread_id": thread_id}} that serves as the primary key for the checkpointer to save and load agent state. The function invokes the agent with the user message and configuration, then returns the agent's response content. The thread_id enables thread isolation, allowing independent conversations with separate memories.

def chat(user_id, thread_id, message):
    config = {"configurable": {"user_id": user_id, "thread_id": thread_id}}
    response = agent.invoke({"messages": [{"role": "user", "content": message}]}, config=config)
    return response['messages'][-1].content