Update alz_companion/agent.py

alz_companion/agent.py

@@ -30,7 +30,7 @@ from .prompts import (
 
 
 # -----------------------------
-# Multimodal Processing Functions (NEW)
+# Multimodal Processing Functions
 # -----------------------------
 
 def _openai_client() -> Optional[OpenAI]:
@@ -68,7 +68,6 @@ def describe_image(image_path: str) -> str:
 # -----------------------------
 # NLU Classification Function
 # -----------------------------
-# (Unchanged from before)
 def detect_tags_from_query(query: str, behavior_options: list, emotion_options: list) -> Dict[str, Optional[str]]:
     """Uses an LLM call to classify the user's query into a behavior and emotion tag."""
     behavior_str = ", ".join(f'"{opt}"' for opt in behavior_options if opt != "None")
@@ -89,7 +88,6 @@ def detect_tags_from_query(query: str, behavior_options: list, emotion_options:
 # -----------------------------
 # Embeddings & VectorStore
 # -----------------------------
-# (Unchanged from before)
 
 def _default_embeddings():
     """Lightweight, widely available model."""
@@ -104,7 +102,6 @@ def build_or_load_vectorstore(docs: List[Document], index_path: str, is_personal
         except Exception:
             pass
     
-    # If it's a new personal vector store with no docs, create a placeholder
     if is_personal and not docs:
         docs = [Document(page_content="(This is the start of the personal memory journal.)", metadata={"source": "placeholder"})]
         
@@ -148,7 +145,6 @@ def bootstrap_vectorstore(sample_paths: List[str] | None = None, index_path: str
 # -----------------------------
 # LLM Call
 # -----------------------------
-# (Unchanged from before)
 def call_llm(messages: List[Dict[str, str]], temperature: float = 0.6) -> str:
     """Call OpenAI Chat Completions if available; else return a fallback."""
     client = _openai_client()
@@ -162,7 +158,7 @@ def call_llm(messages: List[Dict[str, str]], temperature: float = 0.6) -> str:
         return f"[LLM API Error: {e}]"
 
 # -----------------------------
-# Prompting & RAG Chain (HEAVILY MODIFIED)
+# Prompting & RAG Chain
 # -----------------------------
 
 def _format_sources(docs: List[Document]) -> List[str]:
@@ -181,70 +177,56 @@ def make_rag_chain(
 ):
     """Returns a callable that performs the complete, two-tiered RAG process."""
     
-    retriever_general = vs_general.as_retriever(search_kwargs={"k": 3})
-    retriever_personal = vs_personal.as_retriever(search_kwargs={"k": 3})
-
     def _format_docs(docs: List[Document], default_msg: str) -> str:
         if not docs: return default_msg
         return "\n".join([f"- {d.page_content.strip()}" for d in docs])
 
     def _answer_fn(query: str, chat_history: List[Dict[str, str]], scenario_tag: Optional[str] = None, emotion_tag: Optional[str] = None) -> Dict[str, Any]:
         
-        # --- Step 1: Search the Personal Knowledge Base ---
-        personal_docs = retriever_personal.invoke(query)
-        personal_context = _format_docs(personal_docs, "(No relevant personal memories found.)")
-        
-        # --- Step 2: Search the General Knowledge Base with filters ---
+        # Build a dynamic filter that will be used for BOTH knowledge bases
         search_filter = {}
         if scenario_tag and scenario_tag != "None":
             search_filter["behaviors"] = scenario_tag.lower()
         if emotion_tag and emotion_tag != "None":
             search_filter["emotion"] = emotion_tag.lower()
 
+        # Use the filter on both searches if available
         if search_filter:
+            personal_docs = vs_personal.similarity_search(query, k=3, filter=search_filter)
             general_docs = vs_general.similarity_search(query, k=3, filter=search_filter)
         else:
+            # If no filters, perform standard semantic search on both
+            retriever_personal = vs_personal.as_retriever(search_kwargs={"k": 3})
+            retriever_general = vs_general.as_retriever(search_kwargs={"k": 3})
+            personal_docs = retriever_personal.invoke(query)
             general_docs = retriever_general.invoke(query)
+        
+        personal_context = _format_docs(personal_docs, "(No relevant personal memories found.)")
         general_context = _format_docs(general_docs, "(No general guidance found.)")
         
-        # --- Step 3: Determine Emotion for Response Guidelines ---
+        # Determine emotion for the response guidelines
         first_emotion = None
-        # Prioritize emotion from personal memories, then general context
         all_docs = personal_docs + general_docs
         for doc in all_docs:
             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 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)
 
-        # --- Step 4: Assemble and Call the LLM ---
-        is_tagged_scenario = (scenario_tag and scenario_tag != "None") or (emotion_tag and emotion_tag != "None")
+        # Assemble and Call the LLM
+        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
         
-        # Note the new placeholders: general_context and personal_context
         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: # Calm template only uses a single combined context
+            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:
             combined_context = f"General Guidance:\n{general_context}\n\nPersonal Memories:\n{personal_context}"
             user_prompt = template.format(context=combined_context, question=query, language=language)
 
-        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,
-        )
+        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}]
         messages.extend(chat_history)
@@ -261,13 +243,8 @@ def make_rag_chain(
     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": []}
-    
-    chat_history = kwargs.get("chat_history", [])
-    scenario_tag = kwargs.get("scenario_tag")
-    emotion_tag = kwargs.get("emotion_tag")
-
+    if not callable(chain): return {"answer": "[Error: RAG chain is not callable]", "sources": []}
+    chat_history, scenario_tag, emotion_tag = kwargs.get("chat_history", []), kwargs.get("scenario_tag"), kwargs.get("emotion_tag")
     try:
         return chain(question, chat_history=chat_history, scenario_tag=scenario_tag, emotion_tag=emotion_tag)
     except Exception as e:
@@ -277,8 +254,6 @@ def answer_query(chain, question: str, **kwargs) -> Dict[str, Any]:
 # -----------------------------
 # TTS & Transcription
 # -----------------------------
-# (Unchanged)
-
 def synthesize_tts(text: str, lang: str = "en"):
     if not text or gTTS is None: return None
     try:
@@ -294,11 +269,10 @@ def transcribe_audio(filepath: str, lang: str = "en"):
     client = _openai_client()
     if not client:
         return "[Transcription failed: API key not configured]"
-    
     api_args = {"model": "whisper-1"}
     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
+