indent 0.1.26__py3-none-any.whl

exponent/core/remote_execution/truncation.py

@@ -0,0 +1,296 @@
+"""Generalized truncation framework for tool results."""
+
+from abc import ABC, abstractmethod
+from typing import Any, TypeVar, cast
+
+from msgspec.structs import replace
+
+from exponent.core.remote_execution.cli_rpc_types import (
+    BashToolResult,
+    ErrorToolResult,
+    GlobToolResult,
+    GrepToolResult,
+    ListToolResult,
+    ReadToolResult,
+    ToolResult,
+    WriteToolResult,
+)
+from exponent.core.remote_execution.utils import truncate_output
+
+DEFAULT_CHARACTER_LIMIT = 50_000
+DEFAULT_LIST_ITEM_LIMIT = 1000
+DEFAULT_LIST_PREVIEW_ITEMS = 10
+
+
+class TruncationStrategy(ABC):
+    @abstractmethod
+    def should_truncate(self, result: ToolResult) -> bool:
+        pass
+
+    @abstractmethod
+    def truncate(self, result: ToolResult) -> ToolResult:
+        pass
+
+
+class StringFieldTruncation(TruncationStrategy):
+    def __init__(
+        self,
+        field_name: str,
+        character_limit: int = DEFAULT_CHARACTER_LIMIT,
+    ):
+        self.field_name = field_name
+        self.character_limit = character_limit
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if hasattr(result, self.field_name):
+            value = getattr(result, self.field_name)
+            if isinstance(value, str):
+                return len(value) > self.character_limit
+        return False
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        value = getattr(result, self.field_name)
+        if not isinstance(value, str):
+            return result
+
+        truncated_value, was_truncated = truncate_output(value, self.character_limit)
+
+        updates: dict[str, Any] = {self.field_name: truncated_value}
+        if hasattr(result, "truncated") and was_truncated:
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+class ListFieldTruncation(TruncationStrategy):
+    def __init__(
+        self,
+        field_name: str,
+        item_limit: int = DEFAULT_LIST_ITEM_LIMIT,
+        preview_items: int = DEFAULT_LIST_PREVIEW_ITEMS,
+    ):
+        self.field_name = field_name
+        self.item_limit = item_limit
+        self.preview_items = preview_items
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if hasattr(result, self.field_name):
+            value = getattr(result, self.field_name)
+            if isinstance(value, list):
+                return len(value) > self.item_limit
+        return False
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        value = getattr(result, self.field_name)
+        if not isinstance(value, list):
+            return result
+
+        total_items = len(value)
+        if total_items <= self.item_limit:
+            return result
+
+        truncated_count = max(0, total_items - 2 * self.preview_items)
+        truncated_list = (
+            value[: self.preview_items]
+            + [f"... {truncated_count} items truncated ..."]
+            + value[-self.preview_items :]
+        )
+
+        updates: dict[str, Any] = {self.field_name: truncated_list}
+        if hasattr(result, "truncated"):
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+class CompositeTruncation(TruncationStrategy):
+    def __init__(self, strategies: list[TruncationStrategy]):
+        self.strategies = strategies
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        return any(strategy.should_truncate(result) for strategy in self.strategies)
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        for strategy in self.strategies:
+            if strategy.should_truncate(result):
+                result = strategy.truncate(result)
+        return result
+
+
+class TailTruncation(TruncationStrategy):
+    """Truncation strategy that keeps the end of the output (tail) instead of the beginning."""
+
+    def __init__(
+        self,
+        field_name: str,
+        character_limit: int = DEFAULT_CHARACTER_LIMIT,
+    ):
+        self.field_name = field_name
+        self.character_limit = character_limit
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if hasattr(result, self.field_name):
+            value = getattr(result, self.field_name)
+            if isinstance(value, str):
+                return len(value) > self.character_limit
+        return False
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        value = getattr(result, self.field_name)
+        if not isinstance(value, str):
+            return result
+
+        if len(value) <= self.character_limit:
+            return result
+
+        # Keep the last character_limit characters
+        truncated_value = value[-self.character_limit :]
+
+        # Try to start at a newline if possible for cleaner output
+        newline_pos = truncated_value.find("\n")
+        if (
+            newline_pos != -1 and newline_pos < 1000
+        ):  # Only adjust if newline is reasonably close to start
+            truncated_value = truncated_value[newline_pos + 1 :]
+
+        # Add truncation indicator at the beginning
+        truncation_msg = f"... (output truncated, showing last {len(truncated_value)} characters) ...\n"
+        truncated_value = truncation_msg + truncated_value
+
+        updates: dict[str, Any] = {self.field_name: truncated_value}
+        if hasattr(result, "truncated"):
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+class NoOpTruncation(TruncationStrategy):
+    def should_truncate(self, result: ToolResult) -> bool:
+        return False
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        return result
+
+
+class StringListTruncation(TruncationStrategy):
+    """Truncation for lists of strings that limits both number of items and individual string length."""
+
+    def __init__(
+        self,
+        field_name: str,
+        max_items: int = DEFAULT_LIST_ITEM_LIMIT,
+        preview_items: int = DEFAULT_LIST_PREVIEW_ITEMS,
+        max_item_length: int = 1000,
+    ):
+        self.field_name = field_name
+        self.max_items = max_items
+        self.preview_items = preview_items
+        self.max_item_length = max_item_length
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if not hasattr(result, self.field_name):
+            return False
+
+        items = getattr(result, self.field_name)
+        if not isinstance(items, list):
+            return False
+
+        # Check if we need to truncate number of items
+        if len(items) > self.max_items:
+            return True
+
+        # Check if any individual item is too long
+        for item in items:
+            if isinstance(item, str) and len(item) > self.max_item_length:
+                return True
+            # Handle dict items (e.g., with metadata like file path and line number)
+            elif isinstance(item, dict) and "content" in item:
+                if len(item["content"]) > self.max_item_length:
+                    return True
+
+        return False
+
+    def _truncate_item_content(
+        self, item: str | dict[str, Any]
+    ) -> str | dict[str, Any]:
+        """Truncate an individual item's content."""
+        if isinstance(item, str):
+            if len(item) <= self.max_item_length:
+                return item
+            # Truncate string item
+            truncated, _ = truncate_output(item, self.max_item_length)
+            return truncated
+        elif isinstance(item, dict) and "content" in item:
+            # Handle dict-style items (e.g., with metadata like file path and line number)
+            if len(item["content"]) <= self.max_item_length:
+                return item
+            truncated_content, _ = truncate_output(
+                item["content"], self.max_item_length
+            )
+            return {**item, "content": truncated_content}
+        else:
+            return item
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        items = getattr(result, self.field_name)
+        if not isinstance(items, list):
+            return result
+
+        # First, truncate individual item contents
+        truncated_items = [self._truncate_item_content(item) for item in items]
+
+        # Then, limit the number of items if needed
+        total_items = len(truncated_items)
+        if total_items > self.max_items:
+            truncated_count = max(0, total_items - 2 * self.preview_items)
+            final_items = (
+                truncated_items[: self.preview_items]
+                + [f"... {truncated_count} items truncated ..."]
+                + truncated_items[-self.preview_items :]
+            )
+        else:
+            final_items = truncated_items
+
+        updates: dict[str, Any] = {self.field_name: final_items}
+        if hasattr(result, "truncated"):
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+TRUNCATION_REGISTRY: dict[type[ToolResult], TruncationStrategy] = {
+    ReadToolResult: StringFieldTruncation("content"),
+    WriteToolResult: StringFieldTruncation("message"),
+    BashToolResult: TailTruncation("shell_output"),
+    GrepToolResult: StringListTruncation("matches"),
+    GlobToolResult: StringListTruncation("filenames", max_item_length=4096),
+    ListToolResult: StringListTruncation("files", max_item_length=4096),
+}
+
+
+T = TypeVar("T", bound=ToolResult)
+
+
+def truncate_tool_result(result: T) -> T:
+    if isinstance(result, ErrorToolResult):
+        return cast(T, result)
+
+    result_type = type(result)
+    if result_type in TRUNCATION_REGISTRY:
+        strategy = TRUNCATION_REGISTRY[result_type]
+        if strategy.should_truncate(result):
+            return cast(T, strategy.truncate(result))
+
+    return result