Update sessions and stuff

app.py

@@ -1,16 +1,17 @@
 import os
 import chainlit as cl
 from agent_graph import agent_node
-from dotenv import load_dotenv  # Import dotenv
+from dotenv import load_dotenv
 from typing import List, Dict
 import time
+import uuid
+import asyncio
 
 # Load environment variables from .env file
 load_dotenv()
 
 # Ensure your OpenAI API key is set up in environment variables
 openai_api_key = os.getenv("OPENAI_API_KEY")
-
 if not openai_api_key:
     raise ValueError("OpenAI API key is missing in the .env file")
 
@@ -18,22 +19,17 @@ if not openai_api_key:
 chat_histories: Dict[str, List[Dict[str, str]]] = {}
 
 def get_unique_session_id():
-    """Generate a unique session ID combining user session ID and timestamp."""
-    user_id = cl.user_session.get("id", "unknown")
-    timestamp = int(time.time() * 1000)  # Current time in milliseconds
-    return f"{user_id}_{timestamp}"
+    """Generate a unique session ID using UUID."""
+    return str(uuid.uuid4())
 
 @cl.on_chat_start
 async def start_chat():
-    # Generate a unique session ID
-    session_id = get_unique_session_id()
-    # Store the session ID in user session for later use
-    cl.user_session.set("session_id", session_id)
-    
-    # Initialize empty chat history for this session
-    chat_histories[session_id] = []
-    
-    welcome_message = """📈 Welcome to the AI Stock Assistant!
+    try:
+        session_id = get_unique_session_id()
+        cl.user_session.set("session_id", session_id)
+        chat_histories[session_id] = []
+
+        welcome_message = """📈 Welcome to the AI Stock Assistant!
 
 I'm your intelligent stock market companion. Here's what I can do:
 
@@ -56,60 +52,51 @@ Popular stocks to try:
 • Retail: WMT (Walmart), COST (Costco)
 
 What would you like to know about the stock market?"""
-    
-    await cl.Message(content=welcome_message).send()
+
+        await cl.Message(content=welcome_message).send()
+    except Exception as e:
+        print(f"[Error] Failed to start chat: {e}")
+        await cl.Message(content=f"⚠️ Error starting chat: {str(e)}").send()
 
 @cl.on_message
 async def handle_message(message: cl.Message):
     try:
-        # Get the unique session ID
         session_id = cl.user_session.get("session_id")
         if not session_id:
-            # If session ID is missing, generate a new one
             session_id = get_unique_session_id()
             cl.user_session.set("session_id", session_id)
             chat_histories[session_id] = []
-        
-        # Get chat history for this session
+
         history = chat_histories.get(session_id, [])
-        
-        # Add current message to history
         history.append({"role": "user", "content": message.content})
-        
-        # Create state dictionary with history
+
         state = {
             "input": message.content,
             "chat_history": history
         }
-        
-        # Process the message with agent_node
+
         print(f"[Debug] Processing message: {message.content}")
-        response = agent_node(state)
+        response = await asyncio.to_thread(agent_node, state)
         print(f"[Debug] Agent response: {response}")
-        
-        # Send the response back to the user
+
         if isinstance(response, dict) and "output" in response:
-            # Add response to history
             history.append({"role": "assistant", "content": response["output"]})
-            print(f"[Debug] Sending response to user: {response['output']}")
             await cl.Message(content=response["output"]).send()
-            print("[Debug] Response sent successfully")
         else:
-            print(f"[Debug] Invalid response format: {response}")
-            error_msg = "Received an invalid response format from the agent."
-            await cl.Message(content=error_msg).send()
-            
-        # Update history in storage
+            await cl.Message(content="❌ Received an invalid response format from the agent.").send()
+
         chat_histories[session_id] = history
-            
+
     except Exception as e:
-        print(f"[Debug] Error occurred: {e}")
-        error_msg = "Sorry, something went wrong while processing your request."
-        await cl.Message(content=error_msg).send()
+        print(f"[Error] Error in handle_message: {e}")
+        await cl.Message(content=f"⚠️ Error: {str(e)}").send()
 
 @cl.on_chat_end
 async def end_chat():
-    # Clean up chat history when session ends
-    session_id = cl.user_session.get("session_id")
-    if session_id and session_id in chat_histories:
-        del chat_histories[session_id]
+    try:
+        session_id = cl.user_session.get("session_id")
+        if session_id and session_id in chat_histories:
+            del chat_histories[session_id]
+        cl.user_session.clear()
+    except Exception as e:
+        print(f"[Error] Failed to clean up chat history: {e}")