Delete alz_companion/agent.py

alz_companion/agent.py

@@ -1,724 +0,0 @@
-from __future__ import annotations
-import os
-import json
-import base64
-import time
-import tempfile
-import re
-
-from typing import List, Dict, Any, Optional
-
-try:
-    from openai import OpenAI
-except Exception:
-    OpenAI = None
-
-from langchain.schema import Document
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-
-try:
-    from gtts import gTTS
-except Exception:
-    gTTS = None
-
-from .prompts import (
-    SYSTEM_TEMPLATE, ANSWER_TEMPLATE_CALM, 
-    ANSWER_TEMPLATE_ADQ,
-    # --- ADD YOUR NEW PROMPTS HERE ---
-    ANSWER_TEMPLATE_ADQ_MODERATE,
-    ANSWER_TEMPLATE_ADQ_ADVANCED,
-    # --- END OF ADDITION ---
-    SAFETY_GUARDRAILS, RISK_FOOTER, render_emotion_guidelines,
-    NLU_ROUTER_PROMPT, SPECIALIST_CLASSIFIER_PROMPT,
-    ROUTER_PROMPT,
-    ANSWER_TEMPLATE_FACTUAL,
-    ANSWER_TEMPLATE_GENERAL_KNOWLEDGE,
-    ANSWER_TEMPLATE_GENERAL,
-    ANSWER_TEMPLATE_FACTUAL_MULTI,
-    ANSWER_TEMPLATE_SUMMARIZE,
-    QUERY_EXPANSION_PROMPT
-)
-
-_BEHAVIOR_ALIASES = {
-    "repeating questions": "repetitive_questioning", "repetitive questions": "repetitive_questioning",
-    "confusion": "confusion", "wandering": "wandering", "agitation": "agitation",
-    "accusing people": "false_accusations", "false accusations": "false_accusations",
-    "memory loss": "address_memory_loss", "seeing things": "hallucinations_delusions",
-    "hallucinations": "hallucinations_delusions", "delusions": "hallucinations_delusions",
-    "trying to leave": "exit_seeking", "wanting to go home": "exit_seeking",
-    "aphasia": "aphasia", "word finding": "aphasia", "withdrawn": "withdrawal",
-    "apathy": "apathy", "affection": "affection", "sleep problems": "sleep_disturbance",
-    "anxiety": "anxiety", "sadness": "depression_sadness", "depression": "depression_sadness",
-    "checking orientation": "orientation_check", "misidentification": "misidentification",
-    "sundowning": "sundowning_restlessness", "restlessness": "sundowning_restlessness",
-    "losing things": "object_misplacement", "misplacing things": "object_misplacement",
-    "planning": "goal_breakdown", "reminiscing": "reminiscence_prompting",
-    "communication strategy": "caregiver_communication_template",
-}
-
-def _canon_behavior_list(xs: list[str] | None, opts: list[str]) -> list[str]:
-    out = []
-    for x in (xs or []):
-        y = _BEHAVIOR_ALIASES.get(x.strip().lower(), x.strip())
-        if y in opts and y not in out:
-            out.append(y)
-    return out
-    
-_TOPIC_ALIASES = {
-    "home safety": "treatment_option:home_safety", "long-term care": "treatment_option:long_term_care",
-    "music": "treatment_option:music_therapy", "reassure": "treatment_option:reassurance",
-    "routine": "treatment_option:routine_structuring", "validation": "treatment_option:validation_therapy",
-    "caregiving advice": "caregiving_advice", "medical": "medical_fact",
-    "research": "research_update", "story": "personal_story",
-}
-_CONTEXT_ALIASES = {
-    "mild": "disease_stage_mild", "moderate": "disease_stage_moderate", "advanced": "disease_stage_advanced",
-    "care home": "setting_care_home", "hospital": "setting_clinic_or_hospital", "home": "setting_home_or_community",
-    "group": "interaction_mode_group_activity", "1:1": "interaction_mode_one_to_one", "one to one": "interaction_mode_one_to_one",
-    "family": "relationship_family", "spouse": "relationship_spouse", "staff": "relationship_staff_or_caregiver",
-}
-
-def _canon_topic(x: str, opts: list[str]) -> str:
-    if not x: return "None"
-    y = _TOPIC_ALIASES.get(x.strip().lower(), x.strip())
-    return y if y in opts else "None"
-
-def _canon_context_list(xs: list[str] | None, opts: list[str]) -> list[str]:
-    out = []
-    for x in (xs or []):
-        y = _CONTEXT_ALIASES.get(x.strip().lower(), x.strip())
-        if y in opts and y not in out: out.append(y)
-    return out
-
-
-MULTI_HOP_KEYPHRASES = [
-    r"\bcompare\b", r"\bvs\.?\b", r"\bversus\b", r"\bdifference between\b",
-    r"\b(more|less|fewer) (than|visitors|agitated)\b", r"\bchange after\b",
-    r"\bafter.*(vs|before)\b", r"\bbefore.*(vs|after)\b", r"\b(who|which) .*(more|less)\b",
-    # --- START: REVISED & MORE ROBUST PATTERNS ---
-    r"\b(did|was|is)\b .*\b(where|when|who)\b",      # Catches MH1_new ("Did X happen where Y happened?")
-    r"\bconsidering\b .*\bhow long\b",             # Catches MH2_new
-    r"\b(but|and)\b who was the other person\b",    # Catches MH3_new
-    r"what does the journal say about"             # Catches MH4_new
-    # --- END: REVISED & MORE ROBUST PATTERNS ---
-]
-_MH_PATTERNS = [re.compile(p, re.IGNORECASE) for p in MULTI_HOP_KEYPHRASES]
-
-
-# Add this near the top of agent.py with the other keyphrase lists
-SUMMARIZATION_KEYPHRASES = [
-    r"^\b(summarize|summarise|recap)\b", r"^\b(give me a summary|create a short summary)\b"
-]
-_SUM_PATTERNS = [re.compile(p, re.IGNORECASE) for p in SUMMARIZATION_KEYPHRASES]
-
-def _pre_router_summarization(query: str) -> str | None:
-    q = (query or "")
-    for pat in _SUM_PATTERNS:
-        if re.search(pat, q): return "summarization"
-    return None
-    
-
-CARE_KEYPHRASES = [
-    r"\bwhere am i\b", r"\byou('?| ha)ve stolen my\b|\byou'?ve stolen my\b", 
-    r"\bi lost (the )?word\b|\bword-finding\b|\bcan.?t find the word\b", 
-    r"\bshe didn('?| no)t know me\b|\bhe didn('?| no)t know me\b", 
-    r"\bdisorient(?:ed|ation)\b|\bagitation\b|\bconfus(?:ed|ion)\b", 
-    r"\bcare home\b|\bnursing home\b|\bthe.*home\b", 
-    r"\bplaylist\b|\bsongs?\b.*\b(memories?|calm|soothe|familiar)\b", 
-    r"\bi want to keep teaching\b|\bi want to keep driving\b|\bi want to go home\b",
-    r"music therapy",
-    # --- ADD THESE LINES for handle test cases ---
-    r"music therapy" 
-    r"\bremembering the\b", # Catches P7
-    r"\bmissed you so much\b"  # Catches P4
-    r"\b(i forgot my job|what did i work as|do you remember my job)\b" # Catches queries about forgetting profession
-]
-_CARE_PATTERNS = [re.compile(p) for p in CARE_KEYPHRASES]
-
-
-
-_STRIP_PATTERNS = [(r'^\s*(your\s+(final\s+)?answer|your\s+response)\s+in\s+[A-Za-z\-]+\s*:?\s*', ''), (r'\bbased on (?:the |any )?(?:provided )?(?:context|information|details)(?: provided)?(?:,|\.)?\s*', ''), (r'^\s*as an ai\b.*?(?:,|\.)\s*', ''), (r'\b(according to|from)\s+(the\s+)?(sources?|context)\b[:,]?\s*', ''), (r'\bI hope this helps[.!]?\s*$', '')]
-
-def _clean_surface_text(text: str) -> str:
-    # This function remains unchanged from agent_work.py
-    out = text or ""
-    for pat, repl in _STRIP_PATTERNS:
-        out = re.sub(pat, repl, out, flags=re.IGNORECASE)
-    return re.sub(r'\n{3,}', '\n\n', out).strip()
-
-
-
-
-# Utilities
-def _openai_client() -> Optional[OpenAI]:
-    api_key = os.getenv("OPENAI_API_KEY", "").strip()
-    return OpenAI(api_key=api_key) if api_key and OpenAI else None
-
-def describe_image(image_path: str) -> str:
-    # This function remains unchanged from agent_work.py
-    client = _openai_client()
-    if not client: return "(Image description failed: OpenAI API key not configured.)"
-    try:
-        extension = os.path.splitext(image_path)[1].lower()
-        mime_type = f"image/{'jpeg' if extension in ['.jpg', '.jpeg'] else extension.strip('.')}"
-        with open(image_path, "rb") as image_file:
-            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
-        response = client.chat.completions.create(
-            model="gpt-4o",
-            messages=[{"role": "user", "content": [{"type": "text", "text": "Describe this image concisely for a memory journal. Focus on people, places, and key objects. Example: 'A photo of John and Mary smiling on a bench at the park.'"},{"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{base64_image}"}}]}], max_tokens=100)
-        return response.choices[0].message.content or "No description available."
-    except Exception as e:
-        return f"[Image description error: {e}]"
-
-# --- MODIFICATION 1: Use the new, corrected NLU function ---
-def detect_tags_from_query(
-    query: str,
-    nlu_vectorstore: FAISS,
-    behavior_options: list,
-    emotion_options: list,
-    topic_options: list,
-    context_options: list,
-    settings: dict = None
-) -> Dict[str, Any]:
-    """Uses a dynamic two-step NLU process: Route -> Retrieve Examples -> Classify."""
-    result_dict = {"detected_behaviors": [], "detected_emotion": "None", "detected_topics": [], "detected_contexts": []}
-    router_prompt = NLU_ROUTER_PROMPT.format(query=query)
-    primary_goal_raw = call_llm([{"role": "user", "content": router_prompt}], temperature=0.0).strip().lower()
-    goal_for_filter = "practical_planning" if "practical" in primary_goal_raw else "emotional_support"
-    goal_for_prompt = "Practical Planning" if "practical" in primary_goal_raw else "Emotional Support"
-
-    if settings and settings.get("debug_mode"):
-        print(f"\n--- NLU Router ---\nGoal: {goal_for_prompt} (Filter: '{goal_for_filter}')\n------------------\n")
-
-    retriever = nlu_vectorstore.as_retriever(search_kwargs={"k": 2, "filter": {"primary_goal": goal_for_filter}})
-    retrieved_docs = retriever.invoke(query)
-    if not retrieved_docs:
-        retrieved_docs = nlu_vectorstore.as_retriever(search_kwargs={"k": 2}).invoke(query)
-    
-    selected_examples = "\n".join(
-        f"User Query: \"{doc.page_content}\"\n{json.dumps(doc.metadata['classification'], indent=4)}"
-        for doc in retrieved_docs
-    )
-    if not selected_examples:
-        selected_examples = "(No relevant examples found)"
-        if settings and settings.get("debug_mode"):
-             print("WARNING: NLU retriever found no examples for this query.")
-
-    behavior_str = ", ".join(f'"{opt}"' for opt in behavior_options if opt != "None")
-    emotion_str = ", ".join(f'"{opt}"' for opt in emotion_options if opt != "None")
-    topic_str = ", ".join(f'"{opt}"' for opt in topic_options if opt != "None")
-    context_str = ", ".join(f'"{opt}"' for opt in context_options if opt != "None")
-
-    prompt = SPECIALIST_CLASSIFIER_PROMPT.format(
-        primary_goal=goal_for_prompt, examples=selected_examples,
-        behavior_options=behavior_str, emotion_options=emotion_str,
-        topic_options=topic_str, context_options=context_str, query=query
-    )
-
-    messages = [{"role": "system", "content": "You are a helpful NLU classification assistant."}, {"role": "user", "content": prompt}]
-    response_str = call_llm(messages, temperature=0.0, response_format={"type": "json_object"})
-
-    if settings and settings.get("debug_mode"):
-        print(f"\n--- NLU Specialist Full Response ---\n{response_str}\n----------------------------------\n")
-
-    try:
-        start_brace = response_str.find('{')
-        end_brace = response_str.rfind('}')
-        if start_brace == -1 or end_brace <= start_brace:
-            raise json.JSONDecodeError("No valid JSON object found in response.", response_str, 0)
-        
-        json_str = response_str[start_brace : end_brace + 1]
-        result = json.loads(json_str)
-
-        result_dict["detected_emotion"] = result.get("detected_emotion") or "None"
-        
-        behaviors_raw = result.get("detected_behaviors")
-        behaviors_canon = _canon_behavior_list(behaviors_raw, behavior_options)
-        if behaviors_canon:
-            result_dict["detected_behaviors"] = behaviors_canon
-
-        topics_raw = result.get("detected_topics") or result.get("detected_topic")
-        detected_topics = []
-        if isinstance(topics_raw, list):
-            for t in topics_raw:
-                ct = _canon_topic(t, topic_options)
-                if ct != "None": detected_topics.append(ct)
-        elif isinstance(topics_raw, str):
-            ct = _canon_topic(topics_raw, topic_options)
-            if ct != "None": detected_topics.append(ct)
-        result_dict["detected_topics"] = detected_topics
-        
-        contexts_raw = result.get("detected_contexts")
-        contexts_canon = _canon_context_list(contexts_raw, context_options)
-        if contexts_canon:
-            result_dict["detected_contexts"] = contexts_canon
-            
-        return result_dict
-        
-    except (json.JSONDecodeError, AttributeError) as e:
-        print(f"ERROR parsing NLU Specialist JSON: {e}")
-        return result_dict
-
-def _default_embeddings():
-    # This function remains unchanged from agent_work.py
-    model_name = os.getenv("EMBEDDINGS_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
-    return HuggingFaceEmbeddings(model_name=model_name)
-
-def build_or_load_vectorstore(docs: List[Document], index_path: str, is_personal: bool = False) -> FAISS:
-    # This function remains unchanged from agent_work.py
-    os.makedirs(os.path.dirname(index_path), exist_ok=True)
-    if os.path.isdir(index_path) and os.path.exists(os.path.join(index_path, "index.faiss")):
-        try:
-            return FAISS.load_local(index_path, _default_embeddings(), allow_dangerous_deserialization=True)
-        except Exception: pass
-    if is_personal and not docs:
-        docs = [Document(page_content="(This is the start of the personal memory journal.)", metadata={"source": "placeholder"})]
-    vs = FAISS.from_documents(docs, _default_embeddings())
-    vs.save_local(index_path)
-    return vs
-
-def texts_from_jsonl(path: str) -> List[Document]:
-    # This function remains unchanged from agent_work.py
-    out: List[Document] = []
-    try:
-        with open(path, "r", encoding="utf-8") as f:
-            for i, line in enumerate(f):
-                obj = json.loads(line.strip())
-                txt = obj.get("text") or ""
-                if not txt.strip(): continue
-                md = {"source": os.path.basename(path), "chunk": i}
-                for k in ("behaviors", "emotion", "topic_tags", "context_tags"):
-                    if k in obj and obj[k]: md[k] = obj[k]
-                out.append(Document(page_content=txt, metadata=md))
-    except Exception: return []
-    return out
-
-# Some vectorstores might return duplicates.
-# This is useful when top-k cutoff might otherwise include near-duplicates from query expansion
-def dedup_docs(scored_docs):
-    seen = set()
-    unique = []
-    for doc, score in scored_docs:
-        uid = doc.metadata.get("source", "") + "::" + doc.page_content.strip()
-        if uid not in seen:
-            unique.append((doc, score))
-            seen.add(uid)
-    return unique
-
-    
-def bootstrap_vectorstore(sample_paths: List[str] | None = None, index_path: str = "data/faiss_index") -> FAISS:
-    # This function remains unchanged from agent_work.py
-    docs: List[Document] = []
-    for p in (sample_paths or []):
-        try:
-            if p.lower().endswith(".jsonl"):
-                docs.extend(texts_from_jsonl(p))
-            else:
-                with open(p, "r", encoding="utf-8", errors="ignore") as fh:
-                    docs.append(Document(page_content=fh.read(), metadata={"source": os.path.basename(p)}))
-        except Exception: continue
-    if not docs:
-        docs = [Document(page_content="(empty index)", metadata={"source": "placeholder"})]
-    return build_or_load_vectorstore(docs, index_path=index_path)
-
-def call_llm(messages: List[Dict[str, str]], temperature: float = 0.6, stop: Optional[List[str]] = None, response_format: Optional[dict] = None) -> str:
-    # This function remains unchanged from agent_work.py
-    client = _openai_client()
-    if client is None: raise RuntimeError("OpenAI client not configured (missing API key?).")
-    model = os.getenv("OPENAI_CHAT_MODEL", "gpt-4o-mini")
-    api_args = {"model": model, "messages": messages, "temperature": float(temperature if temperature is not None else 0.6)}
-    if stop: api_args["stop"] = stop
-    if response_format: api_args["response_format"] = response_format
-    resp = client.chat.completions.create(**api_args)
-    content = ""
-    try:
-        content = resp.choices[0].message.content or ""
-    except Exception:
-        msg = getattr(resp.choices[0], "message", None)
-        if isinstance(msg, dict): content = msg.get("content") or ""
-    return content.strip()
-
-MULTI_HOP_KEYPHRASES = [r"\bcompare\b", r"\bvs\.?\b", r"\bversus\b", r"\bdifference between\b", r"\b(more|less|fewer) (than|visitors|agitated)\b", r"\bchange after\b", r"\bafter.*(vs|before)\b", r"\bbefore.*(vs|after)\b", r"\b(who|which) .*(more|less)\b"]
-_MH_PATTERNS = [re.compile(p, re.IGNORECASE) for p in MULTI_HOP_KEYPHRASES]
-
-def _pre_router_multi_hop(query: str) -> str | None:
-    # This function remains unchanged from agent_work.py
-    q = (query or "")
-    for pat in _MH_PATTERNS:
-        if re.search(pat, q): return "multi_hop"
-    return None
-
-def _pre_router(query: str) -> str | None:
-    # This function remains unchanged from agent_work.py
-    q = (query or "").lower()
-    for pat in _CARE_PATTERNS:
-        if re.search(pat, q): return "caregiving_scenario"
-    return None
-
-def _llm_route_with_prompt(query: str, temperature: float = 0.0) -> str:
-    # This function remains unchanged from agent_work.py
-    router_messages = [{"role": "user", "content": ROUTER_PROMPT.format(query=query)}]
-    query_type = call_llm(router_messages, temperature=temperature).strip().lower()
-    return query_type
-
-# OLD use this new pre-router and place it in the correct order of priority.
-# OLD def route_query_type(query: str) -> str:
-# NEW the severity override only apply to moderate or advanced stages   
-def route_query_type(query: str, severity: str = "Normal / Unspecified"):
-    # This new, adaptive logic ONLY applies if severity is set to moderate or advanced.
-    if severity in ["Moderate Stage", "Advanced Stage"]:
-        # Check if it's an obvious other type first (e.g., summarization)
-        if not _pre_router_summarization(query) and not _pre_router_multi_hop(query):
-             print(f"Query classified as: caregiving_scenario (severity override)")
-             return "caregiving_scenario"
-    # END
-    
-    # FOR "Normal / Unspecified", THE CODE CONTINUES HERE, USING THE EXISTING LOGIC
-    # This is your original code path.    
-    # Priority 1: Check for specific, structural queries first.
-    mh_hit = _pre_router_multi_hop(query)
-    if mh_hit:
-        print(f"Query classified as: {mh_hit} (multi-hop pre-router)")
-        return mh_hit
-    
-    # Priority 2: Check for explicit commands like "summarize".
-    sum_hit = _pre_router_summarization(query)
-    if sum_hit:
-        print(f"Query classified as: {sum_hit} (summarization pre-router)")
-        return sum_hit
-
-    # Priority 3: Check for general caregiving keywords.
-    care_hit = _pre_router(query)
-    if care_hit:
-        print(f"Query classified as: {care_hit} (caregiving pre-router)")
-        return care_hit
-        
-    # Fallback: If no pre-routers match, use the LLM for nuanced classification.
-    query_type = _llm_route_with_prompt(query, temperature=0.0)
-    print(f"Query classified as: {query_type} (LLM router)")
-    return query_type
-
-
-# helper: put near other small utils in agent.py
-# In agent.py, replace the _source_ids_for_eval function
-
-def _source_ids_for_eval(docs, cap=5):
-    """
-    Return the source identifiers for evaluation.
-    - For jsonl files, it returns the numeric chunk ID or the scene_id if present.
-    - For ANY other source, it returns the generic name "Text Input".
-    - It excludes the 'placeholder' source.
-    """
-    out, seen = [], set()
-    for d in docs or []:
-        md = getattr(d, "metadata", {}) or {}
-        src = str(md.get("source", "")).lower()
-
-        if src == 'placeholder':
-            continue
-        
-        key = None
-        
-        if src.endswith(".jsonl"):
-            # Prioritize 'scene_id' if it exists (for alive_inside.jsonl)
-            if 'scene_id' in md:
-                key = str(md['scene_id'])
-            # Fallback to numeric chunk ID for other jsonl files
-            elif 'chunk' in md and isinstance(md['chunk'], int):
-                key = str(md['chunk'])
-        else:
-            key = "Text Input"
-        
-        if key and key not in seen:
-            seen.add(key)
-            out.append(str(key))
-            if len(out) >= cap:
-                break
-    return out
-
-
-# In agent.py, replace the ENTIRE make_rag_chain function with this one.
-# def make_rag_chain(vs_general: FAISS, vs_personal: FAISS, *, for_evaluation: bool = False, role: str = "patient", temperature: float = 0.6, language: str = "English", patient_name: str = "the patient", caregiver_name: str = "the caregiver", tone: str = "warm"):
-# NEW:  accept the new disease_stage parameter.
-def make_rag_chain(vs_general: FAISS, vs_personal: FAISS, *, for_evaluation: bool = False, role: str = "patient", temperature: float = 0.6, language: str = "English", patient_name: str = "the patient", caregiver_name: str = "the caregiver", tone: str = "warm", disease_stage: str = "Normal / Unspecified"):
-    """Returns a callable that performs the complete RAG process."""
-
-    RELEVANCE_THRESHOLD = 0.85
-    SCORE_MARGIN = 0.10  # Margin to decide if scores are "close enough" to blend.
-
-    def _format_docs(docs: List[Document], default_msg: str) -> str:
-        if not docs: return default_msg
-        unique_docs = {doc.page_content: doc for doc in docs}.values()
-        return "\n".join([f"- {d.page_content.strip()}" for d in unique_docs])
-
-    # def _answer_fn(query: str, query_type: str, chat_history: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
-    # NEW
-    def _answer_fn(query: str, query_type: str, chat_history: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
-
-        # --- ADD THIS LINE FOR VERIFICATION ---
-        print(f"DEBUG: RAG chain received disease_stage = '{disease_stage}'")
-        # --- END OF ADDITION ---
-    
-        # Create a local variable for test_temperature to avoid the UnboundLocalError.
-        test_temperature = temperature
-        
-        p_name = patient_name or "the patient"
-        c_name = caregiver_name or "the caregiver"
-        
-        perspective_line = (f"You are speaking directly to {p_name}, who is the patient...") if role == "patient" else (f"You are communicating with {c_name}, the caregiver, about {p_name}.")
-        system_message = SYSTEM_TEMPLATE.format(tone=tone, language=language, perspective_line=perspective_line, guardrails=SAFETY_GUARDRAILS)
-        messages = [{"role": "system", "content": system_message}]
-        messages.extend(chat_history)
-
-        if "general_knowledge_question" in query_type or "general_conversation" in query_type:
-            template = ANSWER_TEMPLATE_GENERAL_KNOWLEDGE if "general_knowledge" in query_type else ANSWER_TEMPLATE_GENERAL
-            user_prompt = template.format(question=query, language=language)
-            messages.append({"role": "user", "content": user_prompt})
-            raw_answer = call_llm(messages, temperature=test_temperature)
-            answer = _clean_surface_text(raw_answer)
-            sources = ["General Knowledge"] if "general_knowledge" in query_type else []
-            return {"answer": answer, "sources": sources, "source_documents": []}
-
-        expansion_prompt = QUERY_EXPANSION_PROMPT.format(question=query)
-        expansion_response = call_llm([{"role": "user", "content": expansion_prompt}], temperature=0.1)
-        try:
-            search_queries = [query] + json.loads(expansion_response.strip().replace("```json", "").replace("```", ""))
-        except json.JSONDecodeError:
-            search_queries = [query]
-
-        # NEW: Determine sourcing weight
-        if disease_stage in ["Moderate Stage", "Advanced Stage"]:
-            top_k_general = 5
-            top_k_personal = 1
-        else:  # current default
-            top_k_general = 2
-            top_k_personal = 3
-
-        # NEW: pass top_k_personal and top_k_general parameters
-        personal_results_with_scores = [
-            result for q in search_queries for result in vs_personal.similarity_search_with_score(q, k=top_k_personal)
-        ]
-        general_results_with_scores = [
-            result for q in search_queries for result in vs_general.similarity_search_with_score(q, k=top_k_general)
-        ]
-
-        # NEW: Remove duplicates
-        personal_results_with_scores = dedup_docs(personal_results_with_scores)
-        general_results_with_scores = dedup_docs(general_results_with_scores)
-
-        ## BEGIN DEBUGGING 
-        print(f"[DEBUG] Retrieved {len(personal_results_with_scores)} personal, {len(general_results_with_scores)} general results")
-        if personal_results_with_scores:
-            print(f"Top personal score: {max([s for _, s in personal_results_with_scores]):.3f}")
-        if general_results_with_scores:
-            print(f"Top general score: {max([s for _, s in general_results_with_scores]):.3f}")
-        
-        print("\n--- DEBUG: Personal Search Results with Scores (Before Filtering) ---")
-        if personal_results_with_scores:
-            for doc, score in personal_results_with_scores:
-                print(f"  - Score: {score:.4f} | Source: {doc.metadata.get('source', 'N/A')}")
-        else:
-            print("  - No results found.")
-        print("-----------------------------------------------------------------")
-        
-        print("\n--- DEBUG: General Search Results with Scores (Before Filtering) ----")
-        if general_results_with_scores:
-            for doc, score in general_results_with_scores:
-                print(f"  - Score: {score:.4f} | Source: {doc.metadata.get('source', 'N/A')}")
-        else:
-            print("  - No results found.")
-        print("-----------------------------------------------------------------")
-        ## END DEBUGGING 
-        
-        # Return the most relevant doc if not return the best score; and all strip OUT placehoder doc
-        def get_best_docs_with_fallback(results_with_scores: list[tuple[Document, float]]) -> (list[Document], float):
-            valid_results = [res for res in results_with_scores if res[0].metadata.get("source") != "placeholder"]
-            if not valid_results:
-                return [], float('inf')
-            
-            best_score = sorted(valid_results, key=lambda x: x[1])[0][1]
-            filtered_docs = [doc for doc, score in valid_results if score < RELEVANCE_THRESHOLD]
-            
-            if not filtered_docs:
-                return [sorted(valid_results, key=lambda x: x[1])[0][0]], best_score
-            
-            return filtered_docs, best_score
-        # END def get_best_docs_with_fallback
-        
-        if disease_stage in ["Moderate Stage", "Advanced Stage"]:
-            # Use top-k selection (e.g. top 5 for general, top 1 for personal)
-            filtered_general_docs = [doc for doc, score in general_results_with_scores[:top_k_general]]
-            best_general_score = general_results_with_scores[0][1] if general_results_with_scores else 0.0
-
-            filtered_personal_docs = [doc for doc, score in personal_results_with_scores[:top_k_personal]]
-            best_personal_score = personal_results_with_scores[0][1] if personal_results_with_scores else 0.0
-        else:
-            # Use standard fallback-based scoring
-            filtered_personal_docs, best_personal_score = get_best_docs_with_fallback(personal_results_with_scores)
-            filtered_general_docs, best_general_score = get_best_docs_with_fallback(general_results_with_scores)
-        
-        print("\n--- DEBUG: Filtered Personal Docs (After Threshold/Fallback) ---")
-        if filtered_personal_docs:
-            for doc in filtered_personal_docs:
-                print(f"  - Source: {doc.metadata.get('source', 'N/A')}")
-        else:
-            print("  - No documents met the criteria.")
-        print("----------------------------------------------------------------")
-
-        print("\n--- DEBUG: Filtered General Docs (After Threshold/Fallback) ----")
-        if filtered_general_docs:
-            for doc in filtered_general_docs:
-                print(f"  - Source: {doc.metadata.get('source', 'N/A')}")
-        else:
-            print("  - No documents met the criteria.")
-        print("----------------------------------------------------------------")
-        
-        personal_memory_routes = ["factual", "multi_hop", "summarization"]
-        is_personal_route = any(route_keyword in query_type for route_keyword in personal_memory_routes)
-        
-        all_retrieved_docs = []
-        if is_personal_route:
-            # --- MODIFIED AS PER YOUR SPECIFICATION ---
-            # Implements the simple fallback logic for personal routes.
-            # the logic of it always returns a personal doc unless it's not loaded with personal memory
-            if filtered_personal_docs:
-                all_retrieved_docs = filtered_personal_docs
-            else:
-                all_retrieved_docs = filtered_general_docs
-            # --- END OF MODIFICATION ---
-        else: # caregiving_scenario
-            if disease_stage in ["Moderate Stage", "Advanced Stage"]:
-                # --- STAGE-AWARE LOGIC FOR CAREGIVING SCENARIOS ---
-                if filtered_general_docs:
-                    all_retrieved_docs = filtered_general_docs
-                elif filtered_personal_docs:
-                    all_retrieved_docs = filtered_personal_docs
-                else:
-                    all_retrieved_docs = []
-                # --- END STAGE-AWARE BLOCK ---
-            else:
-                # --- NORMAL ROUTING LOGIC ---
-                # Conditional Blending logic for caregiving remains.
-                if abs(best_personal_score - best_general_score) <= SCORE_MARGIN:
-                    all_retrieved_docs = list({doc.page_content: doc for doc in filtered_personal_docs + filtered_general_docs}.values())[:4]
-                elif best_personal_score < best_general_score:
-                    all_retrieved_docs = filtered_personal_docs
-                else:
-                    all_retrieved_docs = filtered_general_docs
-        
-        # --- Prompt Generation and LLM Call ---
-        answer = ""
-        if is_personal_route:
-            personal_context = _format_docs(all_retrieved_docs, "(No relevant personal memories found.)")
-            # New modify for test evaluation, general_context is empty but use general context in live chat
-            general_context = _format_docs([], "") if for_evaluation else _format_docs(filtered_general_docs, "(No general information found.)")
-            # End
-            
-            template = ANSWER_TEMPLATE_SUMMARIZE if "summarization" in query_type else ANSWER_TEMPLATE_FACTUAL
-            user_prompt = ""
-            if "summarization" in query_type:
-                if for_evaluation:  # for evaluation, use only personal
-                    user_prompt = template.format(context=personal_context, question=query, language=language, patient_name=p_name, caregiver_name=c_name, role=role)
-                else:  # for live chat, use more context
-                    combined_context = f"{personal_context}\n{general_context}".strip()
-                    user_prompt = template.format(context=combined_context, question=query, language=language, patient_name=p_name, caregiver_name=c_name, role=role)
-            
-            else: # ANSWER_TEMPLATE_FACTUAL
-                user_prompt = template.format(personal_context=personal_context, general_context=general_context, question=query, language=language, patient_name=p_name, caregiver_name=c_name)
-            
-            messages.append({"role": "user", "content": user_prompt})
-            if for_evaluation: # if evaluation test, set temperature (creativity) low from 0.6 input
-                test_temperature = 0.0 # Modify the local variable
-            raw_answer = call_llm(messages, temperature=test_temperature)
-            answer = _clean_surface_text(raw_answer)
-        
-        else: # caregiving_scenario
-            # --- MODIFICATION START: Integrate the severity-based logic ---
-            # The disease_stage variable is available here from the outer function's scope
-            
-            # 1. Select the appropriate template based on the disease stage setting.
-            if disease_stage == "Advanced Stage":
-                template = ANSWER_TEMPLATE_ADQ_ADVANCED
-            elif disease_stage == "Moderate Stage":
-                template = ANSWER_TEMPLATE_ADQ_MODERATE
-            else: # Normal / Unspecified or Mild Stage
-                template = ANSWER_TEMPLATE_ADQ
-
-            # 2. The rest of the logic remains the same. It will use the 'template' variable
-            # that was just selected above.
-            personal_sources = {'1 Complaints of a Dutiful Daughter.txt', 'Saved Chat', 'Text Input'}
-            personal_context = _format_docs([d for d in all_retrieved_docs if d.metadata.get('source') in personal_sources], "(No relevant personal memories found.)")
-            general_context = _format_docs([d for d in all_retrieved_docs if d.metadata.get('source') not in personal_sources], "(No general guidance found.)")
-            
-            first_emotion = next((d.metadata.get("emotion") for d in all_retrieved_docs if d.metadata.get("emotion")), None)
-            emotions_context = render_emotion_guidelines(first_emotion or kwargs.get("emotion_tag"))
-            
-            # NEW: Add Emotion Tag
-            user_prompt = template.format(general_context=general_context, personal_context=personal_context, 
-                                          question=query, scenario_tag=kwargs.get("scenario_tag"), 
-                                          emotions_context=emotions_context, role=role, language=language, 
-                                          patient_name=p_name, caregiver_name=c_name, 
-                                          emotion_tag=kwargs.get("emotion_tag"))
-            messages.append({"role": "user", "content": user_prompt})
-            # --- MODIFICATION END ---
-            
-            # OLD
-            # template = ANSWER_TEMPLATE_ADQ
-            # user_prompt = template.format(general_context=general_context, personal_context=personal_context, 
-            #                               question=query, scenario_tag=kwargs.get("scenario_tag"), 
-            #                               emotions_context=emotions_context, role=role, language=language, 
-            #                               patient_name=p_name, caregiver_name=c_name)
-            # messages.append({"role": "user", "content": user_prompt})
-            
-            if for_evaluation: # if evaluation test, set temperature (creativity) low from 0.6 input
-                test_temperature = 0.0 # Modify the local variable
-            raw_answer = call_llm(messages, temperature=test_temperature)
-            answer = _clean_surface_text(raw_answer)
-            
-            high_risk_scenarios = ["exit_seeking", "wandering", "elopement"]
-            if kwargs.get("scenario_tag") and kwargs["scenario_tag"].lower() in high_risk_scenarios:
-                answer += f"\n\n---\n{RISK_FOOTER}"
-        
-        if for_evaluation:
-            sources = _source_ids_for_eval(all_retrieved_docs)
-        else:
-            sources = sorted(list(set(d.metadata.get("source", "unknown") for d in all_retrieved_docs if d.metadata.get("source") != "placeholder")))
-        
-        print("DEBUG Sources (After Filtering):", sources)
-        return {"answer": answer, "sources": sources, "source_documents": all_retrieved_docs}
-
-    return _answer_fn
-              
-    # END of make_rag_chain
-
-def answer_query(chain, question: str, **kwargs) -> Dict[str, Any]:
-    # This function remains unchanged from agent_work.py
-    if not callable(chain): return {"answer": "[Error: RAG chain is not callable]", "sources": []}
-    try:
-        return chain(question, **kwargs)
-    except Exception as e:
-        print(f"ERROR in answer_query: {e}")
-        return {"answer": f"[Error executing chain: {e}]", "sources": []}
-
-def synthesize_tts(text: str, lang: str = "en"):
-    # This function remains unchanged from agent_work.py
-    if not text or gTTS is None: return None
-    try:
-        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as fp:
-            tts = gTTS(text=text, lang=(lang or "en"))
-            tts.save(fp.name)
-            return fp.name
-    except Exception:
-        return None
-
-def transcribe_audio(filepath: str, lang: str = "en"):
-    # This function remains unchanged from agent_work.py
-    client = _openai_client()
-    if not client: return "[Transcription failed: API key not configured]"
-    model = os.getenv("TRANSCRIBE_MODEL", "whisper-1")
-    api_args = {"model": model}
-    if lang and lang != "auto": api_args["language"] = lang
-    with open(filepath, "rb") as audio_file:
-        transcription = client.audio.transcriptions.create(file=audio_file, **api_args)
-    return transcription.text