Update alz_companion/agent.py

alz_companion/agent.py

@@ -8,6 +8,7 @@ import re
 import random # for random select songs
 
 from typing import List, Dict, Any, Optional
+from sentence_transformers import CrossEncoder
 
 try:
     from openai import OpenAI
@@ -107,6 +108,25 @@ MULTI_HOP_KEYPHRASES = [
 _MH_PATTERNS = [re.compile(p, re.IGNORECASE) for p in MULTI_HOP_KEYPHRASES]
 
 
+
+FACTUAL_KEYPHRASES = [
+    r"\b(what is|what was) my\b",
+    r"\b(who is|who was) my\b",
+    r"\b(where is|where was) my\b",
+    r"\b(how old am i)\b",
+    # r"\b(when did|what did) the journal say\b"
+]
+_FQ_PATTERNS = [re.compile(p, re.IGNORECASE) for p in FACTUAL_KEYPHRASES]
+
+def _pre_router_factual(query: str) -> str | None:
+    """Checks for patterns common in direct factual questions about personal memory."""
+    q = (query or "")
+    for pat in _FQ_PATTERNS:
+        if re.search(pat, q):
+            return "factual_question"
+    return None
+
+
 # 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"
@@ -259,11 +279,13 @@ def detect_tags_from_query(
         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)
@@ -277,6 +299,23 @@ def build_or_load_vectorstore(docs: List[Document], index_path: str, is_personal
     vs.save_local(index_path)
     return vs
 
+    
+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 texts_from_jsonl(path: str) -> List[Document]:
     # This function remains unchanged from agent_work.py
     out: List[Document] = []
@@ -293,6 +332,36 @@ def texts_from_jsonl(path: str) -> List[Document]:
     except Exception: return []
     return out
 
+
+def rerank_documents(query: str, documents: list[tuple[Document, float]]) -> list[tuple[tuple[Document, float], float]]:
+    """
+    Re-ranks a list of retrieved documents against a query using a CrossEncoder model.
+    Returns the original document tuples along with their new re-ranker score.
+    """
+    if not documents or not query:
+        return []
+
+    model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
+    
+    doc_contents = [doc.page_content for doc, score in documents]
+    query_doc_pairs = [[query, doc_content] for doc_content in doc_contents]
+    
+    scores = model.predict(query_doc_pairs)
+    
+    reranked_results = list(zip(documents, scores))
+    reranked_results.sort(key=lambda x: x[1], reverse=True)
+
+    print(f"\n[DEBUG] Re-ranked Top 3 Sources:")
+    for doc_tuple, score in reranked_results[:3]:
+        doc, _ = doc_tuple
+        # --- MODIFICATION: Add score to debug log ---
+        print(f"  - New Rank | Source: {doc.metadata.get('source')} | Score: {score:.4f}")
+        
+    # --- MODIFICATION: Return the results with scores ---
+    return reranked_results 
+
+    
+
 # 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):
@@ -305,21 +374,6 @@ def dedup_docs(scored_docs):
             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
@@ -398,6 +452,7 @@ def route_query_type(query: str, severity: str = "Normal / Unspecified"):
         print(f"Query classified as: {sum_hit} (summarization pre-router)")
         return sum_hit
 
+    # Priority 4: Check for music requests.
     # NEW Add Music Support before care_hit = _pre_router(query)
     # the general "caregiving" keyword checker (_pre_router) is called before 
     # the specific "play music" checker (_pre_router_music).
@@ -405,8 +460,8 @@ def route_query_type(query: str, severity: str = "Normal / Unspecified"):
     if music_hit:
         print(f"Query classified as: {music_hit} (music re-router)")
         return music_hit
-    
-    # Priority 3: Check for general caregiving keywords.
+
+    # Priority 5: Check for general caregiving keywords.
     care_hit = _pre_router(query)
     if care_hit:
         print(f"Query classified as: {care_hit} (caregiving pre-router)")
@@ -421,7 +476,6 @@ def route_query_type(query: str, severity: str = "Normal / Unspecified"):
 
 # 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.
@@ -584,15 +638,13 @@ def make_rag_chain(vs_general: FAISS, vs_personal: FAISS, *, for_evaluation: boo
         # --- END OF REVISED MUSIC LOGIC ---
         # END --- MUSIC PLAYBACK LOGIC ---
 
-        
         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)
@@ -601,232 +653,109 @@ def make_rag_chain(vs_general: FAISS, vs_personal: FAISS, *, for_evaluation: boo
             answer = _clean_surface_text(raw_answer)
             sources = ["General Knowledge"] if "general_knowledge" in query_type else []
             return {"answer": answer, "sources": sources, "source_documents": []}
+        # --- END: Non-RAG Route Handling ---
 
-        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]
+        all_retrieved_docs = []
+        is_personal_route = "factual" in query_type or "summarization" in query_type or "multi_hop" in query_type
 
-        # 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')}")
+
+        # --- NEW: DEDICATED LOGIC PATHS FOR RETRIEVAL ---
+        if is_personal_route:
+            # For personal queries, semantic search is unreliable. We retrieve ALL personal documents.
+            print("[DEBUG] Personal Memory Route Activated. Retrieving all personal documents.")
+            if vs_personal and vs_personal._collection.count() > 0:
+                personal_docs_data = vs_personal.get(include=["metadatas", "documents"])
+                all_retrieved_docs.extend([
+                    Document(page_content=doc_content, metadata=metadata)
+                    for doc_content, metadata in zip(personal_docs_data['documents'], personal_docs_data['metadatas'])
+                ])
         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]
+            # For caregiving scenarios, use our powerful Multi-Stage Retrieval algorithm.
+            print("[DEBUG] Using Multi-Stage Retrieval for caregiving scenario...")
+            print("[DEBUG] Expanding query...")
+            search_queries = [query]
+            try:
+                expansion_prompt = QUERY_EXPANSION_PROMPT.format(question=query)
+                expansion_messages = [{"role": "user", "content": expansion_prompt}]
+                raw_expansion = call_llm(expansion_messages, temperature=0.0)
+                expanded = json.loads(raw_expansion)
+                if isinstance(expanded, list):
+                    search_queries.extend(expanded)
+            except Exception as e:
+                print(f"[DEBUG] Query expansion failed: {e}")
+
+            scenario_tags = kwargs.get("scenario_tag")
+            if isinstance(scenario_tags, str): scenario_tags = [scenario_tags]
+            primary_behavior = (scenario_tags or [None])[0]
             
-            if not filtered_docs:
-                return [sorted(valid_results, key=lambda x: x[1])[0][0]], best_score
+            candidate_docs = []
+            if primary_behavior and primary_behavior != "None":
+                print(f"  - Stage 1a: High-precision search for behavior: '{primary_behavior}'")
+                for q in search_queries:
+                    candidate_docs.extend(vs_general.similarity_search_with_score(q, k=10, filter={"behaviors": primary_behavior}))
+
+            print("  - Stage 1b: High-recall semantic search (k=20)")
+            for q in search_queries:
+                candidate_docs.extend(vs_general.similarity_search_with_score(q, k=20))
+
+            all_candidate_docs = dedup_docs(candidate_docs)
+            print(f"[DEBUG] Total unique candidates for re-ranking: {len(all_candidate_docs)}")
+            reranked_docs_with_scores = rerank_documents(query, all_candidate_docs) if all_candidate_docs else []
+
+         
+            # --- BEST method code:  Recall 90% and Precision 73%
+            final_docs_with_scores = []
+            if reranked_docs_with_scores:
+                RELATIVE_SCORE_MARGIN = 3.0
+                top_doc_tuple, top_score = reranked_docs_with_scores[0]
+                final_docs_with_scores.append(top_doc_tuple)
+                for doc_tuple, score in reranked_docs_with_scores[1:]:
+                    if score > (top_score - RELATIVE_SCORE_MARGIN):
+                        final_docs_with_scores.append(doc_tuple)
+                    else: break
             
-            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.")
+            limit = 5 if disease_stage in ["Moderate Stage", "Advanced Stage"] else 3
+            final_docs_with_scores = final_docs_with_scores[:limit]
+            all_retrieved_docs = [doc for doc, score in final_docs_with_scores]
+            # BEFORE FINAL PROCESSING (Applies to all RAG routes)
+
+        # --- FINAL PROCESSING (Applies to all RAG routes) ---
+        print("\n--- DEBUG: Final Selected Docs ---")
+        for doc in all_retrieved_docs:
+            print(f"    - Source: {doc.metadata.get('source', 'N/A')}")
         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 = []
+        personal_sources_set = {'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_set], "(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_set], "(No general guidance found.)")
+            
         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:
-                print("[DEBUG] Factual/Sum/Multi: Prioritizing personal docs.")
-                all_retrieved_docs = filtered_personal_docs
-            else:
-                print("[DEBUG] Factual/Sum/Multi: Prioritizing general docs.")
-                all_retrieved_docs = filtered_general_docs
-            # --- END OF MODIFICATION ---
+            template = ANSWER_TEMPLATE_SUMMARIZE if "summarization" in query_type else ANSWER_TEMPLATE_FACTUAL_MULTI if "multi_hop" in query_type 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, context=personal_context, role=role)
+            print("[DEBUG] Personal Route Factual / Sum / Multi PROMPT") 
         else: # caregiving_scenario
-            if disease_stage in ["Moderate Stage", "Advanced Stage"]:
-                # --- STAGE-AWARE LOGIC FOR CAREGIVING SCENARIOS ---
-                if filtered_general_docs:
-                    print("[DEBUG] Moderate / Advanced: Prioritizing general documents.")
-                    all_retrieved_docs = filtered_general_docs
-                elif filtered_personal_docs:
-                    print("[DEBUG] Moderate / Advanced: Falling back to personal documents.")
-                    all_retrieved_docs = filtered_personal_docs
-                else:
-                    print("[DEBUG] Moderate / Advanced: No relevant documents found.")
-                    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:
-                    print("[DEBUG] Caregiving Case: Blending personal and general docs (scores are close).")
-                    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:
-                    print("[DEBUG] Caregiving Case: Prioritizing personal docs (better score).")
-                    all_retrieved_docs = filtered_personal_docs
-                else:
-                    print("[DEBUG] Caregiving Case: Prioritizing general docs (better score).")
-                    all_retrieved_docs = filtered_general_docs
+            if disease_stage == "Advanced Stage": template = ANSWER_TEMPLATE_ADQ_ADVANCED
+            elif disease_stage == "Moderate Stage": template = ANSWER_TEMPLATE_ADQ_MODERATE
+            else: template = ANSWER_TEMPLATE_ADQ
+            emotions_context = render_emotion_guidelines(kwargs.get("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"))
+            print("[DEBUG] Caregiving Scenario PROMPT") 
+        # end
         
-        # --- 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
-
-            print(f"[DEBUG] Personal Context: {personal_context}")
-            print(f"[DEBUG] General Context: {general_context}")
-            
-            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)
-            print("[DEBUG] Factual / Sum / Multi LLM Answer: ", {answer})
+        messages.append({"role": "user", "content": user_prompt})
         
-        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.)")
-
-            print(f"[DEBUG] Personal Context: {personal_context}")
-            print(f"[DEBUG] General Context: {general_context}")
-            
-            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)
-            print("[DEBUG] Caregiving Case LLM Answer: ", {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")))
+        raw_answer = call_llm(messages, temperature=0.0 if for_evaluation else temperature)
+        answer = _clean_surface_text(raw_answer)
+        print("[DEBUG] LLM Answer", {answer}) 
+
+        if (kwargs.get("scenario_tag") or "").lower() in ["exit_seeking", "wandering"]:
+            answer += f"\n\n---\n{RISK_FOOTER}"
         
+        sources = _source_ids_for_eval(all_retrieved_docs) if for_evaluation else 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
-              
+    return _answer_fn              
     # END of make_rag_chain
 
 def answer_query(chain, question: str, **kwargs) -> Dict[str, Any]: