feat: add interface, refactor state nodes and app.py for UX

agent/utils/nodes.py

@@ -175,8 +175,11 @@ async def write_memory(state: AgentState, config: RunnableConfig, store: BaseSto
     # Store the updated memory using async interface
     await store.aput(namespace, "user_memory", {"memory": new_memory.content})
     
-    
-    return state
+    # Update the state with the new memory
+    return {
+        **state,
+        "user_memories": {"memory": new_memory.content}
+    }
 
 # Initialize tool node
 tool_node = ToolNode(tool_belt)

agent/utils/state.py

@@ -1,6 +1,7 @@
-from typing import Annotated, TypedDict
+from typing import Annotated, TypedDict, Optional
 from langgraph.graph.message import add_messages
 
 class AgentState(TypedDict):
     messages: Annotated[list, add_messages]
     context: list  # Store retrieved context
+    user_memories: Optional[dict]  # Store user memory information

interfaces/__init__.py

@@ -0,0 +1 @@
+"""Interfaces package for Chief Joy Officer.""" 

interfaces/chainlit/__init__.py

@@ -0,0 +1 @@
+"""Chainlit interface for Chief Joy Officer.""" 

interfaces/chainlit/app.py

@@ -1,11 +1,19 @@
 import uuid
+import os
+import json
 from langchain_core.messages import HumanMessage, AIMessage, AIMessageChunk
 from langchain.schema.runnable.config import RunnableConfig
 import chainlit as cl
+import sys
+from pathlib import Path
+
+# Add the project root to the Python path
+project_root = str(Path(__file__).parent.parent.parent)
+if project_root not in sys.path:
+    sys.path.append(project_root)
+
 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"
 
@@ -19,9 +27,6 @@ async def on_chat_start():
         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"
@@ -30,21 +35,14 @@ async def on_chat_start():
         
     except Exception as e:
         print(f"Error in chat initialization: {str(e)}")
-        error_message = cl.Message(
+        await cl.Message(
             content="I apologize, but I encountered an error during initialization. Please try refreshing the page.",
             author="System"
-        )
-        await error_message.send()
+        ).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)
-    
     # Initialize response message
     msg = cl.Message(content="")
     
@@ -54,39 +52,24 @@ async def on_message(message: cl.Message):
                 # 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)
+                # Get session ID
+                session_id = cl.user_session.get("session_id")
                 
-                # Create current state
-                current_state = AgentState(
-                    messages=message_history,
-                    context=cl.user_session.get("last_context", [])
-                )
-                
-                # Stream the response
+                # Process through graph with current message
                 async for chunk in graph.astream(
-                    current_state,
-                    config={"configurable": {"thread_id": session_id}},
+                    {"messages": [HumanMessage(content=message.content)]},
+                    {"configurable": {"thread_id": session_id}},
                     stream_mode="messages"
                 ):
-                    # Handle different node outputs
-                    if isinstance(chunk[0], AIMessageChunk):
+                    if chunk[1]["langgraph_node"] == "agent" and isinstance(
+                        chunk[0], (AIMessageChunk, AIMessage)
+                    ):
                         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()