Langmem
- Lets build a graph which is a personal assistant with Postgress store
-
Personal assistant will remember city and favorite drink
-
Sample non working skeleton code
import os
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.store.postgres import PostgresStore
from dotenv import load_dotenv
from utils import get_llm
from langmem import create_memory_manager, create_memory_store_manager
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from pydantic import BaseModel
from langgraph.checkpoint.postgres import PostgresSaver
from langchain_core.runnables.config import RunnableConfig
load_dotenv()
model = get_llm()
from langchain_google_vertexai.embeddings import VertexAIEmbeddings
def get_embedding(name:str = 'text-embedding-005') -> VertexAIEmbeddings:
return VertexAIEmbeddings(model=name)
class UserPreferences(BaseModel):
name: str
city: str
favorite_drink: str
# setup store
with PostgresStore.from_conn_string(
os.getenv('PG_VECTOR_DB_URI'),
index={
"dims": 768,
"embed": get_embedding(),
"fields": ["memory_text"] # specify which fields to embed. Default is the whole serialized value
}
) as store:
store.setup()
memory_store_manager = create_memory_store_manager(
model,
schemas=[UserPreferences],
store=store,
namespace=("user", "{user_id}", "preferences")
)
# setup checkpointer
with PostgresSaver.from_conn_string(os.getenv('PG_VECTOR_DB_URI')) as saver:
saver.setup()
def get_memories_block(preferences):
return ""
def ask_llm(state: MessagesState, config: RunnableConfig):
# lets get the memory_manager_store
preferences = memory_store_manager.search(query="user preferences")
# from this prefernces
# we will create a prompt
# fetch long term memories
memories_block = get_memories_block(preferences)
system_prompt = "You are personal assitant with the following memories {memory_block}"
response = model.invoke(
{
"messages": [SystemMessage(system_prompt)] + state['messages']
}
)
return {
"messages": response
}
stategraph = StateGraph(MessagesState)
stategraph.add_node("ask", ask_llm)
stategraph.set_entry_point("ask")
stategraph.set_finish_point("ask")
with PostgresSaver.from_conn_string(os.getenv('PG_VECTOR_DB_URI')) as saver:
graph = stategraph.compile(checkpointer=saver)
if __name__ == "__main__":
response = graph.invoke({
"messages": [HumanMessage("Im alice and i like coffee and i live in mumbai")]
config= {
"configurable" : {
"thread_id": "t1",
"user_id": "alice@wonderland"
}
}
})
response = graph.invoke({
"messages": [HumanMessage("where do i live ?")]
config= {
"configurable" : {
"thread_id": "t2",
"user_id": "alice@wonderland"
}
}
})
