app_module.py

import os
import torch
import gradio as gr
import time
import re
import codecs
import uuid
import json
import logging
import tempfile
import numpy as np
import scipy.io.wavfile as wavfile
import asyncio
import warnings
from typing import List, Tuple, Generator, Dict
from threading import Thread

# ML / Transformers
import transformers
transformers.utils.logging.set_verbosity_error()
warnings.filterwarnings("ignore", category=UserWarning, module="gradio.components.dropdown")
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer

# --- Logging Setup ---
# Set root logger to ERROR to suppress external library noise
logging.basicConfig(level=logging.ERROR, format='%(name)s [%(levelname)s] %(message)s')

# Specific library suppressions
for lib in ["transformers", "accelerate", "httpx", "gradio", "langchain"]:
    logging.getLogger(lib).setLevel(logging.ERROR)

# Application-level logger
logger = logging.getLogger("app")
logger.setLevel(logging.DEBUG)
logger.propagate = False  # DO NOT propagate to root to avoid double-logging or filtering
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
ch.setFormatter(logging.Formatter('[app] [%(levelname)s] %(message)s'))
logger.addHandler(ch)

# --------------------------------------------------------------------
# Konfiguration & Globale States
# --------------------------------------------------------------------
EMBED_MODEL_ID = "google/embeddinggemma-300m"
LLM_MODEL_ID = "google/gemma-3-4b-it"

EMBEDDING_FUNCTION = None
LLM_MODEL = None
LLM_PROCESSOR = None

# --- UI Premium Aesthetics ---
PREMIUM_CSS = """
.glass-panel {
    background: rgba(255, 255, 255, 0.05) !important;
    backdrop-filter: blur(10px) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 15px !important;
    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37) !important;
}
.sidebar-panel {
    border-right: 1px solid rgba(255, 255, 255, 0.1) !important;
    height: 100vh;
}
    border-bottom: 2px solid #0f3460;
}
.desktop-only { display: block; }
.mobile-only { display: none; }
@media (max-width: 768px) {
    .desktop-only { display: none !important; }
    .mobile-only { display: block !important; }
    .sidebar-panel { display: none !important; }
}
"""

try:
    from pypdf import PdfReader
    from langchain_community.vectorstores import FAISS
    from langchain_huggingface import HuggingFaceEmbeddings
    from langchain_core.documents import Document
    from langchain_text_splitters import RecursiveCharacterTextSplitter
    from mongochain import MongoDBHandler
except ImportError:
    pass

# Spiritual Integration
try:
    from spiritual_bridge import get_oracle_data
except ImportError:
    get_oracle_data = None

# --- Model Loading ---
def get_device() -> torch.device:
    if torch.cuda.is_available(): return torch.device("cuda")
    return torch.device("cpu")

def get_embedding_function():
    global EMBEDDING_FUNCTION
    if EMBEDDING_FUNCTION is None:
        device = get_device()
        logger.debug(f"Initialisiere Embedding-Modell '{EMBED_MODEL_ID}' auf Device '{device}'.")
        EMBEDDING_FUNCTION = HuggingFaceEmbeddings(
            model_name=EMBED_MODEL_ID,
            model_kwargs={'device': device}
        )
        logger.debug("Embedding-Modell erfolgreich initialisiert.")
    return EMBEDDING_FUNCTION

def get_llm():
    global LLM_MODEL, LLM_PROCESSOR
    if LLM_MODEL is None or LLM_PROCESSOR is None:
        device = get_device()
        logger.debug(f"Initialisiere LLM '{LLM_MODEL_ID}' auf Device '{device}'.")
        dtype = torch.bfloat16 if "cuda" in device.type else torch.float32
        LLM_MODEL = Gemma3ForConditionalGeneration.from_pretrained(
            LLM_MODEL_ID,
            dtype=dtype,
            device_map="auto",
        ).eval()
        LLM_PROCESSOR = AutoProcessor.from_pretrained(LLM_MODEL_ID)
        logger.debug("LLM und Prozessor erfolgreich initialisiert.")
    return LLM_MODEL, LLM_PROCESSOR

# --- Language Detection ---
def detect_language(text: str) -> str:
    if not text or len(text) < 3: return "English"
    model, processor = get_llm()
    prompt = f"Detect the language of the following text and return ONLY the language name (e.g., 'English', 'German', 'French'):\n\n\"{text}\""
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(inputs, max_new_tokens=20, do_sample=False)
    raw_output = processor.batch_decode(outputs[:, inputs.shape[1]:], skip_special_tokens=True)[0].strip()
    logger.debug(f"DEBUG: Raw Language Detection Output: '{raw_output}'")
    
    keywords = ["English", "German", "French", "Spanish", "Italian", "Dutch", "Russian", "Chinese", "Japanese"]
    for k in keywords:
        if k.lower() in raw_output.lower():
            logger.debug(f"DEBUG: Detected User Language (Normalized): '{k}'")
            return k
    return "English"

# --- Document Handling ---
def extract_text_from_file(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext in [".txt", ".md", ".markdown"]:
        with open(path, "r", encoding="utf-8", errors="ignore") as f: return f.read()
    if ext == ".pdf":
        text_parts = []
        try:
            reader = PdfReader(path)
            for page in reader.pages:
                page_text = page.extract_text()
                if page_text: text_parts.append(page_text)
            return "\n\n".join(text_parts)
        except Exception as e:
            logger.error(f"Error reading PDF {path}: {e}"); return ""
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f: return f.read()
    except Exception: return ""

def get_text_splitter() -> RecursiveCharacterTextSplitter:
    return RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200, length_function=len)

# --- RAG Core ---
def index_files(file_paths, mongo_uri, db_name, coll_name, use_mongo, vs_state, mh_state, progress=gr.Progress(track_tqdm=True)):
    if not file_paths: return "Keine Dateien zum Indexieren ausgewählt.", vs_state, mh_state
    
    logger.debug(f"Indexierung gestartet für {len(file_paths)} Datei(en).")
    embed_fn = get_embedding_function()
    splitter = get_text_splitter()
    documents = []
    
    for path in progress.tqdm(file_paths, desc="1/2: Dateien verarbeiten"):
        if path is None: continue
        text = extract_text_from_file(path)
        if not text.strip(): continue
        chunks = splitter.split_text(text)
        source_name = os.path.basename(path)
        for c in chunks:
            documents.append(Document(page_content=c, metadata={"source": source_name}))
    
    logger.debug(f"Total chunks created: {len(documents)}")
    if not documents: return "Kein Text zum Indexieren gefunden.", vs_state, mh_state
    
    progress(0.7, desc="2/2: Indexing...")
    new_vs = FAISS.from_documents(documents, embed_fn)
    if vs_state: 
        vs_state.merge_from(new_vs)
    else: 
        vs_state = new_vs
        
    mh_state = None
    if use_mongo:
        try:
            mh_state = MongoDBHandler(uri=mongo_uri, db_name=db_name, collection_name=coll_name)
            mh_state.connect()
            logger.debug(f"Pushe {len(documents)} Chunks nach MongoDB...")
            for doc in documents:
                mh_state.insert_chunk(doc.page_content, doc.metadata)
            logger.debug("MongoDB-Sync abgeschlossen.")
        except Exception as e:
            logger.error(f"Mongo Error: {e}")

    logger.debug(f"Indexierung abgeschlossen. Gesamt: {vs_state.index.ntotal} Chunks.")
    return f"Index aktualisiert: {vs_state.index.ntotal} Chunks insgesamt.", vs_state, mh_state

def clear_index():
    import gc; gc.collect()
    logger.debug("Vektor-Index wurde geleert.")
    return "Index geleert.", None, None

def retrieve_relevant_chunks(query, vs_state, mh_state, top_k=3):
    if not vs_state: return []
    logger.debug(f"Suche in FAISS: '{query}'")
    docs = vs_state.similarity_search(query, k=top_k)
    return [{"content": d.page_content, "source": d.metadata.get("source", "Unknown")} for d in docs]

def build_rag_prompt(user_question: str, retrieved_chunks: List[Dict]) -> str:
    if not retrieved_chunks: context_str = "Kein relevanter Kontext gefunden."
    else:
        context_parts = [f"[{i}] (Quelle: {ch['source']}): \"{ch['content']}\"" for i, ch in enumerate(retrieved_chunks, 1)]
        context_str = "\n\n".join(context_parts)
    return (f"Beantworte die Benutzerfrage nur basierend auf dem Kontext.\n\n"
            f"--- Kontext ---\n{context_str}\n\n"
            f"--- Frage ---\n{user_question}\n\n"
            f"--- Antwort ---")

# --- Agent System ---
def build_agent_prompt(query, context, history, language="English", short_answers=False):
    context_str = "\n".join([f"- {c['content']} (Source: {c['source']})" for i, c in enumerate(context)])
    
    style_instruction = "Be concise." if short_answers else ""
    
    system = f"""You are Sage 6.5, a spiritual AI guide. 
Respond in {language}. {style_instruction}
If you need to use a tool, you MUST use the following JSON format inside <tool_call> tags:
<tool_call>{{"name": "tool_name", "arguments": {{"arg1": "val1"}}}}</tool_call>

Available Tools:
1. oracle_consultation: Consult the archive for deep wisdom. Arguments: {{"topic": "str", "name": "str (Optional. Use ONLY if the user explicitly stated their name, otherwise omit)"}}
"""
    return system + f"\n\nContext:\n{context_str}\n\nUser Question: {query}"

def chat_agent_stream(query, history, vs_state, mh_state, user_lang=None, short_answers=False):
    model, processor = get_llm()
    lang = user_lang if user_lang else detect_language(query)
    context = retrieve_relevant_chunks(query, vs_state, mh_state)
    prompt = build_agent_prompt(query, context, history, language=lang, short_answers=short_answers)
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    
    logger.info(f"[AGENT] 🏁 Starting Agent Loop for Query: '{query}' (Lang: {lang})")
    
def chat_agent_stream(query, history, vs_state, mh_state, user_lang=None, short_answers=False):
    model, processor = get_llm()
    lang = user_lang if user_lang else detect_language(query)
    context = retrieve_relevant_chunks(query, vs_state, mh_state)
    prompt = build_agent_prompt(query, context, history, language=lang, short_answers=short_answers)
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    
    logger.info(f"[AGENT] 🏁 Starting Agent Loop for Query: '{query}' (Lang: {lang})")
    
    max_turns = 3
    for turn in range(max_turns):
        logger.info(f"[AGENT] 🔄 Turn {turn+1}/{max_turns}")
        
        input_ids = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
        
        gen_kwargs = {"input_ids": input_ids, "streamer": streamer, "max_new_tokens": 512, "do_sample": False}
        thread = Thread(target=model.generate, kwargs=gen_kwargs)
        thread.start()
        
        current_turn_text = ""
        # We yield a TUPLE: (accumulated_text_for_THIS_turn, is_final)
        # But wait, the wrapper needs to handle new messages.
        # Strategy: Yield just the text of THIS turn. Wrapper handles appending to a NEW history item each turn.
        
        logger.info("[AGENT] ⏳ Streaming response...")
        for new_text in streamer:
            current_turn_text += new_text
            clean_chunk = re.sub(r"<tool_call>.*?</tool_call>", "", current_turn_text, flags=re.DOTALL)
            yield clean_chunk.strip()
            
        logger.info(f"[AGENT] 🛑 Raw Model Output: {current_turn_text}")
        
        # Tool Detection
        tool_match = re.search(r"<tool_call>(.*?)</tool_call>", current_turn_text, re.DOTALL)
        if tool_match:
            # If tool found, this turn is OVER regarding user output.
            # We yield a special signal to indicate "End of Message, Start Next Logic"?
            # actually, if we yield, the wrapper updates history[-1]. 
            # If we want a NEW message, we need to tell wrapper to append.
            # Simplified: Use a separator? No, wrapper loop is easier.
            
            # For now, let's keep the generator simple. 
            # It yields text updates for the CURRENT turn.
            # Once loop breaks (tool found), we start next turn.
            # BUT: How to tell wrapper "This turn is done, start a new bubble"?
            # Generator yields: {"text": "...", "new_bubble": True/False}
            
            try:
                tool_data = json.loads(tool_match.group(1))
                logger.info(f"[AGENT] 🛠️ Tool Call Detected: {tool_data}")
                
                tool_name = tool_data.get("name")
                tool_args = tool_data.get("arguments", {})
                
                if tool_name == "oracle_consultation":
                    topic = tool_args.get("topic", "")
                    
                    # Name Handling: Use provided name or default to 'Seeker'
                    req_name = tool_args.get("name", "").strip()
                    effective_name = req_name if req_name else "Seeker"
                    
                    logger.info(f"[AGENT] 🔮 Executing Oracle with topic: '{topic}' for '{effective_name}'")
                    if get_oracle_data:
                        try:
                            # Call backend
                            oracle_raw = get_oracle_data(name=effective_name, topic=topic, date_str="")
                            
                            # FILTERING LOGIC (User Request: Only 3 sources, no BOS API/ELS)
                            # We construct a filtered dictionary
                            filtered_result = {
                                "wisdom_nodes": oracle_raw.get("wisdom_nodes", [])
                            }
                            # If wisdom_nodes is missing/empty, maybe keep raw but warn?
                            # Use strict filtering as requested.
                            
                            tool_result = json.dumps(filtered_result, indent=2)
                            logger.info(f"[AGENT] ✅ Oracle Result Obtained (Filtered Size: {len(tool_result)} bytes)")
                        except Exception as e:
                            logger.error(f"[AGENT] ❌ Oracle Backend Error: {e}")
                            tool_result = f"Error executing oracle: {str(e)}"
                    else:
                        logger.warning("[AGENT] ⚠️ Oracle module not available")
                        tool_result = "Oracle module not available."
                else:
                    logger.warning(f"[AGENT] ⚠️ Unknown tool requested: {tool_name}")
                    tool_result = f"Unknown tool: {tool_name}"

                messages.append({"role": "assistant", "content": [{"type": "text", "text": current_turn_text}]})
                
                tool_injection = f"""<tool_result>{tool_result}</tool_result>
Now interpret this result soulfully and poetically for the user. Do not mention JSON.
IMPORTANT: Connect this smoothly to your previous statement. Ensure a fluid, cohesive narrative without abrupt jumps."""
                
                logger.info("[AGENT] 💉 Injecting Tool Result into context for interpretation...")
                messages.append({"role": "user", "content": [{"type": "text", "text": tool_injection}]})
                
                # Yield a special marker to say "Turn Finished"
                yield "__TURN_END__" 
                continue
            except Exception as e:
                logger.error(f"[AGENT] 💥 Tool parsing/logic crash: {e}")
                break
        else:
            logger.info("[AGENT] ✨ No tool calls. Finalizing response.")
            break

# --- Voice Engine ---
async def generate_speech(text: str, lang: str = "English"):
    import edge_tts
    VOICES = {"English": "en-US-GuyNeural", "German": "de-DE-ConradNeural", "French": "fr-FR-HenriNeural"}
    voice = VOICES.get(lang, VOICES["English"])
    logger.debug(f"TRACE: generate_speech() called. Text len: {len(text)}, Lang: {lang}")
    temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(temp_wav.name)
    return temp_wav.name

def transcribe_audio(path: str):
    logger.debug(f"TRACE: transcribe_audio() called with path: {path}")
    return "Transcribed text"

# --- Gradio Wrappers ---
def voice_chat_wrapper(audio_path, history, threads, tid, vs_state, mh_state, short_answers):
    if audio_path is None: yield history, threads, gr.update(), gr.update(), None; return
    text = transcribe_audio(audio_path)
    detected_lang = detect_language(text)
    final_history, final_threads, final_update = history, threads, gr.update()
    if text:
        gen = chat_wrapper(text, history, threads, tid, vs_state, mh_state, short_answers=short_answers, lang=detected_lang)
        for h, t, tr1, tr2, _ in gen:
            final_history, final_threads, final_update = h, t, tr1
            yield h, t, tr1, tr2, None
    import asyncio
    last_msg = final_history[-1]["content"] if final_history else ""
    if last_msg:
        voice_path = asyncio.run(generate_speech(last_msg, lang=detected_lang))
        yield final_history, final_threads, final_update, final_update, voice_path
    else:
        yield final_history, final_threads, final_update, final_update, None

def chat_wrapper(message, history, threads, tid, vs_state, mh_state, short_answers=False, lang=None):
    if not message.strip(): 
        upd = gr.update(choices=[(v["title"], k) for k, v in threads.items()], value=tid)
        yield history, threads, upd, upd, None
        return
    history.append({"role": "user", "content": message})
    yield history, threads, gr.update(), gr.update(), None
    
    # Start first response bubble
    history.append({"role": "assistant", "content": ""})
    
    for response_part in chat_agent_stream(message, history[:-2], vs_state, mh_state, user_lang=lang, short_answers=short_answers):
        if response_part == "__TURN_END__":
            # Start NEW bubble for next turn
            history.append({"role": "assistant", "content": ""})
            yield history, threads, gr.update(), gr.update(), None
        else:
            history[-1]["content"] = response_part
            yield history, threads, gr.update(), gr.update(), None
            
    # Cleanup empty bubble if exists (rare edge case)
    if not history[-1]["content"]: history.pop()
    
    if tid not in threads: threads[tid] = {"title": "Conversation", "history": []}
    threads[tid]["history"] = history
    if len(history) <= 2: 
        threads[tid]["title"] = (message[:25] + "..") if message else "Conversation"
    choices = [(v["title"], k) for k, v in threads.items()]
    upd = gr.update(choices=choices, value=tid)
    yield history, threads, upd, upd, None

def stream_handler(stream, state):
    if stream is None: return state, None
    sr, y = stream
    if y is None or len(y) == 0: return state, None
    y = y.astype(np.float32)
    y = y / np.max(np.abs(y)) if np.max(np.abs(y)) > 0 else y
    rms = np.sqrt(np.mean(y**2))
    SILENCE_THRESHOLD, SILENCE_CHUNKS = 0.01, 20
    if state is None: state = {"buffer": [], "silence_counter": 0, "is_speaking": False}
    state["buffer"].append((sr, stream[1]))
    if rms > SILENCE_THRESHOLD:
        state["is_speaking"], state["silence_counter"] = True, 0
    elif state["is_speaking"]:
        state["silence_counter"] += 1
    if state["is_speaking"] and state["silence_counter"] > SILENCE_CHUNKS:
        full_audio = np.concatenate([c[1] for c in state["buffer"]])
        sr_final = state["buffer"][0][0]
        temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        wavfile.write(temp_wav.name, sr_final, full_audio)
        return {"buffer": [], "silence_counter": 0, "is_speaking": False}, temp_wav.name
    return state, None

# --- INTERNAL CALLBACKS ---
def create_new_thread_callback(threads):
    nid = str(uuid.uuid4())
    threads[nid] = {"title": "New Conversation", "history": []}
    choices = [(v["title"], k) for k, v in threads.items()]
    return threads, nid, gr.update(choices=choices, value=nid), []

def switch_thread(tid, t_state):
    logger.debug(f"TRACE: switch_thread() called for tid: {tid}")
    if isinstance(tid, list):
        if not tid: return [], gr.update(), gr.update(), gr.update()
        tid = tid[0]
    tid = str(tid)
    history = t_state.get(tid, {}).get("history", [])
    choices = [(v["title"], k) for k, v in t_state.items()]
    upd = gr.update(value=tid, choices=choices)
    return history, tid, upd, upd

def session_import_handler(file):
    if not file: return [], {}, None, gr.update(), gr.update()
    with open(file.name, "r") as f: data = json.load(f)
    imported_threads = data.get("threads", {})
    active_id = data.get("active_id", list(imported_threads.keys())[0] if imported_threads else None)
    history = imported_threads.get(active_id, {}).get("history", []) if active_id else []
    choices = [(v["title"], k) for k, v in imported_threads.items()]
    upd = gr.update(choices=choices, value=active_id)
    return history, imported_threads, active_id, upd, upd

def session_export_handler(chatbot_val, threads, active_id):
    export_data = {"threads": threads, "active_id": active_id}
    path = "sage_session_export.json"
    with open(path, "w") as f: json.dump(export_data, f, indent=2)
    return path

def build_demo() -> gr.Blocks:
    initial_thread_id = str(uuid.uuid4())
    with gr.Blocks(title="Gemma 3 Sage v6.5 SP1", theme="soft", css=PREMIUM_CSS) as demo:
        threads_state = gr.State({initial_thread_id: {"title": "New Chat", "history": []}})
        active_thread_id = gr.State(initial_thread_id)
        vector_store_state = gr.State(None)
        mongo_handler_state = gr.State(None)
        
        with gr.Row(elem_classes="header-tray"):
            gr.Markdown("# 🌌 Gemma 3 Sage <small>v6.5 SP1</small>")

        with gr.Row():
            # Desktop Sidebar (Radio List)
            with gr.Column(scale=1, variant="panel", elem_classes="sidebar-panel glass-panel desktop-only"):
                gr.Markdown("### 🕒 Recent Chats")
                # Using Radio as a list selector
                thread_list = gr.Radio(choices=[(f"New Chat", initial_thread_id)], value=initial_thread_id, label=None, interactive=True, container=False)
                new_thread_btn = gr.Button("➕ New Conversation", variant="secondary")
                
            with gr.Column(scale=4):
                with gr.Tabs() as tabs:
                    with gr.Tab("💬 Live Conversation", id=0, elem_classes="glass-panel"):
                        
                        # Mobile Menu (Accordion + Dropdown)
                        with gr.Accordion("🕒 Conversations (Mobile)", open=False, visible=True, elem_classes="mobile-only") as mobile_sessions:
                            m_thread_list = gr.Dropdown(choices=[("New Chat", initial_thread_id)], value=initial_thread_id, label="Select Session")
                            m_new_btn = gr.Button("➕ New Conversation", variant="secondary")
                        
                        chatbot = gr.Chatbot(label="Sage", type="messages", height=600, show_label=False, autoscroll=False)
                        with gr.Row():
                            msg_textbox = gr.Textbox(placeholder="Whisper your heart or type...", label=None, scale=8, container=False)
                            submit_btn = gr.Button("Send", variant="primary", scale=1)
                        # Moved Short Answer checkbox here for visibility
                        with gr.Row():
                            short_ans_cb = gr.Checkbox(label="Short Answers", value=False)
                        with gr.Row():
                            stream_state = gr.State({"buffer": [], "silence_counter": 0, "is_speaking": False})
                            audio_input = gr.Audio(label="Voice", sources="microphone", type="numpy", streaming=True)
                            processed_audio, audio_output = gr.State(None), gr.Audio(label="Sage Voice", autoplay=True, visible=False)
                        with gr.Row(elem_classes="glass-panel"):
                            export_btn = gr.Button("📤 Export Session", variant="secondary", size="sm")
                            import_file = gr.File(label="Import", file_count="single", height=60)
                            export_file = gr.File(label="Download", interactive=False, visible=False)

                    with gr.Tab("📚 Sacred Knowledge", id=1, elem_classes="glass-panel"):
                        file_uploader = gr.File(label="Upload", file_count="multiple", type="filepath")
                        index_button = gr.Button("🔄 Sync Index", variant="primary")
                        index_status = gr.Markdown("Bereit.")
                        with gr.Accordion("⚙️ MongoDB Settings", open=False):
                            mongo_uri = gr.Textbox(label="URI", value="mongodb://localhost:27017/")
                            mongo_db = gr.Textbox(label="DB", value="rag_db")
                            mongo_coll = gr.Textbox(label="Coll", value="gemma_chunks")
                            use_mongo_cb = gr.Checkbox(label="Sync to Mongo", value=True)
                            clear_mongo_btn = gr.Button("🗑️ Clear Mongo")
                        clear_idx_btn = gr.Button("🧹 Clear FAISS", variant="stop")
                        
                        clear_mongo_btn.click(lambda u, d, c: MongoDBHandler(u, d, c).connect() and MongoDBHandler(u, d, c).clear() or "Mongo geleert", [mongo_uri, mongo_db, mongo_coll], index_status)

        audio_input.stream(stream_handler, [audio_input, stream_state], [stream_state, processed_audio])
        processed_audio.change(voice_chat_wrapper, [processed_audio, chatbot, threads_state, active_thread_id, vector_store_state, mongo_handler_state, short_ans_cb], [chatbot, threads_state, thread_list, m_thread_list, audio_output])
        msg_textbox.submit(chat_wrapper, [msg_textbox, chatbot, threads_state, active_thread_id, vector_store_state, mongo_handler_state, short_ans_cb], [chatbot, threads_state, thread_list, m_thread_list, audio_output]).then(lambda: "", None, msg_textbox)
        submit_btn.click(chat_wrapper, [msg_textbox, chatbot, threads_state, active_thread_id, vector_store_state, mongo_handler_state, short_ans_cb], [chatbot, threads_state, thread_list, m_thread_list, audio_output]).then(lambda: "", None, msg_textbox)
        new_thread_btn.click(create_new_thread_callback, [threads_state], [threads_state, active_thread_id, thread_list, chatbot])
        m_new_btn.click(create_new_thread_callback, [threads_state], [threads_state, active_thread_id, m_thread_list, chatbot])
        thread_list.change(switch_thread, [thread_list, threads_state], [chatbot, active_thread_id, thread_list, m_thread_list])
        m_thread_list.change(switch_thread, [m_thread_list, threads_state], [chatbot, active_thread_id, thread_list, m_thread_list])
        index_button.click(index_files, [file_uploader, mongo_uri, mongo_db, mongo_coll, use_mongo_cb, vector_store_state, mongo_handler_state], [index_status, vector_store_state, mongo_handler_state], show_progress="full")
        clear_idx_btn.click(clear_index, outputs=[index_status, vector_store_state, mongo_handler_state], show_progress="full")
        import_file.change(session_import_handler, import_file, [chatbot, threads_state, active_thread_id, thread_list, m_thread_list], show_progress="full")
        export_btn.click(session_export_handler, [chatbot, threads_state, active_thread_id], export_file, show_progress="full").then(lambda: gr.update(visible=True), None, export_file)
    return demo