update rag_search in build_retrievier

agent.py

@@ -8,6 +8,8 @@ from typing import List, TypedDict, Annotated, Optional
 import requests
 from urllib.parse import urlparse
 import shutil
+import io
+from typing import Dict, Any
 
 from langchain.tools import tool, StructuredTool
 from langchain_community.document_loaders import (
@@ -21,6 +23,7 @@ from langchain.schema import BaseMessage, SystemMessage, HumanMessage
 from langgraph.graph.message import add_messages
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.documents import Document
 
 from youtube_transcript_api import YouTubeTranscriptApi
 from PIL import Image
@@ -328,19 +331,32 @@ class MyAgent:
         """
         if not self.docs:
             return
-        self.retriever = BM25Retriever.from_documents(self.docs)
-
-        @tool(name="rag_search")
-        def rag_search(query: str) -> str:
-            """Retrieve top-3 relevant document chunks via BM25."""
-            res = self.retriever.get_relevant_documents(query)
-            if res:
-                return "\n\n".join([doc.page_content for doc in res[:3]])
-            return ""
-
-        # Register RAG tool
-        self.tools.append(rag_search)
-
+            
+        # Build retriever
+        try:
+            self.retriever = BM25Retriever.from_documents(self.docs)
+            
+            # Define tool with proper error handling
+            @tool(name="rag_search")
+            def rag_search(query: str) -> str:
+                """Search loaded documents for relevant information."""
+                try:
+                    if not self.retriever:
+                        return "No documents have been loaded for search."
+                        
+                    res = self.retriever.get_relevant_documents(query)
+                    if res:
+                        return "\n\n".join(f"Document {i+1}:\n{doc.page_content}" 
+                                         for i, doc in enumerate(res[:3]))
+                    return "No relevant information found in loaded documents."
+                except Exception as e:
+                    return f"Error searching documents: {e}"
+            
+            # Remove existing rag_search if present to prevent duplicates
+            self.tools = [t for t in self.tools if t.name != "rag_search"]
+            self.tools.append(rag_search)
+        except Exception as e:
+            print(f"Error building retriever: {e}")
 
     def __call__(
         self,
@@ -382,8 +398,11 @@ class MyAgent:
             builder.add_edge(START, "assistant")
             
             # Always allow the assistant to hand off to the tools node
-            builder.add_edge("assistant", "tools")
-            # And then return from tools back to the assistant
+            builder.add_conditional_edges(
+                "assistant",
+                lambda s: any(t.name in s["messages"][-1].content for t in self.tools),
+                {True: "tools", False: "assistant"}
+            )
             builder.add_edge("tools", "assistant")
             graph = builder.compile()
 
@@ -402,10 +421,22 @@ class MyAgent:
         return self(question, file_paths)
     
     def _assistant_node(self, state: dict) -> dict:
-        # Invoke the chat model with our BaseMessage list
-        resp = self.llm(state["messages"])
-        state["messages"].append(resp)
-        return state
+        """Process messages with the LLM."""
+        try:
+            # Check if messages exist
+            if not state["messages"]:
+                # Add a system message if empty
+                state["messages"].append(SystemMessage(content=SYSTEM_PROMPT))
+                
+            # Invoke the chat model with our BaseMessage list
+            resp = self.llm(state["messages"])
+            state["messages"].append(resp)
+            return state
+        except Exception as e:
+            # Handle errors by adding an error message
+            error_msg = f"Error calling LLM: {str(e)}"
+            print(error_msg)
+            return state