SupplyAgent/supply_agent/tools/base.py

from __future__ import annotations

import inspect
import json
from collections.abc import Callable
from typing import Any, get_type_hints

from pydantic import BaseModel

from supply_agent.types import ToolDefinition


def _python_type_to_json_schema(py_type: type) -> dict[str, Any]:
    """Map Python types to JSON Schema types."""
    origin = getattr(py_type, "__origin__", None)

    if py_type is str or py_type is inspect.Parameter.empty:
        return {"type": "string"}
    if py_type is int:
        return {"type": "integer"}
    if py_type is float:
        return {"type": "number"}
    if py_type is bool:
        return {"type": "boolean"}
    if origin is list:
        args = getattr(py_type, "__args__", (Any,))
        item_type = args[0] if args else Any
        return {"type": "array", "items": _python_type_to_json_schema(item_type)}
    if origin is dict:
        return {"type": "object"}
    return {"type": "string"}


def _build_parameters(func: Callable[..., Any]) -> dict[str, Any]:
    """Build JSON Schema parameters from function signature."""
    sig = inspect.signature(func)
    hints = get_type_hints(func)
    properties: dict[str, Any] = {}
    required: list[str] = []

    for name, param in sig.parameters.items():
        if name in ("self", "cls"):
            continue
        py_type = hints.get(name, str)
        prop: dict[str, Any] = _python_type_to_json_schema(py_type)

        if param.default is not inspect.Parameter.empty:
            if not isinstance(param.default, (BaseModel,)):
                prop["default"] = param.default
        else:
            required.append(name)

        properties[name] = prop

    return {"type": "object", "properties": properties, "required": required}


def tool(
    func: Callable[..., Any] | None = None,
    *,
    name: str | None = None,
    description: str | None = None,
) -> Callable[..., Any]:
    """
    Decorator to register a function as an agent tool.

    Usage:
        @tool
        def search(query: str, limit: int = 10) -> str:
            '''Search the web for information.'''
            ...
    """
    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
        fn._is_tool = True  # type: ignore[attr-defined]
        fn._tool_name = name or fn.__name__  # type: ignore[attr-defined]
        fn._tool_description = description or (fn.__doc__ or "").strip()  # type: ignore[attr-defined]
        fn._tool_parameters = _build_parameters(fn)  # type: ignore[attr-defined]
        return fn

    if func is not None:
        return decorator(func)
    return decorator


class Tool:
    """A callable tool with schema metadata."""

    def __init__(
        self,
        func: Callable[..., Any],
        name: str | None = None,
        description: str | None = None,
    ) -> None:
        self.func = func
        self.name = name or getattr(func, "_tool_name", func.__name__)
        self.description = description or getattr(
            func, "_tool_description", (func.__doc__ or "").strip()
        )
        self.parameters = getattr(func, "_tool_parameters", _build_parameters(func))

    @property
    def definition(self) -> ToolDefinition:
        return ToolDefinition(
            name=self.name,
            description=self.description,
            parameters=self.parameters,
        )

    def __call__(self, **kwargs: Any) -> Any:
        return self.func(**kwargs)

    async def acall(self, **kwargs: Any) -> Any:
        result = self.func(**kwargs)
        if inspect.isawaitable(result):
            return await result
        return result

    def execute(self, arguments: str | dict[str, Any]) -> str:
        """Execute the tool with JSON or dict arguments, returning a string result."""
        if isinstance(arguments, str):
            args = json.loads(arguments) if arguments.strip() else {}
        else:
            args = arguments
        try:
            result = self(**args)
            if inspect.isawaitable(result):
                raise RuntimeError(
                    f"Tool '{self.name}' is async. Use ToolRegistry.aexecute() instead."
                )
            return str(result) if result is not None else ""
        except Exception as e:
            return json.dumps({"error": str(e)}, ensure_ascii=False)

    async def aexecute(self, arguments: str | dict[str, Any]) -> str:
        """Async execute the tool."""
        if isinstance(arguments, str):
            args = json.loads(arguments) if arguments.strip() else {}
        else:
            args = arguments
        try:
            result = await self.acall(**args)
            return str(result) if result is not None else ""
        except Exception as e:
            return json.dumps({"error": str(e)}, ensure_ascii=False)