Coverage for src\funcall\decorators.py: 83%
29 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-28 01:17 +0900
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-28 01:17 +0900
1"""
2decorators.py
4This module contains custom decorators for the funcall package.
5"""
7import asyncio
8import functools
9from collections.abc import Awaitable, Callable
10from typing import Generic, ParamSpec, TypeVar, cast
12P = ParamSpec("P")
13R = TypeVar("R")
16class ToolWrapper(Generic[P, R]):
17 def __init__(self, func: Callable[P, R], *, require_confirm: bool = False, return_direct: bool = False) -> None:
18 functools.update_wrapper(self, func)
19 self._func = func
20 self.require_confirm = require_confirm
21 self.return_direct = return_direct
22 self.is_async = asyncio.iscoroutinefunction(func)
24 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
25 if self.is_async: 25 ↛ 26line 25 didn't jump to line 26 because the condition on line 25 was never true
26 msg = "This tool function is async, use 'await' to call it."
27 raise RuntimeError(msg)
28 return self._func(*args, **kwargs)
30 async def acall(self, *args: P.args, **kwargs: P.kwargs) -> R:
31 if self.is_async: 31 ↛ 33line 31 didn't jump to line 33 because the condition on line 31 was always true
32 return await cast("Callable[P, Awaitable[R]]", self._func)(*args, **kwargs)
33 loop = asyncio.get_running_loop()
34 return await loop.run_in_executor(None, functools.partial(self._func, *args, **kwargs))
37def tool(
38 func: Callable[P, R] | None = None,
39 *,
40 return_immediately: bool = False,
41 require_confirmation: bool = False,
42) -> Callable[[Callable[P, R]], ToolWrapper[P, R]] | ToolWrapper[P, R]:
43 """
44 Decorator: Mark a function as a tool, specifying if it should return immediately and/or require human confirmation.
46 Can be used as @tool or @tool(...)
47 """
49 def decorator(f: Callable[P, R]) -> ToolWrapper[P, R]:
50 return ToolWrapper(f, return_direct=return_immediately, require_confirm=require_confirmation)
52 if func is not None:
53 return decorator(func)
54 return decorator