feat: fix and test chainlit interface for cross thread memory

.gitignore

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

agent/graph.py

@@ -26,7 +26,6 @@ def create_graph_builder():
     # Set entry point
     builder.set_entry_point("agent")
 
-    # builder.add_edge("agent", "write_memory")
 
     # Add conditional edges from agent
     builder.add_conditional_edges(

app.py

@@ -1,151 +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 create_agent_graph, create_agent_graph_without_memory, get_checkpointer
+from agent import create_agent_graph, get_checkpointer
 from agent.utils.state import AgentState
 import os
 import json
 
-# Path to SQLite database for short-term memory
 SHORT_TERM_MEMORY_DB_PATH = "data/short_term.db"
 
-# Ensure the data directory exists
 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 empty message history
-    cl.user_session.set("message_history", [])
-    
-    # Initialize config using stored session ID
-    config = RunnableConfig(
-        configurable={
-            "thread_id": session_id,
-            "session_id": session_id,
-            "checkpoint_ns": session_id
-        }
-    )
-    
-    # Initialize empty state with auth
+    """Initialize the chat session"""
     try:
-        async with get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
-            graph = await create_agent_graph(saver)
-            initial_state = AgentState(
-                messages=[], 
-                context=[]
-            )
-            
-            await graph.ainvoke(initial_state, config=config)
-            
-            # Store initial state
-            cl.user_session.set("last_state", {
-                "messages": [], 
-                "context": []
-            })
+        # 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")
     if not session_id:
         session_id = str(uuid.uuid4())
         cl.user_session.set("session_id", session_id)
     
-    print(f"Session ID: {session_id}")
-    
-    # Get message history
-    message_history = cl.user_session.get("message_history", [])
-    
-    # Add new message to history
-    current_message = HumanMessage(content=message.content)
-    message_history.append(current_message)
-    cl.user_session.set("message_history", message_history)
-    
-    config = RunnableConfig(
-        configurable={
-            "thread_id": session_id,
-            "session_id": session_id,
-            "checkpoint_ns": session_id
-        }
-    )
+    # Initialize response message
+    msg = cl.Message(content="")
     
     try:
-        async with get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
-            # Create graph with memory
-            graph = await create_agent_graph(saver)
-            
-            # Get the last state or create new one
-            last_state_dict = cl.user_session.get("last_state", {"messages": [], "context": []})
-            
-            # Create new state with current message history
-            current_state = AgentState(
-                messages=message_history,
-                context=last_state_dict.get("context", [])
-            )
-            
-            # 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
-            last_state = None  # Track the final state
-            
-            # Stream the response
-            async for chunk in graph.astream(
-                current_state,
-                config=config
-            ):
-                for node, values in chunk.items():
-                    if node == "retrieve":
-                        if loading_msg:
-                            await loading_msg.remove()
-                        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)
-                            # Add AI message to history
-                            message_history.append(last_message)
-                            cl.user_session.set("message_history", message_history)
-                        # Update last state
-                        last_state = values
-            
-            # Update the last state as a serializable dict
-            if last_state:
-                cl.user_session.set("last_state", {
-                    "messages": [msg.content for msg in message_history],
-                    "context": last_state.get("context", [])
-                })
-            await final_answer.send()
-            
+        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 in message handler: {str(e)}")
-        await cl.Message(content="I apologize, but I encountered an error processing your message. Please try again.").send()
+        await cl.Message(
+            content="I apologize, but I encountered an error processing your message. Please try again.",
+            author="System"
+        ).send()