biz-bud/tests/stubs/langchain_core/tools.py

"""Lightweight stub implementations of :mod:`langchain_core.tools`."""

from __future__ import annotations

import asyncio
import inspect
from typing import Any, Callable, Coroutine, Mapping


def _ensure_async(result: Any) -> Coroutine[Any, Any, Any]:
    if inspect.isawaitable(result):
        return result  # type: ignore[return-value]
    async def _wrapper() -> Any:
        return result
    return _wrapper()


def _coerce_kwargs(input_data: Any, kwargs: Mapping[str, Any] | None = None) -> dict[str, Any]:
    if kwargs:
        return dict(kwargs)
    if input_data is None:
        return {}
    if isinstance(input_data, Mapping):
        return dict(input_data)
    return {"input": input_data}


class BaseTool:
    """Minimal approximation of LangChain's ``BaseTool``."""

    name: str = "tool"
    description: str | None = None
    args_schema: Any | None = None
    return_direct: bool = False
    is_single_input: bool = True
    handle_tool_error: Any | None = None

    def __init__(
        self,
        *,
        name: str | None = None,
        description: str | None = None,
        args_schema: Any | None = None,
        return_direct: bool | None = None,
        is_single_input: bool | None = None,
        handle_tool_error: Any | None = None,
    ) -> None:
        if name is not None:
            self.name = name
        if description is not None:
            self.description = description
        if args_schema is not None:
            self.args_schema = args_schema
        if return_direct is not None:
            self.return_direct = return_direct
        if is_single_input is not None:
            self.is_single_input = is_single_input
        if handle_tool_error is not None:
            self.handle_tool_error = handle_tool_error

    # ------------------------------------------------------------------
    # Invocation helpers
    # ------------------------------------------------------------------
    def invoke(self, input: Any | None = None, **kwargs: Any) -> Any:
        return self._run(**_coerce_kwargs(input, kwargs))

    async def ainvoke(self, input: Any | None = None, **kwargs: Any) -> Any:
        return await self._arun(**_coerce_kwargs(input, kwargs))

    # The real BaseTool exposes ``arun``/``run`` wrappers. These helpers are
    # convenience aliases used in a handful of call sites.
    def run(self, *args: Any, **kwargs: Any) -> Any:
        return self.invoke(*args, **kwargs)

    async def arun(self, *args: Any, **kwargs: Any) -> Any:
        return await self.ainvoke(*args, **kwargs)

    # ------------------------------------------------------------------
    # Extension points for subclasses
    # ------------------------------------------------------------------
    def _run(self, **kwargs: Any) -> Any:  # pragma: no cover - override hook
        raise NotImplementedError("BaseTool subclasses must implement _run")

    async def _arun(self, **kwargs: Any) -> Any:  # pragma: no cover - override hook
        raise NotImplementedError("BaseTool subclasses must implement _arun")


class _CallableTool(BaseTool):
    def __init__(
        self,
        func: Callable[..., Any],
        *,
        name: str | None = None,
        description: str | None = None,
        args_schema: Any | None = None,
        return_direct: bool | None = None,
        is_single_input: bool | None = None,
        handle_tool_error: Any | None = None,
    ) -> None:
        super().__init__(
            name=name or func.__name__,
            description=description or (inspect.getdoc(func) or ""),
            args_schema=args_schema,
            return_direct=return_direct,
            is_single_input=is_single_input,
            handle_tool_error=handle_tool_error,
        )
        self._func = func

    def _run(self, **kwargs: Any) -> Any:
        result = self._func(**kwargs)
        if inspect.isawaitable(result):
            loop = asyncio.get_event_loop()
            return loop.run_until_complete(result)
        return result

    async def _arun(self, **kwargs: Any) -> Any:
        result = self._func(**kwargs)
        return await _ensure_async(result)


def tool(
    func: Callable[..., Any] | str | None = None,
    *,
    name: str | None = None,
    description: str | None = None,
    args_schema: Any | None = None,
    return_direct: bool | None = None,
    is_single_input: bool | None = None,
    handle_tool_error: Any | None = None,
    infer_schema: bool | None = None,
    **_: Any,
) -> Callable[[Callable[..., Any]], _CallableTool] | _CallableTool:
    """Decorator returning a lightweight ``BaseTool`` implementation."""

    initial_name = name
    if isinstance(func, str):
        initial_name = func
        func = None

    def decorator(target: Callable[..., Any]) -> _CallableTool:
        return _CallableTool(
            target,
            name=initial_name,
            description=description,
            args_schema=args_schema,
            return_direct=return_direct,
            is_single_input=is_single_input,
            handle_tool_error=handle_tool_error,
        )

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


__all__ = ["BaseTool", "tool"]