Add truncation for messages and tools. Version to 3.0.0.

Author: monoxgas

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rigging"
-version = "3.0.0-rc.4"
+version = "3.0.0"
 description = "LLM Interaction Framework"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 license = "MIT"

rigging/message.py

@@ -30,7 +30,7 @@
 from rigging.model import Model, ModelT
 from rigging.parsing import try_parse_many
 from rigging.tool.api import ApiToolCall
-from rigging.util import AudioFormat, identify_audio_format, truncate_string
+from rigging.util import AudioFormat, identify_audio_format, shorten_string, truncate_string
 
 Role = t.Literal["system", "user", "assistant", "tool"]
 """The role of a message. Can be 'system', 'user', 'assistant', or 'tool'."""
@@ -111,7 +111,7 @@ class ImageUrl(BaseModel):
     """Cache control entry for prompt caching."""
 
     def __str__(self) -> str:
-        return f"<ContentImageUrl url='{truncate_string(self.image_url.url, 50)}'>"
+        return f"<ContentImageUrl url='{shorten_string(self.image_url.url, 50)}'>"
 
     @classmethod
     def from_file(
@@ -241,7 +241,7 @@ def __str__(self) -> str:
         return (
             f"<ContentAudioInput format='{self.input_audio.format}' "
             f"transcript='{self.input_audio.transcript}' "
-            f"data='{truncate_string(self.input_audio.data, 50)}'>"
+            f"data='{shorten_string(self.input_audio.data, 50)}'>"
         )
 
     @classmethod
@@ -655,6 +655,8 @@ def clone(self) -> "Message":
             self.role,
             copy.deepcopy(self.content_parts),
             parts=copy.deepcopy(self.parts),
+            tool_calls=copy.deepcopy(self.tool_calls),
+            tool_call_id=self.tool_call_id,
         )
 
     def cache(self, cache_control: dict[str, str] | bool = True) -> "Message":  # noqa: FBT002
@@ -705,6 +707,20 @@ def apply(self, **kwargs: str) -> "Message":
         new.content = template.safe_substitute(**kwargs)
         return new
 
+    def truncate(self, max_length: int, suffix: str = "\n[truncated]") -> "Message":
+        """
+        Truncates the message content to a maximum length.
+
+        Args:
+            max_length: The maximum length of the message content.
+
+        Returns:
+            The truncated message.
+        """
+        new = self.clone()
+        new.content = truncate_string(new.content, max_length, suf=suffix)
+        return new
+
     def strip(
         self,
         model_type: type[Model],

rigging/tool/base.py

@@ -73,6 +73,8 @@ class Tool(t.Generic[P, R]):
     - `True`: Catch all exceptions.
     - `list[type[Exception]]`: Catch only the specified exceptions.
     """
+    truncate: int | None = None
+    """If set, the maximum number of characters to truncate any tool output to."""
 
     _signature: inspect.Signature | None = field(default=None, init=False, repr=False)
     _type_adapter: TypeAdapter[t.Any] | None = field(
@@ -100,6 +102,7 @@ def from_callable(
         name: str | None = None,
         description: str | None = None,
         catch: bool | t.Iterable[type[Exception]] = False,
+        truncate: int | None = None,
     ) -> te.Self:
         from rigging.prompt import Prompt
 
@@ -194,6 +197,7 @@ def empty_func(*args, **kwargs):  # type: ignore [no-untyped-def] # noqa: ARG001
             parameters_schema=schema,
             fn=fn,
             catch=catch if isinstance(catch, bool) else set(catch),
+            truncate=truncate,
         )
 
         self._signature = signature
@@ -367,6 +371,9 @@ async def handle_tool_call(
         else:
             message.content_parts = [ContentText(text=str(result))]
 
+        if self.truncate:
+            message = message.truncate(self.truncate)
+
         # If this is a native tool call, we should wrap up our
         # result in a NativeToolResult object to provide clarity to the
         # generator. Otherwise we can rely on the `tool` role and associated
@@ -402,6 +409,7 @@ def tool(
     name: str | None = None,
     description: str | None = None,
     catch: bool | t.Iterable[type[Exception]] = False,
+    truncate: int | None = None,
 ) -> t.Callable[[t.Callable[P, R]], Tool[P, R]]:
     ...
 
@@ -421,6 +429,7 @@ def tool(
     name: str | None = None,
     description: str | None = None,
     catch: bool | t.Iterable[type[Exception]] = False,
+    truncate: int | None = None,
 ) -> t.Callable[[t.Callable[P, R]], Tool[P, R]] | Tool[P, R]:
     """
     Decorator for creating a Tool, useful for overriding a name or description.
@@ -433,6 +442,7 @@ def tool(
             - `False`: Do not catch exceptions.
             - `True`: Catch all exceptions.
             - `list[type[Exception]]`: Catch only the specified exceptions.
+        truncate: If set, the maximum number of characters to truncate any tool output to.
 
     Returns:
         The decorated Tool object.
@@ -453,7 +463,13 @@ def make_tool(func: t.Callable[..., t.Any]) -> Tool[P, R]:
                 stacklevel=3,
             )
 
-        return Tool.from_callable(func, name=name, description=description, catch=catch)
+        return Tool.from_callable(
+            func,
+            name=name,
+            description=description,
+            catch=catch,
+            truncate=truncate,
+        )
 
     if func is not None:
         return make_tool(func)
@@ -496,6 +512,7 @@ def tool_method(
     name: str | None = None,
     description: str | None = None,
     catch: bool | t.Iterable[type[Exception]] = False,
+    truncate: int | None = None,
 ) -> t.Callable[[t.Callable[t.Concatenate[t.Any, P], R]], ToolMethod[P, R]]:
     ...
 
@@ -515,6 +532,7 @@ def tool_method(
     name: str | None = None,
     description: str | None = None,
     catch: bool | t.Iterable[type[Exception]] = False,
+    truncate: int | None = None,
 ) -> t.Callable[[t.Callable[t.Concatenate[t.Any, P], R]], ToolMethod[P, R]] | ToolMethod[P, R]:
     """
     Decorator for creating a Tool from a class method.
@@ -570,6 +588,7 @@ def wrapper(self: t.Any, *args: P.args, **kwargs: P.kwargs) -> R:
             name=name,
             description=description,
             catch=catch,
+            truncate=truncate,
         )
 
     if func is not None:

rigging/util.py

@@ -150,16 +150,31 @@ def get_qualified_name(obj: t.Callable[..., t.Any]) -> str:
 # Formatting
 
 
-def truncate_string(content: str, max_length: int, *, sep: str = "...") -> str:
-    """Return a string at most max_length characters long."""
+def shorten_string(content: str, max_length: int, *, sep: str = "...") -> str:
+    """Return a string at most max_length characters long by removing the middle of the string."""
     if len(content) <= max_length:
         return content
 
     remaining = max_length - len(sep)
+    if remaining <= 0:
+        return sep
+
     middle = remaining // 2
     return content[:middle] + sep + content[-middle:]
 
 
+def truncate_string(content: str, max_length: int, *, suf: str = "...") -> str:
+    """Return a string at most max_length characters long by removing the end of the string."""
+    if len(content) <= max_length:
+        return content
+
+    remaining = max_length - len(suf)
+    if remaining <= 0:
+        return suf
+
+    return content[:remaining] + suf
+
+
 # List utilities