edited app.py file

app.py

@@ -1,4 +1,3 @@
-import uuid
 import os
 import gradio as gr
 import pandas as pd
@@ -10,12 +9,12 @@ 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 json
 import re
-from typing import Dict, List, Tuple
+import random
+import sqlite3
+from langgraph.checkpoint.sqlite import SqliteSaver
 
 # === Configuration ===
 genai.configure(api_key=os.environ["GEMINI_API_KEY"])
@@ -39,7 +38,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, session_id):
+def log_response(question, answer, source_ids, knowledge_pairs):
     """
     Log a question, answer, source IDs, and knowledge base question-answer pairs to the Google Sheet.
     
@@ -50,13 +49,13 @@ def log_response(question, answer, source_ids, knowledge_pairs, session_id):
         knowledge_pairs (list): List of tuples containing (question, answer) from the knowledge base.
     """
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    # Prepare row with user question, answer, source IDs, and knowledge base pairs
     knowledge_question_1 = knowledge_pairs[0][0] if len(knowledge_pairs) > 0 else "N/A"
     knowledge_answer_1 = knowledge_pairs[0][1] if len(knowledge_pairs) > 0 else "N/A"
     knowledge_question_2 = knowledge_pairs[1][0] if len(knowledge_pairs) > 1 else "N/A"
     knowledge_answer_2 = knowledge_pairs[1][1] if len(knowledge_pairs) > 1 else "N/A"
     row = [
         timestamp,
-        session_id,
         question,
         answer,
         source_ids,
@@ -73,28 +72,46 @@ def log_response(question, answer, source_ids, knowledge_pairs, session_id):
         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)
+# === Load and Clean Knowledge Base ===
+df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
+df_kb.dropna(subset=['Content'], inplace=True)
+
+def prepare_documents(data):
+    documents, metadatas, ids = [], [], []
+    for item in data:
+        documents.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
+        metadatas.append({
+            "question": item["Question"],
+            "content": item["Content"],
+            "section": item.get("Section", ""),
+            "source": item.get("Source", ""),
+            "owner": item.get("Owner", ""),
+            "tag": item.get("Tag", ""),
+            "id": item["ID"]
+        })
+        ids.append(item["ID"])
+    return documents, metadatas, ids
+
+xeno_data_list = df_kb.to_dict('records')
+documents, metadatas, ids = prepare_documents(xeno_data_list)
+
+# === Setup ChromaDB ===
+try:
+    client = chromadb.PersistentClient(path="/tmp/xeno_db")
+    try:
+        collection = client.get_collection(name=collection_name)
+        print(f"Loaded existing ChromaDB collection: {collection_name}")
+    except:
+        print(f"Creating new ChromaDB collection: {collection_name}")
+        collection = client.create_collection(name=collection_name)
+        collection.add(documents=documents, metadatas=metadatas, ids=ids)
+except Exception as e:
+    print(f"Failed to initialize ChromaDB: {e}")
+    raise
+
+vector_store = Chroma(client=client, collection_name=collection_name)
+retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})
 
-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):
@@ -148,7 +165,6 @@ class IntentClassifier:
         for intent_name, intent_data in self.intent_patterns.items():
             for pattern in intent_data['patterns']:
                 if re.search(pattern, message_lower, re.IGNORECASE):
-                    import random
                     response = random.choice(intent_data['responses'])
                     return intent_name, response
         
@@ -162,45 +178,32 @@ class IntentClassifier:
 # Initialize intent classifier
 intent_classifier = IntentClassifier()
 
-# === Load and Clean Knowledge Base ===
-df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
-df_kb.dropna(subset=['Content'], inplace=True)
-
-def prepare_documents(data):
-    documents, metadatas, ids = [], [], []
-    for item in data:
-        documents.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
-        metadatas.append({
-            "question": item["Question"],
-            "content": item["Content"],
-            "section": item.get("Section", ""),
-            "source": item.get("Source", ""),
-            "owner": item.get("Owner", ""),
-            "tag": item.get("Tag", ""),
-            "id": item["ID"]
-        })
-        ids.append(item["ID"])
-    return documents, metadatas, ids
-
-xeno_data_list = df_kb.to_dict('records')
-documents, metadatas, ids = prepare_documents(xeno_data_list)
-
-# === Setup ChromaDB ===
-try:
-    client = chromadb.PersistentClient(path="/tmp/xeno_db")
-    try:
-        collection = client.get_collection(name=collection_name)
-        print(f"Loaded existing ChromaDB collection: {collection_name}")
-    except:
-        print(f"Creating new ChromaDB collection: {collection_name}")
-        collection = client.create_collection(name=collection_name)
-        collection.add(documents=documents, metadatas=metadatas, ids=ids)
-except Exception as e:
-    print(f"Failed to initialize ChromaDB: {e}")
-    raise
+# === LangGraph Memory Setup ===
+conn = sqlite3.connect("xeno_memory.db", check_same_thread=False)
+memory = SqliteSaver(conn=conn)
 
-vector_store = Chroma(client=client, collection_name=collection_name)
-retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})
+def update_memory(config, user_message, assistant_message):
+    current_checkpoint = memory.get(config) or {
+        "v": 1,
+        "ts": datetime.now().isoformat(),
+        "channel_values": {"messages": []},
+        "channel_versions": {},
+        "versions_seen": {},
+        "pending_sends": [],
+    }
+    messages = current_checkpoint["channel_values"].get("messages", [])
+    
+    messages.append({"role": "user", "content": user_message})
+    messages.append({"role": "assistant", "content": assistant_message})
+    
+    # Update version for messages channel
+    current_version = current_checkpoint["channel_versions"].get("messages", 0) + 1
+    current_checkpoint["channel_versions"]["messages"] = current_version
+    
+    current_checkpoint["channel_values"]["messages"] = messages
+    current_checkpoint["ts"] = datetime.now().isoformat()
+    
+    memory.put(config, current_checkpoint, {})
 
 # === Prompt System ===
 SYSTEM_PROMPT = """You are a friendly XENO Support Assistant, an AI-powered helpful and professional customer service representative.
@@ -229,101 +232,72 @@ def process_context(results, cosine_scores, max_results=2):
         knowledge_pairs.append((question, answer))
     return formatted_context, source_ids, knowledge_pairs
 
-# === LLM Generation (Refactored) ===
-def generate_xeno_response(context, question, chat_history):
-    """Generates a response but does NOT handle memory."""
+# === LLM Generation ===
+def generate_xeno_response(context, question, conversation_history=""):
     model = genai.GenerativeModel(llm_model_name)
-    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}"
-    
+    prompt = f"""{SYSTEM_PROMPT}
+
+### PREVIOUS CONVERSATION ###
+{conversation_history}
+
+### CONTEXT ###
+{context}
+
+### QUESTION ###
+{question}"""
     response = model.generate_content(prompt)
     return response.text.strip()
 
-
-# === Main Interface Logic (Refactored) ===
-def get_context_and_answer(message, history, session_id="default"):
-    """
-    Handles intent classification, RAG, and memory updates in one place.
-    """
-    config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
+# === Main Interface Logic ===
+def get_context_and_answer(message, history):
+    config = {"configurable": {"thread_id": "xeno_default"}}  # Fixed thread ID for simplicity; consider per-session in production
     
-    full_checkpoint = memory.get(config) or {}
-    chat_history = full_checkpoint.get("channel_values", {}).get("messages", [])
-    intent, direct_response = intent_classifier.classify_intent(message)
+    # Classify intent first
+    intent, predefined_response = intent_classifier.classify_intent(message)
     
-    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?"
-        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 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)
+    if predefined_response:
+        # Handle simple intents with predefined responses
+        log_response(message, predefined_response, "N/A", [])
+        update_memory(config, message, predefined_response)
+        return predefined_response
+    
+    # Load conversation history for queries
+    current_checkpoint = memory.get(config) or {"channel_values": {"messages": []}}
+    previous_messages = current_checkpoint["channel_values"].get("messages", [])
+    conversation_history = "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in previous_messages])
+    
+    # For queries, proceed with RAG
+    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)
 
-            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."
+    if max(cosine_scores) < 0.6:
+        answer = "I'm sorry, I couldn't find the specific information you're looking for in my knowledge base."
+        log_response(message, answer, "N/A", [])
+        update_memory(config, message, answer)
+        return answer
 
+    context, source_ids, knowledge_pairs = process_context(queried_results, cosine_scores)
+    answer = generate_xeno_response(context, message, conversation_history)
+    log_response(message, answer, ", ".join(source_ids), knowledge_pairs)
     update_memory(config, message, answer)
-    log_response(message, answer, source_ids, knowledge_pairs, session_id)
-    
     return answer
 
-# === Enhanced Gradio UI ===
-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)
-    
-    config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
-    updated_messages = (memory.get(config) or {}).get("messages", [])
-    
-    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
+# === Gradio UI ===
+iface = gr.ChatInterface(
+    fn=get_context_and_answer,
+    title="ASKXENO",
+    description="Ask anything about XENO's financial services.",
+    theme="soft"
+)
 
 if __name__ == "__main__":
-    iface = create_interface()
-    iface.launch(share=False, server_name="0.0.0.0", server_port=7860)
+    iface.launch(share=False)