agent.py

from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool


class add(Tool):
    name = "add"
    description = """
        This tool adds two integers together and returns an integer.

        Args:
            a: first int
            b: second 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
        """
    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
        """
    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
        """
    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
        """
    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(),
    DuckDuckGoSearchTool(),
    WikipediaSearchTool()
]


def create_agent():
    agent = CodeAgent(
        model = HfApiModel(),
        tools = tools,
        max_steps=10,
        verbosity_level=2
    )

    return agent