Adding memory management

app.py

@@ -1,3 +1,4 @@
+import uuid
 import os
 import gradio as gr
 import pandas as pd
@@ -9,6 +10,8 @@ import chromadb
 from langchain_chroma import Chroma
 import gspread
 from google.oauth2.service_account import Credentials
+from langgraph.checkpoint.sqlite import SqliteSaver
+import sqlite3
 import json
 from datetime import datetime
 import re
@@ -36,7 +39,7 @@ client_gspread = gspread.authorize(get_google_sheets_credentials())
 # Open the Google Sheet
 sheet = client_gspread.open("Response_Log").sheet1
 
-def log_response(question, answer, source_ids, knowledge_pairs):
+def log_response(question, answer, source_ids, knowledge_pairs, session_id):
     """
     Log a question, answer, source IDs, and knowledge base question-answer pairs to the Google Sheet.
     
@@ -53,6 +56,7 @@ def log_response(question, answer, source_ids, knowledge_pairs):
     knowledge_answer_2 = knowledge_pairs[1][1] if len(knowledge_pairs) > 1 else "N/A"
     row = [
         timestamp,
+        session_id,
         question,
         answer,
         source_ids,
@@ -69,6 +73,28 @@ def log_response(question, answer, source_ids, knowledge_pairs):
         with open("/tmp/response_log.txt", "a") as f:
             f.write(f"{timestamp},{question},{answer},{source_ids},{knowledge_question_1},{knowledge_answer_1},{knowledge_question_2},{knowledge_answer_2}\n")
 
+# === LangGraph Memory Setup ===
+conn = sqlite3.connect("xeno_memory.db", check_same_thread=False)
+memory = SqliteSaver(conn=conn)
+
+def update_memory(config, user_message, assistant_message):
+    full_checkpoint = memory.get(config) or {}
+    messages = full_checkpoint.get("channel_values", {}).get("messages", [])
+    
+    messages.append({"role": "user", "content": user_message})
+    messages.append({"role": "assistant", "content": assistant_message})
+    
+    checkpoint_to_save = {
+        "v": 1,
+        "id": str(uuid.uuid4()),
+        "ts": datetime.now().isoformat(),
+        "channel_values": {"messages": messages},
+        "channel_versions": {},
+        "versions_seen": {},   
+    }
+    
+    memory.put(config, checkpoint_to_save, {}, {})
+    
 # === Intent Classification System ===
 class IntentClassifier:
     def __init__(self):
@@ -203,106 +229,101 @@ def process_context(results, cosine_scores, max_results=2):
         knowledge_pairs.append((question, answer))
     return formatted_context, source_ids, knowledge_pairs
 
-# === LLM Generation ===
-def generate_xeno_response(context, question):
+# === LLM Generation (Refactored) ===
+def generate_xeno_response(context, question, chat_history):
+    """Generates a response but does NOT handle memory."""
     model = genai.GenerativeModel(llm_model_name)
-    prompt = f"""{SYSTEM_PROMPT}
-### CONTEXT ###
-{context}
-### QUESTION ###
-{question}"""
+    formatted_history = "\n".join(
+        [f"{msg['role'].capitalize()}: {msg['content']}" for msg in chat_history]
+    ) if chat_history else "None"
+    
+    prompt = f"{SYSTEM_PROMPT}\n### HISTORY ###\n{formatted_history}\n### CONTEXT ###\n{context}\n### QUESTION ###\n{question}"
+    
     response = model.generate_content(prompt)
     return response.text.strip()
 
-# === Enhanced Main Interface Logic with Intent Classification ===
-def get_context_and_answer(message, history):
+
+# === Main Interface Logic (Refactored) ===
+def get_context_and_answer(message, history, session_id="default"):
     """
-    Enhanced pipeline with intent classification
+    Handles intent classification, RAG, and memory updates in one place.
     """
-    # Step 1: Intent Classification
-    intent, direct_response = intent_classifier.classify_intent(message)
+    config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
     
-    # Step 2: Handle simple intents directly
-    if intent_classifier.is_simple_intent(intent) and direct_response:
-        log_response(message, direct_response, "N/A", [])
-        return direct_response
+    full_checkpoint = memory.get(config) or {}
+    chat_history = full_checkpoint.get("channel_values", {}).get("messages", [])
+    intent, direct_response = intent_classifier.classify_intent(message)
     
-    # Step 3: For queries that need RAG processing
-    if intent == 'query':
-        # Check if message is too short or unclear
+    answer = ""
+    source_ids = "N/A"
+    knowledge_pairs = []
+
+    if intent != 'query':
+        answer = direct_response
+    else: 
         if len(message.strip()) < 3:
-            answer = "I'd be happy to help! Could you please provide more details about what you'd like to know about XENO services?"
-            log_response(message, answer, "N/A", [])
-            return answer
-        
-        # Retrieve relevant documents
-        try:
-            queried_results = retriever.invoke(message)
-            query_embedding = genai.embed_content(
-                model=embedding_model,
-                content=message,
-                task_type="retrieval_query"
-            )['embedding']
-            
-            cosine_scores = []
-            for doc in queried_results:
-                doc_embedding = genai.embed_content(
-                    model=embedding_model,
-                    content=doc.page_content,
-                    task_type="retrieval_document"
-                )['embedding']
-                cos_sim = util.cos_sim(
-                    torch.tensor(query_embedding).float(), 
-                    torch.tensor(doc_embedding).float()
-                )[0][0].item()
-                cosine_scores.append(cos_sim)
+            answer = "I'd be happy to help! Could you please provide more details about what you'd like to know?"
+        else:
+            try:
+                queried_results = retriever.invoke(message)
+                query_embedding = genai.embed_content(model=embedding_model, content=message, task_type="retrieval_query")['embedding']
+                
+                doc_embeddings = [genai.embed_content(model=embedding_model, content=doc.page_content, task_type="retrieval_document")['embedding'] for doc in queried_results]
+                
+                cosine_scores = util.cos_sim(torch.tensor(query_embedding).float(), torch.tensor(doc_embeddings).float())[0].tolist()
 
-            # If none of the results have sufficient similarity, fallback
-            if max(cosine_scores) < 0.4:
-                answer = "I'm sorry, I couldn't find the specific information you're looking for in my knowledge base. Could you try rephrasing your question or contact XENO support directly for assistance?"
-                log_response(message, answer, "N/A", [])
-                return answer
+                if max(cosine_scores) < 0.4:
+                    answer = "I'm sorry, I couldn't find specific information for your question. Could you try rephrasing it, or contact XENO support directly?"
+                else:
+                    context, source_ids_list, knowledge_pairs = process_context(queried_results, cosine_scores)
+                    answer = generate_xeno_response(context, message, chat_history)
+                    source_ids = ", ".join(source_ids_list)
 
-            context, source_ids, knowledge_pairs = process_context(queried_results, cosine_scores)
-            answer = generate_xeno_response(context, message)
-            log_response(message, answer, ", ".join(source_ids), knowledge_pairs)
-            return answer
-            
-        except Exception as e:
-            answer = "I apologize, but I'm experiencing a technical issue. Please contact XENO support directly for assistance with your query."
-            log_response(message, answer, "N/A", [])
-            return answer
-    
-    # Handle goodbye intent (not simple, but has direct response)
-    if intent == 'goodbye' and direct_response:
-        log_response(message, direct_response, "N/A", [])
-        return direct_response
+            except Exception as e:
+                print(f"Error during RAG processing: {e}")
+                answer = "I apologize, but I'm having a technical issue. Please try again shortly or contact XENO support."
 
-    # Fallback for any unhandled cases
-    answer = "I'm here to help with XENO financial services. What would you like to know?"
-    log_response(message, answer, "N/A", [])
+    update_memory(config, message, answer)
+    log_response(message, answer, source_ids, knowledge_pairs, session_id)
+    
     return answer
 
 # === Enhanced Gradio UI ===
-def create_interface():
-    """Create the Gradio interface with custom styling"""
+def respond(message, history, session_id):
+    """Gradio's main response function."""
+    if not session_id:
+        session_id = str(uuid.uuid4())
+    
+    response = get_context_and_answer(message, history, session_id)
     
-    iface = gr.ChatInterface(
-        fn=get_context_and_answer,
-        title=" ASKXENO",
-        description="""**Welcome to XENO AI Support!** 
-I can help you with questions about XENO financial services including:
-• Account management and setup
-• Transaction processes and fees  
-• Platform features and troubleshooting
-• General service information
-*Simply type your question below to get started!*""",
-        theme="soft"
-    )
+    config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
+    updated_messages = (memory.get(config) or {}).get("messages", [])
     
-    return iface
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": response})
+    
+    return "", history
+def create_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("""ASKXENO 
+                    
+            **Welcome to XENO AI Support!** 
+            I can help you with questions about XENO financial services including:
+            • Account management and setup
+            • Transaction processes and fees  
+            • Platform features and troubleshooting
+            • General service information
+            *Simply type your question below to get started!*
+                    """)
+        
+        session_id_box = gr.Textbox(label="Session ID", value=str(uuid.uuid4()), interactive=True)
+        
+        chatbot = gr.Chatbot(label="XENO Assistant", bubble_full_width=False, height=500, type="messages")
+        msg = gr.Textbox(label="Your Message", placeholder="Type your question here...")
+        
+        msg.submit(respond, [msg, chatbot, session_id_box], [msg, chatbot])
+    return demo
 
-# === Main Execution ===
 if __name__ == "__main__":
     iface = create_interface()
-    iface.launch(share=False)
+    iface.launch(share=False, server_name="0.0.0.0", server_port=7860)