Update tools.py

tools.py

@@ -8,7 +8,6 @@ from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
 from langchain_community.tools.tavily_search import TavilySearchResults
 from dotenv import load_dotenv
 import os
-import requests
 
 load_dotenv()
 
@@ -18,15 +17,13 @@ load_dotenv()
 VECTORSTORE_DIR = "data/vectorstore"
 os.makedirs(VECTORSTORE_DIR, exist_ok=True)
 
-# ==============================
-# GLOBAL RETRIEVER
-# ==============================
 retriever = None
 
 
 def load_retriever():
     """Load FAISS retriever from disk if available."""
     global retriever
+
     try:
         embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
         index_path = os.path.join(VECTORSTORE_DIR, "index.faiss")
@@ -35,16 +32,22 @@ def load_retriever():
             vectorstore = FAISS.load_local(
                 VECTORSTORE_DIR,
                 embeddings,
-                allow_dangerous_deserialization=True
+                allow_dangerous_deserialization=True,
             )
-            retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-            print("✅ Vectorstore loaded from disk")
+            retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
+            print("✅ Retriever loaded successfully")
+        else:
+            print("⚠️ No vectorstore found yet")
+
     except Exception as e:
-        print("❌ Failed to load vectorstore:", e)
+        print("❌ Retriever load error:", e)
+
+
+# Load on startup
+load_retriever()
 
 
 def build_vectorstore(path: str):
-    """Build FAISS vector store from uploaded PDF."""
     loader = PyPDFLoader(path)
     docs = loader.load()
 
@@ -53,20 +56,18 @@ def build_vectorstore(path: str):
         chunk_overlap=100
     )
 
-    split_docs = splitter.split_documents(docs)
+    chunks = splitter.split_documents(docs)
     embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
 
-    vectorstore = FAISS.from_documents(split_docs, embeddings)
+    vectorstore = FAISS.from_documents(chunks, embeddings)
     vectorstore.save_local(VECTORSTORE_DIR)
 
     return vectorstore
 
 
-def update_retriever(pdf_path: str):
-    """Update retriever after document upload."""
+def update_retriever(path: str):
     global retriever
-    vectorstore = build_vectorstore(pdf_path)
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+    retriever = build_vectorstore(path).as_retriever(search_kwargs={"k": 4})
 
 
 # ==============================
@@ -76,48 +77,41 @@ def create_rag_tool():
 
     @tool
     def rag_search(query: str) -> str:
-        """
-        Retrieve relevant information from uploaded documents.
-        Uses FAISS-based semantic search.
-        """
-        global retriever
+        """Retrieve relevant context from uploaded documents."""
 
-        if retriever is None:
-            load_retriever()
+        global retriever
 
         if retriever is None:
-            return "No document has been uploaded yet."
+            return "No document uploaded yet."
 
         docs = retriever.invoke(query)
 
         if not docs:
-            return "No relevant information found in the uploaded document."
+            return "No relevant information found in the document."
 
         return "\n\n".join(d.page_content for d in docs)
 
     return rag_search
 
 
-# ==============================
-# EXTERNAL TOOLS
-# ==============================
+# -----------------------------
+# External tools (safe)
+# -----------------------------
 
 @tool
-def arxiv_search(query: str) -> dict:
-    """Search academic papers from arXiv."""
+def wikipedia_search(query: str) -> dict:
+    """Search Wikipedia."""
     try:
-        arxiv = ArxivQueryRun(api_wrapper=ArxivAPIWrapper())
-        return {"results": arxiv.run(query)}
+        return {"results": WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()).run(query)}
     except Exception as e:
         return {"error": str(e)}
 
 
 @tool
-def wikipedia_search(query: str) -> dict:
-    """Search Wikipedia for relevant information."""
+def arxiv_search(query: str) -> dict:
+    """Search academic papers on arXiv."""
     try:
-        wiki = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
-        return {"results": wiki.run(query)}
+        return {"results": ArxivQueryRun(api_wrapper=ArxivAPIWrapper()).run(query)}
     except Exception as e:
         return {"error": str(e)}
 
@@ -126,7 +120,6 @@ def wikipedia_search(query: str) -> dict:
 def tavily_search(query: str) -> dict:
     """Search the web using Tavily."""
     try:
-        search = TavilySearchResults(max_results=5)
-        return {"results": search.run(query)}
+        return {"results": TavilySearchResults(max_results=5).run(query)}
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": str(e)}