Update agent.py

agent.py

@@ -6,7 +6,7 @@ from serpapi import GoogleSearch
 import pandas as pd
 import faiss
 from langgraph import Graph, LLM, tool #or other graph library
-
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 
 # ─── 1) Load & embed all documents at startup ───
 # 1a) Read CSV of docs
@@ -95,7 +95,26 @@ def web_search(query: str, num_results: int = 5) -> str:
     results = search.get_dict().get("organic_results", [])
     snippets = [r.get("snippet","")for r in results]
     return "\n".join(f"- {s}" for s in snippets)
-    
+
+@tool
+def wiki_search(query: str) -> str:
+    """
+    Search Wikipedia for up to 2 pages matching 'query',
+    and return their contents.
+    """
+    #load up to 2 pages
+    pages = WikipediaLoader(query=query, load_max_docs=2).load()
+    #format as plain text
+    return "\n\n---\n\n".join(doc.page_content for doc in pages)
+
+@tool
+def arxiv_search(query:str) -> str:
+    """
+    Search ArXiv for up to 3 abstracts matching 'query',
+    and return their first 1000 characters.
+    """
+    papers = ArxivLoader(query=query, load_max_docs=3).load()
+    return "\n\n---\n\n".join(doc.page_content[:1000]for doc in papers)
     
 # 2) Build your graph
 def build_graph(provider: str = "huggingface") -> Graph:
@@ -107,7 +126,9 @@ def build_graph(provider: str = "huggingface") -> Graph:
     tools = [
         calculator,
         retrieve_docs,
-        web_search,  # add more tools here
+        web_search,
+        wiki_search,
+        arxiv_search, # add more tools here
     ]
     llm_with_tools = llm.with_tools(tools) 
 
@@ -118,11 +139,15 @@ def build_graph(provider: str = "huggingface") -> Graph:
     # tool nodes
     graph.add_node("calc", calculator)
     graph.add_node("retrieve", retrieve_docs)
-    graph.add_node("search", web_search)
+    graph.add_node("web_search", web_search)
+    graph.add_node("wiki", wiki_search)
+    graph.add_node("arxiv", arxiv_search)
     # allow the LLM to call any tool:
     graph.add_edge("ask", "calc") # allow ask -> calc
     graph.add_edge("ask", "retrieve")
     graph.add_edge("ask", "web_search")
+    graph.add_edge("ask", "wiki")
+    graph.add_edge("ask", "arxiv")
     
     graph.set.start("ask")
     return graph