updated with more tools

model.py

@@ -2,30 +2,73 @@
 import os
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import AsyncHtmlLoader
 from langchain.tools import tool
 from langchain.prompts import ChatPromptTemplate
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from prompt import get_prompt
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+
 import re
 
+# --- Define Tools ---
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two integers."""
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two integers."""
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract b from a."""
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    """Divide a by b, error on zero."""
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Compute a mod b."""
+    return a % b
+
+@tool
+def wiki_search(query: str) -> dict:
+    """Search Wikipedia and return up to 2 documents."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    results = [f"<Document source=\"{d.metadata['source']}\" page=\"{d.metadata.get('page','')}\"/>\n{d.page_content}" for d in docs]
+    return {"wiki_results": "\n---\n".join(results)}
+
+@tool
+def web_search(query: str) -> dict:
+    """Search Tavily and return up to 3 results."""
+    docs = TavilySearchResults(max_results=3).invoke(query=query)
+    results = [f"<Document source=\"{d.metadata['source']}\" page=\"{d.metadata.get('page','')}\"/>\n{d.page_content}" for d in docs]
+    return {"web_results": "\n---\n".join(results)}
+
+@tool
+def arxiv_search(query: str) -> dict:
+    """Search Arxiv and return up to 3 docs."""
+    docs = ArxivLoader(query=query, load_max_docs=3).load()
+    results = [f"<Document source=\"{d.metadata['source']}\" page=\"{d.metadata.get('page','')}\"/>\n{d.page_content[:1000]}" for d in docs]
+    return {"arxiv_results": "\n---\n".join(results)}
+
 class Model:
     def __init__(self):        
         #load_dotenv(find_dotenv())
         self.token = os.getenv("HF_TOKEN")
-        print(f"token: {self.token}")
         self.system_prompt = get_prompt()
         print(f"system_prompt: {self.system_prompt}")
         self.agent_executor = self.setup_model()
 
-    # Define a tool for the agent to use
-    @tool
-    def scrape_webpage(self,url: str) -> str:
-        """Scrapes a given URL and returns the content."""
-        loader = AsyncHtmlLoader(url)
-        docs = loader.load()
-        return docs[0].page_content# Define the search tool
-
     def get_answer(self, question: str) -> str:
         try:
             result = self.agent_executor.invoke({"input": question})
@@ -52,8 +95,17 @@ class Model:
         search = DuckDuckGoSearchRun()
 
         # # Define a tool for the agent to use
-        tools = [search, self.scrape_webpage]
-
+        tools = [
+            multiply,
+            add,
+            subtract,
+            divide,
+            modulus,
+            wiki_search,
+            search,
+            arxiv_search,
+        ]
+        
         llm = HuggingFaceEndpoint(
             repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
             huggingfacehub_api_token=self.token,