mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +08:00
core/tools: embed truncate tokens in read_file (max_chars)
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from plyngent.agent import ToolTag, mark_lineno_read, tool
|
||||
from plyngent.tools.truncate_token import truncate_with_token
|
||||
from plyngent.tools.workspace import resolve_path
|
||||
|
||||
_LINENO_WIDTH = 6
|
||||
@@ -24,6 +25,7 @@ async def read_file(
|
||||
offset: int = 0,
|
||||
limit: int | None = None,
|
||||
with_lineno: bool = False,
|
||||
max_chars: int | None = None,
|
||||
) -> str:
|
||||
"""Read a text file under the workspace.
|
||||
|
||||
@@ -31,6 +33,9 @@ async def read_file(
|
||||
When ``with_lineno`` is true, each line is prefixed with its 1-based file line
|
||||
number (`` N|…``), matching ``edit_lineno`` numbering, and those lines are
|
||||
marked readable for ``edit_lineno`` this turn.
|
||||
|
||||
``max_chars`` caps the returned slice; a ``TRUNCATE_TOKEN`` is appended so
|
||||
``get_truncated`` can continue reading the rest without a new request.
|
||||
"""
|
||||
target = resolve_path(path)
|
||||
if not target.is_file():
|
||||
@@ -46,5 +51,18 @@ async def read_file(
|
||||
slice_lines = lines[start:end]
|
||||
if with_lineno:
|
||||
mark_lineno_read(str(target), set(range(start + 1, end + 1)))
|
||||
return _format_with_lineno(slice_lines, start_lineno=start + 1)
|
||||
return "".join(slice_lines)
|
||||
out = _format_with_lineno(slice_lines, start_lineno=start + 1)
|
||||
else:
|
||||
out = "".join(slice_lines)
|
||||
if max_chars is not None and max_chars >= 1:
|
||||
char_start = len("".join(lines[:start]))
|
||||
out, _ = truncate_with_token(
|
||||
out,
|
||||
max_chars,
|
||||
kind="file",
|
||||
location=path, # arg form: short token, resolves to the same file
|
||||
offset=char_start,
|
||||
limit=max_chars,
|
||||
total_len=len(text),
|
||||
)
|
||||
return out
|
||||
|
||||
@@ -1,8 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
|
||||
from plyngent.tools.file import edit_replace, listdir, read_file, write_file
|
||||
from tests.test_tools.helpers import call_sync
|
||||
|
||||
get_truncated_module = importlib.import_module("plyngent.tools.file.get_truncated")
|
||||
|
||||
|
||||
def _extract_token(out: str) -> str | None:
|
||||
for line in out.splitlines():
|
||||
if "TRUNCATE_TOKEN:" in line:
|
||||
return line.split("[TRUNCATE_TOKEN: ", 1)[1].rstrip("]")
|
||||
return None
|
||||
|
||||
|
||||
def test_write_read_listdir_edit(workspace: object) -> None:
|
||||
del workspace
|
||||
@@ -16,6 +28,28 @@ def test_write_read_listdir_edit(workspace: object) -> None:
|
||||
assert call_sync(read_file, "notes/a.txt") == "hello there"
|
||||
|
||||
|
||||
def test_read_file_max_chars_embeds_token(workspace: Path) -> None:
|
||||
(workspace / "big.txt").write_text("a" * 1000, encoding="utf-8")
|
||||
out = call_sync(read_file, "big.txt", max_chars=200)
|
||||
content, _, marker = out.partition("\n...[")
|
||||
assert content == "a" * len(content)
|
||||
assert len(content) < 200
|
||||
assert "TRUNCATE_TOKEN:" in marker
|
||||
token = _extract_token(out)
|
||||
assert token is not None
|
||||
parsed = get_truncated_module.decode_truncate_token(token)
|
||||
assert parsed is not None
|
||||
assert parsed.kind == "file"
|
||||
assert parsed.offset == len(content) # cursor lands exactly where content stops
|
||||
|
||||
|
||||
def test_read_file_max_chars_small_file_untouched(workspace: Path) -> None:
|
||||
(workspace / "small.txt").write_text("tiny", encoding="utf-8")
|
||||
out = call_sync(read_file, "small.txt", max_chars=200)
|
||||
assert out == "tiny"
|
||||
assert "TRUNCATE_TOKEN:" not in out
|
||||
|
||||
|
||||
def test_edit_replace_first_only(workspace: object) -> None:
|
||||
del workspace
|
||||
_ = call_sync(write_file, "t.txt", "aa aa")
|
||||
|
||||
Reference in New Issue
Block a user