Update app.py

app.py

@@ -2,21 +2,33 @@ 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
+    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
+    )
     from alz_companion.prompts import BEHAVIOUR_TAGS, EMOTION_STYLES
+    from langchain.schema import Document
     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 make_rag_chain(vs, **kwargs): return lambda q, **k: {"answer": f"(Demo) You asked: {q}", "sources": []}
+    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
     BEHAVIOUR_TAGS = {"None": []}
     EMOTION_STYLES = {"None": {}}
     print(f"WARNING: Could not import from alz_companion ({e}). Running in UI-only demo mode.")
@@ -34,9 +46,11 @@ CONFIG = {
 # --- 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:
@@ -83,7 +97,7 @@ def seed_files_into_theme(theme: str):
             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]
@@ -98,9 +112,55 @@ 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."
+    docs_to_add = []
+    if text_input and text_input.strip():
+        docs_to_add.append(Document(page_content=text_input.strip(), metadata={"source": "Text Input"}))
+    if file_input:
+        transcribed_text = transcribe_audio(file_input.name)
+        docs_to_add.append(Document(page_content=transcribed_text, metadata={"source": os.path.basename(file_input.name)}))
+    if image_input:
+        described_text = describe_image(image_input.name)
+        docs_to_add.append(Document(page_content=described_text, metadata={"source": "Image Input"}))
+    if not docs_to_add:
+        return "No processable content found to add."
+    if personal_vectorstore is None:
+        personal_vectorstore = build_or_load_vectorstore(docs_to_add, PERSONAL_INDEX_PATH, is_personal=True)
+    else:
+        personal_vectorstore.add_documents(docs_to_add)
+    personal_vectorstore.save_local(PERSONAL_INDEX_PATH)
+    return f"Successfully added {len(docs_to_add)} new item(s) to personal knowledge base."
+
+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 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")
@@ -110,71 +170,43 @@ def chat_fn(user_text, audio_file, settings, chat_history):
             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)
-
+    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"),
+        "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)
-
+    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
-    )
-
+    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
 
@@ -206,27 +238,9 @@ def auto_setup_on_load(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; }
-"""
+# --- 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({})
@@ -234,34 +248,45 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
     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"):
+        gr.Markdown("## Add to Personal Knowledge Base")
+        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.Tab("Settings"):
         with gr.Group():
             gr.Markdown("## Conversation & Persona Settings")
             with gr.Row():
-                role = gr.Radio(CONFIG["roles"], value="patient", label="Your Role")
+                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")
-                emotion_tag = gr.Dropdown(CONFIG["emotion_tags"], value="None", label="Emotion Filter")
-
+                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 ## Knowledge Base Management")
+        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):
@@ -281,17 +306,37 @@ with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
         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])
+    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])
     
     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])
+    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])
 
-# --- Launch (Must be the last part of the script) ---
+# --- 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() # <-- ADD THIS LINE to pre-load everything
+    pre_load_indexes()
     demo.queue().launch(debug=True)
-