File size: 4,766 Bytes
4dbe519
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
64462d2
 
 
4dbe519
 
 
 
64462d2
 
 
 
 
4dbe519
 
 
 
 
 
 
 
 
 
 
 
64462d2
 
 
 
 
 
4dbe519
 
 
 
 
 
 
 
64462d2
 
 
4dbe519
 
 
 
 
 
 
 
 
 
 
64462d2
 
 
4dbe519
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
64462d2
 
 
4dbe519
 
 
 
 
 
 
 
 
 
 
 
64462d2
 
 
 
 
4dbe519
 
 
 
 
 
 
 
64462d2
 
 
4dbe519
 
 
 
64462d2
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
"""Tool system for the agent framework."""

from abc import ABC, abstractmethod
from typing import Dict, Any, Callable
import inspect
from .models import ExecutionContext
from .utils import function_to_input_schema, format_tool_definition


class BaseTool(ABC):
    """Abstract base class for all tools."""
    
    def __init__(
        self, 
        name: str = None, 
        description: str = None, 
        tool_definition: Dict[str, Any] = None,
        # Confirmation support
        requires_confirmation: bool = False,
        confirmation_message_template: str = None
    ):
        self.name = name or self.__class__.__name__
        self.description = description or self.__doc__ or ""
        self._tool_definition = tool_definition
        self.requires_confirmation = requires_confirmation
        self.confirmation_message_template = confirmation_message_template or (
            "The agent wants to execute '{name}' with arguments: {arguments}. "
            "Do you approve?"
        )
    
    @property
    def tool_definition(self) -> Dict[str, Any] | None:
        return self._tool_definition
    
    @abstractmethod
    async def execute(self, context: ExecutionContext, **kwargs) -> Any:
        pass
    
    async def __call__(self, context: ExecutionContext, **kwargs) -> Any:
        return await self.execute(context, **kwargs)

    def get_confirmation_message(self, arguments: dict[str, Any]) -> str:
        """Generate a confirmation message for this tool call."""
        return self.confirmation_message_template.format(
            name=self.name,
            arguments=arguments
        )
class FunctionTool(BaseTool):
    """Wraps a Python function as a BaseTool."""
    
    def __init__(
        self, 
        func: Callable, 
        name: str = None, 
        description: str = None,
        tool_definition: Dict[str, Any] = None,
        requires_confirmation: bool = False,
        confirmation_message_template: str = None
    ):
        self.func = func
        self.needs_context = 'context' in inspect.signature(func).parameters
        
        self.name = name or func.__name__
        self.description = description or (func.__doc__ or "").strip()
        tool_definition = tool_definition or self._generate_definition()
        
        super().__init__(
            name=self.name, 
            description=self.description, 
            tool_definition=tool_definition,
            requires_confirmation=requires_confirmation,
            confirmation_message_template=confirmation_message_template
        )
    
    async def execute(self, context: ExecutionContext = None, **kwargs) -> Any:
        """Execute the wrapped function.
        
        Context is only required if the wrapped function has a 'context' parameter.
        """
        if self.needs_context:
            if context is None:
                raise ValueError(
                    f"Tool '{self.name}' requires a context parameter. "
                    f"Please provide an ExecutionContext instance."
                )
            result = self.func(context=context, **kwargs)
        else:
            result = self.func(**kwargs)
        
        # Handle both sync and async functions
        if inspect.iscoroutine(result):
            return await result
        return result
    
    def _generate_definition(self) -> Dict[str, Any]:
        """Generate tool definition from function signature."""
        parameters = function_to_input_schema(self.func)
        return format_tool_definition(self.name, self.description, parameters)


def tool(
    func: Callable = None,
    *,
    name: str = None,
    description: str = None,
    tool_definition: Dict[str, Any] = None,
    requires_confirmation: bool = False,
    confirmation_message: str = None
):
    """Decorator to convert a function into a FunctionTool.
    
    Usage:
        @tool
        def my_function(x: int) -> int:
            return x * 2
        
        # Or with parameters:
        @tool(name="custom_name", description="Custom description")
        def my_function(x: int) -> int:
            return x * 2
        
        # With confirmation:
        @tool(requires_confirmation=True, confirmation_message="Delete file?")
        def delete_file(filename: str) -> str:
            ...
    """
    from typing import Union
    
    def decorator(f: Callable) -> FunctionTool:
        return FunctionTool(
            func=f,
            name=name,
            description=description,
            tool_definition=tool_definition,
            requires_confirmation=requires_confirmation,
            confirmation_message_template=confirmation_message
        )
    
    if func is not None:
        return decorator(func)
    return decorator