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

1""" 

2decorators.py 

3 

4This module contains custom decorators for the funcall package. 

5""" 

6 

7import asyncio 

8import functools 

9from collections.abc import Awaitable, Callable 

10from typing import Generic, ParamSpec, TypeVar, cast 

11 

12P = ParamSpec("P") 

13R = TypeVar("R") 

14 

15 

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) 

23 

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) 

29 

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)) 

35 

36 

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. 

45 

46 Can be used as @tool or @tool(...) 

47 """ 

48 

49 def decorator(f: Callable[P, R]) -> ToolWrapper[P, R]: 

50 return ToolWrapper(f, return_direct=return_immediately, require_confirm=require_confirmation) 

51 

52 if func is not None: 

53 return decorator(func) 

54 return decorator