Update app.py

app.py

@@ -6,7 +6,9 @@ from langchain.agents import AgentExecutor, create_react_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.tools import tool
 from langchain.prompts import PromptTemplate
-from langchain_openai import ChatOpenAI  # OpenAI-compatible for Groq API
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+#from langchain_openai import ChatOpenAI  # OpenAI-compatible for Groq API
 from ddgs import DDGS  # Updated DuckDuckGo Search
 from dotenv import load_dotenv
 import re
@@ -48,6 +50,70 @@ def duckduckgo_search(query: str) -> str:
         return json.dumps([{"title": r["title"], "snippet": r["body"]} for r in results])
     except Exception as e:
         return f"Search error: {str(e)}"
+
+@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 wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    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 {"wiki_results": formatted_search_docs}
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
         
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
@@ -55,7 +121,8 @@ class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-        tools = [duckduckgo_search, python_code_executor, download_file]
+        #tools = [duckduckgo_search, python_code_executor, download_file]
+        tools = [duckduckgo_search, python_code_executor, download_file,arvix_search,wiki_search,multiply,add,subtract,divide]
         # Prompt for short, exact answers (GAIA-style) with required ReAct variables
         prompt_template = PromptTemplate.from_template("""
         You are an AI agent solving GAIA benchmark tasks. Use tools if needed (web search, code execution, file download).