Update agent.py

agent.py

@@ -286,17 +286,27 @@ def build_graph(provider: str = "google"):
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
     
     def retriever(state: MessagesState):
-        """Retriever node"""
-        similar_question = vector_store.similarity_search(state["messages"][0].content)
+        user_query = state["messages"][0].content
+        similar_docs = vector_store.similarity_search(user_query)
     
-        if not similar_question:
-            example_msg = HumanMessage(content="No similar question found.")
+        if not similar_docs:
+            print("No similar docs found in FAISS. Using wiki_search.")
+            wiki_result = wiki_search.invoke(user_query)
+            return {
+                "messages": [
+                    sys_msg,
+                    state["messages"][0],
+                    HumanMessage(content=f"Using Wikipedia search:\n\n{wiki_result['wiki_results']}")
+                ]
+            }
         else:
-            example_msg = HumanMessage(
-                content=f"Here I provide a similar question and answer for reference:\n\n{similar_question[0].page_content}",
-            )
-    
-        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
+            return {
+                "messages": [
+                    sys_msg,
+                    state["messages"][0],
+                    HumanMessage(content=f"Reference question:\n\n{similar_docs[0].page_content}")
+                ]
+            }