python-codex 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

pycodex/utils/truncation.py

@@ -1,9 +1,10 @@
 """Shared truncation helpers for model-visible tool output."""
 
 import math
+import typing
+from dataclasses import replace
 
 from ..protocol import JSONValue, ToolResult
-import typing
 
 DEFAULT_MAX_OUTPUT_TOKENS = 10_000
 TRUNCATION_SERIALIZATION_BUDGET_MULTIPLIER = 1.2
@@ -13,7 +14,7 @@ HISTORY_TOOL_OUTPUT_TOKENS = int(
 APPROX_BYTES_PER_TOKEN = 4
 
 
-def approx_token_count(text: 'str') -> 'int':
+def approx_token_count(text: "str") -> "int":
     """Estimate token count using the upstream Codex 4-bytes-per-token rule."""
     if not text:
         return 0
@@ -24,7 +25,7 @@ def approx_token_count(text: 'str') -> 'int':
     )
 
 
-def formatted_truncate_text(text: 'str', max_tokens: 'int') -> 'str':
+def formatted_truncate_text(text: "str", max_tokens: "int") -> "str":
     """Format a direct tool response with line count plus middle truncation."""
     byte_budget = _approx_bytes_for_tokens(max_tokens)
     if len(text.encode("utf-8")) <= byte_budget:
@@ -35,8 +36,8 @@ def formatted_truncate_text(text: 'str', max_tokens: 'int') -> 'str':
 
 
 def truncate_tool_result_for_history(
-    result: 'ToolResult',
-) -> 'ToolResult':
+    result: "ToolResult",
+) -> "ToolResult":
     """Truncate model-visible ToolResult content before storing it in history."""
     if result.content_items is not None:
         truncated_content_items = _truncate_content_items(
@@ -45,41 +46,16 @@ def truncate_tool_result_for_history(
         )
         if truncated_content_items == result.content_items:
             return result
-        return ToolResult(
-            call_id=result.call_id,
-            name=result.name,
-            output=result.output,
-            content_items=truncated_content_items,
-            success=result.success,
-            is_error=result.is_error,
-            tool_type=result.tool_type,
-        )
+        return replace(result, content_items=truncated_content_items)
 
     output_text = _tool_output_text(result.output)
     truncated_output = _truncate_text(output_text, HISTORY_TOOL_OUTPUT_TOKENS)
     if truncated_output == output_text:
         return result
-    return ToolResult(
-        call_id=result.call_id,
-        name=result.name,
-        output=truncated_output,
-        success=result.success,
-        is_error=result.is_error,
-        tool_type=result.tool_type,
-    )
-
-
-def truncate_tool_results_for_history(
-    results: 'typing.Iterable[ToolResult]',
-) -> 'typing.List[ToolResult]':
-    """Apply history-layer truncation to a batch of completed tool results."""
-    return [
-        truncate_tool_result_for_history(result)
-        for result in results
-    ]
+    return replace(result, output=truncated_output)
 
 
-def _tool_output_text(output: 'JSONValue') -> 'str':
+def _tool_output_text(output: "JSONValue") -> "str":
     if isinstance(output, str):
         return output
 
@@ -93,10 +69,10 @@ def _tool_output_text(output: 'JSONValue') -> 'str':
 
 
 def _truncate_content_items(
-    content_items: 'typing.Tuple[typing.Dict[str, typing.Any], ...]',
-    token_limit: 'int',
-) -> 'typing.Tuple[typing.Dict[str, typing.Any], ...]':
-    output: 'typing.List[typing.Dict[str, typing.Any]]' = []
+    content_items: "typing.Tuple[typing.Dict[str, typing.Any], ...]",
+    token_limit: "int",
+) -> "typing.Tuple[typing.Dict[str, typing.Any], ...]":
+    output: "typing.List[typing.Dict[str, typing.Any]]" = []
     remaining_budget = token_limit
     omitted_text_items = 0
 
@@ -135,26 +111,26 @@ def _truncate_content_items(
     return tuple(output)
 
 
-def _approx_tokens_from_byte_count(byte_count: 'int') -> 'int':
+def _approx_tokens_from_byte_count(byte_count: "int") -> "int":
     if byte_count <= 0:
         return 0
     return (byte_count + APPROX_BYTES_PER_TOKEN - 1) // APPROX_BYTES_PER_TOKEN
 
 
-def _approx_bytes_for_tokens(token_count: 'int') -> 'int':
+def _approx_bytes_for_tokens(token_count: "int") -> "int":
     return max(token_count, 0) * APPROX_BYTES_PER_TOKEN
 
 
-def _split_budget(byte_budget: 'int') -> 'typing.Tuple[int, int]':
+def _split_budget(byte_budget: "int") -> "typing.Tuple[int, int]":
     left_budget = byte_budget // 2
     return left_budget, byte_budget - left_budget
 
 
 def _split_string(
-    text: 'str',
-    beginning_bytes: 'int',
-    end_bytes: 'int',
-) -> 'typing.Tuple[str, str]':
+    text: "str",
+    beginning_bytes: "int",
+    end_bytes: "int",
+) -> "typing.Tuple[str, str]":
     if not text:
         return "", ""
 
@@ -187,7 +163,7 @@ def _split_string(
     return text[:prefix_end], text[suffix_start:]
 
 
-def _truncate_text(text: 'str', max_tokens: 'int') -> 'str':
+def _truncate_text(text: "str", max_tokens: "int") -> "str":
     if not text:
         return ""