import os import json import shutil import gradio as gr from typing import List, Dict, Any, Optional # --- Agent Imports & Safe Fallbacks --- try: from alz_companion.agent import bootstrap_vectorstore, make_rag_chain, answer_query, synthesize_tts, transcribe_audio, detect_tags_from_query from alz_companion.prompts import BEHAVIOUR_TAGS, EMOTION_STYLES AGENT_OK = True except Exception as e: AGENT_OK = False def bootstrap_vectorstore(sample_paths=None, index_path="data/"): return object() def make_rag_chain(vs, **kwargs): return lambda q, **k: {"answer": f"(Demo) You asked: {q}", "sources": []} def answer_query(chain, q, **kwargs): return chain(q, **kwargs) def synthesize_tts(text: str, lang: str = "en"): return None def transcribe_audio(filepath: str, lang: str = "en"): return "This is a transcribed message." def detect_tags_from_query(query: str, behavior_options: list, emotion_options: list): return {"detected_behavior": "None", "detected_emotion": "None"} BEHAVIOUR_TAGS = {"None": []} EMOTION_STYLES = {"None": {}} print(f"WARNING: Could not import from alz_companion ({e}). Running in UI-only demo mode.") # --- Centralized Configuration --- CONFIG = { "themes": ["All", "The Father", "Still Alice", "Away from Her", "General Caregiving"], "roles": ["patient", "caregiver"], "behavior_tags": ["None"] + list(BEHAVIOUR_TAGS.keys()), "emotion_tags": ["None"] + list(EMOTION_STYLES.keys()), "languages": {"English": "en", "Chinese": "zh", "Malay": "ms", "French": "fr", "Spanish": "es"}, "tones": ["warm", "neutral", "formal", "playful"] } # --- File Management & Vector Store Logic --- INDEX_BASE = os.getenv('INDEX_BASE', 'data') UPLOADS_BASE = os.path.join(INDEX_BASE, "uploads") os.makedirs(UPLOADS_BASE, exist_ok=True) THEME_PATHS = {t: os.path.join(INDEX_BASE, f"faiss_index_{t.replace(' ', '').lower()}") for t in CONFIG["themes"]} vectorstores = {} def canonical_theme(tk: str) -> str: return tk if tk in CONFIG["themes"] else "All" def theme_upload_dir(theme: str) -> str: p = os.path.join(UPLOADS_BASE, f"theme_{canonical_theme(theme).replace(' ', '').lower()}") os.makedirs(p, exist_ok=True) return p def load_manifest(theme: str) -> Dict[str, Any]: p = os.path.join(theme_upload_dir(theme), "manifest.json") if os.path.exists(p): try: with open(p, "r", encoding="utf-8") as f: return json.load(f) except Exception: pass return {"files": {}} def save_manifest(theme: str, man: Dict[str, Any]): with open(os.path.join(theme_upload_dir(theme), "manifest.json"), "w", encoding="utf-8") as f: json.dump(man, f, indent=2) def list_theme_files(theme: str) -> List[tuple[str, bool]]: man = load_manifest(theme) base = theme_upload_dir(theme) found = [(n, bool(e)) for n, e in man.get("files", {}).items() if os.path.exists(os.path.join(base, n))] existing = {n for n, e in found} for name in sorted(os.listdir(base)): if name not in existing and os.path.isfile(os.path.join(base, name)): found.append((name, False)) man["files"] = dict(found) save_manifest(theme, man) return found def copy_into_theme(theme: str, src_path: str) -> str: fname = os.path.basename(src_path) dest = os.path.join(theme_upload_dir(theme), fname) shutil.copy2(src_path, dest) return dest def seed_files_into_theme(theme: str): SEED_FILES = [ ("sample_data/caregiving_tips.txt", True), ("sample_data/the_father_segments_tagged_with_emotion_hybrid.jsonl", True), ("sample_data/still_alice_segments_tagged_with_emotion_hybrid.jsonl", True), ("sample_data/away_from_her_segments_tagged_with_emotion_hybrid.jsonl", True) ] man, changed = load_manifest(theme), False for path, enable in SEED_FILES: if not os.path.exists(path): continue fname = os.path.basename(path) if not os.path.exists(os.path.join(theme_upload_dir(theme), fname)): copy_into_theme(theme, path) man["files"][fname] = bool(enable) changed = True if changed: save_manifest(theme, man) # this line is not needed -> return changed def ensure_index(theme='All'): theme = canonical_theme(theme) if theme in vectorstores: return vectorstores[theme] upload_dir = theme_upload_dir(theme) enabled_files = [os.path.join(upload_dir, n) for n, enabled in list_theme_files(theme) if enabled] index_path = THEME_PATHS.get(theme) vectorstores[theme] = bootstrap_vectorstore(sample_paths=enabled_files, index_path=index_path) return vectorstores[theme] # --- Gradio Callbacks --- def collect_settings(*args): keys = ["role", "patient_name", "caregiver_name", "tone", "language", "tts_lang", "temperature", "behaviour_tag", "emotion_tag", "active_theme", "tts_on", "debug_mode"] return dict(zip(keys, args)) def chat_fn(user_text, audio_file, settings, chat_history): question = (user_text or "").strip() if audio_file and not question: try: voice_lang_name = settings.get("tts_lang", "English") voice_lang_code = CONFIG["languages"].get(voice_lang_name, "en") question = transcribe_audio(audio_file, lang=voice_lang_code) except Exception as e: err_msg = f"Audio Error: {e}" if settings.get("debug_mode") else "Sorry, I couldn't understand the audio." chat_history.append({"role": "assistant", "content": err_msg}) return "", None, chat_history if not question: return "", None, chat_history chat_history.append({"role": "user", "content": question}) # --- NLU Classification Logic --- manual_behavior_tag = settings.get("behaviour_tag") manual_emotion_tag = settings.get("emotion_tag") if manual_behavior_tag not in [None, "None"] or manual_emotion_tag not in [None, "None"]: scenario_tag = manual_behavior_tag emotion_tag = manual_emotion_tag print("Using manual tags from UI.") else: print("No manual tags set, performing auto-detection...") behavior_options = CONFIG.get("behavior_tags", []) emotion_options = CONFIG.get("emotion_tags", []) detected_tags = detect_tags_from_query(question, behavior_options=behavior_options, emotion_options=emotion_options) scenario_tag = detected_tags.get("detected_behavior") emotion_tag = detected_tags.get("detected_emotion") if (scenario_tag and scenario_tag != "None") or (emotion_tag and emotion_tag != "None"): detected_msg = f"*(Auto-detected context: Behavior=`{scenario_tag}`, Emotion=`{emotion_tag}`)*" chat_history.append({"role": "assistant", "content": detected_msg}) # --- END OF NLU LOGIC --- active_theme = settings.get("active_theme", "All") vs = ensure_index(active_theme) rag_chain_settings = { "role": settings.get("role"), "temperature": settings.get("temperature"), "language": settings.get("language"), "patient_name": settings.get("patient_name"), "caregiver_name": settings.get("caregiver_name"), "tone": settings.get("tone"), } chain = make_rag_chain(vs, **rag_chain_settings) if scenario_tag == "None": scenario_tag = None if emotion_tag == "None": emotion_tag = None simple_history = chat_history[:-1] response = answer_query( chain, question, chat_history=simple_history, scenario_tag=scenario_tag, emotion_tag=emotion_tag ) answer = response.get("answer", "[No answer found]") chat_history.append({"role": "assistant", "content": answer}) audio_out = None if settings.get("tts_on") and answer: tts_lang_code = CONFIG["languages"].get(settings.get("tts_lang"), "en") audio_out = synthesize_tts(answer, lang=tts_lang_code) from gradio import update return "", (update(value=audio_out, visible=bool(audio_out))), chat_history def upload_knowledge(files, current_theme): if not files: return "No files were selected to upload." added = 0 for f in files: try: copy_into_theme(current_theme, f.name); added += 1 except Exception as e: print(f"Error uploading file {f.name}: {e}") if added > 0 and current_theme in vectorstores: del vectorstores[current_theme] return f"Uploaded {added} file(s). Refreshing file list..." def save_file_selection(current_theme, enabled_files): man = load_manifest(current_theme) for fname in man['files']: man['files'][fname] = fname in enabled_files save_manifest(current_theme, man) if current_theme in vectorstores: del vectorstores[current_theme] return f"Settings saved. Index for theme '{current_theme}' will rebuild on the next query." def refresh_file_list_ui(current_theme): files = list_theme_files(current_theme) enabled = [f for f, en in files if en] msg = f"Found {len(files)} file(s). {len(enabled)} enabled." return gr.update(choices=[f for f, _ in files], value=enabled), msg def auto_setup_on_load(current_theme): theme_dir = theme_upload_dir(current_theme) if not os.listdir(theme_dir): print("First-time setup: Auto-seeding sample data...") seed_files_into_theme(current_theme) all_settings = collect_settings("patient", "", "", "warm", "English", "English", 0.7, "None", "None", "All", True, False) files_ui, status_msg = refresh_file_list_ui(current_theme) return all_settings, files_ui, status_msg # In app.py, add this new function def pre_load_indexes(): """Pre-builds the FAISS index for all themes at startup.""" print("Pre-loading knowledge base indexes at startup...") for theme in CONFIG["themes"]: print(f" - Loading index for theme: '{theme}'") try: ensure_index(theme) print(f" ...'{theme}' theme loaded successfully.") except Exception as e: print(f" ...Error loading theme '{theme}': {e}") print("All indexes loaded. Application is ready.") # --- UI Definition (Must be at the end of the file) ---------------- CSS = """ .gradio-container { font-size: 14px; } #chatbot { min-height: 250px; } #audio_out audio { max-height: 20px; } #audio_in audio { max-height: 20px; padding: 0; } """ with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo: settings_state = gr.State({}) with gr.Tab("Chat"): user_text = gr.Textbox(show_label=False, placeholder="Type your message here...") audio_in = gr.Audio(sources=["microphone"], type="filepath", label="Voice Input", elem_id="audio_in") with gr.Row(): submit_btn = gr.Button("Send", variant="primary") clear_btn = gr.Button("Clear") audio_out = gr.Audio(label="Response Audio", autoplay=True, visible=True, elem_id="audio_out") chatbot = gr.Chatbot(elem_id="chatbot", label="Conversation", type="messages") with gr.Tab("Settings"): with gr.Group(): gr.Markdown("## Conversation & Persona Settings") with gr.Row(): role = gr.Radio(CONFIG["roles"], value="patient", label="Your Role") temperature = gr.Slider(0.0, 1.2, value=0.7, step=0.1, label="Creativity") tone = gr.Dropdown(CONFIG["tones"], value="warm", label="Response Tone") with gr.Row(): patient_name = gr.Textbox(label="Patient's Name", placeholder="e.g., 'Dad' or 'John'") caregiver_name = gr.Textbox(label="Caregiver's Name", placeholder="e.g., 'me' or 'Jane'") behaviour_tag = gr.Dropdown(CONFIG["behavior_tags"], value="None", label="Behaviour Filter") emotion_tag = gr.Dropdown(CONFIG["emotion_tags"], value="None", label="Emotion Filter") with gr.Accordion("Language, Voice & Debugging", open=False): language = gr.Dropdown(list(CONFIG["languages"].keys()), value="English", label="Response Language") tts_lang = gr.Dropdown(list(CONFIG["languages"].keys()), value="English", label="Voice Language") tts_on = gr.Checkbox(True, label="Enable Voice Response (TTS)") debug_mode = gr.Checkbox(False, label="Show Debug Info") gr.Markdown("--- \n ## Knowledge Base Management") active_theme = gr.Radio(CONFIG["themes"], value="All", label="Active Knowledge Theme") with gr.Row(): with gr.Column(scale=1): files_in = gr.File(file_count="multiple", file_types=[".jsonl", ".txt"], label="Upload Knowledge Files") upload_btn = gr.Button("Upload to Theme", variant="secondary") seed_btn = gr.Button("Import Sample Data", variant="secondary") with gr.Column(scale=2): mgmt_status = gr.Markdown() files_box = gr.CheckboxGroup(choices=[], label="Enable Files for the Selected Theme") with gr.Row(): save_files_btn = gr.Button("Save Selection", variant="primary") refresh_btn = gr.Button("Refresh List") # --- Event Wiring --- all_settings_components = [role, patient_name, caregiver_name, tone, language, tts_lang, temperature, behaviour_tag, emotion_tag, active_theme, tts_on, debug_mode] for component in all_settings_components: component.change(fn=collect_settings, inputs=all_settings_components, outputs=settings_state) submit_btn.click(fn=chat_fn, inputs=[user_text, audio_in, settings_state, chatbot], outputs=[user_text, audio_out, chatbot]) clear_btn.click(lambda: (None, None, [], None, ""), outputs=[user_text, audio_out, chatbot, audio_in, user_text]) upload_btn.click(upload_knowledge, inputs=[files_in, active_theme], outputs=[mgmt_status]).then(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status]) save_files_btn.click(save_file_selection, inputs=[active_theme, files_box], outputs=[mgmt_status]) seed_btn.click(seed_files_into_theme, inputs=[active_theme], outputs=[]).then(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status]) refresh_btn.click(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status]) active_theme.change(refresh_file_list_ui, inputs=[active_theme], outputs=[files_box, mgmt_status]) demo.load(auto_setup_on_load, inputs=[active_theme], outputs=[settings_state, files_box, mgmt_status]) # --- Launch (Must be the last part of the script) --- if __name__ == "__main__": pre_load_indexes() # <-- ADD THIS LINE to pre-load everything demo.queue().launch(debug=True)