Update alz_companion/agent.py

alz_companion/agent.py

@@ -201,10 +201,14 @@ def texts_from_jsonl(path: str) -> List[Document]:
                 obj = json.loads(line)
                 txt = obj.get("text") or ""
                 if not isinstance(txt, str) or not txt.strip(): continue
+
+                # fix bugs by adding tags for topic and context
                 md = {"source": os.path.basename(path), "chunk": i}
-                for k in ("behaviors", "emotion"):
-                    if k in obj: md[k] = obj[k]
+                for k in ("behaviors", "emotion", "topic_tags", "context_tags"):
+                    if k in obj and obj[k]: # Ensure the key exists and is not empty
+                        md[k] = obj[k]
                 out.append(Document(page_content=txt, metadata=md))
+                
     except Exception:
         return []
     return out
@@ -340,21 +344,40 @@ def make_rag_chain(
                 search_filter["behaviors"] = scenario_tag.lower()
             if emotion_tag and emotion_tag != "None":
                 search_filter["emotion"] = emotion_tag.lower()
-
+            # fix bug by adding topic tag and context tag
+            if topic_tag and topic_tag != "None": # <-- ADD THESE TWO LINES
+                search_filter["topic_tags"] = topic_tag.lower()
+            if context_tags: # <-- ADD THESE TWO LINES
+                search_filter["context_tags"] = {"in": [tag.lower() for tag in context_tags]}
+
+            # --- Robust Search Strategy ---
+            # 1. Start with a general, unfiltered search to always get text-based matches.
+            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)
+            
+            # 2. If filters exist, perform a second, more specific search and add the results.
             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:
-                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)
+                print(f"Performing additional search with filter: {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))
 
-            personal_context = _format_docs(personal_docs, "(No relevant personal memories found.)")
-            general_context = _format_docs(general_docs, "(No general guidance found.)")
+            # 3. Combine and de-duplicate the results to get the best of both searches.
+            all_personal_docs = list({doc.page_content: doc for doc in personal_docs}.values())
+            all_general_docs = list({doc.page_content: doc for doc in general_docs}.values())
+
+            # 4. Define the context variables based on the new, combined results.
+            personal_context = _format_docs(all_personal_docs, "(No relevant personal memories found.)")
+            general_context = _format_docs(all_general_docs, "(No general guidance found.)")
 
             first_emotion = None
-            all_docs_care = personal_docs + general_docs
+            all_docs_care = all_personal_docs + all_general_docs
+
+            # -- end of Robust Search Strategy
+
+            
             for doc in all_docs_care:
                 if "emotion" in doc.metadata and doc.metadata["emotion"]:
                     emotion_data = doc.metadata["emotion"]
@@ -384,15 +407,21 @@ def make_rag_chain(
     return _answer_fn
 
 
+# Fix bug by adding topic tag  ... how about context tag??
 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, scenario_tag, emotion_tag = kwargs.get("chat_history", []), kwargs.get("scenario_tag"), kwargs.get("emotion_tag")
+    chat_history = kwargs.get("chat_history", [])
+    scenario_tag = kwargs.get("scenario_tag")
+    emotion_tag = kwargs.get("emotion_tag")
+    topic_tag = kwargs.get("topic_tag") # <-- ADD THIS LINE
+    context_tags = kwargs.get("context_tags") # <-- ADD THIS LINE
     try:
-        return chain(question, chat_history=chat_history, scenario_tag=scenario_tag, emotion_tag=emotion_tag)
+        return chain(question, chat_history=chat_history, scenario_tag=scenario_tag, emotion_tag=emotion_tag, topic_tag=topic_tag) # <-- ADD topic_tag
     except Exception as e:
         print(f"ERROR in answer_query: {e}")
         return {"answer": f"[Error executing chain: {e}]", "sources": []}
 
+
 # -----------------------------
 # TTS & Transcription
 # -----------------------------