Update app.py

app.py

@@ -1,15 +1,3 @@
-# ==========================================
-# CRITICAL FIX FOR CHROMADB / SQLITE
-# This block must be at the very top before other imports
-try:
-    __import__('pysqlite3')
-    import sys
-    sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
-    print("Successfully patched sqlite3 for ChromaDB")
-except ImportError:
-    print("Warning: pysqlite3-binary not installed. ChromaDB might fail if system sqlite is old.")
-# ==========================================
-
 import uuid
 import os
 import gradio as gr
@@ -30,7 +18,7 @@ import re
 from typing import Dict, List, Tuple
 import time
 from contextlib import contextmanager
-import threading
+import threading  # Required for background logging
 import logging
 import traceback
 import sys
@@ -56,12 +44,15 @@ class PipelineTimer:
         self.reset()
     
     def reset(self):
+        """Reset all timing data for a new request"""
         self.start_time = time.time()
         self.step_times = {}
+        self.step_start = None
         self.current_step = None
     
     @contextmanager
     def time_step(self, step_name: str):
+        """Context manager to time a specific step"""
         step_start = time.time()
         self.current_step = step_name
         try:
@@ -71,9 +62,12 @@ class PipelineTimer:
             self.step_times[step_name] = round((step_end - step_start) * 1000, 2)
             self.current_step = None
     
+    def get_total_time(self):
+        return round((time.time() - self.start_time) * 1000, 2)
+    
     def get_timing_summary(self):
         return {
-            'total_time_ms': round((time.time() - self.start_time) * 1000, 2),
+            'total_time_ms': self.get_total_time(),
             'step_times': self.step_times,
             'timestamp': datetime.now().isoformat()
         }
@@ -81,11 +75,10 @@ class PipelineTimer:
 timer = PipelineTimer()
 
 # === Configuration ===
-api_key = os.environ.get("GEMINI_API_KEY")
-if not api_key:
-    print("WARNING: GEMINI_API_KEY not set. App may crash.")
-
-genai.configure(api_key=api_key)
+if "GEMINI_API_KEY" not in os.environ:
+    print("WARNING: GEMINI_API_KEY environment variable not found.")
+    
+genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
 embedding_model = "models/embedding-001"
 llm_model_name = "models/gemma-3-4b-it"
 collection_name = "xeno_collection"
@@ -94,156 +87,210 @@ collection_name = "xeno_collection"
 def get_google_sheets_credentials():
     credentials_json = os.environ.get("GOOGLE_SHEETS_CREDENTIALS")
     if not credentials_json:
-        # Return None to handle gracefully later instead of crashing
-        print("WARNING: GOOGLE_SHEETS_CREDENTIALS not set.")
-        return None
-    
-    try:
-        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)
-    except Exception as e:
-        print(f"Error parsing Google Credentials: {e}")
-        return None
-
-# Initialize Sheets with Robust Failover
-spreadsheet = None
-response_sheet = None
-timing_sheet = None
-feedback_sheet = None
+        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"]
+    creds = Credentials.from_service_account_info(credentials_dict, scopes=scope)
+    return creds
 
+# Initialize Sheets
 try:
-    creds = get_google_sheets_credentials()
-    if creds:
-        client_gspread = gspread.authorize(creds)
-        try:
-            spreadsheet = client_gspread.open("Response_Log")
-            response_sheet = spreadsheet.sheet1
-        except Exception as e:
-            print(f"Could not open spreadsheet: {e}")
-            
-        if spreadsheet:
-            # Init Timing Sheet
-            try:
-                timing_sheet = spreadsheet.worksheet("Timing_Log")
-            except:
-                try:
-                    timing_sheet = spreadsheet.add_worksheet(title="Timing_Log", rows="1000", cols="15")
-                    timing_sheet.append_row(["Timestamp", "Session_ID", "Question", "Total_Time_MS", "Details"])
-                except: pass
-                
-            # Init Feedback Sheet
-            try:
-                feedback_sheet = spreadsheet.worksheet("Feedback_Log")
-            except:
-                try:
-                    feedback_sheet = spreadsheet.add_worksheet(title="Feedback_Log", rows="1000", cols="6")
-                    feedback_sheet.append_row(["Timestamp", "Session_ID", "User_Message", "Bot_Response", "Rating", "Reason"])
-                except: pass
+    client_gspread = gspread.authorize(get_google_sheets_credentials())
+    spreadsheet = client_gspread.open("Response_Log")
+    response_sheet = spreadsheet.sheet1
 except Exception as e:
-    print(f"Google Sheets init failed: {e}")
+    print(f"Error connecting to Google Sheets: {e}")
+    # Dummy classes for dev/fallback
+    class DummySheet:
+        def append_row(self, *args, **kwargs): pass
+        def worksheet(self, *args): return self
+        def add_worksheet(self, *args, **kwargs): return self
+    spreadsheet = DummySheet()
+    response_sheet = DummySheet()
+
+# Timing Sheet
+try:
+    timing_sheet = spreadsheet.worksheet("Timing_Log")
+except:
+    try:
+        timing_sheet = spreadsheet.add_worksheet(title="Timing_Log", rows="1000", cols="15")
+        headers = [
+            "Timestamp", "Session_ID", "Question", "Total_Time_MS",
+            "Intent_Classification_MS", "Memory_Retrieval_MS", "RAG_Retrieval_MS", 
+            "Embedding_Generation_MS", "Similarity_Calculation_MS", "Context_Processing_MS",
+            "LLM_Generation_MS", "Memory_Update_MS", "Logging_MS", "Error_Step", "Notes"
+        ]
+        timing_sheet.append_row(headers)
+    except:
+        timing_sheet = None
+
+# Feedback Sheet
+try:
+    feedback_sheet = spreadsheet.worksheet("Feedback_Log")
+except:
+    try:
+        feedback_sheet = spreadsheet.add_worksheet(title="Feedback_Log", rows="1000", cols="6")
+        headers = ["Timestamp", "Session_ID", "User_Message", "Bot_Response", "Rating", "Flag_Reason"]
+        feedback_sheet.append_row(headers)
+    except:
+        feedback_sheet = None
+
+# === Logging Functions ===
 
-# === Logging Helper Functions ===
 def log_response(question, answer, source_ids, knowledge_pairs, session_id):
-    if not response_sheet: return
+    """Log the main chat interaction"""
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    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,
+        knowledge_question_1, knowledge_answer_1, knowledge_question_2, knowledge_answer_2
+    ]
     try:
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        kq1 = knowledge_pairs[0][0] if len(knowledge_pairs) > 0 else ""
-        ka1 = knowledge_pairs[0][1] if len(knowledge_pairs) > 0 else ""
-        kq2 = knowledge_pairs[1][0] if len(knowledge_pairs) > 1 else ""
-        ka2 = knowledge_pairs[1][1] if len(knowledge_pairs) > 1 else ""
-        
-        row = [timestamp, session_id, question, answer, source_ids, kq1, ka1, kq2, ka2]
         response_sheet.append_row(row)
+        print(f"Logged response: {question} | Sources: {source_ids}")
     except Exception as e:
-        print(f"Log response failed: {e}")
+        print(f"Failed to log response: {e}")
 
 def log_timing_data(question, session_id, timing_summary, error_step=None, notes=None):
-    if not timing_sheet: return
+    """Log performance metrics"""
+    if timing_sheet is None: return
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    step_times = timing_summary['step_times']
+    row = [
+        timestamp, session_id, question[:100], timing_summary['total_time_ms'],
+        step_times.get('intent_classification', 0), step_times.get('memory_retrieval', 0),
+        step_times.get('rag_retrieval', 0), step_times.get('embedding_generation', 0),
+        step_times.get('similarity_calculation', 0), step_times.get('context_processing', 0),
+        step_times.get('llm_generation', 0), step_times.get('memory_update', 0),
+        step_times.get('response_logging', 0), error_step or "", notes or ""
+    ]
     try:
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        row = [
-            timestamp, session_id, question[:50], 
-            timing_summary['total_time_ms'], 
-            json.dumps(timing_summary['step_times']),
-            error_step, notes
-        ]
         timing_sheet.append_row(row)
-    except Exception: pass
+    except Exception as e:
+        print(f"Failed to log timing: {e}")
+
+# === Feedback Functions ===
 
 def _log_feedback_background(row):
+    """Background worker to send feedback to Google Sheets"""
     try:
-        if feedback_sheet: feedback_sheet.append_row(row)
-    except Exception as e: print(f"Feedback log failed: {e}")
+        if feedback_sheet:
+            feedback_sheet.append_row(row)
+            print("Feedback logged successfully.")
+        else:
+            print("Feedback sheet not available.")
+    except Exception as e:
+        print(f"Failed to log feedback: {e}")
 
-# === Feedback Logic ===
 def handle_vote(data: gr.LikeData, history, session_id):
+    """
+    Handles the Google AI Studio style Thumbs Up/Down events.
+    Triggered when user clicks the icon on the chat bubble.
+    """
     if not history: return
+    
     try:
+        # Determine rating
         rating = "Positive" if data.liked else "Negative"
-        idx = data.index
-        if idx < len(history):
-            interaction = history[idx]
-            row = [datetime.now().strftime("%Y-%m-%d %H:%M:%S"), session_id, interaction[0], interaction[1], rating, "Quick Vote"]
+        
+        # Get the interaction from history using data.index
+        # history is a list of [user_msg, bot_msg]
+        interaction_index = data.index
+        
+        # Safety check on index
+        if interaction_index < len(history):
+            interaction = history[interaction_index]
+            user_msg = interaction[0]
+            bot_msg = interaction[1]
+            
+            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            row = [timestamp, session_id, user_msg, bot_msg, rating, "Quick Vote (Icon Click)"]
+            
+            # Run in background thread
             threading.Thread(target=_log_feedback_background, args=(row,)).start()
-    except Exception as e: print(f"Vote error: {e}")
+            print(f"Vote registered: {rating}")
+            
+    except Exception as e:
+        print(f"Error handling vote: {e}")
 
 def submit_manual_flag(reason, history, session_id):
-    if not history: return "No history."
+    """Handles the manual text feedback submission"""
+    if not history: return "No conversation to flag."
+    
     try:
-        interaction = history[-1]
-        row = [datetime.now().strftime("%Y-%m-%d %H:%M:%S"), session_id, interaction[0], interaction[1], "Negative", reason]
+        last_interaction = history[-1]
+        user_msg = last_interaction[0]
+        bot_msg = last_interaction[1]
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        
+        row = [timestamp, session_id, user_msg, bot_msg, "Negative", reason]
         threading.Thread(target=_log_feedback_background, args=(row,)).start()
-        return "Report submitted."
-    except Exception as e: return f"Error: {e}"
+        
+        return "Report submitted. Thank you."
+    except Exception as e:
+        return f"Error submitting report: {str(e)}"
 
-# === Core Logic ===
-# Use a file-based DB that persists in the container's storage
-db_path = "xeno_memory.db"
-conn = sqlite3.connect(db_path, check_same_thread=False)
+# === Core Logic & Classes ===
+conn = sqlite3.connect("xeno_memory.db", check_same_thread=False)
 memory = SqliteSaver(conn=conn)
 
-def update_memory(config, user_msg, bot_msg):
-    checkpoint = {
-        "v": 1, "id": str(uuid.uuid4()), "ts": datetime.now().isoformat(),
-        "channel_values": {"messages": [{"role": "user", "content": user_msg}, {"role": "assistant", "content": bot_msg}]},
-        "channel_versions": {}, "versions_seen": {}
-    }
-    memory.put(config, checkpoint, {}, {})
+def update_memory(config, user_message, assistant_message):
+    with timer.time_step("memory_update"):
+        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 = {
+            "v": 1, "id": str(uuid.uuid4()), "ts": datetime.now().isoformat(),
+            "channel_values": {"messages": messages},
+            "channel_versions": {}, "versions_seen": {},
+        }
+        memory.put(config, checkpoint, {}, {})
 
 def retrieve_memory(config):
-    res = memory.get(config)
-    return res.get("channel_values", {}).get("messages", []) if res else []
+    with timer.time_step("memory_retrieval"):
+        full_checkpoint = memory.get(config) or {}
+        return full_checkpoint.get("channel_values", {}).get("messages", [])
 
 class IntentClassifier:
-    def classify_intent(self, msg):
-        msg = msg.lower()
-        if re.search(r'\b(hi|hello|hey)\b', msg): return 'greeting', "Hello! How can I help with XENO?"
-        if re.search(r'\b(thanks|thank)\b', msg): return 'thanks', "You're welcome!"
+    def __init__(self):
+        self.intent_patterns = {
+            'greeting': {
+                'patterns': [r'\b(hi|hello|hey|greetings)\b', r'^(hi|hello)[\s!.]*$'],
+                'responses': ["Hello! I'm XENO Assistant. How can I help you with XENO financial services?"]
+            },
+            'thanks': {
+                'patterns': [r'\b(thank|thanks)\b'],
+                'responses': ["You're welcome! Let me know if you need anything else."]
+            }
+        }
+    
+    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, re.IGNORECASE):
+                    return intent_name, intent_data['responses'][0]
         return 'query', ''
 
 intent_classifier = IntentClassifier()
 
-# === Knowledge Base ===
-# Wrapped in try-except to prevent crash on file read error
-documents, metadatas, ids = [], [], []
+# === Knowledge Base & ChromaDB ===
 try:
-    if os.path.exists("XENO_Uganda_KnowledgeBase_Advisory.json"):
-        df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
-        df_kb.dropna(subset=['Content'], inplace=True)
-        for r in df_kb.to_dict('records'):
-            documents.append(f"Q: {r['Question']}\nA: {r['Content']}")
-            metadatas.append({"question": r["Question"], "content": r["Content"], "id": str(r["ID"])})
-            ids.append(str(r["ID"]))
-    else:
-        print("Warning: Knowledge base JSON file not found.")
-except Exception as e:
-    print(f"Error loading KB: {e}")
+    df_kb = pd.read_json("XENO_Uganda_KnowledgeBase_Advisory.json")
+    df_kb.dropna(subset=['Content'], inplace=True)
+    xeno_data_list = df_kb.to_dict('records')
+    
+    documents, metadatas, ids = [], [], []
+    for item in xeno_data_list:
+        documents.append(f"Question: {item['Question']}\nAnswer: {item['Content']}")
+        metadatas.append({"question": item["Question"], "content": item["Content"], "id": str(item["ID"])})
+        ids.append(str(item["ID"]))
 
-# === ChromaDB ===
-try:
-    # Use a persistent path that is writable in most containers
-    client = chromadb.PersistentClient(path="./xeno_db_storage")
+    client = chromadb.PersistentClient(path="/tmp/xeno_db")
     try:
         collection = client.get_collection(name=collection_name)
     except:
@@ -253,83 +300,155 @@ try:
     vector_store = Chroma(client=client, collection_name=collection_name)
     retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})
 except Exception as e:
-    print(f"ChromaDB Fatal Error: {e}")
-    # Fallback to prevent crash
+    print(f"DB Init Error: {e}")
+    # Define dummy retriever to allow UI to load even if DB fails
     class DummyRetriever:
-        def invoke(self, x): return []
+        def invoke(self, *args): return []
     retriever = DummyRetriever()
 
-# === Generation ===
-def generate_response(context, question, history):
-    try:
+# === Prompt & Generation ===
+SYSTEM_PROMPT = """You are a friendly XENO Support Assistant. 
+Use only the information provided in the context to answer. 
+If context is missing, apologize and say you cannot assist. Do not hallucinate."""
+
+def process_context(results, cosine_scores, max_results=2):
+    with timer.time_step("context_processing"):
+        if not results: return "", [], []
+        sorted_indices = np.argsort(cosine_scores)[::-1][:max_results]
+        formatted_context = ""
+        source_ids = []
+        knowledge_pairs = []
+        for i, idx in enumerate(sorted_indices, 1):
+            if idx < len(results):
+                result = results[idx]
+                question = result.metadata.get('question', 'N/A')
+                answer = result.metadata.get('content', 'N/A')
+                formatted_context += f"Info {i}: Q: {question}\n A: {answer}\n---\n"
+                source_ids.append(str(result.metadata.get('id', 'N/A')))
+                knowledge_pairs.append((question, answer))
+        return formatted_context, source_ids, knowledge_pairs
+
+def generate_xeno_response(context, question, chat_history):
+    with timer.time_step("llm_generation"):
         model = genai.GenerativeModel(llm_model_name)
-        hist_str = "\n".join([f"{m['role']}: {m['content']}" for m in history])
-        prompt = f"System: You are XENO Support.\nHistory:{hist_str}\nContext:{context}\nUser:{question}"
-        return model.generate_content(prompt).text.strip()
-    except Exception as e:
-        return f"I'm having trouble connecting to my brain right now. ({str(e)})"
+        hist_text = "\n".join([f"{m['role']}: {m['content']}" for m in chat_history]) if chat_history else ""
+        prompt = f"{SYSTEM_PROMPT}\nHistory:\n{hist_text}\nContext:\n{context}\nQuestion:\n{question}"
+        response = model.generate_content(prompt)
+        return response.text.strip()
 
 # === Main Pipeline ===
-def process_message(message, history, session_id):
+def get_context_and_answer(message, history, session_id):
     timer.reset()
-    if not session_id: session_id = str(uuid.uuid4())
-    config = {"configurable": {"thread_id": str(session_id)}}
+    error_step = None
+    notes = []
     
     try:
-        with timer.time_step("intent"):
-            intent, direct_resp = intent_classifier.classify_intent(message)
+        config = {"configurable": {"thread_id": str(session_id), "checkpoint_ns": ""}}
+        
+        with timer.time_step("intent_classification"):
+            intent, direct_response = intent_classifier.classify_intent(message)
         
+        chat_history = retrieve_memory(config)
+        answer, source_ids, knowledge_pairs = "", "N/A", []
+
         if intent != 'query':
-            resp = direct_resp
+            answer = direct_response
+            notes.append(f"Intent: {intent}")
         else:
-            with timer.time_step("retrieval"):
-                docs = retriever.invoke(message)
-            
-            context = "\n".join([d.page_content for d in docs])
-            chat_hist = retrieve_memory(config)
-            
-            with timer.time_step("generation"):
-                resp = generate_response(context, message, chat_hist)
+            try:
+                with timer.time_step("rag_retrieval"):
+                    queried_results = retriever.invoke(message)
+                
+                with timer.time_step("embedding_generation"):
+                    q_embed = genai.embed_content(model=embedding_model, content=message, task_type="retrieval_query")['embedding']
+                    d_embeds = [genai.embed_content(model=embedding_model, content=d.page_content, task_type="retrieval_document")['embedding'] for d in queried_results]
+                
+                with timer.time_step("similarity_calculation"):
+                    if d_embeds:
+                        cosine_scores = util.cos_sim(torch.tensor(q_embed).float(), torch.tensor(d_embeds).float())[0].tolist()
+                        max_score = max(cosine_scores)
+                    else:
+                        cosine_scores, max_score = [], 0
+
+                if max_score < 0.4:
+                    answer = "I'm sorry, I couldn't find specific information for your question."
+                    notes.append(f"Low score: {max_score}")
+                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)
+                    notes.append(f"Score: {max_score:.2f}")
+
+            except Exception as e:
+                error_step = "rag_pipeline"
+                answer = "I apologize, but I'm having a technical issue."
+                print(f"RAG Error: {e}")
+
+        update_memory(config, message, answer)
         
-        update_memory(config, message, resp)
-        log_timing_data(message, session_id, timer.get_timing_summary())
-        return resp
-    
+        with timer.time_step("response_logging"):
+            log_response(message, answer, source_ids, knowledge_pairs, session_id)
+        
+        log_timing_data(message, session_id, timer.get_timing_summary(), error_step, "; ".join(notes))
+        return answer
+
     except Exception as e:
-        print(f"Pipeline Error: {e}")
-        traceback.print_exc()
-        return "I encountered a system error. Please try again."
-
-# === UI ===
-def respond(msg, hist, sid):
-    if not sid: sid = str(uuid.uuid4())
-    resp = process_message(msg, hist, sid)
-    hist.append([msg, resp])
-    return "", hist
-
-def create_demo():
+        log_timing_data(message, session_id, timer.get_timing_summary(), "pipeline_crash", str(e))
+        return "System Error. Please try again."
+
+# === UI Logic ===
+def respond(message, history, session_id):
+    if not session_id: session_id = str(uuid.uuid4())
+    bot_response = get_context_and_answer(message, history, session_id)
+    history.append([message, bot_response])
+    return "", history
+
+def create_interface():
+    # 'fill_height=True' is key for the modern full-screen chat look
     with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
-        gr.Markdown("## ASKXENO")
-        sid = gr.Textbox(value=str(uuid.uuid4()), visible=False)
-        cb = gr.Chatbot(scale=1, likeable=True, show_copy_button=True, bubble_full_width=False)
+        gr.Markdown("## ASKXENO Support")
+        
+        session_id_box = gr.Textbox(label="Session ID", value=str(uuid.uuid4()), visible=False)
+        
+        # likeable=True adds the Thumbs Up/Down icons to bubbles
+        chatbot = gr.Chatbot(
+            label="XENO Assistant",
+            scale=1,
+            likeable=True,  
+            show_copy_button=True,
+            bubble_full_width=False
+        )
         
         with gr.Row(variant="compact"):
-            txt = gr.Textbox(placeholder="Ask XENO...", scale=6, container=False, autofocus=True)
-            btn = gr.Button("Send", scale=1)
-            
-        with gr.Accordion("Report Issue", open=False):
+            msg = gr.Textbox(
+                placeholder="Ask about XENO services...",
+                scale=6,
+                lines=1,
+                show_label=False,
+                autofocus=True,
+                container=False
+            )
+            send_btn = gr.Button("Send", variant="primary", scale=1, min_width=80)
+
+        # Collapsible Flagging Section
+        with gr.Accordion("Report an Issue", open=False):
             with gr.Row():
-                ftxt = gr.Textbox(placeholder="Issue details", show_label=False, scale=4)
-                fbtn = gr.Button("Submit", scale=1)
-            flbl = gr.Label(show_label=False)
-
-        txt.submit(respond, [txt, cb, sid], [txt, cb])
-        btn.click(respond, [txt, cb, sid], [txt, cb])
-        cb.like(handle_vote, [cb, sid], None)
-        fbtn.click(submit_manual_flag, [ftxt, cb, sid], [flbl])
+                flag_reason = gr.Textbox(placeholder="Describe the issue (e.g. incorrect fees)", show_label=False, scale=4)
+                flag_btn = gr.Button("Submit Report", scale=1)
+            flag_status = gr.Label(value="", show_label=False)
+
+        # Event Wiring
+        msg.submit(respond, [msg, chatbot, session_id_box], [msg, chatbot])
+        send_btn.click(respond, [msg, chatbot, session_id_box], [msg, chatbot])
         
+        # Handle the native Google AI Studio style likes
+        chatbot.like(handle_vote, [chatbot, session_id_box], None)
+        
+        # Handle manual text flagging
+        flag_btn.click(submit_manual_flag, [flag_reason, chatbot, session_id_box], [flag_status])
+
     return demo
 
 if __name__ == "__main__":
-    demo = create_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+    iface = create_interface()
+    iface.launch(share=False, server_name="0.0.0.0", server_port=7860, ssr_mode=False)