Update app.py

app.py

@@ -11,7 +11,9 @@ from typing import TypedDict, List, Dict, Any, Optional,Callable
 from langgraph.graph import StateGraph, END
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
+from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
 from youtube_transcript_api import YouTubeTranscriptApi
 
 # (Keep Constants as is)
@@ -87,19 +89,45 @@ def divide(a:int,b:int)->float:
         raise ValueError("Division by zero is not allowed.")
     return a / b
 
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
 search_tool = DuckDuckGoSearchTool()
 
 @tool
 def web_search(query: str) -> str:
-    """
-    Performs a web search for the given query.
+    """Search Tavily for a query and return maximum 3 results.
+    
     Args:
-        query (str): The search query.
-    Returns:
-        str: The search results as a string.
-    """
-    result=search_tool(query)
-    return f"Search results for '{query}' : {result}."
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_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}
 
 @tool
 def wikipedia_tool(query: str) -> str:
@@ -163,7 +191,7 @@ class BasicAgent:
         token=os.environ["OPENROUTER_API_KEY"]  
         self.system_prompt= """You are an expert assistant with knowledge about everything.
         You will be given a task to do and you need to do it with the help of tools provied.
-        if you are asked about a person or anything related to a work you can use the web_search tool or the wikipedia_tool to come up with the most suitable answer.
+        if you are asked about a person or anything related to a work you can use the web_search tool or the wikipedia_tool or arvix_search tool to come up with the most suitable answer.
         If you are given any mathematical tasks you can use the add,subtract,multiply and divide tools provided.
         If you are given a youtube video link use the youtube_transcript_tool to transcibe the audio and to find the answer of the query being asked.
         If you are asked something related to how many ,use the web_search tool to directly find the answer ,if not available use wikipedia_tool and use your senses to find out the answer yourself and always return a number to such queries.
@@ -182,7 +210,7 @@ class BasicAgent:
     provider="openrouter"                                 # Explicitly set to openrouter
 )
         self.agent= CodeAgent(
-            tools = [add, subtract, multiply, divide, web_search, image_generation_tool,youtube_transcript_tool, wikipedia_tool],
+            tools = [add, subtract, multiply, divide,modulus,arvix_search, web_search, image_generation_tool,youtube_transcript_tool, wikipedia_tool],
             model=model,
         )
     def __call__(self, question: str, context: str = "") -> str: