now using smolagents instead of langchain

gaia_agent.py

@@ -1,185 +1,164 @@
 import os
 import requests
-from langchain.agents import create_react_agent, AgentExecutor
-from langchain.tools import tool
-from langchain_core.prompts import PromptTemplate
+from smolagents import CodeAgent, Tool, HfApiModel
 from dotenv import load_dotenv
-from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
-from duckduckgo_search import DDGS
 from bs4 import BeautifulSoup
 
 # Load environment variables
 # load_dotenv()
 
 # --- Agent Setup ---
-openrouter_key = os.getenv("OPENROUTER_API_KEY")
-if not openrouter_key:
-    raise RuntimeError("Set OPENROUTER_API_KEY in your .env (OpenRouter API key)")
-
-from langchain_openai import ChatOpenAI
-
-model = ChatOpenAI(
-    api_key=openrouter_key,
-    base_url="https://openrouter.ai/api/v1",
-    model="openai/gpt-4o-mini",
-    max_tokens=10000,
-    temperature=0
-)
-
-# --- Tools Definition ---
-@tool
-def multiply(a: float, b: float) -> float:
-    """Multiply two numbers.
-    Args:
-        a: first number
-        b: second number
-    """
-    return a * b
-
-@tool
-def add(a: float, b: float) -> float:
-    """Add two numbers.
-    Args:
-        a: first number
-        b: second number
-    """
-    return a + b
-
-@tool
-def subtract(a: float, b: float) -> float:
-    """Subtract two numbers.
-    Args:
-        a: first number
-        b: second number
-    """
-    return a - b
-
-@tool
-def divide(a: float, b: float) -> float:
-    """Divide two numbers.
-    Args:
-        a: first number
-        b: second number
-    """
-    if b == 0:
-        raise ValueError("Cannot divide by zero.")
-    return a / b
-
-@tool
-def modulus(a: float, b: float) -> float:
-    """Get the modulus of two numbers.
-    Args:
-        a: first number
-        b: second number
-    """
-    return a % b
-
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results."""
-    try:
-        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'\n{doc.page_content}\n'
-                for doc in search_docs
-            ]
-        )
-        return formatted_search_docs
-    except Exception as e:
-        return f"Error searching Wikipedia: {str(e)}"
-
-@tool
-def web_search(query: str) -> str:
-    """Search the web for a query and return maximum 3 results."""
-    try:
-        search_docs = DDGS().text(query, max_results=3)
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'Title:{doc["title"]}\nContent:{doc["body"]}\n--\n'
-                for doc in search_docs
-            ]
-        )
-        return formatted_search_docs
-    except Exception as e:
-        return f"Error searching web: {str(e)}"
-
-@tool
-def arxiv_search(query: str) -> str:
-    """Search arXiv for a query and return maximum 3 results."""
-    try:
-        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'\n{doc.page_content[:1000]}\n'
-                for doc in search_docs
-            ]
-        )
-        return formatted_search_docs
-    except Exception as e:
-        return f"Error searching arXiv: {str(e)}"
-
-@tool
-def fetch_url_content(url: str) -> str:
-    """Fetch and return the text content from a webpage URL."""
-    try:
-        response = requests.get(url, timeout=10)
-        response.raise_for_status()
-        soup = BeautifulSoup(response.text, 'html.parser')
-        for script in soup(["script", "style"]):
-            script.decompose()
-        text = soup.get_text(separator='\n', strip=True)
-        return text[:3000] + ("..." if len(text) > 3000 else "")
-    except Exception as e:
-        return f"Error fetching URL: {str(e)}"
-
-# Tools list
-tools = [
-    multiply, add, subtract, divide, modulus,
-    wiki_search, web_search, arxiv_search,
-    fetch_url_content,
-]
-
-# React prompt template
-react_prompt = PromptTemplate.from_template("""You are a helpful assistant that answers questions accurately and concisely.
-
-Answer the following questions as best you can. You have access to the following tools:
-
-{tools}
+hf_token = os.getenv("HF_TOKEN")
+if not hf_token:
+    raise RuntimeError("Set HF_TOKEN in your Space secrets")
+
+# Initialize the model
+model = HfApiModel(token=hf_token, model_id="Qwen/Qwen2.5-72B-Instruct")
+
+# --- Custom Tools Definition ---
+class MultiplyTool(Tool):
+    name = "multiply"
+    description = "Multiplies two numbers together. Takes two arguments: a (first number) and b (second number)."
+    inputs = {"a": {"type": "number", "description": "First number"}, 
+              "b": {"type": "number", "description": "Second number"}}
+    output_type = "number"
+
+    def forward(self, a: float, b: float) -> float:
+        return a * b
+
+class AddTool(Tool):
+    name = "add"
+    description = "Adds two numbers together. Takes two arguments: a (first number) and b (second number)."
+    inputs = {"a": {"type": "number", "description": "First number"}, 
+              "b": {"type": "number", "description": "Second number"}}
+    output_type = "number"
+
+    def forward(self, a: float, b: float) -> float:
+        return a + b
+
+class SubtractTool(Tool):
+    name = "subtract"
+    description = "Subtracts b from a. Takes two arguments: a (first number) and b (second number)."
+    inputs = {"a": {"type": "number", "description": "First number"}, 
+              "b": {"type": "number", "description": "Second number"}}
+    output_type = "number"
+
+    def forward(self, a: float, b: float) -> float:
+        return a - b
+
+class DivideTool(Tool):
+    name = "divide"
+    description = "Divides a by b. Takes two arguments: a (first number) and b (second number)."
+    inputs = {"a": {"type": "number", "description": "First number"}, 
+              "b": {"type": "number", "description": "Second number"}}
+    output_type = "number"
+
+    def forward(self, a: float, b: float) -> float:
+        if b == 0:
+            raise ValueError("Cannot divide by zero.")
+        return a / b
+
+class ModulusTool(Tool):
+    name = "modulus"
+    description = "Calculates the modulus (remainder) of a divided by b. Takes two arguments: a (first number) and b (second number)."
+    inputs = {"a": {"type": "number", "description": "First number"}, 
+              "b": {"type": "number", "description": "Second number"}}
+    output_type = "number"
+
+    def forward(self, a: float, b: float) -> float:
+        return a % b
+
+class WikipediaSearchTool(Tool):
+    name = "wikipedia_search"
+    description = "Searches Wikipedia for a given query and returns relevant information."
+    inputs = {"query": {"type": "string", "description": "The search query"}}
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        try:
+            import wikipedia
+            wikipedia.set_lang("en")
+            results = wikipedia.search(query, results=2)
+            if not results:
+                return "No results found."
+            
+            summaries = []
+            for result in results[:2]:
+                try:
+                    page = wikipedia.page(result, auto_suggest=False)
+                    summaries.append(f"Title: {page.title}\n{page.summary[:500]}")
+                except:
+                    continue
+            
+            return "\n\n---\n\n".join(summaries) if summaries else "No content found."
+        except Exception as e:
+            return f"Error searching Wikipedia: {str(e)}"
 
-Use the following format:
+class WebSearchTool(Tool):
+    name = "web_search"
+    description = "Searches the web using DuckDuckGo and returns relevant results."
+    inputs = {"query": {"type": "string", "description": "The search query"}}
+    output_type = "string"
 
-Question: the input question you must answer
-Thought: you should always think about what to do
-Action: the action to take, should be one of [{tool_names}]
-Action Input: the input to the action
-Observation: the result of the action
-... (this Thought/Action/Action Input/Observation can repeat N times)
-Thought: I now know the final answer
-Final Answer: the final answer to the original input question
+    def forward(self, query: str) -> str:
+        try:
+            from duckduckgo_search import DDGS
+            results = DDGS().text(query, max_results=3)
+            
+            formatted_results = []
+            for result in results:
+                formatted_results.append(
+                    f"Title: {result.get('title', 'N/A')}\n"
+                    f"Content: {result.get('body', 'N/A')}\n"
+                    f"URL: {result.get('href', 'N/A')}"
+                )
+            
+            return "\n\n---\n\n".join(formatted_results) if formatted_results else "No results found."
+        except Exception as e:
+            return f"Error searching web: {str(e)}"
 
-IMPORTANT: Your Final Answer must be:
-- Short and direct (just the answer, no extra explanation)
-- A single value or short phrase
-- No formatting, no bullet points, no extra text
-- Just the factual answer
+class FetchURLTool(Tool):
+    name = "fetch_url_content"
+    description = "Fetches and returns the text content from a given webpage URL."
+    inputs = {"url": {"type": "string", "description": "The URL to fetch content from"}}
+    output_type = "string"
 
-Begin!
+    def forward(self, url: str) -> str:
+        try:
+            response = requests.get(url, timeout=10, headers={'User-Agent': 'Mozilla/5.0'})
+            response.raise_for_status()
+            soup = BeautifulSoup(response.text, 'html.parser')
+            
+            # Remove script and style elements
+            for script in soup(["script", "style", "meta", "link"]):
+                script.decompose()
+            
+            text = soup.get_text(separator='\n', strip=True)
+            # Limit to 3000 characters
+            return text[:3000] + ("..." if len(text) > 3000 else "")
+        except Exception as e:
+            return f"Error fetching URL: {str(e)}"
 
-Question: {input}
-Thought:{agent_scratchpad}""")
+# Initialize tools
+tools = [
+    MultiplyTool(),
+    AddTool(),
+    SubtractTool(),
+    DivideTool(),
+    ModulusTool(),
+    WikipediaSearchTool(),
+    WebSearchTool(),
+    FetchURLTool(),
+]
 
 class GAIAAgent:
     def __init__(self):
-        # create internal agent with React agent
         try:
-            agent = create_react_agent(model, tools, react_prompt)
-            self.agent_executor = AgentExecutor(
-                agent=agent,
+            self.agent = CodeAgent(
                 tools=tools,
-                verbose=True,
-                handle_parsing_errors=True,
-                max_iterations=15
+                model=model,
+                max_steps=10,
+                verbosity_level=2
             )
         except Exception as e:
             print(f"Error creating agent: {e}")
@@ -187,13 +166,12 @@ class GAIAAgent:
 
     def __call__(self, question: str) -> str:
         try:
-            result = self.agent_executor.invoke({"input": question})
-            answer = result.get("output", "")
+            result = self.agent.run(question)
             
-            # Clean up the answer - remove any extra formatting
-            answer = answer.strip()
+            # Clean up the answer
+            answer = str(result).strip()
             
-            # Remove common prefixes that might be added
+            # Remove common prefixes
             prefixes_to_remove = [
                 "The answer is:",
                 "The final answer is:",