AgentCoreService/pqai_agent/toolkit/base.py

from typing import List, Optional
import functools
import threading
from pqai_agent.toolkit.function_tool import FunctionTool


class ToolServiceError(Exception):

    def __init__(self,
                 exception: Optional[Exception] = None,
                 code: Optional[str] = None,
                 message: Optional[str] = None,
                 extra: Optional[dict] = None):
        if exception is not None:
            super().__init__(exception)
        else:
            super().__init__(f'\nError code: {code}. Error message: {message}')
        self.exception = exception
        self.code = code
        self.message = message
        self.extra = extra

def with_timeout(timeout=None):
    r"""Decorator that adds timeout functionality to functions.

    Executes functions with a specified timeout value. Returns a timeout
    message if execution time is exceeded.

    Args:
        timeout (float, optional): The timeout duration in seconds. If None,
            will try to get timeout from the instance's timeout attribute.
            (default: :obj:`None`)

    Example:
        >>> @with_timeout(5)
        ... def my_function():
        ...     return "Success"
        >>> my_function()

        >>> class MyClass:
        ...     timeout = 5
        ...     @with_timeout()
        ...     def my_method(self):
        ...         return "Success"
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Determine the effective timeout value
            effective_timeout = timeout
            if effective_timeout is None and args:
                effective_timeout = getattr(args[0], 'timeout', None)

            # If no timeout value is provided, execute function normally
            if effective_timeout is None:
                return func(*args, **kwargs)

            # Container to hold the result of the function call
            result_container = []

            def target():
                result_container.append(func(*args, **kwargs))

            # Start the function in a new thread
            thread = threading.Thread(target=target)
            thread.start()
            thread.join(effective_timeout)

            # Check if the thread is still alive after the timeout
            if thread.is_alive():
                return (
                    f"Function `{func.__name__}` execution timed out, "
                    f"exceeded {effective_timeout} seconds."
                )
            else:
                return result_container[0]

        return wrapper

    # Handle both @with_timeout and @with_timeout() usage
    if callable(timeout):
        # If timeout is passed as a function, apply it to the decorator
        func, timeout = timeout, None
        return decorator(func)

    return decorator

class BaseToolkit:
    r"""Base class for toolkits.

    Args:
        timeout (Optional[float]): The timeout for the toolkit.
    """

    timeout: Optional[float] = None

    def __init__(self, timeout: Optional[float] = None):
        # check if timeout is a positive number
        if timeout is not None and timeout <= 0:
            raise ValueError("Timeout must be a positive number.")
        self.timeout = timeout

    # Add timeout to all callable methods in the toolkit
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        for attr_name, attr_value in cls.__dict__.items():
            if callable(attr_value) and not attr_name.startswith("__"):
                setattr(cls, attr_name, with_timeout(attr_value))

    def get_tools(self) -> List[FunctionTool]:
        r"""Returns a list of FunctionTool objects representing the
        functions in the toolkit.

        Returns:
            List[FunctionTool]: A list of FunctionTool objects
                representing the functions in the toolkit.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    def get_tool(self, name: str) -> FunctionTool:
        r"""Returns a FunctionTool object by name.

        Args:
            name (str): The name of the tool to retrieve.

        Returns:
            Optional[FunctionTool]: The FunctionTool object if found, else None.
        """
        tools = self.get_tools()
        for tool in tools:
            if tool.name == name:
                return tool
        raise NotImplementedError("Tool not found in the toolkit.")