SupplyAgent/supply_agent/tools/base.py

from __future__ import annotations

import asyncio
import inspect
import json
import types
from collections.abc import Callable
from typing import Any, Union, get_args, get_origin, get_type_hints

from pydantic import BaseModel

from supply_agent.types import ToolDefinition


def _unwrap_optional(py_type: Any) -> Any:
    """若为 Optional[T] / T | None，返回 T；否则原样返回。"""
    origin = get_origin(py_type)
    if origin is Union or origin is types.UnionType:
        args = [a for a in get_args(py_type) if a is not type(None)]
        if len(args) == 1:
            return args[0]
    return py_type


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

    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 = get_args(py_type) or (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 _coerce_args(func: Callable[..., Any], args: dict[str, Any]) -> dict[str, Any]:
    """按函数类型注解做轻量转换，兼容模型把整数写成字符串等情况。"""
    hints = get_type_hints(func)
    coerced: dict[str, Any] = {}
    for name, value in args.items():
        if name not in hints or value is None:
            coerced[name] = value
            continue

        target = _unwrap_optional(hints[name])
        try:
            if target is int and not isinstance(value, bool):
                coerced[name] = int(value)
            elif target is float and not isinstance(value, bool):
                coerced[name] = float(value)
            elif target is bool and isinstance(value, str):
                coerced[name] = value.strip().lower() in {"1", "true", "yes", "y"}
            else:
                coerced[name] = value
        except (TypeError, ValueError):
            coerced[name] = value
    return coerced


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:
        if asyncio.iscoroutinefunction(self.func):
            return await self.func(**kwargs)
        # 同步函数（如阻塞式 DB 调用）丢到线程池执行，避免整体卡住事件循环——
        # 否则一旦某次同步调用阻塞，async 调用方设置的超时（如 asyncio.wait_for）
        # 完全无法取消/中断它，会导致 Agent 无限期卡死。
        result = await asyncio.to_thread(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."""
        try:
            if isinstance(arguments, str):
                args = json.loads(arguments) if arguments.strip() else {}
            else:
                args = arguments
            args = _coerce_args(self.func, args)
            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 (json.JSONDecodeError, TypeError) as e:
            return json.dumps(
                {"error": str(e), "input_error": True},
                ensure_ascii=False,
            )
        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."""
        try:
            if isinstance(arguments, str):
                args = json.loads(arguments) if arguments.strip() else {}
            else:
                args = arguments
            args = _coerce_args(self.func, args)
            result = await self.acall(**args)
            return str(result) if result is not None else ""
        except (json.JSONDecodeError, TypeError) as e:
            return json.dumps(
                {"error": str(e), "input_error": True},
                ensure_ascii=False,
            )
        except Exception as e:
            return json.dumps({"error": str(e)}, ensure_ascii=False)