Update agent.py

agent.py

@@ -1,68 +1,135 @@
-from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, tool
+from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool
 
 
-@tool
-def add(a: int, b: int) -> int:
-    """
+class add(Tool):
+    name = "add"
+    description = """
         This tool adds two integers together and returns an integer.
 
         Args:
             a: first int
             b: second int
-    """
-    return a + b
-
-@tool
-def subtract(a: int, b: int) -> int:
-    """
+        """
+    inputs = {
+        "a":{
+            "type":"int",
+            "description":"first int"
+        },
+        "b":{
+            "type":"int",
+            "description":"second int"
+        }
+    }
+    output_type = "int"
+
+    def forward(self, a: int, b: int):
+        return a + b
+
+
+class subtract(Tool):
+    name = "subtract"
+    description = """
         This tool subtracts two integers and returns an integer.
 
         Args:
             a: first int
             b: second int
-    """
-    return a - b
-
-@tool
-def multiply(a: int, b: int) -> int:
-    """
-        This tool multiplies two integers together and returns an integer.
+        """
+    inputs = {
+        "a":{
+            "type":"int",
+            "description":"first int"
+        },
+        "b":{
+            "type":"int",
+            "description":"second int"
+        }
+    }
+    output_type = "int"
+
+    def forward(self, a: int, b: int):
+        return a - b
+
+class multiply(Tool):
+    name = "multiply"
+    description = """
+        This tool multiplies two integers and returns an integer.
 
         Args:
             a: first int
             b: second int
-    """
-    return a * b
-
-@tool
-def divide(a: int, b: int) -> int:
-    """
-        This tool divides two integers together and returns a float.
+        """
+    inputs = {
+        "a":{
+            "type":"int",
+            "description":"first int"
+        },
+        "b":{
+            "type":"int",
+            "description":"second int"
+        }
+    }
+    output_type = "int"
+
+    def forward(self, a: int, b: int):
+        return a * b
+
+class divide(Tool):
+    name = "divide"
+    description = """
+        This tool divides two integers and returns an integer.
 
         Args:
             a: first int
             b: second int
-    """
-    if b == 0:
-        raise ValueError("Cannot divide by zero.")
-    return a / b
-
-def modulo(a: int, b: int) -> float:
-    """
+        """
+    inputs = {
+        "a":{
+            "type":"int",
+            "description":"first int"
+        },
+        "b":{
+            "type":"int",
+            "description":"second int"
+        }
+    }
+    output_type = "int"
+
+    def forward(self, a: int, b: int):
+        if b == 0:
+            raise ValueError("Cannot divide by zero.")
+        return a / b
+
+class modulo(Tool):
+    name = "modulo"
+    description = """
         This tool returns the modulus of two integers
 
         Args:
             a: first int
             b: second int
-    """
-    return a % b
+        """
+    inputs = {
+        "a":{
+            "type":"int",
+            "description":"first int"
+        },
+        "b":{
+            "type":"int",
+            "description":"second int"
+        }
+    }
+    output_type = "int"
+
+    def forward(self, a: int, b: int):
+        return a % b
 
 tools=[
-    add,
-    subtract,
-    multiply,
-    divide,
-    modulo,
+    add(),
+    subtract(),
+    multiply(),
+    divide(),
+    modulo(),
     DuckDuckGoSearchTool(),
     WikipediaSearchTool()
 ]