tools/timer.py

from functools import wraps
import time
from typing import Any, Optional
from smolagents.tools import Tool
from smolagents import tool

class Timer:
    def __init__(self):
        self.start, self.end = 0,0
    
    def start_timer(self) -> None:
        '''Start timer that measures the time taken to answer the question.'''
        self.start = time.time()

    def stop_timer(self) -> float:
        '''Stop timer that measures the time taken to answer the question.'''
        self.end = time.time()
        return self.end - self.start
    

def timer_tools() -> list:
    timer = Timer()
    def wrap(method):
        @wraps(method)
        def wrapper(*args, **kwargs):
            return method(*args, **kwargs)
        return wrapper
    
    start_timer = tool(wrap(timer.start_timer))
    stop_timer = tool(wrap(timer.stop_timer))
    return start_timer, stop_timer