diff --git a/interpreter/core/utils/truncate_output.py b/interpreter/core/utils/truncate_output.py index 4de2869b66..2454855090 100644 --- a/interpreter/core/utils/truncate_output.py +++ b/interpreter/core/utils/truncate_output.py @@ -1,3 +1,30 @@ +import re +import tempfile +from pathlib import Path + + +ANSI_ESCAPE_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]") +TRUNCATION_PREFIX = "Output truncated (" + + +def _strip_ansi_codes(data): + return ANSI_ESCAPE_RE.sub("", data) + + +def _latest_output_path(): + return Path(tempfile.gettempdir()) / "oi-output-latest.txt" + + +def _build_truncation_message(total_chars, chars_per_end, output_path): + return ( + f"Output truncated ({total_chars:,} visible characters total). " + f"Showing {chars_per_end:,} visible characters from start/end. " + f"Full output saved to `{output_path}`. " + "To inspect it, use tools like `head`, `tail`, or `grep`, " + "or rerun the command in smaller steps.\n\n" + ) + + def truncate_output(data, max_output_chars=2800, add_scrollbars=False): # if "@@@DO_NOT_TRUNCATE@@@" in data: # return data @@ -6,33 +33,37 @@ def truncate_output(data, max_output_chars=2800, add_scrollbars=False): # Calculate how much to show from start and end chars_per_end = max_output_chars // 2 - - message = (f"Output truncated ({len(data):,} characters total). " - f"Showing {chars_per_end:,} characters from start/end. " - "To handle large outputs, store result in python var first " - "`result = command()` then `computer.ai.summarize(result)` for " - "a summary, search with `result.find('text')`, " - "repeat shell commands with wc/grep/sed, etc. or break it down " - "into smaller steps.\n\n") + visible_data = _strip_ansi_codes(data) # This won't work because truncated code is stored in interpreter.messages :/ # If the full code was stored, we could do this: if add_scrollbars: - message = ( - message.strip() - + f" Run `get_last_output()[0:{max_output_chars}]` to see the first page.\n\n" + extra_scrollbar_message = ( + f" Run `get_last_output()[0:{max_output_chars}]` to see the first page." ) + else: + extra_scrollbar_message = "" # Then we have code in `terminal.py` which makes that function work. It should be a computer tool though to just access messages IMO. Or like, self.messages. # Remove previous truncation message if it exists - if data.startswith(message): - data = data[len(message) :] + if data.startswith(TRUNCATION_PREFIX): + _, separator, remainder = data.partition("\n\n") + if separator: + data = remainder + visible_data = _strip_ansi_codes(data) needs_truncation = True # If data exceeds max length, truncate it and add message - if len(data) > max_output_chars or needs_truncation: - first_part = data[:chars_per_end] - last_part = data[-chars_per_end:] + if len(visible_data) > max_output_chars or needs_truncation: + output_path = _latest_output_path() + output_path.write_text(data) + message = _build_truncation_message( + len(visible_data), chars_per_end, output_path + ) + if extra_scrollbar_message: + message = message.strip() + extra_scrollbar_message + "\n\n" + first_part = visible_data[:chars_per_end] + last_part = visible_data[-chars_per_end:] data = message + first_part + "\n[...]\n" + last_part return data diff --git a/tests/core/utils/test_truncate_output.py b/tests/core/utils/test_truncate_output.py new file mode 100644 index 0000000000..5c55ea651b --- /dev/null +++ b/tests/core/utils/test_truncate_output.py @@ -0,0 +1,43 @@ +import importlib.util +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +MODULE_PATH = Path(__file__).resolve().parents[3] / "interpreter/core/utils/truncate_output.py" +SPEC = importlib.util.spec_from_file_location("truncate_output_module", MODULE_PATH) +truncate_output_module = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(truncate_output_module) +truncate_output = truncate_output_module.truncate_output + + +class TestTruncateOutput(unittest.TestCase): + def test_does_not_truncate_when_only_ansi_codes_exceed_limit(self): + data = "\x1b[31mhello\x1b[0m" + + result = truncate_output(data, max_output_chars=5) + + self.assertEqual(result, data) + + def test_truncation_message_uses_real_recovery_paths(self): + data = "abcdefghijklmnopqrstuvwxyz" + + with tempfile.TemporaryDirectory() as temp_dir: + with mock.patch.object( + truncate_output_module.tempfile, + "gettempdir", + return_value=temp_dir, + ): + result = truncate_output(data, max_output_chars=10) + + saved_output = Path(temp_dir) / "oi-output-latest.txt" + self.assertTrue(saved_output.exists()) + self.assertEqual(saved_output.read_text(), data) + + self.assertIn("head", result) + self.assertIn("tail", result) + self.assertIn("grep", result) + self.assertNotIn("computer.ai.summarize(result)", result) +