Update llm_handler.py

llm_handler.py

@@ -1,5 +1,6 @@
 import os
-import uuid
+import random
+import time
 from typing import Dict, List
 from openai import OpenAI
 
@@ -28,8 +29,18 @@ def initialize_llm():
 
 def create_chat_session() -> str:
     """Creates a new chat session and returns the session ID."""
-    session_id = str(uuid.uuid4())
+    # Generate a unique session ID using timestamp + random number
+    timestamp = int(time.time() * 1000)  # milliseconds
+    random_num = random.randint(1000, 9999)
+    session_id = f"{timestamp}_{random_num}"
+    
+    # Ensure uniqueness (very unlikely to collide, but just in case)
+    while session_id in chat_sessions:
+        random_num = random.randint(1000, 9999)
+        session_id = f"{timestamp}_{random_num}"
+    
     chat_sessions[session_id] = []
+    print(f"🆕 Created new chat session: {session_id}")
     return session_id
 
 def clear_chat_session(session_id: str) -> bool:
@@ -50,6 +61,16 @@ def get_chat_history(session_id: str) -> List[Dict[str, str]]:
     """Gets the chat history for a specific session."""
     return chat_sessions.get(session_id, [])
 
+def cleanup_old_sessions():
+    """Clean up old sessions - can be called periodically."""
+    # Keep only 15 most recent sessions to save memory
+    if len(chat_sessions) > 15:
+        # Keep only the most recent 10 sessions when cleanup is triggered
+        session_items = list(chat_sessions.items())
+        chat_sessions.clear()
+        chat_sessions.update(dict(session_items[-10:]))
+        print(f"🧹 Cleaned up old chat sessions. Current count: {len(chat_sessions)}")
+
 def add_to_chat_history(session_id: str, role: str, content: str):
     """Adds a message to the chat history."""
     if session_id not in chat_sessions:
@@ -60,19 +81,40 @@ def add_to_chat_history(session_id: str, role: str, content: str):
         "content": content
     })
     
-    # Keep only the last 20 messages to prevent memory overflow
+    # Keep only the last 20 messages per session to prevent memory overflow
     # (10 user messages + 10 assistant responses)
     if len(chat_sessions[session_id]) > 20:
         chat_sessions[session_id] = chat_sessions[session_id][-20:]
+    
+    # Trigger cleanup if we have too many sessions
+    if len(chat_sessions) > 15:
+        cleanup_old_sessions()
+
+def get_chat_session_count() -> int:
+    """Returns the number of active chat sessions."""
+    return len(chat_sessions)
+
+def clear_all_chat_sessions() -> int:
+    """Clears all chat sessions and returns the count of sessions that were cleared."""
+    session_count = len(chat_sessions)
+    chat_sessions.clear()
+    print(f"🧹 All chat sessions cleared. Removed {session_count} sessions.")
+    return session_count
 
-def get_rag_response(query: str, session_id: str = None) -> str:
+def get_rag_response(query: str, session_id: str = None) -> tuple[str, str]:
     """Generates a response using Retrieval-Augmented Generation with chat memory."""
     if not all([encoder, chroma_collection, openrouter_client]):
-        return "Chatbot is not ready. Models or clients are not loaded."
+        return "Chatbot is not ready. Models or clients are not loaded.", session_id or create_chat_session()
 
     # Create a new session if none provided
     if session_id is None:
         session_id = create_chat_session()
+        print(f"🆕 Created new chat session: {session_id}")
+    
+    # Validate session exists, create if it doesn't
+    if session_id not in chat_sessions:
+        chat_sessions[session_id] = []
+        print(f"🔄 Session {session_id} not found, created new one")
 
     # 1. Retrieve relevant documents from ChromaDB
     query_embedding = encoder.encode([query])[0].tolist()
@@ -129,19 +171,4 @@ If the context doesn't have the answer, use your own general knowledge to provid
         
     except Exception as e:
         print(f"❌ Error calling OpenRouter API: {e}")
-        return "Sorry, I encountered an error while processing your request.", session_id
-
-def get_chat_session_count() -> int:
-    """Returns the number of active chat sessions."""
-    return len(chat_sessions)
-
-def cleanup_old_sessions():
-    """Clean up old sessions - can be called periodically."""
-    # This is a simple cleanup - in production you might want to track timestamps
-    # and clean up sessions older than a certain time
-    if len(chat_sessions) > 1000:  # If too many sessions
-        # Keep only the most recent 500 sessions
-        session_items = list(chat_sessions.items())
-        chat_sessions.clear()
-        chat_sessions.update(dict(session_items[-500:]))
-        print(f"🧹 Cleaned up old chat sessions. Current count: {len(chat_sessions)}")
+        return "Sorry, I encountered an error while processing your request.", session_id