message and session management

ai_chat.py

@@ -2,52 +2,96 @@ from langgraph.graph import StateGraph, END
 from langchain_core.runnables import RunnableLambda
 from typing import TypedDict
 from config import llm, OPENAI_API_KEY
+from database import get_session_messages, save_message
 
-def chat_without_memory(state):
-    """Chat function without memory - stateless responses"""
+def chat_with_session_memory(state):
+    """Chat function with session-based memory"""
     user_message = state["user_message"]
     user_info = state.get("user_info", {})
+    session_id = state.get("session_id")
+    wa_id = state.get("wa_id")
+    wamid = state.get("wamid")
+    
+    # Get conversation history from database
+    session_messages = []
+    if session_id:
+        # This will be populated by the async wrapper
+        session_messages = state.get("session_messages", [])
     
     # Add system message with user context
     system_message = "You are a helpful and concise assistant."
     if user_info.get("name") and user_info["name"] != "Unknown":
         system_message += f" The user's name is {user_info['name']}."
     
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": user_message}
-    ]
+    # Build messages array with history
+    messages = [{"role": "system", "content": system_message}]
+    
+    # Add conversation history (last 10 messages)
+    for msg in session_messages[-30:]:
+        messages.append({"role": msg["role"], "content": msg["content"]})
+    
+    # Add current user message
+    messages.append({"role": "user", "content": user_message})
 
     try:
         if not OPENAI_API_KEY:
             return {"response": "Sorry, AI chat is not available. Please check your OpenAI API key configuration."}
         
         response = llm.invoke(messages)
-        return {"response": response.content}
+        ai_response = response.content
+        
+        # Save messages to database (this will be handled by the async wrapper)
+        return {
+            "response": ai_response,
+            "user_message": user_message,
+            "ai_response": ai_response,
+            "session_id": session_id,
+            "wa_id": wa_id,
+            "wamid": wamid
+        }
     except Exception as e:
-        print(f"Error in chat_without_memory: {e}")
+        print(f"Error in chat_with_session_memory: {e}")
         return {"response": "Sorry, something went wrong: " + str(e)}
 
 class ChatState(TypedDict):
     user_message: str
     response: str
     user_info: dict
+    session_id: str
+    wa_id: str
+    wamid: str
+    session_messages: list
 
 # --- Build LangGraph ---
 graph = StateGraph(ChatState)
-graph.add_node("chat", RunnableLambda(chat_without_memory))
+graph.add_node("chat", RunnableLambda(chat_with_session_memory))
 graph.set_entry_point("chat")
 graph.add_edge("chat", END)
 chat_graph = graph.compile()
 
-async def process_message(user_message: str, user_info: dict = None):
-    """Process a message through the AI chat system"""
+async def process_message(user_message: str, user_info: dict = None, session_id: str = None, wa_id: str = None, wamid: str = None):
+    """Process a message through the AI chat system with session memory"""
     if user_info is None:
         user_info = {}
     
+    # Get session messages for context
+    session_messages = []
+    if session_id:
+        session_messages = await get_session_messages(session_id, limit=10)
+    
+    # Process with AI
     result = await chat_graph.ainvoke({
         "user_message": user_message,
-        "user_info": user_info
+        "user_info": user_info,
+        "session_id": session_id,
+        "wa_id": wa_id,
+        "wamid": wamid,
+        "session_messages": session_messages
     })
     
+    # Save messages to database
+    if session_id and wa_id and wamid:
+        await save_message(session_id, wa_id, wamid, "user", user_message)
+        await save_message(session_id, wa_id, f"{wamid}_ai", "assistant", result["response"])
+    
     return result["response"] 

database.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from config import supabase
 
 async def get_or_create_user(wa_id: str, name: str = None) -> dict:
@@ -42,6 +42,107 @@ async def update_user_activity(wa_id: str):
     except Exception as e:
         print(f"Error updating user activity: {e}")
 
+# --- Session Management ---
+async def get_active_session(wa_id: str) -> dict:
+    """Get the active session for a user (not ended and within 3 hours)"""
+    if not supabase:
+        return None
+    
+    try:
+        # Calculate 3 hours ago
+        three_hours_ago = datetime.utcnow() - timedelta(hours=3)
+        
+        response = supabase.table("chat_sessions").select("*").eq("wa_id", wa_id).eq("ended", False).gte("last_activity", three_hours_ago.isoformat()).order("last_activity", desc=True).limit(1).execute()
+        
+        return response.data[0] if response.data else None
+    except Exception as e:
+        print(f"Error getting active session: {e}")
+        return None
+
+async def create_new_session(wa_id: str) -> dict:
+    """Create a new chat session for a user"""
+    if not supabase:
+        return {"id": "local_session", "wa_id": wa_id, "started_at": datetime.utcnow().isoformat(), "last_activity": datetime.utcnow().isoformat(), "ended": False}
+    
+    try:
+        new_session = {
+            "wa_id": wa_id,
+            "started_at": datetime.utcnow().isoformat(),
+            "last_activity": datetime.utcnow().isoformat(),
+            "ended": False
+        }
+        response = supabase.table("chat_sessions").insert(new_session).execute()
+        return response.data[0] if response.data else new_session
+    except Exception as e:
+        print(f"Error creating new session: {e}")
+        return {"id": "local_session", "wa_id": wa_id, "started_at": datetime.utcnow().isoformat(), "last_activity": datetime.utcnow().isoformat(), "ended": False}
+
+async def update_session_activity(session_id: str):
+    """Update session's last activity timestamp"""
+    if not supabase or session_id == "local_session":
+        return
+    
+    try:
+        supabase.table("chat_sessions").update({"last_activity": datetime.utcnow().isoformat()}).eq("id", session_id).execute()
+    except Exception as e:
+        print(f"Error updating session activity: {e}")
+
+async def end_session(session_id: str):
+    """Mark a session as ended"""
+    if not supabase or session_id == "local_session":
+        return
+    
+    try:
+        supabase.table("chat_sessions").update({"ended": True}).eq("id", session_id).execute()
+    except Exception as e:
+        print(f"Error ending session: {e}")
+
+# --- Message Management ---
+async def save_message(session_id: str, wa_id: str, wamid: str, role: str, content: str):
+    """Save a message to the database"""
+    if not supabase or session_id == "local_session":
+        return
+    
+    try:
+        message = {
+            "session_id": session_id,
+            "wa_id": wa_id,
+            "wamid": wamid,
+            "role": role,
+            "content": content,
+            "created_at": datetime.utcnow().isoformat()
+        }
+        supabase.table("messages").insert(message).execute()
+    except Exception as e:
+        print(f"Error saving message: {e}")
+
+async def get_session_messages(session_id: str, limit: int = 10) -> list:
+    """Get recent messages for a session (for memory context)"""
+    if not supabase or session_id == "local_session":
+        return []
+    
+    try:
+        response = supabase.table("messages").select("*").eq("session_id", session_id).order("created_at", desc=True).limit(limit).execute()
+        # Reverse to get chronological order
+        messages = response.data[::-1] if response.data else []
+        return messages
+    except Exception as e:
+        print(f"Error getting session messages: {e}")
+        return []
+
+# --- Session Management Helper ---
+async def get_or_create_active_session(wa_id: str) -> dict:
+    """Get active session or create new one if needed"""
+    active_session = await get_active_session(wa_id)
+    
+    if active_session:
+        # Update activity
+        await update_session_activity(active_session["id"])
+        return active_session
+    else:
+        # Create new session
+        return await create_new_session(wa_id)
+
 async def get_user(wa_id: str):
     """Get user by WhatsApp ID"""
     if not supabase:

main.py

@@ -3,7 +3,7 @@ from fastapi.responses import JSONResponse
 
 # Import our modular components
 from config import supabase
-from database import get_or_create_user, update_user_activity
+from database import get_or_create_user, update_user_activity, get_or_create_active_session
 from whatsapp import send_whatsapp_message
 from ai_chat import process_message
 from api_routes import router
@@ -35,6 +35,7 @@ async def receive_message(req: Request):
 
         if msg.get("type") == "text":
             user_message = msg.get("text", {}).get("body", "")
+            wamid = msg.get("id")  # Meta message ID
             
             # Get or create user in database
             user_info = await get_or_create_user(wa_id, contact_name)
@@ -42,13 +43,26 @@ async def receive_message(req: Request):
             # Update user activity
             await update_user_activity(wa_id)
             
-            # Process with AI including user context
-            ai_response = await process_message(user_message, user_info)
+            # Get or create active session
+            session = await get_or_create_active_session(wa_id)
+            
+            # Process with AI including session memory
+            ai_response = await process_message(
+                user_message=user_message,
+                user_info=user_info,
+                session_id=session["id"],
+                wa_id=wa_id,
+                wamid=wamid
+            )
             
             # Send response back to WhatsApp
             await send_whatsapp_message(wa_id, ai_response)
             
-            return JSONResponse({"status": "replied", "user_id": wa_id})
+            return JSONResponse({
+                "status": "replied", 
+                "user_id": wa_id,
+                "session_id": session["id"]
+            })
 
     except Exception as e:
         print("Error processing message:", e)

supabase_setup.sql

@@ -6,17 +6,52 @@ CREATE TABLE IF NOT EXISTS users (
     updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
 );
 
--- Create index for better query performance
+-- Create chat_sessions table
+CREATE TABLE IF NOT EXISTS chat_sessions (
+    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    wa_id TEXT NOT NULL REFERENCES users(wa_id) ON DELETE CASCADE,
+    started_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
+    last_activity TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
+    ended BOOLEAN DEFAULT FALSE
+);
+
+-- Create messages table
+CREATE TABLE IF NOT EXISTS messages (
+    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    session_id UUID NOT NULL REFERENCES chat_sessions(id) ON DELETE CASCADE,
+    wa_id TEXT NOT NULL REFERENCES users(wa_id) ON DELETE CASCADE,
+    wamid TEXT NOT NULL, -- Meta message ID
+    role TEXT NOT NULL CHECK (role IN ('user', 'assistant')),
+    content TEXT NOT NULL,
+    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
+);
+
+-- Create indexes for better query performance
 CREATE INDEX IF NOT EXISTS idx_users_created_at ON users(created_at DESC);
 CREATE INDEX IF NOT EXISTS idx_users_updated_at ON users(updated_at DESC);
+CREATE INDEX IF NOT EXISTS idx_chat_sessions_wa_id ON chat_sessions(wa_id);
+CREATE INDEX IF NOT EXISTS idx_chat_sessions_last_activity ON chat_sessions(last_activity DESC);
+CREATE INDEX IF NOT EXISTS idx_chat_sessions_ended ON chat_sessions(ended);
+CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id);
+CREATE INDEX IF NOT EXISTS idx_messages_wa_id ON messages(wa_id);
+CREATE INDEX IF NOT EXISTS idx_messages_created_at ON messages(created_at DESC);
+CREATE INDEX IF NOT EXISTS idx_messages_wamid ON messages(wamid);
 
 -- Enable Row Level Security (RLS) - optional but recommended
 ALTER TABLE users ENABLE ROW LEVEL SECURITY;
+ALTER TABLE chat_sessions ENABLE ROW LEVEL SECURITY;
+ALTER TABLE messages ENABLE ROW LEVEL SECURITY;
 
--- Create policy to allow all operations (you can restrict this based on your needs)
+-- Create policies to allow all operations (you can restrict this based on your needs)
 CREATE POLICY "Allow all operations on users" ON users
     FOR ALL USING (true);
 
+CREATE POLICY "Allow all operations on chat_sessions" ON chat_sessions
+    FOR ALL USING (true);
+
+CREATE POLICY "Allow all operations on messages" ON messages
+    FOR ALL USING (true);
+
 -- Optional: Create a function to automatically update the updated_at timestamp
 CREATE OR REPLACE FUNCTION update_updated_at_column()
 RETURNS TRIGGER AS $$
@@ -26,8 +61,22 @@ BEGIN
 END;
 $$ language 'plpgsql';
 
--- Create trigger to automatically update updated_at
+-- Create trigger to automatically update updated_at for users
 CREATE TRIGGER update_users_updated_at 
     BEFORE UPDATE ON users 
     FOR EACH ROW 
-    EXECUTE FUNCTION update_updated_at_column(); 
+    EXECUTE FUNCTION update_updated_at_column();
+
+-- Create trigger to automatically update last_activity for chat_sessions
+CREATE OR REPLACE FUNCTION update_last_activity_column()
+RETURNS TRIGGER AS $$
+BEGIN
+    NEW.last_activity = NOW();
+    RETURN NEW;
+END;
+$$ language 'plpgsql';
+
+CREATE TRIGGER update_chat_sessions_last_activity 
+    BEFORE UPDATE ON chat_sessions 
+    FOR EACH ROW 
+    EXECUTE FUNCTION update_last_activity_column();