Update agent.py

agent.py

@@ -1,31 +1,70 @@
 from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel
-
-# Initialize a model (use your preferred model_id)
-model = InferenceClientModel(model_id="meta-llama/Llama-2-7b-chat-hf")  # Or another model
-
-# Define custom tools if needed (optional)
 from smolagents import tool
+import os
+
+# Initialize model
+model = InferenceClientModel(model_id="meta-llama/Llama-2-70b-chat-hf")
 
+# Custom tools with corrected docstrings
 @tool
 def add(a: int, b: int) -> int:
-    """Add two numbers."""
+    """Add two numbers.
+    
+    Args:
+        a: The first number to add
+        b: The second number to add
+    """
     return a + b
 
 @tool
 def multiply(a: int, b: int) -> int:
-    """Multiply two numbers."""
+    """Multiply two numbers.
+    
+    Args:
+        a: The first number to multiply
+        b: The second number to multiply
+    """
     return a * b
 
-# Create the agent with web search and custom tools
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    
+    Args:
+        a: The number to subtract from
+        b: The number to subtract
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+    
+    Args:
+        a: The number to divide
+        b: The divisor
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    
+    Args:
+        a: The number to get modulus of
+        b: The modulus divisor
+    """
+    return a % b
+
+# Create agent with tools
 agent = CodeAgent(
-    tools=[DuckDuckGoSearchTool(), add, multiply],
+    tools=[DuckDuckGoSearchTool(), add, multiply, subtract, divide, modulus],
     model=model,
 )
 
 def answer_question(question: str) -> str:
-    """Run the agent on a question and return the answer."""
+    """Run agent on question and format output"""
     result = agent.run(question)
-    # Format as per system_prompt.txt
-    if "FINAL ANSWER:" not in result:
-        result = f"FINAL ANSWER: {result}"
-    return result
+    return f"FINAL ANSWER: {result}" if "FINAL ANSWER:" not in result else result