added arxiv tool

agent.py

@@ -1,13 +1,6 @@
-'''
-TODO tools
-
-Facultative
-- wiki_search
-- arxiv_search
-'''
 import os
 from smolagents import CodeAgent, tool, DuckDuckGoSearchTool, OpenAIServerModel, VisitWebpageTool, Tool
-from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 
 @tool
 def add(a:int, b:int) -> int:
@@ -102,6 +95,21 @@ def wiki_search(query: str) -> str:
     return {"wiki_results": formatted_search_docs}
 
 
+@tool
+def arvix_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+    
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
+
+
 def get_agent() -> CodeAgent:
   search_tool = DuckDuckGoSearchTool()
   web_page_tool = VisitWebpageTool()
@@ -113,4 +121,4 @@ def get_agent() -> CodeAgent:
     api_base="https://codestral.mistral.ai/v1/",
     api_key=api_key)
 
-  return CodeAgent(tools=[add, subtract, multiply, divide, modulus, rounder, search_tool, web_page_tool, wiki_search], model=model)
+  return CodeAgent(tools=[add, subtract, multiply, divide, modulus, rounder, search_tool, web_page_tool, wiki_search, arvix_search], model=model)

requirements.txt

@@ -5,4 +5,6 @@ langchain_community
 wikipedia
 duckduckgo_search
 requests
-markdownify
+markdownify
+arxiv
+pymupdf