Basic working implementation of cross thread memory

agent/__init__.py

@@ -1,3 +1,3 @@
-from agent.graph import create_agent_graph, create_agent_graph_without_memory, get_checkpointer
+from agent.graph import create_agent_graph, create_agent_graph_without_memory, get_checkpointer, langgraph_studio_graph
 
-__all__ = ["create_agent_graph", "create_agent_graph_without_memory", "get_checkpointer"] 
+__all__ = ["create_agent_graph", "create_agent_graph_without_memory", "get_checkpointer", "langgraph_studio_graph"] 

agent/graph.py

@@ -1,29 +1,54 @@
-from langgraph.graph import StateGraph
+from langgraph.graph import StateGraph, END
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from langgraph.store.memory import InMemoryStore
 import aiosqlite
 from types import TracebackType
 from typing import Optional, Type
 
 from agent.utils.state import AgentState
-from agent.utils.nodes import call_model, tool_node, should_continue
+from agent.utils.nodes import (
+    call_model,
+    tool_node,
+    read_memory,
+    write_memory,
+    should_continue
+)
 
 def create_graph_builder():
     """Create a base graph builder with nodes and edges configured."""
-    # Create the graph
     builder = StateGraph(AgentState)
+    
 
     # Add nodes
     builder.add_node("agent", call_model)
     builder.add_node("action", tool_node)
+    builder.add_node("read_memory", read_memory)
+    builder.add_node("write_memory", write_memory)
 
-    # Update edges
+    # Set entry point
     builder.set_entry_point("agent")
+
+    builder.add_edge("agent", "write_memory")
+    builder.add_edge("write_memory", END)
+
+    # Add conditional edges from agent
     builder.add_conditional_edges(
         "agent",
         should_continue,
+        {
+            "action": "action",
+            "read_memory": "read_memory",
+            "write_memory": "write_memory",
+            END: END
+        }
     )
+    
+    # Connect action back to agent
     builder.add_edge("action", "agent")
     
+    # Memory operations should end after completion
+    builder.add_edge("read_memory", "agent")
+    
     return builder
 
 def create_agent_graph_without_memory():
@@ -59,11 +84,21 @@ def get_checkpointer(db_path: str = "data/short_term.db") -> SQLiteCheckpointer:
     """Create and return a SQLiteCheckpointer instance."""
     return SQLiteCheckpointer(db_path)
 
+# Initialize store for across-thread memory
+across_thread_memory = InMemoryStore()
+
 async def create_agent_graph(checkpointer: AsyncSqliteSaver):
-    """Create an agent graph with SQLite-based memory persistence."""
+    """Create an agent graph with memory persistence."""
     builder = create_graph_builder()
-    graph = builder.compile(checkpointer=checkpointer)
+    # Compile with both SQLite checkpointer for within-thread memory
+    # and InMemoryStore for across-thread memory
+    graph = builder.compile(
+        checkpointer=checkpointer,
+        store=across_thread_memory
+    )
     return graph
 
+langgraph_studio_graph = create_agent_graph_without_memory()
+
 # Export the graph builder functions
 __all__ = ["create_agent_graph", "create_agent_graph_without_memory", "get_checkpointer"] 

agent/utils/nodes.py

@@ -1,21 +1,30 @@
 from langchain_openai import ChatOpenAI
-from langchain_core.messages import SystemMessage
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langgraph.graph import END
 from langgraph.prebuilt import ToolNode
+from langchain.memory import ConversationBufferMemory
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.store.base import BaseStore
+from typing import Literal
+# from chainlit.logger import logger
+
 
 from agent.utils.tools import tool_belt
 from agent.utils.state import AgentState
 
-# Initialize LLM
-llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
-model = llm.bind_tools(tool_belt)
+# Initialize LLM for memory operations
+model = ChatOpenAI(model="gpt-4", temperature=0)
+
+# Define system prompt with memory
+SYSTEM_PROMPT = """You are a Chief Joy Officer, an AI assistant focused on helping people find fun and enriching activities in London.
+You have access to memory about the user's preferences and past interactions.
 
-# Define system prompt
-SYSTEM_PROMPT = SystemMessage(content="""You are a Chief Joy Officer, an AI assistant focused on helping people find fun and enriching activities in London.
+Here is what you remember about this user:
+{memory}
 
 Your core objectives are to:
 1. Understand and remember user preferences and interests
-2. Provide personalized activity recommendations
+2. Provide personalized activity recommendations based on their interests
 3. Be engaging and enthusiastic while maintaining professionalism
 4. Give clear, actionable suggestions
 
@@ -23,35 +32,210 @@ Key tools at your disposal:
 - retrieve_context: For finding specific information about events and activities
 - tavily_search: For general web searches about London activities
 
-Always aim to provide value while being mindful of the user's time and interests.""")
+Always aim to provide value while being mindful of the user's time and interests."""
 
-# Define memory prompt
-MEMORY_PROMPT = """Here is the conversation history and relevant information about the user:
+# Define memory creation/update prompt
+MEMORY_UPDATE_PROMPT = """You are analyzing the conversation to update the user's profile and preferences.
 
+CURRENT USER INFORMATION:
 {memory}
 
-Please use this context to provide more personalized responses. When appropriate, reference past interactions and demonstrated preferences to make your suggestions more relevant.
+INSTRUCTIONS:
+1. Review the chat history carefully
+2. Identify new information about the user, such as:
+   - Activity preferences (indoor/outdoor, cultural/sports, etc.)
+   - Specific interests (art, music, food, etc.)
+   - Location preferences in London
+   - Time/schedule constraints
+   - Past experiences with activities
+   - Budget considerations
+3. Merge new information with existing memory
+4. Format as a clear, bulleted list
+5. If new information conflicts with existing memory, keep the most recent
 
-Remember to:
-1. Acknowledge previously mentioned interests
-2. Build upon past recommendations
-3. Avoid repeating suggestions already discussed
-4. Note any changes in preferences
+Remember: Only include factual information directly stated by the user. Do not make assumptions.
 
-Current conversation:
-{conversation}"""
+Based on the conversation, please update the user information:"""
 
-def call_model(state: AgentState):
-    messages = [SYSTEM_PROMPT] + state["messages"]
+def get_last_human_message(state: AgentState):
+    """Get the last human message from the state."""
+    for message in reversed(state["messages"]):
+        if isinstance(message, HumanMessage):
+            return message
+    return None
+
+def call_model(state: AgentState, config: RunnableConfig, store: BaseStore):
+    """Process messages using memory from the store."""
+    # Get the user ID from the config
+    user_id = config["configurable"].get("session_id", "default")
+    
+    # Retrieve memory from the store
+    namespace = ("memory", user_id)
+    existing_memory = store.get(namespace, "user_memory")
+    
+    # Extract memory content or use default
+    memory_content = existing_memory.value.get('memory') if existing_memory else "No previous information about this user."
+    
+    # Create messages list with system prompt including memory
+    messages = [
+        SystemMessage(content=SYSTEM_PROMPT.format(memory=memory_content))
+    ] + state["messages"]
+    
     response = model.invoke(messages)
     return {"messages": [response]}
 
+def update_memory(state: AgentState, config: RunnableConfig, store: BaseStore):
+    """Update user memory based on conversation."""
+    user_id = config["configurable"].get("session_id", "default")
+    namespace = ("memory", user_id)
+    existing_memory = store.get(namespace, "user_memory")
+    
+    memory_content = existing_memory.value.get('memory') if existing_memory else "No previous information about this user."
+    
+    update_prompt = MEMORY_UPDATE_PROMPT.format(memory=memory_content)
+    new_memory = model.invoke([
+        SystemMessage(content=update_prompt)
+    ] + state["messages"])
+    
+    store.put(namespace, "user_memory", {"memory": new_memory.content})
+    return state
+
+def should_continue(state: AgentState) -> Literal["action", "read_memory", "write_memory", END]:
+    """Determine the next node in the graph."""
+    if not state["messages"]:
+        return END
+        
+    last_message = state["messages"][-1]
+    if isinstance(last_message, list):
+        last_message = last_message[-1]
+        
+    last_human_message = get_last_human_message(state)
+
+    # Handle tool calls
+    if hasattr(last_message, "additional_kwargs") and last_message.additional_kwargs.get("tool_calls"):
+        return "action"
+    
+    # Handle memory operations for human messages
+    if last_human_message:
+            
+        # Write memory for longer messages that might contain personal information
+        if len(last_human_message.content.split()) > 3:
+            return "write_memory"
+        # Read memory for short queries to ensure personalized responses
+        else:
+            return "read_memory"
+            
+    return END
+
+def read_memory(state: AgentState, config: RunnableConfig, store: BaseStore):
+    """Read and apply memory context without updating it."""
+    user_id = config["configurable"].get("session_id", "default")
+    namespace = ("memory", user_id)
+    existing_memory = store.get(namespace, "user_memory")
+    
+    if existing_memory:
+        memory_content = existing_memory.value.get('memory')
+        # Add memory context to state for next model call
+        state["memory_context"] = memory_content
+    
+    return state
+
+# Define the memory creation prompt
+MEMORY_CREATION_PROMPT = """"You are collecting information about the user to personalize your responses.
+
+CURRENT USER INFORMATION:
+{memory}
+
+INSTRUCTIONS:
+1. Review the chat history below carefully
+2. Identify new information about the user, such as:
+   - Personal details (name, location)
+   - Preferences (likes, dislikes)
+   - Interests and hobbies
+   - Past experiences
+   - Goals or future plans
+3. Merge any new information with existing memory
+4. Format the memory as a clear, bulleted list
+5. If new information conflicts with existing memory, keep the most recent version
+
+Remember: Only include factual information directly stated by the user. Do not make assumptions or inferences.
+
+Based on the chat history below, please update the user information:"""
+
+async def write_memory(state: AgentState, config: RunnableConfig, store: BaseStore) -> AgentState:
+    """Reflect on the chat history and save a memory to the store."""
+
+    # Get the session ID from config
+    session_id = config["configurable"].get("session_id", "default")
+    
+    # Define the namespace for this user's memory
+    namespace = ("memory", session_id)
+    
+    # Get existing memory using async interface
+    existing_memory = await store.aget(namespace, "user_memory")
+    memory_content = existing_memory.value.get('memory') if existing_memory else "No previous information about this user."
+    
+    # Create system message with memory context
+    system_msg = SystemMessage(content=MEMORY_CREATION_PROMPT.format(memory=memory_content))
+    
+    # Get messages and ensure we're working with the correct format
+    messages = state.get("messages", [])
+    if not messages:
+        return state
+        
+    # Create memory using the model
+    new_memory = await model.ainvoke([system_msg] + messages)
+    
+    # Store the updated memory using async interface
+    await store.aput(namespace, "user_memory", {"memory": new_memory.content})
+    
+    
+    return state
+
 # Initialize tool node
 tool_node = ToolNode(tool_belt)
 
-# Simple flow control - always go to final
-def should_continue(state):
+def should_call_memory(state: AgentState) -> Literal["update_memory", "end"]:
+    """
+    Determine if we should update memory based on the conversation state.
+    
+    Rules for updating memory:
+    1. Only update after human messages (not tool responses)
+    2. Update if the message might contain personal information
+    3. Don't update for simple queries or acknowledgments
+    """
+    if not state["messages"]:
+        return "end"
+        
     last_message = state["messages"][-1]
-    if last_message.tool_calls:
-        return "action"
-    return END 
+    
+    # Skip memory update for tool calls
+    if hasattr(last_message, "additional_kwargs") and last_message.additional_kwargs.get("tool_calls"):
+        return "agent"
+    
+    # Skip memory update for very short messages (likely acknowledgments)
+    if isinstance(last_message, HumanMessage) and len(last_message.content.split()) <= 3:
+        return "agent"
+    
+    # Update memory for human messages that might contain personal information
+    if isinstance(last_message, HumanMessage):
+        return "update_memory"
+    
+    return "agent"
+
+# def route_message(state: MessagesState, config: RunnableConfig, store: BaseStore) -> Literal[END, "update_todos", "update_instructions", "update_profile"]:
+    
+#     """Reflect on the memories and chat history to decide whether to update the memory collection."""
+#     message = state['messages'][-1]
+#     if len(message.tool_calls) ==0:
+#         return END
+#     else:
+#         tool_call = message.tool_calls[0]
+#         if tool_call['args']['update_type'] == "user":
+#             return "update_profile"
+#         elif tool_call['args']['update_type'] == "todo":
+#             return "update_todos"
+#         elif tool_call['args']['update_type'] == "instructions":
+#             return "update_instructions"
+#         else:
+#             raise ValueError

agent/utils/state.py

@@ -3,4 +3,4 @@ from langgraph.graph.message import add_messages
 
 class AgentState(TypedDict):
     messages: Annotated[list, add_messages]
-    context: list  # Store retrieved context 
+    context: list  # Store retrieved context

app.py

@@ -26,7 +26,8 @@ async def on_chat_start():
     config = RunnableConfig(
         configurable={
             "thread_id": session_id,
-            "sessionId": session_id
+            "session_id": session_id,
+            "checkpoint_ns": session_id
         }
     )
     
@@ -34,10 +35,18 @@ async def on_chat_start():
     try:
         async with get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
             graph = await create_agent_graph(saver)
-            initial_state = AgentState(messages=[], context=[])
+            initial_state = AgentState(
+                messages=[], 
+                context=[]
+            )
+            
             await graph.ainvoke(initial_state, config=config)
-            # Store initial state as a serializable dict
-            cl.user_session.set("last_state", {"messages": [], "context": []})
+            
+            # Store initial state
+            cl.user_session.set("last_state", {
+                "messages": [], 
+                "context": []
+            })
     except Exception as e:
         print(f"Error initializing state: {str(e)}")
     
@@ -67,14 +76,14 @@ async def on_message(message: cl.Message):
     config = RunnableConfig(
         configurable={
             "thread_id": session_id,
-            "checkpoint_ns": session_id,  # Use session_id as namespace
-            "sessionId": session_id
+            "session_id": session_id,
+            "checkpoint_ns": session_id
         }
     )
     
     try:
         async with get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
-            # Create graph with SQLite memory
+            # Create graph with memory
             graph = await create_agent_graph(saver)
             
             # Get the last state or create new one

example.json

@@ -0,0 +1,4 @@
+{
+  "role": "user",
+  "content": "I like sailing a lot, tell me about some activities I can do and remember this fact about me"
+}

langgraph.json

@@ -1,7 +1,7 @@
 {
   "dependencies": ".",
   "graphs": {
-    "agent": "agent:graph"
+    "agent": "agent:langgraph_studio_graph"
   },
   "env": ".env"
 }

tools.py

@@ -1,28 +0,0 @@
-from langchain_core.tools import tool
-from langchain_community.tools.tavily_search import TavilySearchResults
-from rag import create_rag_pipeline, add_urls_to_vectorstore
-
-# Initialize RAG pipeline
-rag_components = create_rag_pipeline(collection_name="london_events")
-
-# Add some initial URLs to the vector store
-urls = [
-    "https://www.timeout.com/london/things-to-do-in-london-this-weekend",
-    "https://www.timeout.com/london/london-events-in-march"
-]
-add_urls_to_vectorstore(
-    rag_components["vector_store"],
-    rag_components["text_splitter"],
-    urls
-)
-
-@tool
-def retrieve_context(query: str) -> list[str]:
-    """Searches the knowledge base for relevant information about events and activities. Use this when you need specific details about events."""
-    return [doc.page_content for doc in rag_components["retriever"].get_relevant_documents(query)]
-
-# Initialize Tavily search tool
-tavily_tool = TavilySearchResults(max_results=5)
-
-# Create tool belt
-tool_belt = [tavily_tool, retrieve_context]