improve rag_search

agent.py

@@ -21,7 +21,7 @@ from langchain.agents import initialize_agent, AgentType
 from langchain_community.retrievers import BM25Retriever
 from langchain.schema import BaseMessage, SystemMessage, HumanMessage
 from langgraph.graph.message import add_messages
-from langgraph.graph import START, StateGraph
+from langgraph.graph import START, END, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
 from langchain_core.documents import Document
 
@@ -33,6 +33,7 @@ from dotenv import load_dotenv
 from contextlib import redirect_stdout
 from langchain_community.tools import TavilySearchResults
 from tavily import TavilyClient
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 # Load environment variables from .env file
 # in HF Spaces, the .env file is saved in Variables and secrets in settings
@@ -274,12 +275,29 @@ def download_file(url_or_path: str, save_dir: str = "./downloads") -> str:
         return f"Error downloading/copying file: {e}"
 
 
+@tool
+def extract_table(file_path: str, query: str = "") -> str:
+    """Extract relevant rows from a CSV or Excel file based on a query."""
+    import pandas as pd
+    ext = Path(file_path).suffix.lower()
+    if ext in [".csv"]:
+        df = pd.read_csv(file_path)
+    elif ext in [".xlsx", ".xls"]:
+        df = pd.read_excel(file_path)
+    else:
+        return "Unsupported file type."
+    # Simple filter: return all if no query, else filter columns containing query
+    if query:
+        mask = df.apply(lambda row: row.astype(str).str.contains(query, case=False).any(), axis=1)
+        df = df[mask]
+    return df.head(10).to_csv(index=False)
+
 # Update tools list
 tools: List[StructuredTool] = [
-    calculate, web_search, wikipedia_search, image_recognition,
+    calculate, tavily_search, wikipedia_search, image_recognition,
     read_pdf, read_csv, read_spreadsheet, transcribe_audio,
     youtube_transcript_tool, youtube_transcript_api, read_jsonl,
-    python_interpreter, download_file, tavily_search  # Add tavily_search here
+    python_interpreter, download_file, extract_table   # Add tavily_search here
 ]
 
 class AgentState(TypedDict):
@@ -352,6 +370,15 @@ class MyAgent:
             except Exception as e:
                 print(f"Error loading {path}: {e}")
                 continue
+        # After loading each doc:
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
+        for doc in loaded_docs:
+            chunks = text_splitter.split_text(doc.page_content)
+            for i, chunk in enumerate(chunks):
+                self.docs.append(Document(
+                    page_content=chunk,
+                    metadata={**doc.metadata, "chunk": i, "source": path}
+                ))
 
     def build_retriever(self):
         """
@@ -368,18 +395,13 @@ class MyAgent:
             @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}"
-            
+                if not self.retriever:
+                    return "No documents loaded."
+                docs = self.retriever.get_relevant_documents(query)
+                if not docs:
+                    return "No relevant information found."
+                return "\n\n".join(f"{doc.metadata.get('source', '')}: {doc.page_content[:500]}" for doc in docs[:3])
+
             # 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)
@@ -426,10 +448,17 @@ class MyAgent:
             builder.add_edge(START, "assistant")
             
             # Fix conditional edges with better check
+            def _should_use_tools(state):
+                # If there are loaded docs, always use rag_search first
+                if state.get("input_file"):
+                    return "tools"
+                # Otherwise, let the assistant try to answer
+                return "assistant"
+
             builder.add_conditional_edges(
                 "assistant",
-                tools_condition,  # Use built-in tools_condition
-                "tools"
+                _should_use_tools,
+                {"tools": "tools", "assistant": END}
             )
             builder.add_edge("tools", "assistant")
             
@@ -441,8 +470,10 @@ class MyAgent:
             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()
+            import re
+            match = re.search(r"FINAL ANSWER[:\s]*([^\n]*)", last_message, re.IGNORECASE)
+            if match:
+                return match.group(1).strip()
             return last_message.strip()
         except Exception as e:
             return f"Error processing question: {e}"