add chat history

app.py

@@ -2,6 +2,12 @@ import chainlit as cl
 import pandas as pd
 import time
 from typing import Dict, Any
+from chainlit.session import get_session
+import os
+import json
+from datetime import datetime
+import hashlib
+import uuid
 
 from agents.table_selection import table_selection_agent
 from agents.data_retrieval import sample_data_retrieval_agent
@@ -11,42 +17,172 @@ from agents.execution import execution_agent
 from utils.bigquery_utils import init_bigquery_connection
 from utils.feedback_utils import save_feedback_to_bigquery
 
+# File to store chat history
+HISTORY_DIR = "chat_history"
+os.makedirs(HISTORY_DIR, exist_ok=True)
+
+# Function to get or create a user ID
+def get_user_id(session):
+    # Check if user already has an ID
+    user_id = session.user_session.get("user_id")
+    
+    if not user_id:
+        # Generate a new user ID if none exists
+        # In a real app, you might use authentication or cookies
+        user_id = str(uuid.uuid4())
+        session.user_session.set("user_id", user_id)
+        
+        # Create a user directory for this user
+        user_dir = os.path.join(HISTORY_DIR, user_id)
+        os.makedirs(user_dir, exist_ok=True)
+    
+    return user_id
+
+# Function to save chat history
+def save_chat_history(user_id, conversation_id, user_message, assistant_response):
+    user_dir = os.path.join(HISTORY_DIR, user_id)
+    os.makedirs(user_dir, exist_ok=True)
+    
+    history_file = os.path.join(user_dir, f"{conversation_id}.json")
+    
+    # Load existing history if it exists
+    if os.path.exists(history_file):
+        with open(history_file, 'r') as f:
+            history = json.load(f)
+    else:
+        history = {
+            "conversation_id": conversation_id,
+            "created_at": datetime.now().isoformat(),
+            "messages": []
+        }
+    
+    # Add new messages
+    timestamp = datetime.now().isoformat()
+    history["messages"].append({
+        "timestamp": timestamp,
+        "user": user_message,
+        "assistant": assistant_response
+    })
+    
+    # Save updated history
+    with open(history_file, 'w') as f:
+        json.dump(history, f, indent=2)
+
+# Function to load chat history for a user
+def load_chat_history(user_id, conversation_id=None):
+    user_dir = os.path.join(HISTORY_DIR, user_id)
+    
+    if not os.path.exists(user_dir):
+        return []
+    
+    if conversation_id:
+        # Load specific conversation
+        history_file = os.path.join(user_dir, f"{conversation_id}.json")
+        if os.path.exists(history_file):
+            with open(history_file, 'r') as f:
+                return json.load(f)
+        return None
+    else:
+        # Load all conversations (just metadata)
+        conversations = []
+        for filename in os.listdir(user_dir):
+            if filename.endswith('.json'):
+                with open(os.path.join(user_dir, filename), 'r') as f:
+                    data = json.load(f)
+                    conversations.append({
+                        "conversation_id": data["conversation_id"],
+                        "created_at": data["created_at"],
+                        "message_count": len(data["messages"])
+                    })
+        return sorted(conversations, key=lambda x: x["created_at"], reverse=True)
+
 @cl.on_chat_start
-async def on_chat_start():
-    """Initialize the chat session."""
-    # Initialize BigQuery client
-    client = init_bigquery_connection()
+async def start():
+    # Get the current session
+    session = get_session()
     
-    # Store the client in the user session
-    cl.user_session.set("client", client)
+    # Get or create user ID
+    user_id = get_user_id(session)
     
-    # Send a welcome message
-    await cl.Message(
-        content="👋 Welcome to the Natural Language to SQL Query Assistant! Ask me any question about your e-commerce data.",
-        author="SQL Assistant"
-    ).send()
+    # Create a new conversation ID for this session
+    conversation_id = str(uuid.uuid4())
+    session.user_session.set("conversation_id", conversation_id)
     
-    # Add some example questions without using actions
-    await cl.Message(
-        content="Here are some example questions you can ask:",
-        author="SQL Assistant"
-    ).send()
+    # Load previous conversations for this user
+    conversations = load_chat_history(user_id)
+    
+    if conversations:
+        # Display a welcome back message
+        await cl.Message(
+            content=f"Welcome back! You have {len(conversations)} previous conversations.",
+        ).send()
+        
+        # Create a dropdown to select previous conversations
+        options = [{"value": conv["conversation_id"], "label": f"Conversation from {conv['created_at'][:10]} ({conv['message_count']} messages)"} for conv in conversations]
+        options.insert(0, {"value": "new", "label": "Start a new conversation"})
+        
+        res = await cl.AskUserMessage(
+            content="Would you like to continue a previous conversation or start a new one?",
+            timeout=30,
+            select={"options": options, "value": "new"}
+        ).send()
+        
+        if res and res.get("value") != "new":
+            # User selected a previous conversation
+            selected_conv_id = res.get("value")
+            session.user_session.set("conversation_id", selected_conv_id)
+            
+            # Load the selected conversation
+            conversation = load_chat_history(user_id, selected_conv_id)
+            if conversation:
+                # Display previous messages
+                await cl.Message(content="Loading your previous conversation...").send()
+                for msg in conversation["messages"]:
+                    await cl.Message(content=msg["user"], author="You").send()
+                    await cl.Message(content=msg["assistant"]).send()
+    else:
+        # First time user
+        await cl.Message(
+            content="Welcome to the E-commerce Analytics Assistant! How can I help you today?",
+        ).send()
     
-    examples = [
-        "What are the top 5 products by revenue?",
-        "How many orders were placed in the last month?",
-        "Which customers spent the most in 2023?",
-        "What is the average order value by product category?"
-    ]
+    # Add option to clear history
+    await cl.Action(name="clear_history", label="Clear Chat History", description="Erase all previous conversation history").send()
+
+@cl.on_message
+async def on_message(message: cl.Message):
+    # Get the current session
+    session = get_session()
+    user_id = session.user_session.get("user_id")
+    conversation_id = session.user_session.get("conversation_id")
     
-    # Display all examples in a single message
-    examples_text = "\n\n".join([f"• {example}" for example in examples])
-    examples_text += "\n\n(You can copy and paste any of these examples to try them out)"
+    # Process the message and generate a response
+    # Replace this with your actual processing logic
+    response = f"You said: {message.content}"
     
-    await cl.Message(
-        content=examples_text,
-        author="SQL Assistant"
-    ).send()
+    # Send the response
+    await cl.Message(content=response).send()
+    
+    # Save to chat history
+    save_chat_history(user_id, conversation_id, message.content, response)
+
+@cl.action_callback("clear_history")
+async def clear_history_callback(action):
+    session = get_session()
+    user_id = session.user_session.get("user_id")
+    conversation_id = session.user_session.get("conversation_id")
+    
+    # Clear specific conversation or all conversations
+    user_dir = os.path.join(HISTORY_DIR, user_id)
+    history_file = os.path.join(user_dir, f"{conversation_id}.json")
+    
+    if os.path.exists(history_file):
+        os.remove(history_file)
+        await cl.Message(content="Your current conversation history has been cleared.").send()
+    
+    # Create a new conversation ID
+    new_conversation_id = str(uuid.uuid4())
+    session.user_session.set("conversation_id", new_conversation_id)
 
 @cl.on_message
 async def on_message(message: cl.Message):