Add parameter to tools. Make tools callable directly.

Author: monoxgas

rigging/chat.py

@@ -364,7 +364,7 @@ def inject_system_content(self, content: str) -> "Chat":
             self.messages[0].content += "\n\n" + content
         return self
 
-    def inject_tool_prompt(self, tools: t.Sequence[Tool], mode: ToolMode) -> "Chat":
+    def inject_tool_prompt(self, tools: t.Sequence[Tool[..., t.Any]], mode: ToolMode) -> "Chat":
         """
         Injects a default tool use prompt into the system prompt.
 
@@ -695,7 +695,7 @@ def __init__(
         """How to handle cache_control entries on messages."""
 
         self.until_types: list[type[Model]] = []
-        self.tools: list[Tool] = []
+        self.tools: list[Tool[..., t.Any]] = []
         self.tool_mode: ToolMode = "auto"
         self.api_tool_choice: ApiToolChoice | None = None
         self.inject_tool_prompt = True
@@ -1066,7 +1066,9 @@ def wrap(self, func: t.Callable[[CallableT], CallableT]) -> "ChatPipeline":
 
     def using(
         self,
-        *tools: Tool | t.Callable[..., t.Any] | t.Sequence[Tool | t.Callable[..., t.Any]],
+        *tools: Tool[..., t.Any]
+        | t.Callable[..., t.Any]
+        | t.Sequence[Tool[..., t.Any] | t.Callable[..., t.Any]],
         mode: ToolMode | None = None,
         choice: ApiToolChoice | None = None,
         max_depth: int = DEFAULT_MAX_DEPTH,

rigging/prompt.py

@@ -460,7 +460,7 @@ class Prompt(t.Generic[P, R]):
 
     params: GenerateParams | None = None
     """The parameters to be used when generating chats for this prompt."""
-    tools: list[Tool] = dataclasses.field(default_factory=list)
+    tools: list[Tool[..., t.Any]] = dataclasses.field(default_factory=list)
     """The API tools to be made available when generating chats for this prompt."""
     system_prompt: str | None = None
     """A system prompt fragment to be injected into the messages before generation."""
@@ -1078,7 +1078,7 @@ def prompt(
     pipeline: ChatPipeline | None = None,
     generator: Generator | None = None,
     generator_id: str | None = None,
-    tools: list[Tool | t.Callable[..., t.Any]] | None = None,
+    tools: list[Tool[..., t.Any] | t.Callable[..., t.Any]] | None = None,
     system_prompt: str | None = None,
 ) -> t.Callable[[t.Callable[P, t.Coroutine[t.Any, t.Any, R]] | t.Callable[P, R]], Prompt[P, R]]:
     ...
@@ -1092,7 +1092,7 @@ def prompt(
     pipeline: ChatPipeline | None = None,
     generator: Generator | None = None,
     generator_id: str | None = None,
-    tools: list[Tool | t.Callable[..., t.Any]] | None = None,
+    tools: list[Tool[..., t.Any] | t.Callable[..., t.Any]] | None = None,
     system_prompt: str | None = None,
 ) -> Prompt[P, R]:
     ...
@@ -1106,7 +1106,7 @@ def prompt(
     pipeline: ChatPipeline | None = None,
     generator: Generator | None = None,
     generator_id: str | None = None,
-    tools: list[Tool | t.Callable[..., t.Any]] | None = None,
+    tools: list[Tool[..., t.Any] | t.Callable[..., t.Any]] | None = None,
     system_prompt: str | None = None,
 ) -> Prompt[P, R]:
     ...
@@ -1119,7 +1119,7 @@ def prompt(
     pipeline: ChatPipeline | None = None,
     generator: Generator | None = None,
     generator_id: str | None = None,
-    tools: list[Tool | t.Callable[..., t.Any]] | None = None,
+    tools: list[Tool[..., t.Any] | t.Callable[..., t.Any]] | None = None,
     system_prompt: str | None = None,
 ) -> (
     t.Callable[[t.Callable[P, t.Coroutine[t.Any, t.Any, R]] | t.Callable[P, R]], Prompt[P, R]]

rigging/tool/base.py

@@ -27,6 +27,9 @@
 if t.TYPE_CHECKING:
     from rigging.message import Message
 
+P = t.ParamSpec("P")
+R = t.TypeVar("R")
+
 ToolMode = t.Literal["auto", "api", "xml", "json-in-xml"]
 """
 How tool calls are handled.
@@ -39,7 +42,7 @@
 
 
 @dataclass
-class Tool:
+class Tool(t.Generic[P, R]):
     """Base class for representing a tool to a generator."""
 
     name: str
@@ -48,8 +51,16 @@ class Tool:
     """A description of the tool."""
     parameters_schema: dict[str, t.Any]
     """The JSON schema for the tool's parameters."""
-    fn: t.Callable[..., t.Any]
+    fn: t.Callable[P, R]
     """The function to call."""
+    catch: bool | set[type[Exception]] = False
+    """
+    Whether to catch exceptions and return them as messages.
+
+    - `False`: Do not catch exceptions.
+    - `True`: Catch all exceptions.
+    - `list[type[Exception]]`: Catch only the specified exceptions.
+    """
 
     _signature: inspect.Signature | None = field(default=None, init=False, repr=False)
     _type_adapter: TypeAdapter[t.Any] | None = field(default=None, init=False, repr=False)
@@ -68,11 +79,12 @@ class Tool:
     @classmethod
     def from_callable(
         cls,
-        fn: t.Callable[..., t.Any],
+        fn: t.Callable[P, R],
         *,
         name: str | None = None,
         description: str | None = None,
-    ) -> "Tool":
+        catch: bool | t.Iterable[type[Exception]] = False,
+    ) -> "Tool[P, R]":
         from rigging.prompt import Prompt
 
         fn_for_signature = fn
@@ -84,7 +96,7 @@ def from_callable(
 
         if isinstance(fn, Prompt):
             fn_for_signature = fn.func  # type: ignore [assignment]
-            fn = fn.run
+            fn = fn.run  # type: ignore [assignment]
 
         # In the case that we are recieving a bound function which is tracking
         # an originating prompt, unwrap the prompt and use it's function for
@@ -162,9 +174,11 @@ def empty_func(*args, **kwargs):  # type: ignore [no-untyped-def] # noqa: ARG001
             description=description or fn_for_signature.__doc__ or "",
             parameters_schema=schema,
             fn=fn,
+            catch=catch if isinstance(catch, bool) else set(catch),
         )
 
         self._signature = signature
+        self.__signature__ = signature  # type: ignore [attr-defined]
 
         # For handling API calls, we'll use the type adapter to validate
         # the arguments before calling the function
@@ -226,7 +240,7 @@ async def handle_tool_call(
             tool_call: The tool call to handle.
 
         Returns:
-            The message to send back to the generator or None if tool calling should not proceed.
+            The message to send back to the generator or `None` if iterative tool calling should not proceed any further.
         """
 
         from rigging.message import ContentText, ContentTypes, Message
@@ -248,12 +262,12 @@ async def handle_tool_call(
 
             # Load + validate arguments
 
-            args: dict[str, t.Any]
+            kwargs: dict[str, t.Any]
             if isinstance(tool_call, ApiToolCall | JsonInXmlToolCall):
-                args = json.loads(tool_call_parameters)
+                kwargs = json.loads(tool_call_parameters)
 
                 if self._type_adapter is not None:
-                    args = self._type_adapter.validate_python(args)
+                    kwargs = self._type_adapter.validate_python(kwargs)
 
             elif isinstance(tool_call, XmlToolCall):
                 parsed = self.model.from_text(
@@ -268,18 +282,26 @@ async def handle_tool_call(
                 # argument object instances. We'll just flatten the
                 # model into a dictionary for the function call.
 
-                args = {
+                kwargs = {
                     field_name: getattr(parameters, field_name, None)
                     for field_name in self.model.model_fields
                 }
 
-            span.set_attribute("arguments", args)
+            span.set_attribute("arguments", kwargs)
 
             # Call the function
 
-            result = self.fn(**args)
-            if inspect.isawaitable(result):
-                result = await result
+            try:
+                result: t.Any = self.fn(**kwargs)  # type: ignore [call-arg]
+                if inspect.isawaitable(result):
+                    result = await result
+            except Exception as e:  # noqa: BLE001
+                if self.catch is True or (
+                    not isinstance(self.catch, bool) and isinstance(e, tuple(self.catch))
+                ):
+                    result = f'<error type="{e.__class__.__name__}">:{e}</error>'
+                else:
+                    raise
 
             span.set_attribute("result", result)
 
@@ -331,6 +353,9 @@ async def handle_tool_call(
 
         return message, True
 
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        return self.fn(*args, **kwargs)
+
 
 # Decorator
 
@@ -342,42 +367,49 @@ def tool(
     *,
     name: str | None = None,
     description: str | None = None,
-) -> t.Callable[[t.Callable[..., t.Any]], Tool]:
+    catch: bool | t.Iterable[type[Exception]] = False,
+) -> t.Callable[[t.Callable[P, R]], Tool[P, R]]:
     ...
 
 
 @t.overload
 def tool(
-    func: t.Callable[..., t.Any],
+    func: t.Callable[P, R],
     /,
     *,
     name: str | None = None,
     description: str | None = None,
-) -> Tool:
+    catch: bool | t.Iterable[type[Exception]] = False,
+) -> Tool[P, R]:
     ...
 
 
 def tool(
-    func: t.Callable[..., t.Any] | None = None,
+    func: t.Callable[P, R] | None = None,
     /,
     *,
     name: str | None = None,
     description: str | None = None,
-) -> t.Callable[[t.Callable[..., t.Any]], Tool] | Tool:
+    catch: bool | t.Iterable[type[Exception]] = False,
+) -> t.Callable[[t.Callable[P, R]], Tool[P, R]] | Tool[P, R]:
     """
     Decorator for creating a Tool, useful for overriding a name or description.
 
     Args:
         func: The function to wrap.
         name: The name of the tool.
         description: The description of the tool.
+        catch: Whether to catch exceptions and return them as messages.
+            - `False`: Do not catch exceptions.
+            - `True`: Catch all exceptions.
+            - `list[type[Exception]]`: Catch only the specified exceptions.
 
     Returns:
         The decorated Tool object.
     """
 
-    def make_tool(func: t.Callable[..., t.Any]) -> Tool:
-        return Tool.from_callable(func, name=name, description=description)
+    def make_tool(func: t.Callable[..., t.Any]) -> Tool[P, R]:
+        return Tool.from_callable(func, name=name, description=description, catch=catch)
 
     if func is not None:
         return make_tool(func)

rigging/tool/mcp.py

@@ -69,7 +69,7 @@ class MCPClient:
     """The transport to use"""
     connection: StdioConnection | SSEConnection
     """Connection configuration"""
-    tools: list[Tool]
+    tools: list[Tool[..., t.Any]]
     """A list of tools available on the server"""
 
     def __init__(self, transport: Transport, connection: StdioConnection | SSEConnection) -> None:

rigging/tool/robopages.py

@@ -44,7 +44,7 @@ async def execute_on_server(**kwargs: t.Any) -> t.Any:
     return execute_on_server
 
 
-def robopages(url: str, *, name_filter: str | None = None) -> list[Tool]:
+def robopages(url: str, *, name_filter: str | None = None) -> list[Tool[..., t.Any]]:
     """
     Create a list of tools from a Robopages server.
 
@@ -83,7 +83,7 @@ def robopages(url: str, *, name_filter: str | None = None) -> list[Tool]:
 
     logger.info(f"Fetched {len(tool_definitions)} functions from Robopages ({url})")
 
-    tools: list[Tool] = []
+    tools: list[Tool[..., t.Any]] = []
     for definition in tool_definitions:
         function = definition.function