Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import shutil | |
| import gradio as gr | |
| from datetime import datetime | |
| 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, describe_image, build_or_load_vectorstore, | |
| _default_embeddings | |
| ) | |
| from alz_companion.prompts import BEHAVIOUR_TAGS, EMOTION_STYLES | |
| from langchain.schema import Document | |
| from langchain_community.vectorstores import FAISS | |
| AGENT_OK = True | |
| except Exception as e: | |
| AGENT_OK = False | |
| # Define all fallback functions and classes | |
| def bootstrap_vectorstore(sample_paths=None, index_path="data/"): return object() | |
| def build_or_load_vectorstore(docs, index_path, is_personal=False): return object() | |
| def make_rag_chain(vs_general, vs_personal, **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"} | |
| def describe_image(image_path: str): return "This is a description of an image." | |
| class Document: | |
| def __init__(self, page_content, metadata): | |
| self.page_content = page_content | |
| self.metadata = metadata | |
| class FAISS: | |
| def __init__(self): | |
| self.docstore = type('obj', (object,), {'_dict': {}})() | |
| 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") | |
| PERSONAL_INDEX_PATH = os.path.join(INDEX_BASE, "personal_faiss_index") | |
| 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 = {} | |
| personal_vectorstore = None | |
| 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) | |
| 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 add_personal_knowledge(text_input, file_input, image_input): | |
| global personal_vectorstore | |
| if not any([text_input, file_input, image_input]): | |
| return "Please provide text, a file, or an image to add." | |
| content_text, content_source = "", "" | |
| if text_input and text_input.strip(): | |
| content_text, content_source = text_input.strip(), "Text Input" | |
| elif file_input: | |
| content_text, content_source = transcribe_audio(file_input.name), os.path.basename(file_input.name) | |
| elif image_input: | |
| content_text, content_source = describe_image(image_input.name), "Image Input" | |
| if not content_text: | |
| return "Could not extract any text content to add." | |
| print("Auto-tagging personal memory...") | |
| behavior_options = CONFIG.get("behavior_tags", []) | |
| emotion_options = CONFIG.get("emotion_tags", []) | |
| detected_tags = detect_tags_from_query(content_text, behavior_options=behavior_options, emotion_options=emotion_options) | |
| detected_behavior = detected_tags.get("detected_behavior") | |
| detected_emotion = detected_tags.get("detected_emotion") | |
| print(f" ...Detected Behavior: {detected_behavior}, Emotion: {detected_emotion}") | |
| metadata = {"source": content_source} | |
| if detected_behavior and detected_behavior != "None": | |
| metadata["behaviors"] = [detected_behavior.lower()] | |
| if detected_emotion and detected_emotion != "None": | |
| metadata["emotion"] = detected_emotion.lower() | |
| doc_to_add = Document(page_content=content_text, metadata=metadata) | |
| if personal_vectorstore is None: | |
| personal_vectorstore = build_or_load_vectorstore([doc_to_add], PERSONAL_INDEX_PATH, is_personal=True) | |
| else: | |
| personal_vectorstore.add_documents([doc_to_add]) | |
| personal_vectorstore.save_local(PERSONAL_INDEX_PATH) | |
| return f"Successfully added memory with tags (Behavior: {detected_behavior}, Emotion: {detected_emotion})" | |
| def save_chat_to_memory(chat_history): | |
| global personal_vectorstore | |
| if not chat_history: | |
| return "Nothing to save." | |
| formatted_chat = [] | |
| for message in chat_history: | |
| role = "User" if message["role"] == "user" else "Assistant" | |
| content = message["content"].strip() | |
| if content.startswith("*(Auto-detected context:"): | |
| continue | |
| formatted_chat.append(f"{role}: {content}") | |
| conversation_text = "\n".join(formatted_chat) | |
| if not conversation_text: | |
| return "No conversation content to save." | |
| timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| doc_to_add = Document(page_content=conversation_text, metadata={"source": f"Conversation saved on {timestamp}"}) | |
| if personal_vectorstore is None: | |
| personal_vectorstore = build_or_load_vectorstore([doc_to_add], PERSONAL_INDEX_PATH, is_personal=True) | |
| else: | |
| personal_vectorstore.add_documents([doc_to_add]) | |
| personal_vectorstore.save_local(PERSONAL_INDEX_PATH) | |
| print(f"Saved conversation to long-term memory.") | |
| return f"Conversation from {timestamp} saved successfully to long-term memory!" | |
| def list_personal_memories(): | |
| global personal_vectorstore | |
| if personal_vectorstore is None or not hasattr(personal_vectorstore.docstore, '_dict') or not personal_vectorstore.docstore._dict: | |
| return gr.update(value=[["No memories to display", ""]]), gr.update(choices=["No memories to select"], value=None) | |
| docs = list(personal_vectorstore.docstore._dict.values()) | |
| dataframe_data = [[doc.metadata.get('source', 'Unknown'), doc.page_content] for doc in docs] | |
| dropdown_choices = [doc.page_content for doc in docs] | |
| return gr.update(value=dataframe_data), gr.update(choices=dropdown_choices) | |
| def delete_personal_memory(memory_to_delete): | |
| global personal_vectorstore | |
| if personal_vectorstore is None or not memory_to_delete: | |
| return "Knowledge base is empty or no memory selected." | |
| all_docs = list(personal_vectorstore.docstore._dict.values()) | |
| docs_to_keep = [doc for doc in all_docs if doc.page_content != memory_to_delete] | |
| if len(all_docs) == len(docs_to_keep): | |
| return "Error: Could not find the selected memory to delete." | |
| print(f"Deleting memory. {len(docs_to_keep)} memories remaining.") | |
| if not docs_to_keep: | |
| if os.path.isdir(PERSONAL_INDEX_PATH): | |
| shutil.rmtree(PERSONAL_INDEX_PATH) | |
| personal_vectorstore = build_or_load_vectorstore([], PERSONAL_INDEX_PATH, is_personal=True) | |
| else: | |
| # Rebuild and save the index | |
| new_vs = FAISS.from_documents(docs_to_keep, _default_embeddings()) | |
| new_vs.save_local(PERSONAL_INDEX_PATH) | |
| personal_vectorstore = new_vs | |
| return "Successfully deleted memory. The list will now refresh." | |
| def chat_fn(user_text, audio_file, settings, chat_history): | |
| global personal_vectorstore | |
| 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}) | |
| 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, emotion_tag = manual_behavior_tag, manual_emotion_tag | |
| else: | |
| 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, emotion_tag = detected_tags.get("detected_behavior"), 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}) | |
| active_theme = settings.get("active_theme", "All") | |
| vs_general = ensure_index(active_theme) | |
| if personal_vectorstore is None: | |
| personal_vectorstore = build_or_load_vectorstore([], PERSONAL_INDEX_PATH, is_personal=True) | |
| 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_general, personal_vectorstore, **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 | |
| # --- UI Definition --- | |
| CSS = ".gradio-container { font-size: 14px; } #chatbot { min-height: 250px; } #audio_out audio { max-height: 40px; } #audio_in audio { max-height: 40px; 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") | |
| save_btn = gr.Button("Save to Memory") | |
| clear_btn = gr.Button("Clear") | |
| chat_status = gr.Markdown() | |
| 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("Personalize"): | |
| with gr.Accordion("Add to Personal Knowledge Base", open=True): | |
| gr.Markdown("Add personal notes, memories, or descriptions of people and places. You can also upload audio/video notes or images.") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| personal_text = gr.Textbox(lines=5, label="Text Input", placeholder="e.g., 'My father's name is John. He loves listening to Frank Sinatra music.'") | |
| with gr.Column(scale=1): | |
| personal_file = gr.File(label="Upload Audio/Video File") | |
| personal_image = gr.Image(type="filepath", label="Upload Image") | |
| with gr.Row(): | |
| personal_add_btn = gr.Button("Add Knowledge to Memory", variant="primary") | |
| personal_status = gr.Markdown() | |
| with gr.Accordion("Manage Personal Knowledge", open=False): | |
| personal_memory_display = gr.DataFrame(headers=["Source", "Content"], label="Saved Personal Memories", interactive=False, row_count=(5, "dynamic")) | |
| with gr.Row(): | |
| personal_refresh_btn = gr.Button("Refresh Memories") | |
| with gr.Row(): | |
| personal_delete_selector = gr.Dropdown(label="Select a memory to delete", scale=3, interactive=True) | |
| personal_delete_btn = gr.Button("Delete Selected Memory", variant="stop", scale=1) | |
| personal_delete_status = gr.Markdown() | |
| with gr.Tab("Settings"): | |
| with gr.Group(): | |
| gr.Markdown("## Conversation & Persona Settings") | |
| with gr.Row(): | |
| role = gr.Radio(CONFIG["roles"], value="caregiver", 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 (Manual Override)") | |
| emotion_tag = gr.Dropdown(CONFIG["emotion_tags"], value="None", label="Emotion Filter (Manual Override)") | |
| 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 ## General 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]) | |
| save_btn.click(fn=save_chat_to_memory, inputs=[chatbot], outputs=[chat_status]) | |
| clear_btn.click(lambda: (None, None, [], None, "", ""), outputs=[user_text, audio_out, chatbot, audio_in, user_text, chat_status]) | |
| personal_add_btn.click(fn=add_personal_knowledge, inputs=[personal_text, personal_file, personal_image], outputs=[personal_status]).then(lambda: (None, None, None), outputs=[personal_text, personal_file, personal_image]) | |
| personal_refresh_btn.click(fn=list_personal_memories, inputs=None, outputs=[personal_memory_display, personal_delete_selector]) | |
| personal_delete_btn.click(fn=delete_personal_memory, inputs=[personal_delete_selector], outputs=[personal_delete_status]).then(fn=list_personal_memories, inputs=None, outputs=[personal_memory_display, personal_delete_selector]) | |
| 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]).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]) | |
| # --- Startup Logic --- | |
| def pre_load_indexes(): | |
| global personal_vectorstore | |
| print("Pre-loading all knowledge base indexes at startup...") | |
| for theme in CONFIG["themes"]: | |
| print(f" - Loading general 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(" - Loading personal knowledge index...") | |
| try: | |
| personal_vectorstore = build_or_load_vectorstore([], PERSONAL_INDEX_PATH, is_personal=True) | |
| print(" ...Personal knowledge loaded successfully.") | |
| except Exception as e: | |
| print(f" ...Error loading personal knowledge: {e}") | |
| print("All indexes loaded. Application is ready.") | |
| if __name__ == "__main__": | |
| pre_load_indexes() | |
| demo.queue().launch(debug=True) | |