mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
core/tools: token-ify per-stream caps and resume memory remainders
This commit is contained in:
@@ -8,24 +8,46 @@ from plyngent.agent import ToolTag, tool
|
|||||||
from plyngent.agent.budget import DEFAULT_TOOL_RESULT_MAX_CHARS
|
from plyngent.agent.budget import DEFAULT_TOOL_RESULT_MAX_CHARS
|
||||||
from plyngent.tools.file.read import line_range_label, read_raw_text
|
from plyngent.tools.file.read import line_range_label, read_raw_text
|
||||||
from plyngent.tools.net.fetch import fetch
|
from plyngent.tools.net.fetch import fetch
|
||||||
from plyngent.tools.truncate_token import decode_truncate_token, truncate_with_token
|
from plyngent.tools.truncate_token import (
|
||||||
|
decode_truncate_token,
|
||||||
|
get_remainder,
|
||||||
|
truncate_with_token,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE)
|
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE)
|
||||||
async def get_truncated(token: str, *, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS) -> str:
|
async def get_truncated(token: str, *, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS) -> str:
|
||||||
"""Fetch the next chunk of a truncated result using its ``TRUNCATE_TOKEN``.
|
"""Fetch the next chunk of a truncated result using its ``truncate_token``.
|
||||||
|
|
||||||
Pass the token shown at the end of a truncated ``read_file``, ``fetch``, or
|
Pass the token shown at the end of a truncated ``read_file``, ``fetch``,
|
||||||
earlier ``get_truncated`` result to keep reading without re-requesting the
|
``run_command``, or earlier ``get_truncated`` result to keep reading without
|
||||||
whole source or raising limits. Truncated output carries a fresh token, so
|
re-requesting the whole source or raising limits. Truncated output carries a
|
||||||
chunks chain indefinitely. File chunks start with a 1-based ``L{begin}-{end}``
|
fresh token, so chunks chain indefinitely. File chunks start with a 1-based
|
||||||
line range like ``read_file``.
|
``L{begin}-{end}`` line range like ``read_file``; memory chunks (generic
|
||||||
|
tool output) are served from a short-lived in-memory store that is forgotten
|
||||||
|
when the agent exits.
|
||||||
"""
|
"""
|
||||||
parsed = decode_truncate_token(token)
|
parsed = decode_truncate_token(token)
|
||||||
if parsed is None:
|
if parsed is None:
|
||||||
return "error: invalid truncate token"
|
return "error: invalid truncate token"
|
||||||
if parsed.kind == "http":
|
if parsed.kind == "http":
|
||||||
return await fetch.handler(parsed.location, offset=parsed.offset, max_chars=parsed.limit)
|
return await fetch.handler(parsed.location, offset=parsed.offset, max_chars=parsed.limit)
|
||||||
|
if parsed.kind == "memory":
|
||||||
|
remainder = get_remainder(parsed.location)
|
||||||
|
if remainder is None:
|
||||||
|
return "error: truncate token expired (no longer in memory); re-run the tool"
|
||||||
|
start = parsed.offset
|
||||||
|
segment = remainder[start : start + max_chars]
|
||||||
|
chunk, _ = truncate_with_token(
|
||||||
|
segment,
|
||||||
|
max_chars,
|
||||||
|
kind="memory",
|
||||||
|
location=parsed.location,
|
||||||
|
offset=start,
|
||||||
|
limit=max_chars,
|
||||||
|
total_len=len(remainder),
|
||||||
|
)
|
||||||
|
return chunk
|
||||||
text, err = read_raw_text(parsed.location)
|
text, err = read_raw_text(parsed.location)
|
||||||
if err:
|
if err:
|
||||||
return err
|
return err
|
||||||
@@ -41,6 +63,6 @@ async def get_truncated(token: str, *, max_chars: int = DEFAULT_TOOL_RESULT_MAX_
|
|||||||
limit=max_chars,
|
limit=max_chars,
|
||||||
total_len=len(text),
|
total_len=len(text),
|
||||||
)
|
)
|
||||||
content_len = len(chunk.split("\n...[", 1)[0]) if "\n...[" in chunk else len(chunk)
|
content_len = len(chunk.split("\n[Truncated", 1)[0]) if "\n[Truncated" in chunk else len(chunk)
|
||||||
header = line_range_label(text, start, start + content_len)
|
header = line_range_label(text, start, start + content_len)
|
||||||
return f"{header}\n{chunk}"
|
return f"{header}\n{chunk}"
|
||||||
|
|||||||
@@ -22,9 +22,19 @@ DEFAULT_MAX_BATCH_STEPS = 20
|
|||||||
|
|
||||||
|
|
||||||
def truncate_output(text: str, label: str, max_chars: int = DEFAULT_MAX_OUTPUT_CHARS) -> str:
|
def truncate_output(text: str, label: str, max_chars: int = DEFAULT_MAX_OUTPUT_CHARS) -> str:
|
||||||
if len(text) <= max_chars:
|
"""Cap a per-stream output block; embed a truncate token when cut.
|
||||||
|
|
||||||
|
``label`` (stdout/stderr/batch) is kept for call-site clarity but the
|
||||||
|
marker itself is the uniform ``[Truncated (N chars max.; M omitted).
|
||||||
|
truncate_token=...]`` (no label), matching every other truncation site.
|
||||||
|
Returns ``text`` unchanged when it fits.
|
||||||
|
"""
|
||||||
|
del label
|
||||||
|
if max_chars < 1 or len(text) <= max_chars:
|
||||||
return text
|
return text
|
||||||
return text[:max_chars] + f"\n...[{label} truncated]"
|
from plyngent.tools.truncate_token import truncate_generic
|
||||||
|
|
||||||
|
return truncate_generic(text, max_chars)
|
||||||
|
|
||||||
|
|
||||||
def format_command_result(
|
def format_command_result(
|
||||||
|
|||||||
@@ -246,9 +246,9 @@ async def test_fetch_char_truncation_embeds_token(workspace: Path, http_server:
|
|||||||
out = await call_async(fetch, f"{http_server}/big", max_chars=1000)
|
out = await call_async(fetch, f"{http_server}/big", max_chars=1000)
|
||||||
assert "truncated=true" in out
|
assert "truncated=true" in out
|
||||||
body = out.split("--- body ---", 1)[-1]
|
body = out.split("--- body ---", 1)[-1]
|
||||||
assert "TRUNCATE_TOKEN:" in body
|
assert "truncate_token=" in body
|
||||||
token_line = next(ln for ln in body.splitlines() if "TRUNCATE_TOKEN:" in ln)
|
token_line = next(ln for ln in body.splitlines() if "truncate_token=" in ln)
|
||||||
token_str = token_line.split("[TRUNCATE_TOKEN: ", 1)[1].rstrip("]")
|
token_str = token_line.split("truncate_token=", 1)[1].rstrip("]")
|
||||||
token = decode_truncate_token(token_str)
|
token = decode_truncate_token(token_str)
|
||||||
assert token is not None
|
assert token is not None
|
||||||
assert token.kind == "http"
|
assert token.kind == "http"
|
||||||
@@ -269,7 +269,7 @@ async def test_fetch_offset_skips_body(workspace: Path, http_server: str) -> Non
|
|||||||
out = await call_async(fetch, f"{http_server}/big", max_chars=200, offset=4900)
|
out = await call_async(fetch, f"{http_server}/big", max_chars=200, offset=4900)
|
||||||
body = out.split("--- body ---", 1)[-1]
|
body = out.split("--- body ---", 1)[-1]
|
||||||
assert body.strip() == "x" * 100
|
assert body.strip() == "x" * 100
|
||||||
assert "TRUNCATE_TOKEN:" not in out
|
assert "truncate_token=" not in out
|
||||||
|
|
||||||
|
|
||||||
async def test_fetch_private_denied_without_grant(workspace: Path, http_server: str) -> None:
|
async def test_fetch_private_denied_without_grant(workspace: Path, http_server: str) -> None:
|
||||||
|
|||||||
@@ -11,8 +11,8 @@ get_truncated_module = importlib.import_module("plyngent.tools.file.get_truncate
|
|||||||
|
|
||||||
def _extract_token(out: str) -> str | None:
|
def _extract_token(out: str) -> str | None:
|
||||||
for line in out.splitlines():
|
for line in out.splitlines():
|
||||||
if "TRUNCATE_TOKEN:" in line:
|
if "truncate_token=" in line:
|
||||||
return line.split("[TRUNCATE_TOKEN: ", 1)[1].rstrip("]")
|
return line.split("truncate_token=", 1)[1].rstrip("]")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
@@ -33,10 +33,10 @@ def test_read_file_max_chars_embeds_token(workspace: Path) -> None:
|
|||||||
out = call_sync(read_file, "big.txt", max_chars=200)
|
out = call_sync(read_file, "big.txt", max_chars=200)
|
||||||
header, _, rest = out.partition("\n")
|
header, _, rest = out.partition("\n")
|
||||||
assert header.startswith("L1-") # 1-based range of the raw slice read
|
assert header.startswith("L1-") # 1-based range of the raw slice read
|
||||||
content, _, marker = rest.partition("\n...[")
|
content, _, marker = rest.partition("\n[Truncated")
|
||||||
assert content == "a" * len(content)
|
assert content == "a" * len(content)
|
||||||
assert len(content) < 200
|
assert len(content) < 200
|
||||||
assert "TRUNCATE_TOKEN:" in marker
|
assert "truncate_token=" in marker
|
||||||
token = _extract_token(out)
|
token = _extract_token(out)
|
||||||
assert token is not None
|
assert token is not None
|
||||||
parsed = get_truncated_module.decode_truncate_token(token)
|
parsed = get_truncated_module.decode_truncate_token(token)
|
||||||
@@ -49,7 +49,7 @@ def test_read_file_max_chars_small_file_untouched(workspace: Path) -> None:
|
|||||||
(workspace / "small.txt").write_text("tiny", encoding="utf-8")
|
(workspace / "small.txt").write_text("tiny", encoding="utf-8")
|
||||||
out = call_sync(read_file, "small.txt", max_chars=200)
|
out = call_sync(read_file, "small.txt", max_chars=200)
|
||||||
assert out == "L1-1\ntiny"
|
assert out == "L1-1\ntiny"
|
||||||
assert "TRUNCATE_TOKEN:" not in out
|
assert "truncate_token=" not in out
|
||||||
|
|
||||||
|
|
||||||
def test_edit_replace_first_only(workspace: object) -> None:
|
def test_edit_replace_first_only(workspace: object) -> None:
|
||||||
|
|||||||
@@ -12,8 +12,8 @@ get_truncated = get_truncated_module.get_truncated
|
|||||||
|
|
||||||
def _extract_token(out: str) -> str | None:
|
def _extract_token(out: str) -> str | None:
|
||||||
for line in out.splitlines():
|
for line in out.splitlines():
|
||||||
if "TRUNCATE_TOKEN:" in line:
|
if "truncate_token=" in line:
|
||||||
return line.split("[TRUNCATE_TOKEN: ", 1)[1].rstrip("]")
|
return line.split("truncate_token=", 1)[1].rstrip("]")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
@@ -29,7 +29,7 @@ async def test_get_truncated_file_chains_without_gap(workspace: Path) -> None:
|
|||||||
total_read = 0
|
total_read = 0
|
||||||
for _ in range(20):
|
for _ in range(20):
|
||||||
out = await call_async(get_truncated, token, max_chars=200)
|
out = await call_async(get_truncated, token, max_chars=200)
|
||||||
body, _, _ = out.partition("\n...[")
|
body, _, _ = out.partition("\n[Truncated")
|
||||||
header, _, content = body.partition("\n")
|
header, _, content = body.partition("\n")
|
||||||
assert header.startswith("L") and "-" in header
|
assert header.startswith("L") and "-" in header
|
||||||
assert content == "a" * len(content)
|
assert content == "a" * len(content)
|
||||||
@@ -62,3 +62,42 @@ async def test_get_truncated_http_passthrough(monkeypatch) -> None:
|
|||||||
out = await call_async(get_truncated, token, max_chars=200)
|
out = await call_async(get_truncated, token, max_chars=200)
|
||||||
assert out == "body-from-500"
|
assert out == "body-from-500"
|
||||||
assert calls == [("https://example.com/x", 500, 400)]
|
assert calls == [("https://example.com/x", 500, 400)]
|
||||||
|
|
||||||
|
|
||||||
|
async def test_get_truncated_memory_chains() -> None:
|
||||||
|
from plyngent.tools.truncate_token import truncate_generic
|
||||||
|
|
||||||
|
bounded = truncate_generic("m" * 1000, 200)
|
||||||
|
assert "truncate_token=" in bounded
|
||||||
|
token = bounded.split("truncate_token=", 1)[1].rstrip("]")
|
||||||
|
parsed = decode_truncate_token(token)
|
||||||
|
assert parsed is not None
|
||||||
|
assert parsed.kind == "memory"
|
||||||
|
cursor = parsed.offset # chars already delivered by truncate_generic
|
||||||
|
delivered = cursor
|
||||||
|
for _ in range(20):
|
||||||
|
out = await call_async(get_truncated, token, max_chars=200)
|
||||||
|
content, _, _ = out.partition("\n[Truncated")
|
||||||
|
assert content == "m" * len(content)
|
||||||
|
delivered += len(content)
|
||||||
|
next_token = _extract_token(out)
|
||||||
|
if next_token is None:
|
||||||
|
break # source exhausted → no token, no marker
|
||||||
|
parsed = decode_truncate_token(next_token)
|
||||||
|
assert parsed is not None
|
||||||
|
assert parsed.kind == "memory"
|
||||||
|
# No gap: the next cursor is exactly where this chunk ended.
|
||||||
|
assert parsed.offset == delivered
|
||||||
|
token = next_token
|
||||||
|
assert delivered == 1000
|
||||||
|
|
||||||
|
|
||||||
|
async def test_get_truncated_memory_expired() -> None:
|
||||||
|
from plyngent.tools.truncate_token import _MEMORY_REMAINDERS, store_remainder
|
||||||
|
|
||||||
|
key = store_remainder("x" * 500)
|
||||||
|
del _MEMORY_REMAINDERS[key]
|
||||||
|
token = encode_truncate_token(TruncateToken(kind="memory", location=key, offset=0, limit=200))
|
||||||
|
out = await call_async(get_truncated, token, max_chars=200)
|
||||||
|
assert out.startswith("error:")
|
||||||
|
assert "expired" in out
|
||||||
|
|||||||
Reference in New Issue
Block a user