Update alz_companion/agent.py

alz_companion/agent.py

@@ -233,7 +233,6 @@ def call_llm(messages: List[Dict[str, str]], temperature: float = 0.6, stop: Opt
 # -----------------------------
 # Prompting & RAG Chain
 # -----------------------------
-
 def make_rag_chain(
     vs_general: FAISS,
     vs_personal: FAISS,
@@ -245,68 +244,107 @@ def make_rag_chain(
     caregiver_name: str = "the caregiver",
     tone: str = "warm",
 ):
+    """Returns a callable that performs the complete, intelligent RAG process."""
+
     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, chat_history: List[Dict[str, str]], scenario_tag: Optional[str] = None, emotion_tag: Optional[str] = None, topic_tag: Optional[str] = None, context_tags: Optional[List[str]] = None) -> Dict[str, Any]:
+
         router_messages = [{"role": "user", "content": ROUTER_PROMPT.format(query=query)}]
         query_type = call_llm(router_messages, temperature=0.0).strip().lower()
         print(f"Query classified as: {query_type}")
 
         system_message = SYSTEM_TEMPLATE.format(tone=tone, language=language, patient_name=patient_name or "the patient", caregiver_name=caregiver_name or "the caregiver", guardrails=SAFETY_GUARDRAILS)
-        messages = [{"role": "system", "content": system_message}, *chat_history]
+        messages = [{"role": "system", "content": system_message}]
+        messages.extend(chat_history)
 
         if "general_knowledge_question" in query_type:
             user_prompt = ANSWER_TEMPLATE_GENERAL_KNOWLEDGE.format(question=query, language=language)
             messages.append({"role": "user", "content": user_prompt})
-            return {"answer": call_llm(messages, temperature=temperature), "sources": ["General Knowledge"]}
+            answer = call_llm(messages, temperature=temperature)
+            return {"answer": answer, "sources": ["General Knowledge"]}
 
         elif "factual_question" in query_type:
+            print(f"Performing query expansion for: '{query}'")
             expansion_prompt = QUERY_EXPANSION_PROMPT.format(question=query)
             expansion_response = call_llm([{"role": "user", "content": expansion_prompt}], temperature=0.1)
             try:
-                expanded_queries = json.loads(expansion_response.strip().replace("```json", "").replace("```", ""))
+                clean_response = expansion_response.strip().replace("```json", "").replace("```", "")
+                expanded_queries = json.loads(clean_response)
                 search_queries = [query] + expanded_queries
             except json.JSONDecodeError:
                 search_queries = [query]
-
+            print(f"Searching with queries: {search_queries}")
             all_docs = []
             for q in search_queries:
                 all_docs.extend(vs_personal.similarity_search(q, k=2))
                 all_docs.extend(vs_general.similarity_search(q, k=2))
-            context = _format_docs(all_docs, "(No relevant information found.)")
+            context = _format_docs(all_docs, "(No relevant information found in the memory journal.)")
             user_prompt = ANSWER_TEMPLATE_FACTUAL.format(context=context, question=query, language=language)
             messages.append({"role": "user", "content": user_prompt})
-            return {"answer": call_llm(messages, temperature=temperature), "sources": list(set(d.metadata.get("source", "unknown") for d in all_docs))}
+            answer = call_llm(messages, temperature=temperature)
+            sources = list(set(d.metadata.get("source", "unknown") for d in all_docs))
+            return {"answer": answer, "sources": sources}
 
         elif "general_conversation" in query_type:
             user_prompt = ANSWER_TEMPLATE_GENERAL.format(question=query, language=language)
             messages.append({"role": "user", "content": user_prompt})
-            return {"answer": call_llm(messages, temperature=temperature), "sources": []}
-
-        else: # Default to caregiving logic
-            search_filter = {}
-            if scenario_tag: search_filter["behaviors"] = scenario_tag.lower()
-            if emotion_tag: search_filter["emotion"] = emotion_tag.lower()
-            if topic_tag: search_filter["topic_tags"] = topic_tag.lower()
-            if context_tags: search_filter["context_tags"] = {"in": [tag.lower() for tag in context_tags]}
+            answer = call_llm(messages, temperature=temperature)
+            return {"answer": answer, "sources": []}
 
+        else: # Default to the original caregiving logic
+            # --- Reworked search strategy to handle filters correctly ---
+            
+            # 1. Start with a general, unfiltered search to always get text-based matches.
             personal_docs = vs_personal.similarity_search(query, k=3)
             general_docs = vs_general.similarity_search(query, k=3)
-            if search_filter:
-                personal_docs.extend(vs_personal.similarity_search(query, k=3, filter=search_filter))
-                general_docs.extend(vs_general.similarity_search(query, k=3, filter=search_filter))
-
+            
+            # 2. Build a filter for simple equality checks (FAISS supported).
+            simple_search_filter = {}
+            if scenario_tag and scenario_tag != "None":
+                simple_search_filter["behaviors"] = scenario_tag.lower()
+            if emotion_tag and emotion_tag != "None":
+                simple_search_filter["emotion"] = emotion_tag.lower()
+            if topic_tag and topic_tag != "None":
+                simple_search_filter["topic_tags"] = topic_tag.lower()
+
+            # 3. If simple filters exist, perform a second, more specific search.
+            if simple_search_filter:
+                print(f"Performing additional search with filter: {simple_search_filter}")
+                personal_docs.extend(vs_personal.similarity_search(query, k=2, filter=simple_search_filter))
+                general_docs.extend(vs_general.similarity_search(query, k=2, filter=simple_search_filter))
+            
+            # 4. If context_tags exist (unsupported by 'in'), loop through them and perform separate searches.
+            if context_tags:
+                print(f"Performing looped context tag search for: {context_tags}")
+                for tag in context_tags:
+                    context_filter = {"context_tags": tag.lower()}
+                    personal_docs.extend(vs_personal.similarity_search(query, k=1, filter=context_filter))
+                    general_docs.extend(vs_general.similarity_search(query, k=1, filter=context_filter))
+
+            # 5. Combine and de-duplicate all results.
             all_docs_care = list({doc.page_content: doc for doc in personal_docs + general_docs}.values())
+            
+            # --- End of reworked search strategy ---
+
             personal_context = _format_docs([d for d in all_docs_care if d in personal_docs], "(No relevant personal memories found.)")
             general_context = _format_docs([d for d in all_docs_care if d in general_docs], "(No general guidance found.)")
 
-            first_emotion = next((d.metadata.get("emotion") for d in all_docs_care if d.metadata.get("emotion")), None)
+            first_emotion = None
+            for doc in all_docs_care:
+                if "emotion" in doc.metadata and doc.metadata["emotion"]:
+                    emotion_data = doc.metadata["emotion"]
+                    if isinstance(emotion_data, list): first_emotion = emotion_data[0]
+                    else: first_emotion = emotion_data
+                    if first_emotion: break
+
             emotions_context = render_emotion_guidelines(first_emotion or emotion_tag)
+            is_tagged_scenario = (scenario_tag and scenario_tag != "None") or (emotion_tag and emotion_tag != "None") or (first_emotion is not None)
+            template = ANSWER_TEMPLATE_ADQ if is_tagged_scenario else ANSWER_TEMPLATE_CALM
 
-            template = ANSWER_TEMPLATE_ADQ if any([scenario_tag, emotion_tag, first_emotion]) else ANSWER_TEMPLATE_CALM
             if template == ANSWER_TEMPLATE_ADQ:
                 user_prompt = template.format(general_context=general_context, personal_context=personal_context, question=query, scenario_tag=scenario_tag, emotions_context=emotions_context, role=role, language=language)
             else:
@@ -316,12 +354,16 @@ def make_rag_chain(
             messages.append({"role": "user", "content": user_prompt})
             answer = call_llm(messages, temperature=temperature)
 
-            if scenario_tag and scenario_tag.lower() in ["exit_seeking", "wandering"]:
+            high_risk_scenarios = ["exit_seeking", "wandering", "elopement"]
+            if scenario_tag and scenario_tag.lower() in high_risk_scenarios:
                 answer += f"\n\n---\n{RISK_FOOTER}"
 
-            return {"answer": answer, "sources": list(set(d.metadata.get("source", "unknown") for d in all_docs_care))}
+            sources = list(set(d.metadata.get("source", "unknown") for d in all_docs_care))
+            return {"answer": answer, "sources": sources}
+
     return _answer_fn
 
+
 def answer_query(chain, question: str, **kwargs) -> Dict[str, Any]:
     if not callable(chain): return {"answer": "[Error: RAG chain is not callable]", "sources": []}
     try: