updating the ui take 2

app.py

@@ -1,23 +1,21 @@
-import os
 import uuid
-import re
-import json
-import sqlite3
-import random
+import os
+import gradio as gr
+import pandas as pd
 import torch
 import numpy as np
-import pandas as pd
-import gradio as gr
-from datetime import datetime
 from sentence_transformers import util
-from typing import Tuple
-
 import google.generativeai as genai
 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
+from typing import Tuple
 
 # === Configuration ===
 genai.configure(api_key=os.environ["GEMINI_API_KEY"])
@@ -31,96 +29,126 @@ def get_google_sheets_credentials():
     if not credentials_json:
         raise ValueError("GOOGLE_SHEETS_CREDENTIALS environment variable not set.")
     credentials_dict = json.loads(credentials_json)
-    scope = [
-        "https://spreadsheets.google.com/feeds",
-        "https://www.googleapis.com/auth/drive"
-    ]
-    return Credentials.from_service_account_info(credentials_dict, scopes=scope)
+    scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
+    creds = Credentials.from_service_account_info(credentials_dict, scopes=scope)
+    return creds
 
 client_gspread = gspread.authorize(get_google_sheets_credentials())
 sheet = client_gspread.open("Response_Log").sheet1
 
 def log_response(question, answer, source_ids, knowledge_pairs, session_id):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    kq1, ka1 = knowledge_pairs[0] if len(knowledge_pairs) > 0 else ("N/A", "N/A")
-    kq2, ka2 = knowledge_pairs[1] if len(knowledge_pairs) > 1 else ("N/A", "N/A")
-
+    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, kq1, ka1, kq2, ka2
+        timestamp,
+        session_id,
+        question,
+        answer,
+        source_ids,
+        knowledge_question_1,
+        knowledge_answer_1,
+        knowledge_question_2,
+        knowledge_answer_2
     ]
     try:
         sheet.append_row(row)
+        print(f"Logged: {question} | Source IDs: {source_ids}")
     except Exception as e:
+        print(f"Failed to log to Google Sheet: {e}")
         with open("/tmp/response_log.txt", "a") as f:
-            f.write(",".join(map(str, row)) + "\n")
+            f.write(f"{timestamp},{question},{answer},{source_ids},{knowledge_question_1},{knowledge_answer_1},{knowledge_question_2},{knowledge_answer_2}\n")
 
-# === Memory Setup ===
+# === LangGraph Memory ===
 conn = sqlite3.connect("xeno_memory.db", check_same_thread=False)
 memory = SqliteSaver(conn=conn)
 
 def update_memory(config, user_message, assistant_message):
-    checkpoint = memory.get(config) or {}
-    messages = checkpoint.get("channel_values", {}).get("messages", [])
+    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})
-    memory.put(config, {
+    
+    checkpoint_to_save = {
         "v": 1,
         "id": str(uuid.uuid4()),
         "ts": datetime.now().isoformat(),
         "channel_values": {"messages": messages},
         "channel_versions": {},
-        "versions_seen": {}
-    }, {}, {})
+        "versions_seen": {},   
+    }
+    
+    memory.put(config, checkpoint_to_save, {}, {})
 
 # === Intent Classifier ===
 class IntentClassifier:
     def __init__(self):
         self.intent_patterns = {
-            "greeting": {
-                "patterns": [r"\b(hi|hello|hey|good morning|good afternoon|good evening)\b"],
-                "responses": [
-                    "Hello! I'm XENO Assistant. How can I help you today?",
-                    "Hi there! What can I assist you with?",
-                    "Good day! How may I assist you?"
+            'greeting': {
+                'patterns': [
+                    r'\b(hi|hello|hey|good morning|good afternoon|good evening|greetings)\b',
+                    r'^(hi|hello|hey)[\s!.]*$',
+                    r'\b(how are you|how do you do)\b'
+                ],
+                'responses': [
+                    "Hello! I'm XENO Assistant. How can I help you with XENO financial services today?",
+                    "Hi there! I'm here to assist you with any questions about XENO services. What can I help you with?",
+                    "Good day! Welcome to XENO Support. How may I assist you today?"
                 ]
             },
-            "thanks": {
-                "patterns": [r"\b(thank you|thanks|appreciate)\b"],
-                "responses": [
-                    "You're welcome! Anything else I can help you with?",
-                    "Happy to help!",
-                    "Glad I could assist!"
+            'thanks': {
+                'patterns': [
+                    r'\b(thank you|thanks|thank u|thx|appreciate|grateful)\b',
+                    r'^(thanks|thank you)[\s!.]*$',
+                    r'\b(much appreciated|thanks a lot|thank you so much)\b'
+                ],
+                'responses': [
+                    "You're welcome! Is there anything else I can help you with regarding XENO services?",
+                    "Happy to help! Feel free to ask if you have any other questions about XENO.",
+                    "Glad I could assist you! Let me know if you need help with anything else."
                 ]
             },
-            "goodbye": {
-                "patterns": [r"\b(bye|goodbye|see you)\b"],
-                "responses": [
-                    "Goodbye! Have a great day!",
-                    "Take care! Come back anytime.",
-                    "See you later!"
+            'goodbye': {
+                'patterns': [
+                    r'\b(bye|goodbye|see you|farewell|take care|have a good day)\b',
+                    r'^(bye|goodbye)[\s!.]*$',
+                    r'\b(talk to you later|see you later|until next time)\b'
+                ],
+                'responses': [
+                    "Goodbye! Thank you for using XENO services. Have a great day!",
+                    "Take care! Feel free to return anytime you need help with XENO services.",
+                    "Have a wonderful day! Don't hesitate to reach out if you need assistance with XENO."
                 ]
             }
         }
-
+    
     def classify_intent(self, message: str) -> Tuple[str, str]:
         message_lower = message.lower().strip()
         for intent_name, intent_data in self.intent_patterns.items():
-            for pattern in intent_data["patterns"]:
-                if re.search(pattern, message_lower):
-                    return intent_name, random.choice(intent_data["responses"])
-        return "query", ""
+            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
+        return 'query', ''
+    
+    def is_simple_intent(self, intent: str) -> bool:
+        return intent in ['greeting', 'thanks']
 
 intent_classifier = IntentClassifier()
 
 # === Load Knowledge Base ===
 df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
-df_kb.dropna(subset=["Content"], inplace=True)
+df_kb.dropna(subset=['Content'], inplace=True)
 
 def prepare_documents(data):
-    docs, metas, ids = [], [], []
+    documents, metadatas, ids = [], [], []
     for item in data:
-        docs.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
-        metas.append({
+        documents.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
+        metadatas.append({
             "question": item["Question"],
             "content": item["Content"],
             "section": item.get("Section", ""),
@@ -130,85 +158,102 @@ def prepare_documents(data):
             "id": item["ID"]
         })
         ids.append(item["ID"])
-    return docs, metas, ids
+    return documents, metadatas, ids
 
-docs, metas, ids = prepare_documents(df_kb.to_dict("records"))
+xeno_data_list = df_kb.to_dict('records')
+documents, metadatas, ids = prepare_documents(xeno_data_list)
 
-# === ChromaDB ===
-client = chromadb.PersistentClient(path="/tmp/xeno_db")
+# === ChromaDB Setup ===
 try:
-    collection = client.get_collection(name=collection_name)
-except:
-    collection = client.create_collection(name=collection_name)
-    collection.add(documents=docs, metadatas=metas, ids=ids)
+    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
 
-retriever = Chroma(client=client, collection_name=collection_name) \
-    .as_retriever(search_type="similarity", search_kwargs={"k": 4})
+vector_store = Chroma(client=client, collection_name=collection_name)
+retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})
 
 # === Prompt System ===
-SYSTEM_PROMPT = """You are a friendly XENO Support Assistant.
-Use only the knowledge base context to answer questions.
-Do not hallucinate. If no relevant info, politely decline.
-Remember previous conversations."""
+SYSTEM_PROMPT = """You are a friendly XENO Support Assistant...
+"""
 
-# === Context Processor ===
+# === Context Processing ===
 def process_context(results, cosine_scores, max_results=2):
-    sorted_idx = np.argsort(cosine_scores)[::-1][:max_results]
-    ctx, src_ids, kpairs = "", [], []
-    for i, idx in enumerate(sorted_idx, 1):
-        q = results[idx].metadata.get("question", "N/A")
-        a = results[idx].metadata.get("content", "N/A")
-        ctx += f"Knowledge Entry {i}:\nQ: {q}\nA: {a}\n" + "-" * 40 + "\n"
-        src_ids.append(results[idx].metadata.get("id", "N/A"))
-        kpairs.append((q, a))
-    return ctx, src_ids, kpairs
-
-# === LLM Generation ===
-def generate_xeno_response(context, question, history):
+    sorted_indices = np.argsort(cosine_scores)[::-1][:max_results]
+    formatted_context = ""
+    source_ids = []
+    knowledge_pairs = []
+    for i, idx in enumerate(sorted_indices, 1):
+        result = results[idx]
+        question = result.metadata.get('question', 'N/A')
+        answer = result.metadata.get('content', 'N/A')
+        formatted_context += f"Knowledge Entry {i}:\nQ: {question}\nA: {answer}\n" + "-" * 40 + "\n"
+        source_ids.append(result.metadata.get('id', 'N/A'))
+        knowledge_pairs.append((question, answer))
+    return formatted_context, source_ids, knowledge_pairs
+
+# === LLM Response ===
+def generate_xeno_response(context, question, chat_history):
     model = genai.GenerativeModel(llm_model_name)
-    hist_text = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in history]) if history else "None"
-    prompt = f"{SYSTEM_PROMPT}\n### HISTORY ###\n{hist_text}\n### CONTEXT ###\n{context}\n### QUESTION ###\n{question}"
-    return model.generate_content(prompt).text.strip()
-
-# === Chat Handler ===
-def chat_handler(message, history):
-    session_id = "default"  # could be made dynamic if needed
-    config = {"configurable": {"thread_id": session_id, "checkpoint_ns": ""}}
-    checkpoint = memory.get(config) or {}
-    chat_history = checkpoint.get("channel_values", {}).get("messages", [])
-
-    intent, quick_reply = intent_classifier.classify_intent(message)
-    answer, src_ids, kpairs = "", "N/A", []
-
-    if intent != "query":
-        answer = quick_reply
-    else:
-        try:
-            results = retriever.invoke(message)
-            query_emb = genai.embed_content(model=embedding_model, content=message, task_type="retrieval_query")['embedding']
-            doc_embs = [genai.embed_content(model=embedding_model, content=doc.page_content, task_type="retrieval_document")['embedding'] for doc in results]
-            cos_scores = util.cos_sim(torch.tensor(query_emb).float(), torch.tensor(doc_embs).float())[0].tolist()
-
-            if max(cos_scores) < 0.4:
-                answer = "I'm sorry, I couldn't find specific information for your question."
-            else:
-                ctx, src_ids_list, kpairs = process_context(results, cos_scores)
-                answer = generate_xeno_response(ctx, message, chat_history)
-                src_ids = ", ".join(src_ids_list)
-        except Exception as e:
-            answer = "I’m having a technical issue. Please try again later."
+    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()
+
+# === Main Chat Logic ===
+def get_context_and_answer(message, history, session_id=None):
+    if session_id is None:
+        session_id = str(uuid.uuid4())
+    config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
+    full_checkpoint = memory.get(config) or {}
+    chat_history = full_checkpoint.get("channel_values", {}).get("messages", [])
+    
+    intent, direct_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?"
+        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."
+                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)
+            except Exception as e:
+                print(f"Error during RAG processing: {e}")
+                answer = "I apologize, but I'm having a technical issue."
 
     update_memory(config, message, answer)
-    log_response(message, answer, src_ids, kpairs, session_id)
+    log_response(message, answer, source_ids, knowledge_pairs, session_id)
     return answer
 
 # === Clean ChatInterface UI ===
-iface = gr.ChatInterface(
-    fn=chat_handler,
-    title="ASKXENO",
-    description="Ask anything about XENO's financial services.",
-    theme="soft"
-)
-
 if __name__ == "__main__":
+    iface = gr.ChatInterface(
+        fn=lambda message, history: (get_context_and_answer(message, history, session_id=str(uuid.uuid4())), history),
+        title="ASKXENO",
+        description="Ask anything about XENO's financial services.",
+        theme="soft"
+    )
     iface.launch(share=False, server_name="0.0.0.0", server_port=7860)