agent.py

'''
TODO tools
- web_search

Facultative
- wiki_search
- arxiv_search
'''

from smolagents import CodeAgent, HfApiModel, tool, DuckDuckGoSearchTool, MLXModel


@tool
def add(a:int, b:int) -> int:
  """
  This tool returns the sum of two numbers.

  Args:
    a: first number
    b: second number
  """

  return a+b


@tool
def subtract(a:int, b:int) -> int:
  """
  This tool returns the difference between two numbers.

  Args:
    a: first number
    b: second number
  """

  return a-b


@tool
def multiply(a:int, b:int) -> int:
  """
  This tool multiplies two numbers.

  Args:
    a: first number
    b: second number
  """

  return a*b


@tool
def divide(a:int, b:int) -> float:
  """
  This tool divides 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:int, b:int) -> int:
  """
  This tool returns the modulus of two numbers.

  Args:
    a: first number
    b: second number
  """

  return a%b


@tool
def rounder(a:float, n:int) -> float:
  """
  This tool return a number rounded to a certain number of decimals.

  Args:
    a: number to be rounded
    n: number of decimals to use when rounding the number
  """

  return round(a,n)


def get_agent() -> CodeAgent:
  search_tool = DuckDuckGoSearchTool()

  model = MLXModel(
    "mlx-community/Qwen2.5-Coder-32B-Instruct-4bit",
    {
        "temperature": 0.7,
        "top_k": 20,  
        "top_p": 0.8, 
        "min_p": 0.05,
        "num_ctx": 32768,
    },
)

  return CodeAgent(tools=[add, subtract, multiply, divide, modulus, rounder, search_tool], model=model)