tools/base.py

"""
Base tool class with error handling, retry, and caching support
"""

from typing import Any, Dict, Optional, Callable
from smolagents import Tool
import functools
import time
import logging

logger = logging.getLogger(__name__)


def retry_with_backoff(max_retries: int = 3, base_delay: float = 1.0):
    """
    Decorator for retrying with exponential backoff.

    Args:
        max_retries: Maximum retry attempts
        base_delay: Base delay in seconds (doubles each retry)
    """
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_retries - 1:
                        delay = base_delay * (2 ** attempt)
                        logger.warning(
                            f"{func.__name__} failed (attempt {attempt + 1}/{max_retries}): {e}. "
                            f"Retrying in {delay}s..."
                        )
                        time.sleep(delay)
                    else:
                        logger.error(f"{func.__name__} failed after {max_retries} attempts: {e}")

            if last_exception:
                raise last_exception
            else:
                raise RuntimeError(f"{func.__name__} failed without exception")

        return wrapper
    return decorator


class BaseAgentTool(Tool):
    """
    Base class for all agent tools with error handling and caching.

    Attributes:
        cacheable: Whether tool results can be cached
        cache_ttl: Cache time-to-live in seconds
        retryable: Whether tool should retry on failure
        max_retries: Maximum retry attempts
    """

    cacheable: bool = False
    cache_ttl: int = 3600
    retryable: bool = True
    max_retries: int = 3

    def _format_error(
        self,
        error: Exception,
        recovery_hint: Optional[str] = None,
        fallback_action: Optional[str] = None
    ) -> str:
        """
        Format error as structured response for agent.

        Args:
            error: Exception that occurred
            recovery_hint: Hint for how to recover
            fallback_action: Alternative action to try

        Returns:
            Formatted error message
        """
        error_dict = {
            "success": False,
            "isError": True,
            "error": str(error),
            "error_type": type(error).__name__,
        }

        if recovery_hint:
            error_dict["recovery_hint"] = recovery_hint

        if fallback_action:
            error_dict["fallback_action"] = fallback_action

        import json
        return json.dumps(error_dict, indent=2)

    def _format_success(self, result: Any, metadata: Optional[Dict[str, Any]] = None) -> str:
        """
        Format success response as structured output.

        Args:
            result: Tool result
            metadata: Optional metadata to include

        Returns:
            Formatted success message
        """
        response = {
            "success": True,
            "result": result
        }

        if metadata:
            response["metadata"] = metadata

        import json
        return json.dumps(response, indent=2)

    def _validate_output(self, result: Any) -> str:
        """
        Validate tool output against ToolOutput schema.

        Optional helper method for subclasses that want strict validation.
        Wraps the result in Pydantic validation and auto-repairs malformed JSON.

        Args:
            result: Raw tool result

        Returns:
            JSON string conforming to ToolOutput schema
        """
        from core.validation import validate_tool_output

        validated = validate_tool_output(result)
        return validated.model_dump_json(indent=2)