Merge pull request #1 from ric9176/memory-store

Add long term memory store, refactor chainlit interface, re-structure project

.gitignore

@@ -1,4 +1,6 @@
 __pycache__/
 .chainlit/
 .venv/ 
-.env
+.env
+.langgraph_api/
+data/

agent/__init__.py

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

agent/agent.py

@@ -1,50 +0,0 @@
-from langgraph.graph import StateGraph
-from langgraph.checkpoint.memory import MemorySaver
-
-from agent.utils.state import AgentState
-from agent.utils.nodes import call_model, tool_node, should_continue
-
-def create_agent_graph():
-    # Create the graph
-    builder = StateGraph(AgentState)
-
-    # Add nodes
-    builder.add_node("agent", call_model)
-    builder.add_node("action", tool_node)
-
-    # Update edges
-    builder.set_entry_point("agent")
-    builder.add_conditional_edges(
-        "agent",
-        should_continue,
-    )
-    builder.add_edge("action", "agent")
-
-    # Initialize memory saver for conversation persistence
-    memory = MemorySaver()
-
-    # Compile the graph with memory
-    return builder.compile(checkpointer=memory)
-
-def create_agent_graph_without_memory():
-    # Create the graph
-    builder = StateGraph(AgentState)
-
-    # Add nodes
-    builder.add_node("agent", call_model)
-    builder.add_node("action", tool_node)
-
-    # Update edges
-    builder.set_entry_point("agent")
-    builder.add_conditional_edges(
-        "agent",
-        should_continue,
-    )
-    builder.add_edge("action", "agent")
-
-    # Compile the graph without memory
-    return builder.compile()
-
-# Create both graph variants
-graph_with_memory = create_agent_graph()
-graph = create_agent_graph_without_memory() 

agent/graph.py

@@ -0,0 +1,96 @@
+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,
+    write_memory,
+    should_continue
+)
+
+def create_graph_builder():
+    """Create a base graph builder with nodes and edges configured."""
+    builder = StateGraph(AgentState)
+    
+
+    # Add nodes
+    builder.add_node("agent", call_model)
+    builder.add_node("action", tool_node)
+    builder.add_node("write_memory", write_memory)
+
+    # Set entry point
+    builder.set_entry_point("agent")
+
+
+    # Add conditional edges from agent
+    builder.add_conditional_edges(
+        "agent",
+        should_continue,
+        {
+            "action": "action",
+            "write_memory": "write_memory",
+        }
+    )
+    
+    # Connect action back to agent
+    builder.add_edge("action", "agent")
+    builder.add_edge("write_memory", END)
+    
+    return builder
+
+def create_agent_graph_without_memory():
+    """Create an agent graph without memory persistence."""
+    builder = create_graph_builder()
+    return builder.compile()
+
+class SQLiteCheckpointer:
+    """Context manager for SQLite checkpointing."""
+    
+    def __init__(self, db_path: str):
+        self.db_path = db_path
+        self.saver: Optional[AsyncSqliteSaver] = None
+    
+    async def __aenter__(self) -> AsyncSqliteSaver:
+        """Initialize and return the AsyncSqliteSaver."""
+        conn = await aiosqlite.connect(self.db_path)
+        self.saver = AsyncSqliteSaver(conn)
+        return self.saver
+    
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Clean up the SQLite connection."""
+        if self.saver and hasattr(self.saver, 'conn'):
+            await self.saver.conn.close()
+            self.saver = None
+
+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 memory persistence."""
+    builder = create_graph_builder()
+    # 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,35 +1,199 @@
 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)
-
-# Define system prompt
-SYSTEM_PROMPT = SystemMessage(content="""
-You are a helpful AI assistant that answers questions clearly and concisely.
-If you don't know something, simply say you don't know.
-Be engaging and professional in your responses.
-Use the retrieve_context tool when you need specific information about events and activities.
-Use the tavily_search tool for general web searches.
-""")
-
-def call_model(state: AgentState):
-    messages = [SYSTEM_PROMPT] + state["messages"]
-    response = model.invoke(messages)
+# 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.
+
+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 based on their interests
+3. Be engaging and enthusiastic while maintaining professionalism
+4. Give clear, actionable suggestions
+
+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."""
+
+# 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}
+
+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: Only include factual information directly stated by the user. Do not make assumptions.
+
+Based on the conversation, please update the user information:"""
+
+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"]
+    tool_calling_model = model.bind_tools(tool_belt)
+    response = tool_calling_model.invoke(messages)
     return {"messages": [response]}
 
-# Initialize tool node
-tool_node = ToolNode(tool_belt)
+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
 
-# Simple flow control - always go to final
-def should_continue(state):
+def should_continue(state: AgentState) -> Literal["action", "write_memory"]:
+    """Determine the next node in the graph."""
+    if not state["messages"]:
+        return END
+        
     last_message = state["messages"][-1]
-    if last_message.tool_calls:
+    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"
-    return END 
+
+    return "write_memory"
+    
+    # # 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"
+            
+    # return END
+
+# 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)
+
+# 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

@@ -1,108 +1,99 @@
 import uuid
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage, AIMessage, AIMessageChunk
 from langchain.schema.runnable.config import RunnableConfig
 import chainlit as cl
-from agent import graph_with_memory as graph
+from agent import create_agent_graph, get_checkpointer
 from agent.utils.state import AgentState
+import os
+import json
+
+SHORT_TERM_MEMORY_DB_PATH = "data/short_term.db"
+
+os.makedirs(os.path.dirname(SHORT_TERM_MEMORY_DB_PATH), exist_ok=True)
 
 @cl.on_chat_start
 async def on_chat_start():
-    # Generate and store a session ID
-    session_id = str(uuid.uuid4())
-    cl.user_session.set("session_id", session_id)
-    
-    # Initialize the conversation state with proper auth
-    cl.user_session.set("messages", [])
-    
-    # Initialize config using stored session ID
-    config = RunnableConfig(
-        configurable={
-            "thread_id": session_id,
-            "sessionId": session_id
-        }
-    )
-    
-    # Initialize empty state with auth
+    """Initialize the chat session"""
     try:
-        await graph.ainvoke(
-            AgentState(messages=[], context=[]),
-            config=config
+        # Generate and store a session ID
+        session_id = str(uuid.uuid4())
+        cl.user_session.set("session_id", session_id)
+        
+        # Initialize empty message history
+        cl.user_session.set("message_history", [])
+        
+        welcome_message = cl.Message(
+            content="Hello! I'm your chief joy officer, here to help you with finding fun things to do in London!",
+            author="Assistant"
         )
+        await welcome_message.send()
+        
     except Exception as e:
-        print(f"Error initializing state: {str(e)}")
-    
-    await cl.Message(
-        content="Hello! I'm your chief joy officer, here to help you with finding fun things to do in London!",
-        author="Assistant"
-    ).send()
+        print(f"Error in chat initialization: {str(e)}")
+        error_message = cl.Message(
+            content="I apologize, but I encountered an error during initialization. Please try refreshing the page.",
+            author="System"
+        )
+        await error_message.send()
 
 @cl.on_message
 async def on_message(message: cl.Message):
+    """Handle incoming messages and stream responses"""
+    # Get or create session ID
     session_id = cl.user_session.get("session_id")
-    print(f"Session ID: {session_id}")
     if not session_id:
         session_id = str(uuid.uuid4())
         cl.user_session.set("session_id", session_id)
     
-    config = RunnableConfig(
-        configurable={
-            "thread_id": session_id,
-            "checkpoint_ns": "default_namespace",
-            "sessionId": session_id
-        }
-    )
+    # Initialize response message
+    msg = cl.Message(content="")
     
-    # Try to retrieve previous conversation state
     try:
-        previous_state = await graph.aget_state(config)
-        if previous_state and previous_state.values:
-            previous_messages = previous_state.values.get('messages', [])
-            print("Found previous state with messages:", len(previous_messages))
-        else:
-            print("Previous state empty or invalid")
-            previous_messages = []
-        current_messages = previous_messages + [HumanMessage(content=message.content)]
+        async with cl.Step(type="run"):
+            async with get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
+                # Create graph with memory
+                graph = await create_agent_graph(saver)
+                
+                # Get message history and add current message
+                message_history = cl.user_session.get("message_history", [])
+                current_message = HumanMessage(content=message.content)
+                message_history.append(current_message)
+                
+                # Create current state
+                current_state = AgentState(
+                    messages=message_history,
+                    context=cl.user_session.get("last_context", [])
+                )
+                
+                # Stream the response
+                async for chunk in graph.astream(
+                    current_state,
+                    config={"configurable": {"thread_id": session_id}},
+                    stream_mode="messages"
+                ):
+                    # Handle different node outputs
+                    if isinstance(chunk[0], AIMessageChunk):
+                        await msg.stream_token(chunk[0].content)
+                    elif isinstance(chunk[0], AIMessage):
+                        if chunk[0] not in message_history:
+                            message_history.append(chunk[0])
+                
+                # Get final state
+                final_state = await graph.aget_state(
+                    config={"configurable": {"thread_id": session_id}}
+                )
+                
+                # Update session state
+                if final_state:
+                    cl.user_session.set("message_history", message_history)
+                    cl.user_session.set("last_context", final_state.values.get("context", []))
+        
+        # Send the final message
+        await msg.send()
+                
     except Exception as e:
-        print(f"Error retrieving previous state: {str(e)}")
-        current_messages = [HumanMessage(content=message.content)]
-    
-    # Setup callback handler and final answer message
-    cb = cl.LangchainCallbackHandler()
-    final_answer = cl.Message(content="")
-    await final_answer.send()
-    
-    loading_msg = None  # Initialize reference to loading message
-    
-    # Stream the response
-    async for chunk in graph.astream(
-        AgentState(messages=current_messages, context=[]),
-        config=RunnableConfig(
-            configurable={
-                "thread_id": session_id,
-            }
-        )
-    ):
-        for node, values in chunk.items():
-            if node == "retrieve":
-                loading_msg = cl.Message(content="🔍 Searching knowledge base...", author="System")
-                await loading_msg.send()
-            elif values.get("messages"):
-                last_message = values["messages"][-1]
-                # Check for tool calls in additional_kwargs
-                if hasattr(last_message, "additional_kwargs") and last_message.additional_kwargs.get("tool_calls"):
-                    tool_name = last_message.additional_kwargs["tool_calls"][0]["function"]["name"]
-                    if loading_msg:
-                        await loading_msg.remove()
-                    loading_msg = cl.Message(
-                        content=f"🔍 Using {tool_name}...",
-                        author="Tool"
-                    )
-                    await loading_msg.send()
-                # Only stream AI messages, skip tool outputs
-                elif isinstance(last_message, AIMessage):
-                    if loading_msg:
-                        await loading_msg.remove()
-                        loading_msg = None
-                    await final_answer.stream_token(last_message.content)
-
-    await final_answer.send()
+        print(f"Error in message handler: {str(e)}")
+        await cl.Message(
+            content="I apologize, but I encountered an error processing your message. Please try again.",
+            author="System"
+        ).send()

docker-compose.yml

@@ -0,0 +1,23 @@
+version: "3.8"
+
+services:
+  qdrant:
+    image: qdrant/qdrant:latest
+    ports:
+      - "6333:6333"
+    volumes:
+      - ./data/long_term_memory:/qdrant/storage
+    restart: unless-stopped
+
+  app:
+    build: .
+    ports:
+      - "8000:8000"
+    volumes:
+      - .:/app
+      - ./data/short_term_memory:/app/data
+    environment:
+      - QDRANT_URL=http://qdrant:6333
+      - OPENAI_API_KEY=${OPENAI_API_KEY}
+    depends_on:
+      - qdrant

langgraph.json

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

studio_example_message.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"
+}

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]