memory-demo.py

memory-demo.py

from langgraph.store.mongodb import MongoDBStore
from pymongo.collection import Collection
from pymongo.database import Database
from pymongo import MongoClient
import os
from langgraph.checkpoint.mongodb import MongoDBSaver
from dotenv import load_dotenv
from langchain_openai import AzureOpenAIEmbeddings
from langchain_openai import AzureChatOpenAI

load_dotenv()

embedding_model = AzureOpenAIEmbeddings(
    model="text-embedding-ada-002",
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    openai_api_version="2025-01-01-preview",
)
llm = AzureChatOpenAI(
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    api_version="2025-01-01-preview",
    model="gpt-4o",
)
# Mongo Client
mongo_conn_str = "mongodb://localhost:27017?retryWrites=true&w=majority&directConnection=true"
mongo_client = MongoClient(mongo_conn_str)

## Init checkpointer - used to save short term conversation state
mdb_checkpointer = MongoDBSaver(mongo_client)

## Init Store - used for Long Term Memory to save/retrieve memories
mdb_store = MongoDBStore(
    collection=Collection(
        database=Database(name="long-term-memory", client=mongo_client),
        name="my-memories",
    ),
    index_config={
        "embed": embedding_model,  # Your embedding model
        "dims": len(embedding_model.embed_query("")),  # Embedding dimensions
        "fields": ["$"],  # Need to find docs on why this works
        "filters": None,
    },
)


from langgraph.prebuilt import create_react_agent
from langmem import create_manage_memory_tool, create_search_memory_tool

## Init Agent

tools = [
    create_manage_memory_tool(namespace="memories"),
    create_search_memory_tool(namespace="memories"),
]

agent = create_react_agent(
    llm,  # compliant LLM class
    tools=tools,
    store=mdb_store,
    checkpointer=mdb_checkpointer,
)

# Start conversation
config = {"configurable": {"thread_id": "1"}}
response = agent.invoke({"messages": "MongoDB is my favorite database"}, config=config)
for message in response["messages"]:
    message.pretty_print()


"""
If you are using Python 3.10 upgrade, 
the langmem library requires a feature available only in Python 3.11 and newer.
"""