update rag_search and build_retriever

agent.py

@@ -17,7 +17,7 @@ from langchain_community.document_loaders import AssemblyAIAudioTranscriptLoader
 from langchain.chat_models import init_chat_model
 from langchain.agents import initialize_agent, AgentType
 from langchain_community.retrievers import BM25Retriever
-from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
+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
@@ -260,8 +260,8 @@ tools: List[StructuredTool] = [
 
 class AgentState(TypedDict):
     # The document provided
-    input_file: Optional[str]  # Contains file path (PDF/PNG)
-    messages: Annotated[list[AnyMessage], add_messages]
+    input_file: Optional[List[str]]  # Contains file path (PDF/PNG)
+    messages: Annotated[List[BaseMessage], add_messages]
 
 # === Agent Class ===
 class MyAgent:
@@ -333,7 +333,7 @@ class MyAgent:
         @tool(name="rag_search")
         def rag_search(query: str) -> str:
             """Retrieve top-3 relevant document chunks via BM25."""
-            res = self.retriever.invoke(query)
+            res = self.retriever.get_relevant_documents(query)
             if res:
                 return "\n\n".join([doc.page_content for doc in res[:3]])
             return ""
@@ -347,63 +347,63 @@ class MyAgent:
         question: str,
         file_paths: Optional[List[str]] = None
     ) -> str:
-        # Prepare state graph
-        state: Dict[str, Any] = {"messages": [], "input_file": None}
+        try:
+            # Prepare state graph
+            state: Dict[str, Any] = {"messages": [], "input_file": None}
 
-        # Use structured tool attributes
-        tool_desc = "\n".join(f"{t.name}: {t.description}" for t in self.tools)
-        
-        # Enhanced system prompt with RAG guidance
-        rag_prompt = """
-        If the question seems to be about any loaded documents, ALWAYS:
-        1. Use the rag_search tool first to find relevant information
-        2. Base your answer on the retrieved content
-        3. If no relevant content is found, say so
-        """
-        
-        sys_msg = SystemMessage(content=f"{SYSTEM_PROMPT}\n\n{rag_prompt if file_paths else ''}\n\nTools:\n{tool_desc}")
-        state["messages"].append(sys_msg)
-
-        # Optionally load RAG docs
-        if file_paths:
-            self.add_files(file_paths)
-            self.build_retriever()
-
-        # Add user question
-        state["messages"].append(HumanMessage(content=question))
-        if file_paths:
-            state["input_file"] = file_paths
-
-        # Build graph
-        builder = StateGraph(dict)
-        builder.add_node("assistant", self._assistant_node)
-        builder.add_node("tools", ToolNode(self.tools))
-        builder.add_edge(START, "assistant")
-        
-        # Updated tool detection logic
-        builder.add_conditional_edges(
-            "assistant",
-            lambda s: any(t.name in s["messages"][-1].content for t in self.tools),
-            "tools"
-        )
-        builder.add_edge("tools", "assistant")
-        graph = builder.compile()
-
-        # Use invoke() instead of run()
-        out = graph.invoke(state)
-        last_message = out["messages"][-1].content
-        
-        # Extract only the FINAL ANSWER part
-        if "FINAL ANSWER:" in last_message:
-            return last_message.split("FINAL ANSWER:")[-1].strip()
-        return last_message.strip()
+            # Use structured tool attributes
+            tool_desc = "\n".join(f"{t.name}: {t.description}" for t in self.tools)
+            
+            # Enhanced system prompt with RAG guidance
+            rag_prompt = """
+            If the question seems to be about any loaded documents, ALWAYS:
+            1. Use the rag_search tool first to find relevant information
+            2. Base your answer on the retrieved content
+            3. If no relevant content is found, say so
+            """
+            
+            sys_msg = SystemMessage(content=f"{SYSTEM_PROMPT}\n\n{rag_prompt if file_paths else ''}\n\nTools:\n{tool_desc}")
+            state["messages"].append(sys_msg)
+
+            # Optionally load RAG docs
+            if file_paths:
+                self.add_files(file_paths)
+                self.build_retriever()
+
+            # Add user question
+            state["messages"].append(HumanMessage(content=question))
+            if file_paths:
+                state["input_file"] = file_paths
+
+            # Build graph
+            builder = StateGraph(dict)
+            builder.add_node("assistant", self._assistant_node)
+            builder.add_node("tools", ToolNode(self.tools))
+            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_edge("tools", "assistant")
+            graph = builder.compile()
+
+            # Use invoke() instead of run()
+            out = graph.invoke(state)
+            last_message = out["messages"][-1].content
+            
+            # Extract only the FINAL ANSWER part
+            if "FINAL ANSWER:" in last_message:
+                return last_message.split("FINAL ANSWER:")[-1].strip()
+            return last_message.strip()
+        except Exception as e:
+            return f"Error processing question: {e}"
     
     def run(self, question: str, file_paths: Optional[List[str]] = None) -> str:
         return self(question, file_paths)
     
     def _assistant_node(self, state: dict) -> dict:
-        # Invoke LLM on current messages
-        resp = self.llm.invoke(state["messages"])
+        # Invoke the chat model with our BaseMessage list
+        resp = self.llm(state["messages"])
         state["messages"].append(resp)
         return state