fix: Tool Error Handling (#222)

* Reworking error handling for pipeline tools/parsing to give more control

* Fixing tests

Author: monoxgas

docs/api/chat.mdx

@@ -1268,13 +1268,20 @@ def __init__(
     self.scorers: list[Scorer[Chat]] = []
     """List of dreadnode scorers to evaluate the generated chat upon completion."""
 
-    self.until_types: list[type[Model]] = []
+    self.until_parsed_as_types: list[type[Model]] = []
+    self.until_parsed_as_catch: bool = True
     self.tools: list[Tool[..., t.Any]] = []
     self.tool_mode: ToolMode = "auto"
     self.inject_tool_prompt = True
     self.add_tool_stop_token = True
-    self.then_callbacks: list[tuple[ThenChatCallback, int, bool]] = []
-    self.map_callbacks: list[tuple[MapChatCallback, int, bool]] = []
+    self.then_callbacks: list[
+        # callback, max_depth, as_task
+        tuple[ThenChatCallback, int, bool]
+    ] = []
+    self.map_callbacks: list[
+        # callback, max_depth, as_task
+        tuple[MapChatCallback, int, bool]
+    ] = []
     self.watch_callbacks: list[WatchChatCallback] = watch_callbacks or []
     self.transforms: list[Transform] = []
 ```
@@ -1726,7 +1733,7 @@ def clone(
     )
     new.chat = (chat or self.chat).clone()
     if not only_messages:
-        new.until_types = self.until_types.copy()
+        new.until_parsed_as_types = self.until_parsed_as_types.copy()
         new.tools = self.tools.copy()
         new.tool_mode = self.tool_mode
         new.metadata = deepcopy(self.metadata)
@@ -3084,6 +3091,7 @@ def transform(
 until_parsed_as(
     *types: type[ModelT],
     max_depth: int = DEFAULT_MAX_DEPTH,
+    catch: bool | None = None,
     attempt_recovery: bool | None = None,
     drop_dialog: bool | None = None,
     max_rounds: int | None = None,
@@ -3105,6 +3113,11 @@ before the generation process completes.
   `DEFAULT_MAX_DEPTH`
   )
   –The maximum depth to re-attempt parsing using recursive pipelines (this is shared between all types).
+* **`catch`**
+  (`bool | None`, default:
+  `None`
+  )
+  –Whether to catch exceptions and return them as messages automatically, otherwise raise them to the pipeline.
 * **`attempt_recovery`**
   (`bool | None`, default:
   `None`
@@ -3132,6 +3145,7 @@ def until_parsed_as(
     self,
     *types: type[ModelT],
     max_depth: int = DEFAULT_MAX_DEPTH,
+    catch: bool | None = None,
     # deprecated
     attempt_recovery: bool | None = None,
     drop_dialog: bool | None = None,
@@ -3144,6 +3158,7 @@ def until_parsed_as(
     Args:
         *types: The type or types of models to wait for.
         max_depth: The maximum depth to re-attempt parsing using recursive pipelines (this is shared between all types).
+        catch: Whether to catch exceptions and return them as messages automatically, otherwise raise them to the pipeline.
         attempt_recovery: deprecated, recovery is always attempted.
         drop_dialog: deprecated, the full dialog is always returned.
         max_rounds: deprecated, use `max_depth` instead.
@@ -3170,7 +3185,8 @@ def until_parsed_as(
             stacklevel=2,
         )
 
-    self.until_types = list(types)
+    self.until_parsed_as_types = list(types)
+    self.until_parsed_as_catch = catch or self.until_parsed_as_catch
 
     max_depth = max_rounds or max_depth
     self.then_callbacks = [
@@ -3197,6 +3213,7 @@ using(
     choice: ToolChoice | None = None,
     max_depth: int = DEFAULT_MAX_DEPTH,
     add_stop_token: bool | None = None,
+    catch: bool | Iterable[type[Exception]] | None = None,
 ) -> ChatPipeline
 ```
 
@@ -3240,6 +3257,13 @@ wrapped in xml tags.
   )
   –When using "xml" tool transforms, use stop tokens to
   immediately process a tool call when observed.
+* **`catch`**
+  (`bool | Iterable[type[Exception]] | None`, default:
+  `None`
+  )
+  –Override the catch setting for all incoming tools, or leave `None` to use the tool's default.
+  By default, catches `json.JSONDecodeError` and `ValidationError`. Set to `{}` to let the pipeline
+  handle all tool exceptions.
 
 **Returns:**
 
@@ -3274,6 +3298,7 @@ def using(
     choice: ToolChoice | None = None,
     max_depth: int = DEFAULT_MAX_DEPTH,
     add_stop_token: bool | None = None,
+    catch: bool | t.Iterable[type[Exception]] | None = None,
 ) -> "ChatPipeline":
     """
     Adds a tool or a sequence of tools to participate in the generation process.
@@ -3294,6 +3319,9 @@ def using(
         max_depth: The maximum depth for recursive tool calls (this is shared between all tools).
         add_stop_token: When using "xml" tool transforms, use stop tokens to
             immediately process a tool call when observed.
+        catch: Override the catch setting for all incoming tools, or leave `None` to use the tool's default.
+            By default, catches `json.JSONDecodeError` and `ValidationError`. Set to `{}` to let the pipeline
+            handle all tool exceptions.
 
     Returns:
         The updated pipeline.
@@ -3332,6 +3360,9 @@ def using(
         else:
             _tools.append(tool)
 
+    if catch is not None:
+        _tools = [tool.with_(catch=catch) for tool in _tools]
+
     existing_names = {tool.name for tool in self.tools}
     new_names = {tool.name for tool in _tools}
     for name in existing_names & new_names:

docs/api/message.mdx

@@ -1175,6 +1175,8 @@ from_model(
     models: Model | Sequence[Model],
     role: Role = "user",
     suffix: str | None = None,
+    tool_call_id: str | None = None,
+    metadata: dict[str, Any] | None = None,
 ) -> Message
 ```
 
@@ -1195,6 +1197,16 @@ Create a Message object from one or more Model objects.
   `None`
   )
   –A suffix to append to the content.
+* **`metadata`**
+  (`dict[str, Any] | None`, default:
+  `None`
+  )
+  –Additional metadata for the Message.
+* **`tool_call_id`**
+  (`str | None`, default:
+  `None`
+  )
+  –The ID of the tool call associated with this message.
 
 **Returns:**
 
@@ -1209,6 +1221,8 @@ def from_model(
     models: Model | t.Sequence[Model],
     role: Role = "user",
     suffix: str | None = None,
+    tool_call_id: str | None = None,
+    metadata: dict[str, t.Any] | None = None,
 ) -> "Message":
     """
     Create a Message object from one or more Model objects.
@@ -1217,6 +1231,8 @@ def from_model(
         models: The Model object(s) to convert to a Message.
         role: The role of the Message.
         suffix: A suffix to append to the content.
+        metadata: Additional metadata for the Message.
+        tool_call_id: The ID of the tool call associated with this message.
 
     Returns:
         The created Message object.
@@ -1240,7 +1256,18 @@ def from_model(
     if suffix is not None:
         content += f"\n{suffix}"
 
-    return cls(role=role, content=content, slices=slices_)
+    # If we building this message from an error, add
+    # the error content to the metadata
+    if isinstance(models, ErrorModel):
+        metadata = {"error": models.content}
+
+    return cls(
+        role=role,
+        content=content,
+        slices=slices_,
+        tool_call_id=tool_call_id,
+        metadata=metadata or {},
+    )
 ```