complete refactor to use AsyncSqliteSaver

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
 
-__all__ = ["graph", "graph_with_memory"] 
+__all__ = ["create_agent_graph", "create_agent_graph_without_memory", "get_checkpointer"] 

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,69 @@
+from langgraph.graph import StateGraph
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+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
+
+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)
+
+    # Update edges
+    builder.set_entry_point("agent")
+    builder.add_conditional_edges(
+        "agent",
+        should_continue,
+    )
+    builder.add_edge("action", "agent")
+    
+    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)
+
+async def create_agent_graph(checkpointer: AsyncSqliteSaver):
+    """Create an agent graph with SQLite-based memory persistence."""
+    builder = create_graph_builder()
+    graph = builder.compile(checkpointer=checkpointer)
+    return graph
+
+# Export the graph builder functions
+__all__ = ["create_agent_graph", "create_agent_graph_without_memory", "get_checkpointer"] 

app.py

@@ -2,8 +2,16 @@ import uuid
 from langchain_core.messages import HumanMessage, AIMessage
 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, create_agent_graph_without_memory, 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():
@@ -11,8 +19,8 @@ async def on_chat_start():
     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 empty message history
+    cl.user_session.set("message_history", [])
     
     # Initialize config using stored session ID
     config = RunnableConfig(
@@ -24,10 +32,12 @@ async def on_chat_start():
     
     # Initialize empty state with auth
     try:
-        await graph.ainvoke(
-            AgentState(messages=[], context=[]),
-            config=config
-        )
+        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 as a serializable dict
+            cl.user_session.set("last_state", {"messages": [], "context": []})
     except Exception as e:
         print(f"Error initializing state: {str(e)}")
     
@@ -38,71 +48,95 @@ async def on_chat_start():
 
 @cl.on_message
 async def on_message(message: cl.Message):
+    # 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)
     
+    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,
-            "checkpoint_ns": "default_namespace",
+            "checkpoint_ns": session_id,  # Use session_id as namespace
             "sessionId": session_id
         }
     )
     
-    # 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 get_checkpointer(SHORT_TERM_MEMORY_DB_PATH) as saver:
+            # Create graph with SQLite 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()
+            
     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.").send()