mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
refactor(truncation): improve file truncation logic (#184)
* refactor(file_store): simplify ChromaDB client initialization and improve file truncation logic - Remove shutil import and _create_chroma_client method from chroma_file_store.py - Directly initialize ChromaDB PersistentClient in start method without retry logic - Reduce DEFAULT_MAX_BYTES from 100KB to 50KB in file_utils.py - Update truncation notice format to provide clearer continuation instructions - Add _truncate_fresh and _retruncate functions for better text truncation handling - Replace inline truncation logic with dedicated function calls in file_utils.py - Rename skills_tool_ids to md_file_tool_ids in tool_result_compactor.py - Update file detection logic to identify any .md files instead of only skill.md - Create comprehensive unit tests for truncation functionality in test_truncate_text_output.py * chore(version): bump version to 0.3.1.6 - Update __version__ from 0.3.1.5 to 0.3.1.6 in __init__.py
This commit is contained in:
jinliyl
GitHub
parent
ff49a77f18
commit
37628ba524
6 changed files with 538 additions and 96 deletions
|
|
@ -6,7 +6,7 @@ from . import extension
|
|||
from . import memory
|
||||
from .reme import ReMe
|
||||
|
||||
__version__ = "0.3.1.5"
|
||||
__version__ = "0.3.1.6"
|
||||
|
||||
__all__ = [
|
||||
"config",
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
|
||||
import json
|
||||
import random
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -109,9 +108,13 @@ class ChromaFileStore(BaseFileStore):
|
|||
except Exception as e:
|
||||
logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}")
|
||||
|
||||
def _create_chroma_client(self):
|
||||
"""Create a ChromaDB PersistentClient instance."""
|
||||
return chromadb.PersistentClient(
|
||||
async def start(self) -> None:
|
||||
"""Initialize ChromaDB client and collection."""
|
||||
if self.client is not None:
|
||||
return
|
||||
|
||||
# Initialize persistent ChromaDB client
|
||||
self.client = chromadb.PersistentClient(
|
||||
path=str(self.db_path),
|
||||
settings=Settings(
|
||||
anonymized_telemetry=False,
|
||||
|
|
@ -119,23 +122,6 @@ class ChromaFileStore(BaseFileStore):
|
|||
),
|
||||
)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Initialize ChromaDB client and collection."""
|
||||
if self.client is not None:
|
||||
return
|
||||
|
||||
# Initialize persistent ChromaDB client, retry once after wiping db_path on failure
|
||||
try:
|
||||
self.client = self._create_chroma_client()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"ChromaDB failed to initialize at {self.db_path} ({e}). " f"Deleting corrupted database and retrying.",
|
||||
)
|
||||
if self.db_path.exists():
|
||||
shutil.rmtree(self.db_path)
|
||||
logger.info(f"Deleted ChromaDB directory: {self.db_path}")
|
||||
self.client = self._create_chroma_client()
|
||||
|
||||
# Get or create the chunks collection
|
||||
# ChromaDB uses cosine distance by default for similarity
|
||||
self.chunks_collection = self.client.get_or_create_collection(
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ class ToolResultCompactor(BaseOp):
|
|||
recent_n += 1
|
||||
split_index = max(0, len(messages) - max(recent_n, self.recent_n))
|
||||
|
||||
skills_tool_ids = set()
|
||||
md_file_tool_ids = set()
|
||||
try:
|
||||
for msg in messages:
|
||||
if not isinstance(msg.content, list):
|
||||
|
|
@ -105,12 +105,12 @@ class ToolResultCompactor(BaseOp):
|
|||
|
||||
if (
|
||||
block.get("name", "").lower() == "read_file"
|
||||
and "skill.md" in (block.get("raw_input") or "").lower()
|
||||
and ".md" in (block.get("raw_input") or "").lower()
|
||||
):
|
||||
skills_tool_ids.add(tool_id)
|
||||
md_file_tool_ids.add(tool_id)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to detect skill tool ids: %s", e)
|
||||
logger.info(f"skills_tool_ids: {skills_tool_ids}")
|
||||
logger.warning("Failed to detect md file tool ids: %s", e)
|
||||
logger.info(f"md_file_tool_ids: {md_file_tool_ids}")
|
||||
|
||||
for idx, msg in enumerate(messages):
|
||||
if not isinstance(msg.content, list):
|
||||
|
|
@ -120,7 +120,7 @@ class ToolResultCompactor(BaseOp):
|
|||
for block in msg.content:
|
||||
if isinstance(block, dict) and block.get("type") == "tool_result" and block.get("output"):
|
||||
tool_use_id = block.get("id", "")
|
||||
if tool_use_id in skills_tool_ids:
|
||||
if tool_use_id in md_file_tool_ids:
|
||||
effective_max_bytes = self.recent_max_bytes
|
||||
else:
|
||||
effective_max_bytes = max_bytes
|
||||
|
|
|
|||
|
|
@ -153,10 +153,13 @@ class FileIO:
|
|||
if text == selected_content and e < total:
|
||||
content_bytes = len(text.encode("utf-8"))
|
||||
notice = (
|
||||
TRUNCATION_NOTICE_MARKER
|
||||
+ f"\nFile: {file_path}\nStarting at start_line={s}, next {content_bytes} bytes."
|
||||
f"\nTotal lines: {total}"
|
||||
f"\nUse start_line={e + 1} to continue."
|
||||
TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated."
|
||||
f"\nThe full content is saved to the file "
|
||||
f"and contains {total} lines in total."
|
||||
f"\nThis excerpt starts at line {s} and "
|
||||
f"covers the next {content_bytes} bytes."
|
||||
"\nIf the current content is not enough, "
|
||||
f"call `read_file` with file_path={file_path} start_line={e + 1} to read more."
|
||||
)
|
||||
text = text + notice
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from ....core.utils import get_logger
|
|||
logger = get_logger()
|
||||
|
||||
# Default truncation limit
|
||||
DEFAULT_MAX_BYTES = 100 * 1024
|
||||
DEFAULT_MAX_BYTES = 50 * 1024
|
||||
|
||||
# Maximum file size to read into memory (1GB)
|
||||
MAX_FILE_READ_BYTES = 1024 * 1024 * 1024
|
||||
|
|
@ -16,17 +16,136 @@ MAX_FILE_READ_BYTES = 1024 * 1024 * 1024
|
|||
# Marker prepended to every truncation notice.
|
||||
# Format:
|
||||
# <<<TRUNCATED>>>
|
||||
# File: <path>
|
||||
# Starting at start_line=X, next N bytes.
|
||||
# Total lines: Z
|
||||
# Use start_line=Y to continue.
|
||||
# The output above was truncated.
|
||||
# The full content is saved to the file and contains Z lines in total.
|
||||
# This excerpt starts at line X and covers the next N bytes.
|
||||
# If the current content is not enough, call `read_file` with file_path=<path> start_line=Y to read more.
|
||||
#
|
||||
# Split output on this marker to recover the original (untruncated) portion:
|
||||
# original = output.split(TRUNCATION_NOTICE_MARKER)[0]
|
||||
TRUNCATION_NOTICE_MARKER = "<<<TRUNCATED>>>"
|
||||
|
||||
|
||||
# pylint: disable=too-many-return-statements
|
||||
def _truncate_fresh(
|
||||
text: str,
|
||||
start_line: int,
|
||||
total_lines: int,
|
||||
max_bytes: int,
|
||||
file_path: str | None,
|
||||
encoding: str,
|
||||
) -> str:
|
||||
"""Truncate fresh text (no prior truncation marker) by bytes with line integrity.
|
||||
|
||||
Slices at the byte boundary and appends a truncation notice with a continuation
|
||||
hint so callers know which line to read next.
|
||||
|
||||
Returns the original text unchanged when it fits within max_bytes, or when the
|
||||
last line itself exceeds max_bytes (unhandled edge case).
|
||||
"""
|
||||
text_bytes = text.encode(encoding)
|
||||
|
||||
# Under the byte limit — return as-is without any modification.
|
||||
if len(text_bytes) <= max_bytes:
|
||||
return text
|
||||
|
||||
# Slice at the byte boundary.
|
||||
# Assuming every single line is shorter than DEFAULT_MAX_BYTES, this cut always
|
||||
# lands mid-line, guaranteeing at least one complete line before the boundary.
|
||||
# Lines that exceed DEFAULT_MAX_BYTES are not handled and may be skipped entirely.
|
||||
truncated = text_bytes[:max_bytes]
|
||||
# Decode back to str; errors="ignore" drops any split multi-byte character
|
||||
# at the cut boundary without raising an exception.
|
||||
result = truncated.decode(encoding, errors="ignore")
|
||||
|
||||
# Count '\n' characters to determine how many complete lines are included.
|
||||
# The tail after the final '\n' is a partial line that will be covered by
|
||||
# the next read starting at next_line.
|
||||
newline_count = result.count("\n")
|
||||
|
||||
# Compute the first line number not yet fully included in this chunk.
|
||||
# max(1, ...) prevents next_line from equaling start_line when a single line
|
||||
# exceeds max_bytes (newline_count == 0), which would make the caller retry
|
||||
# the same range indefinitely.
|
||||
next_line = start_line + max(1, newline_count)
|
||||
|
||||
if next_line <= total_lines:
|
||||
# Truncation fell before the last line — continue reading from next_line.
|
||||
read_from = next_line
|
||||
elif start_line < total_lines:
|
||||
# next_line overshot total_lines, meaning the cut landed inside the last line.
|
||||
# Re-read from the start of the last line so the caller gets it in full.
|
||||
read_from = total_lines
|
||||
else:
|
||||
# start_line == total_lines: the last line itself exceeds DEFAULT_MAX_BYTES.
|
||||
# This case is outside our handled range — return without a truncation notice.
|
||||
return result
|
||||
|
||||
notice = (
|
||||
TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated."
|
||||
f"\nThe full content is saved to the file and contains {total_lines} lines in total."
|
||||
f"\nThis excerpt starts at line {start_line} and covers the next {max_bytes} bytes."
|
||||
f"\nIf the current content is not enough, call `read_file` with file_path={file_path or ''} "
|
||||
f"start_line={read_from} to read more."
|
||||
)
|
||||
|
||||
return result + notice
|
||||
|
||||
|
||||
def _retruncate(
|
||||
text: str,
|
||||
max_bytes: int,
|
||||
encoding: str,
|
||||
) -> str:
|
||||
"""Re-truncate text that was previously truncated (contains TRUNCATION_NOTICE_MARKER).
|
||||
|
||||
Extracts the original content before the marker, applies the new byte limit, and
|
||||
updates the embedded notice (byte count and continuation line number) via regex.
|
||||
|
||||
Returns the original text unchanged when:
|
||||
- the content already fits within max_bytes (with a small slack);
|
||||
- required metadata fields cannot be parsed from the existing notice.
|
||||
"""
|
||||
parts = text.split(TRUNCATION_NOTICE_MARKER, 1)
|
||||
original_content = parts[0]
|
||||
old_notice = parts[1]
|
||||
|
||||
text_bytes = original_content.encode(encoding)
|
||||
|
||||
# Allow a small slack to avoid unnecessary re-truncation when content is just
|
||||
# barely over the limit (e.g. due to minor encoding differences).
|
||||
if len(text_bytes) <= max_bytes + 100:
|
||||
return text
|
||||
|
||||
# Parse start_line from notice; return text unchanged if not found
|
||||
start_match = re.search(r"starts at line (\d+)", old_notice)
|
||||
if not start_match:
|
||||
return text
|
||||
start_line_parsed = int(start_match.group(1))
|
||||
|
||||
# Re-slice to the new byte limit.
|
||||
# Because every line is assumed to be shorter than DEFAULT_MAX_BYTES, the cut
|
||||
# always falls somewhere mid-line, so at least one complete line is preserved.
|
||||
truncated_bytes = text_bytes[:max_bytes]
|
||||
# errors="ignore" silently drops any incomplete multi-byte character at the cut boundary.
|
||||
result = truncated_bytes.decode(encoding, errors="ignore")
|
||||
# Each '\n' in result corresponds to one fully-included line;
|
||||
# anything after the last '\n' is a partial line that was cut off.
|
||||
newline_count = result.count("\n")
|
||||
|
||||
# The next read should start at the line immediately after all complete lines.
|
||||
# max(1, ...) guards against the theoretical zero-newline case
|
||||
# (impossible when every line is shorter than DEFAULT_MAX_BYTES).
|
||||
next_line = start_line_parsed + max(1, newline_count)
|
||||
|
||||
if not re.search(r"covers the next \d+ bytes", old_notice):
|
||||
return text
|
||||
# _truncate_fresh always includes a continuation hint, so both fields are always present.
|
||||
new_notice = re.sub(r"covers the next \d+ bytes", f"covers the next {max_bytes} bytes", old_notice)
|
||||
new_notice = re.sub(r"start_line=\d+ to read more", f"start_line={next_line} to read more", new_notice)
|
||||
|
||||
return result + TRUNCATION_NOTICE_MARKER + new_notice
|
||||
|
||||
|
||||
def truncate_text_output(
|
||||
text: str,
|
||||
start_line: int = 1,
|
||||
|
|
@ -41,9 +160,9 @@ def truncate_text_output(
|
|||
If over limit, truncate at the last complete line that fits,
|
||||
allowing the next read to start from a fresh line.
|
||||
|
||||
If TRUNCATION_NOTICE_MARKER is already in text (previously truncated),
|
||||
extract the original content, re-truncate it, and update the
|
||||
truncation notice using regex.
|
||||
Dispatches to :func:`_truncate_fresh` for text seen for the first time, or to
|
||||
:func:`_retruncate` when the text already contains a TRUNCATION_NOTICE_MARKER
|
||||
from a previous pass.
|
||||
|
||||
Args:
|
||||
text: The output text to truncate.
|
||||
|
|
@ -65,63 +184,16 @@ def truncate_text_output(
|
|||
|
||||
try:
|
||||
if TRUNCATION_NOTICE_MARKER in text:
|
||||
parts = text.split(TRUNCATION_NOTICE_MARKER, 1)
|
||||
original_content = parts[0]
|
||||
old_notice = parts[1]
|
||||
|
||||
text_bytes = original_content.encode(encoding)
|
||||
|
||||
# Allow a small slack to avoid re-truncating near-limit content
|
||||
if len(text_bytes) <= max_bytes + 100:
|
||||
return text
|
||||
|
||||
# Parse start_line and total_lines from notice; return text unchanged if not found
|
||||
start_match = re.search(r"Starting at start_line=(\d+)", old_notice)
|
||||
total_match = re.search(r"Total lines: (\d+)", old_notice)
|
||||
if not start_match or not total_match:
|
||||
return text
|
||||
start_line_parsed = int(start_match.group(1))
|
||||
total_lines_parsed = int(total_match.group(1))
|
||||
|
||||
truncated_bytes = text_bytes[:max_bytes]
|
||||
result = truncated_bytes.decode(encoding, errors="ignore")
|
||||
newline_count = result.count("\n")
|
||||
|
||||
next_line = start_line_parsed + max(1, newline_count)
|
||||
|
||||
if not re.search(r"next \d+ bytes", old_notice):
|
||||
return text
|
||||
has_continuation = bool(re.search(r"Use start_line=\d+", old_notice))
|
||||
new_notice = re.sub(r"next \d+ bytes", f"next {max_bytes} bytes", old_notice)
|
||||
if has_continuation:
|
||||
new_notice = re.sub(r"Use start_line=\d+", f"Use start_line={next_line}", new_notice)
|
||||
elif next_line <= total_lines_parsed:
|
||||
new_notice = re.sub(r"(Total lines: \d+)", f"\\1\nUse start_line={next_line} to continue.", new_notice)
|
||||
|
||||
return result + TRUNCATION_NOTICE_MARKER + new_notice
|
||||
|
||||
return _retruncate(text, max_bytes=max_bytes, encoding=encoding)
|
||||
else:
|
||||
text_bytes = text.encode(encoding)
|
||||
|
||||
if len(text_bytes) <= max_bytes:
|
||||
return text
|
||||
|
||||
truncated = text_bytes[:max_bytes]
|
||||
result = truncated.decode(encoding, errors="ignore")
|
||||
|
||||
newline_count = result.count("\n")
|
||||
|
||||
next_line = start_line + max(1, newline_count)
|
||||
|
||||
continuation = f"\nUse start_line={next_line} to continue." if next_line <= total_lines else ""
|
||||
notice = (
|
||||
TRUNCATION_NOTICE_MARKER
|
||||
+ f"\nFile: {file_path or ''}\nStarting at start_line={start_line}, next {max_bytes} bytes."
|
||||
f"\nTotal lines: {total_lines}{continuation}"
|
||||
return _truncate_fresh(
|
||||
text,
|
||||
start_line=start_line,
|
||||
total_lines=total_lines,
|
||||
max_bytes=max_bytes,
|
||||
file_path=file_path,
|
||||
encoding=encoding,
|
||||
)
|
||||
|
||||
return result + notice
|
||||
|
||||
except Exception:
|
||||
logger.warning("truncate_text_output failed, returning original text", exc_info=True)
|
||||
return text
|
||||
|
|
|
|||
381
tests/light/test_truncate_text_output.py
Normal file
381
tests/light/test_truncate_text_output.py
Normal file
|
|
@ -0,0 +1,381 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
# pylint: disable=missing-function-docstring
|
||||
"""Unit tests for _truncate_fresh, _retruncate, and truncate_text_output.
|
||||
|
||||
Assumptions that mirror the production code:
|
||||
- Every single line is shorter than DEFAULT_MAX_BYTES.
|
||||
- Lines that exceed DEFAULT_MAX_BYTES are explicitly ignored / not fully read.
|
||||
- _truncate_fresh always includes a continuation hint in the notice, so
|
||||
_retruncate can assume the hint is always present.
|
||||
|
||||
Run:
|
||||
cd tests/light && python test_truncate_text_output.py
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
|
||||
from reme.memory.file_based.utils.file_utils import (
|
||||
TRUNCATION_NOTICE_MARKER,
|
||||
_truncate_fresh,
|
||||
_retruncate,
|
||||
truncate_text_output,
|
||||
)
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
LINE_BYTES = 20 # every test line is exactly this many bytes
|
||||
ENC = "utf-8"
|
||||
|
||||
|
||||
def make_line(i: int) -> str:
|
||||
"""Return a line that is exactly LINE_BYTES bytes in UTF-8.
|
||||
|
||||
Format: "L{i}" padded with underscores, terminated with "\\n".
|
||||
Works for i up to 999.
|
||||
"""
|
||||
prefix = f"L{i}"
|
||||
return prefix + "_" * (LINE_BYTES - 1 - len(prefix)) + "\n"
|
||||
|
||||
|
||||
def make_text(n: int, start: int = 1) -> str:
|
||||
"""Build n lines starting from line number `start`."""
|
||||
return "".join(make_line(i) for i in range(start, start + n))
|
||||
|
||||
|
||||
def content_of(result: str) -> str:
|
||||
"""Return the portion before the truncation marker."""
|
||||
return result.split(TRUNCATION_NOTICE_MARKER)[0]
|
||||
|
||||
|
||||
def notice_of(result: str) -> str:
|
||||
"""Return the portion after the truncation marker, or '' if absent."""
|
||||
parts = result.split(TRUNCATION_NOTICE_MARKER, 1)
|
||||
return parts[1] if len(parts) > 1 else ""
|
||||
|
||||
|
||||
def parse_next_line(result: str) -> int:
|
||||
"""Extract 'start_line=X' from the notice, or -1 if absent."""
|
||||
m = re.search(r"start_line=(\d+) to read more", notice_of(result))
|
||||
return int(m.group(1)) if m else -1
|
||||
|
||||
|
||||
def parse_covers_bytes(result: str) -> int:
|
||||
"""Extract 'covers the next X bytes' from the notice."""
|
||||
m = re.search(r"covers the next (\d+) bytes", notice_of(result))
|
||||
return int(m.group(1)) if m else -1
|
||||
|
||||
|
||||
def assert_eq(a, b, msg=""):
|
||||
assert a == b, f"{msg}: expected {b!r}, got {a!r}"
|
||||
|
||||
|
||||
def assert_in(needle, haystack, msg=""):
|
||||
assert needle in haystack, f"{msg}: {needle!r} not found in {haystack!r}"
|
||||
|
||||
|
||||
def assert_not_in(needle, haystack, msg=""):
|
||||
assert needle not in haystack, f"{msg}: {needle!r} unexpectedly found in {haystack!r}"
|
||||
|
||||
|
||||
# ── _truncate_fresh tests ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_fresh_no_truncation_when_under_limit():
|
||||
text = make_text(3) # 60 bytes
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=3, max_bytes=100, file_path=None, encoding=ENC)
|
||||
assert_eq(result, text, "under limit: no change")
|
||||
assert_not_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
|
||||
|
||||
def test_fresh_no_truncation_at_exact_limit():
|
||||
text = make_text(3) # 60 bytes
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=3, max_bytes=60, file_path=None, encoding=ENC)
|
||||
assert_eq(result, text, "at exact limit: no change")
|
||||
assert_not_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
|
||||
|
||||
def test_fresh_mid_file_correct_next_line():
|
||||
# 10 lines × 20 bytes = 200 bytes; max_bytes=95
|
||||
# 95 bytes → 4 complete lines (80 bytes) + 15 bytes into line 5
|
||||
# newline_count=4, next_line=1+4=5, 5<=10 → read_from=5
|
||||
text = make_text(10)
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=10, max_bytes=95, file_path=None, encoding=ENC)
|
||||
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_next_line(result), 5, "next_line after mid-file cut")
|
||||
assert_eq(parse_covers_bytes(result), 95)
|
||||
assert_in("L4", content_of(result), "line 4 fully included")
|
||||
assert_not_in("L5\n", content_of(result), "line 5 not fully included")
|
||||
|
||||
|
||||
def test_fresh_notice_contains_total_lines_and_start_line():
|
||||
text = make_text(10, start=3)
|
||||
result = _truncate_fresh(text, start_line=3, total_lines=12, max_bytes=95, file_path="/tmp/foo.txt", encoding=ENC)
|
||||
notice = notice_of(result)
|
||||
assert_in("contains 12 lines in total", notice)
|
||||
assert_in("starts at line 3", notice)
|
||||
assert_in("file_path=/tmp/foo.txt", notice)
|
||||
|
||||
|
||||
def test_fresh_next_line_equals_total_lines_reads_from_last():
|
||||
# 5 lines × 20 = 100 bytes; max_bytes=85
|
||||
# 85 bytes → 4 complete lines (80 bytes) + 5 bytes into line 5
|
||||
# newline_count=4, next_line=1+4=5 = total_lines → read_from=5
|
||||
text = make_text(5)
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=85, file_path=None, encoding=ENC)
|
||||
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_next_line(result), 5, "should re-read the last line")
|
||||
|
||||
|
||||
def test_fresh_next_line_overshoots_total_lines():
|
||||
# max_bytes=115 → 5 complete lines (100 bytes) + 15 bytes into line 6
|
||||
# newline_count=5, next_line=1+5=6 > total_lines(5), start_line(1) < 5
|
||||
# → read_from = total_lines = 5
|
||||
text = make_text(10)
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=115, file_path=None, encoding=ENC)
|
||||
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_next_line(result), 5, "overshoot: fall back to total_lines")
|
||||
|
||||
|
||||
def test_fresh_long_first_line_skips_to_next_line():
|
||||
# Line 1 is 200 bytes (> max_bytes=100), no '\n' in truncated result.
|
||||
# newline_count=0, next_line=1+max(1,0)=2, 2<=5 → read_from=2
|
||||
long_line = "A" * 199 + "\n" # 200 bytes
|
||||
text = long_line + make_text(4, start=2)
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=100, file_path=None, encoding=ENC)
|
||||
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_next_line(result), 2, "skip to line 2 after long line 1")
|
||||
assert_not_in("A" * 101, content_of(result))
|
||||
|
||||
|
||||
def test_fresh_long_middle_line_skips_to_next_line():
|
||||
# Lines 1-2 normal (40 bytes); line 3 is 200 bytes; lines 4-5 normal.
|
||||
# max_bytes=100: fits lines 1-2 (40 bytes) + 60 bytes of line 3 (no '\n')
|
||||
# newline_count=2, next_line=1+2=3, 3<=5 → read_from=3
|
||||
text = make_text(2, start=1) + "B" * 199 + "\n" + make_text(2, start=4)
|
||||
result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=100, file_path=None, encoding=ENC)
|
||||
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_next_line(result), 3, "restart from long line 3")
|
||||
|
||||
|
||||
def test_fresh_long_last_line_start_equals_total_no_notice():
|
||||
# start_line == total_lines, single line too long → unhandled case, no notice.
|
||||
# newline_count=0, next_line=5+1=6 > 5, start_line(5)==total_lines(5)
|
||||
long_line = "C" * 199 + "\n" # 200 bytes
|
||||
result = _truncate_fresh(long_line, start_line=5, total_lines=5, max_bytes=100, file_path=None, encoding=ENC)
|
||||
|
||||
assert_not_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(len(result), 100, "only max_bytes bytes returned")
|
||||
|
||||
|
||||
def test_fresh_long_last_line_arrived_from_previous_chunk_no_notice():
|
||||
long_line = "D" * 199 + "\n"
|
||||
result = _truncate_fresh(long_line, start_line=5, total_lines=5, max_bytes=80, file_path=None, encoding=ENC)
|
||||
|
||||
assert_not_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
|
||||
|
||||
# ── _retruncate tests ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _first_pass(n_lines: int = 50, max_bytes: int = 490) -> str:
|
||||
"""Produce a first-truncated text via _truncate_fresh."""
|
||||
return _truncate_fresh(
|
||||
make_text(n_lines),
|
||||
start_line=1,
|
||||
total_lines=n_lines,
|
||||
max_bytes=max_bytes,
|
||||
file_path="/file.txt",
|
||||
encoding=ENC,
|
||||
)
|
||||
|
||||
|
||||
def test_retruncate_within_slack_returns_unchanged():
|
||||
# content ≈ 490 bytes; re-truncate with max_bytes=400.
|
||||
# 490 <= 400+100=500 → slack hit, return unchanged.
|
||||
pass1 = _first_pass(n_lines=50, max_bytes=490)
|
||||
result = _retruncate(pass1, max_bytes=400, encoding=ENC)
|
||||
assert_eq(result, pass1, "within slack: unchanged")
|
||||
|
||||
|
||||
def test_retruncate_updates_byte_count_in_notice():
|
||||
pass1 = _first_pass(n_lines=50, max_bytes=490)
|
||||
result = _retruncate(pass1, max_bytes=195, encoding=ENC)
|
||||
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_covers_bytes(result), 195, "notice byte count updated")
|
||||
|
||||
|
||||
def test_retruncate_updates_next_line():
|
||||
# pass1 content ≈ 490 bytes (24 complete lines), start_line=1.
|
||||
# re-truncate at 195 bytes → 9 complete lines, next_line=1+9=10.
|
||||
pass1 = _first_pass(n_lines=50, max_bytes=490)
|
||||
result = _retruncate(pass1, max_bytes=195, encoding=ENC)
|
||||
|
||||
assert_eq(parse_next_line(result), 10, "next_line updated to 10")
|
||||
|
||||
|
||||
def test_retruncate_content_is_smaller():
|
||||
pass1 = _first_pass(n_lines=50, max_bytes=490)
|
||||
result = _retruncate(pass1, max_bytes=195, encoding=ENC)
|
||||
|
||||
assert len(content_of(result)) < len(content_of(pass1)), "content should shrink"
|
||||
|
||||
|
||||
def test_retruncate_missing_starts_at_line_returns_unchanged():
|
||||
# Notice missing "starts at line X" → return unchanged.
|
||||
# Use large content (600 bytes) to bypass the 100-byte slack.
|
||||
large_content = make_text(30) # 600 bytes
|
||||
broken = (
|
||||
large_content + TRUNCATION_NOTICE_MARKER + "\nThe output above was truncated."
|
||||
"\nThe full content is saved to the file and contains 30 lines in total."
|
||||
"\nThis excerpt covers the next 600 bytes." # no "starts at line X"
|
||||
"\nIf the current content is not enough, call `read_file` with file_path=/f.txt start_line=5 to read more."
|
||||
)
|
||||
result = _retruncate(broken, max_bytes=10, encoding=ENC)
|
||||
assert_eq(result, broken, "missing 'starts at line': unchanged")
|
||||
|
||||
|
||||
def test_retruncate_missing_covers_next_bytes_returns_unchanged():
|
||||
# Notice has malformed "covers ??? bytes" → return unchanged.
|
||||
large_content = "X" * 500 + "\n"
|
||||
broken = (
|
||||
large_content + TRUNCATION_NOTICE_MARKER + "\nThe output above was truncated."
|
||||
"\nThe full content is saved to the file and contains 10 lines in total."
|
||||
"\nThis excerpt starts at line 1 and covers ??? bytes."
|
||||
"\nIf the current content is not enough, call `read_file` with file_path=/f.txt start_line=5 to read more."
|
||||
)
|
||||
result = _retruncate(broken, max_bytes=10, encoding=ENC)
|
||||
assert_eq(result, broken, "missing 'covers the next N bytes': unchanged")
|
||||
|
||||
|
||||
# ── truncate_text_output dispatch / guard tests ───────────────────────────────
|
||||
|
||||
|
||||
def test_dispatch_empty_string():
|
||||
result = truncate_text_output("", start_line=1, total_lines=0, max_bytes=10)
|
||||
assert_eq(result, "", "empty string bypassed")
|
||||
|
||||
|
||||
def test_dispatch_max_bytes_zero():
|
||||
text = make_text(5)
|
||||
result = truncate_text_output(text, max_bytes=0)
|
||||
assert_eq(result, text, "max_bytes=0 bypassed")
|
||||
|
||||
|
||||
def test_dispatch_routes_to_fresh_when_no_marker():
|
||||
text = make_text(10)
|
||||
result = truncate_text_output(text, start_line=1, total_lines=10, max_bytes=95)
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, result)
|
||||
assert_eq(parse_next_line(result), 5)
|
||||
|
||||
|
||||
def test_dispatch_routes_to_retruncate_when_marker_present():
|
||||
pass1 = _first_pass(n_lines=50, max_bytes=490)
|
||||
pass2 = truncate_text_output(pass1, max_bytes=195)
|
||||
assert_in(TRUNCATION_NOTICE_MARKER, pass2)
|
||||
assert_eq(parse_covers_bytes(pass2), 195)
|
||||
|
||||
|
||||
# ── multi-pass integration tests ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_three_pass_decreasing_truncation():
|
||||
"""Three successive truncations with shrinking max_bytes."""
|
||||
n_lines = 50
|
||||
text = make_text(n_lines) # 1000 bytes
|
||||
|
||||
pass1 = truncate_text_output(text, start_line=1, total_lines=n_lines, max_bytes=490, file_path="/f.txt")
|
||||
assert TRUNCATION_NOTICE_MARKER in pass1
|
||||
assert parse_covers_bytes(pass1) == 490
|
||||
assert parse_next_line(pass1) > 1
|
||||
|
||||
# 490 > 195+100=295 → re-truncation proceeds
|
||||
pass2 = truncate_text_output(pass1, max_bytes=195)
|
||||
assert TRUNCATION_NOTICE_MARKER in pass2
|
||||
assert parse_covers_bytes(pass2) == 195
|
||||
assert parse_next_line(pass2) < parse_next_line(pass1), "next_line regresses"
|
||||
assert len(content_of(pass2)) < len(content_of(pass1))
|
||||
|
||||
# content_of(pass2) ≈ 195 bytes; 195 > 90+100=190 → re-truncation proceeds
|
||||
pass3 = truncate_text_output(pass2, max_bytes=90)
|
||||
assert TRUNCATION_NOTICE_MARKER in pass3
|
||||
assert parse_covers_bytes(pass3) == 90
|
||||
assert parse_next_line(pass3) < parse_next_line(pass2), "next_line regresses further"
|
||||
assert len(content_of(pass3)) < len(content_of(pass2))
|
||||
|
||||
|
||||
def test_three_pass_next_lines_are_consistent():
|
||||
"""next_line values should monotonically decrease with each re-truncation."""
|
||||
n_lines = 50
|
||||
text = make_text(n_lines)
|
||||
|
||||
pass1 = truncate_text_output(text, start_line=1, total_lines=n_lines, max_bytes=490, file_path="/f.txt")
|
||||
pass2 = truncate_text_output(pass1, max_bytes=195)
|
||||
pass3 = truncate_text_output(pass2, max_bytes=90)
|
||||
|
||||
n1 = parse_next_line(pass1)
|
||||
n2 = parse_next_line(pass2)
|
||||
n3 = parse_next_line(pass3)
|
||||
|
||||
assert n1 > n2 > n3 > 1, f"Expected n1 > n2 > n3 > 1, got {n1} > {n2} > {n3}"
|
||||
|
||||
|
||||
# ── runner ────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def run_all():
|
||||
tests = [
|
||||
# _truncate_fresh
|
||||
test_fresh_no_truncation_when_under_limit,
|
||||
test_fresh_no_truncation_at_exact_limit,
|
||||
test_fresh_mid_file_correct_next_line,
|
||||
test_fresh_notice_contains_total_lines_and_start_line,
|
||||
test_fresh_next_line_equals_total_lines_reads_from_last,
|
||||
test_fresh_next_line_overshoots_total_lines,
|
||||
test_fresh_long_first_line_skips_to_next_line,
|
||||
test_fresh_long_middle_line_skips_to_next_line,
|
||||
test_fresh_long_last_line_start_equals_total_no_notice,
|
||||
test_fresh_long_last_line_arrived_from_previous_chunk_no_notice,
|
||||
# _retruncate
|
||||
test_retruncate_within_slack_returns_unchanged,
|
||||
test_retruncate_updates_byte_count_in_notice,
|
||||
test_retruncate_updates_next_line,
|
||||
test_retruncate_content_is_smaller,
|
||||
test_retruncate_missing_starts_at_line_returns_unchanged,
|
||||
test_retruncate_missing_covers_next_bytes_returns_unchanged,
|
||||
# truncate_text_output dispatch / guard
|
||||
test_dispatch_empty_string,
|
||||
test_dispatch_max_bytes_zero,
|
||||
test_dispatch_routes_to_fresh_when_no_marker,
|
||||
test_dispatch_routes_to_retruncate_when_marker_present,
|
||||
# multi-pass integration
|
||||
test_three_pass_decreasing_truncation,
|
||||
test_three_pass_next_lines_are_consistent,
|
||||
]
|
||||
|
||||
passed = 0
|
||||
failed = 0
|
||||
for t in tests:
|
||||
try:
|
||||
t()
|
||||
print(f" PASS {t.__name__}")
|
||||
passed += 1
|
||||
except Exception as e:
|
||||
print(f" FAIL {t.__name__}: {e}")
|
||||
failed += 1
|
||||
|
||||
print(f"\n{passed} passed, {failed} failed out of {len(tests)} tests.")
|
||||
return failed == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
ok = run_all()
|
||||
sys.exit(0 if ok else 1)
|
||||
Loading…
Add table
Reference in a new issue