test RAG in agent

agent.py

@@ -9,8 +9,18 @@ from langchain.tools import tool
 from langchain_community.document_loaders import (
     CSVLoader,
     YoutubeLoader,
+    PyPDFLoader
+)
+from langchain_community.document_loaders.blob_loaders.youtube_audio import (
+    YoutubeAudioLoader,
+)
+from langchain_community.document_loaders.generic import GenericLoader
+from langchain_community.document_loaders.parsers.audio import (
+    OpenAIWhisperParser,
+    OpenAIWhisperParserLocal,
 )
 
+#from langchain_community.document_loaders import AudioLoader, WhisperLoader
 from langchain.chat_models import init_chat_model
 from langchain.agents import initialize_agent, AgentType
 from langchain_community.retrievers import BM25Retriever
@@ -165,7 +175,7 @@ class AgentState(TypedDict):
     # The document provided
     input_file: Optional[str]  # Contains file path (PDF/PNG)
     messages: Annotated[list[AnyMessage], add_messages]
-    
+
 # === Agent Class ===
 class MyAgent:
     def __init__(
@@ -175,51 +185,103 @@ class MyAgent:
     ):
         # Initialize LLM
         self.llm = init_chat_model(model_name, temperature=temperature)
+        # Base tools
+        self.tools = tools
+        # RAG components
+        self.docs: List[Any] = []
+        self.retriever: Optional[BM25Retriever] = None
 
-        # Base tools: use provided tools or default list
-        self.tools = tools 
-
-        # Human-readable tool descriptions
-        self.textual_tool_desc = "\n".join(t.__doc__.strip() for t in self.tools)
-
-        # Define assistant node
-        def assistant_node(state: AgentState) -> dict:
-            sys_msg = SystemMessage(
-                content="\n".join([
-                    SYSTEM_PROMPT,
-                    "\nTools available:\n" + self.textual_tool_desc
-                ])
-            )
-            msgs = [sys_msg] + state["messages"]
-            response = self.llm(msgs)
-            return {"messages": state["messages"] + [response], "input_file": state.get("input_file")}
-
-        # Condition to invoke tools: check if last LLM message mentions a tool invocation
-        def needs_tool(state: AgentState) -> bool:
-            last = state["messages"][-1].content.lower()
-            return any(f"{t.__name__.lower()}(" in last for t in self.tools)
-
-        # Build the state graph
-        builder = StateGraph(AgentState)
-        builder.add_node("assistant", assistant_node)
-        builder.add_node("tools", ToolNode(self.tools))
-        builder.add_edge(START, "assistant")
-        builder.add_conditional_edges("assistant", needs_tool)
-        builder.add_edge("tools", "assistant")
+    def add_files(self, file_paths: List[str]):
+        """
+        Load and index documents for RAG based on file extensions or URLs.
+        Supports: PDF, CSV, audio (mp3/wav), and YouTube URLs.
+        """
+        for path in file_paths:
+            ext = Path(path).suffix.lower()
+            if ext == ".csv":
+                loader = CSVLoader(path)
+                self.docs.extend(loader.load())
+            elif ext == ".pdf":
+                loader = PyPDFLoader(path)
+                self.docs.extend(loader.load())
+            elif ext in [".mp3", ".wav"]:
+                audio_docs = AudioLoader(path).load()
+                self.docs.extend(WhisperLoader().load(audio_docs))
+            elif "youtube" in path:
+                loader = YoutubeLoader.from_youtube_url(path)
+                self.docs.extend(loader.load())
+            else:
+                continue
+
+    def build_retriever(self):
+        """
+        Create BM25Retriever over the loaded documents and register rag_search tool.
+        """
+        if not self.docs:
+            return
+        self.retriever = BM25Retriever.from_documents(self.docs)
+
+        @tool
+        def rag_search(query: str) -> str:
+            """
+            Retrieve top-3 relevant document chunks via BM25.
+            """
+            res = self.retriever.invoke(query)
+            if res:
+                return "\n\n".join([doc.page_content for doc in res[:3]])
+            return ""
+
+        # Register RAG tool
+        self.tools.append(rag_search)
 
-        self.react_graph = builder.compile()
 
     def __call__(
         self,
-        user_input: str,
-        input_file: Optional[str] = None,
+        question: str,
+        file_paths: Optional[List[str]] = None
     ) -> str:
-        state = AgentState()
-        state["messages"] = [HumanMessage(content=user_input)]
-        state["input_file"] = input_file
-        out = self.react_graph(state)
-        # Return only the final LLM message content
-        return out["messages"][-1].content.strip()
+        # Prepare state graph
+        state: Dict[str, Any] = {"messages": [], "input_file": None}
+
+        # Add system message
+        tool_desc = "\n".join(f"{tool_func.__name__}: {tool_func.__doc__.strip()}" \
+                               for tool_func in self.tools)
+        sys_msg = SystemMessage(content=f"{SYSTEM_PROMPT}\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")
+        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()
+
+        # Run graph until completion
+        out = graph.run(state)
+        return out["messages"][-1].content
+
+    def _assistant_node(self, state: dict) -> dict:
+        # Invoke LLM on current messages
+        resp = self.llm.invoke(state["messages"])
+        state["messages"].append(resp)
+        return state
+
 
 # CLI entrypoint
 if __name__ == "__main__":