arxiv ret

agent_utils.py

@@ -1,5 +1,5 @@
 from langchain_community.document_loaders import WikipediaLoader
-
+import arxiv
 
 def wiki_ret(question: str) -> str:
     """ Retrieve docs from wikipedia """
@@ -18,3 +18,22 @@ def wiki_ret(question: str) -> str:
 
     return {"context": formatted_search_docs}
 
+
+
+def arxiv_ret(query: str, max_results: int = 3) -> str:
+    """Search arXiv for a given query and return a summary of top results."""
+    search = arxiv.Search(
+        query=query,
+        max_results=max_results,
+        sort_by=arxiv.SortCriterion.Relevance
+    )
+
+    results = []
+    for result in search.results():
+        summary = result.summary.replace("\n", " ")
+        results.append(f"Title: {result.title}\nAuthors: {', '.join(a.name for a in result.authors)}\nSummary: {summary}\nURL: {result.entry_id}\n")
+
+    if not results:
+        return "No relevant papers found."
+    
+    return "\n\n".join(results)

app.py

@@ -17,7 +17,7 @@ from langchain_core.messages import HumanMessage
 from langgraph.graph import START, END, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
 
-from agent_utils import wiki_ret
+from agent_utils import wiki_ret, arxiv_ret
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -82,7 +82,7 @@ class BasicAgent:
     def create_graph(self):
         builder = StateGraph(GraphState)
         builder.add_node("agent", self.agent)
-        builder.add_node("tools", ToolNode(tools = [wiki_ret]))
+        builder.add_node("tools", ToolNode(tools = [wiki_ret, arxiv_ret]))
 
         builder.add_edge(START, "agent")
         builder.add_conditional_edges("agent", tools_condition)

requirements.txt

@@ -4,4 +4,5 @@ langchain-groq
 langchain-community
 langchain-core
 wikipedia
-langgraph
+langgraph
+arxiv