Coverage for src\funcall\decorators.py: 0%

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], *, return_immediately: bool, require_confirmation: bool) -> None: 

18 functools.update_wrapper(self, func) 

19 self._func = func 

20 self.return_immediately = return_immediately 

21 self.require_confirmation = require_confirmation 

22 self._is_async = asyncio.iscoroutinefunction(func) 

23 

24 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: 

25 if self._is_async: 

26 msg = "This tool function is async, use 'await acall' 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: 

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(*, return_immediately: bool = False, require_confirmation: bool = False) -> Callable[[Callable[P, R]], ToolWrapper[P, R]]: 

38 """ 

39 Decorator: Mark a function as a tool, specifying if it should return immediately and/or require human confirmation. 

40 

41 Args: 

42 return_immediately (bool): Whether the function should return immediately. 

43 require_confirmation (bool): Whether the function requires human confirmation. 

44 """ 

45 

46 def decorator(func: Callable[P, R]) -> ToolWrapper[P, R]: 

47 return ToolWrapper(func, return_immediately=return_immediately, require_confirmation=require_confirmation) 

48 

49 return decorator