mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
core/tools: resume fetch bodies from memory instead of refetching
This commit is contained in:
@@ -7,7 +7,6 @@ from typing import cast
|
|||||||
from plyngent.agent import ToolTag, tool
|
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_file, read_raw_text
|
from plyngent.tools.file.read import line_range_label, read_file, read_raw_text
|
||||||
from plyngent.tools.net.fetch import fetch
|
|
||||||
from plyngent.tools.truncate_token import (
|
from plyngent.tools.truncate_token import (
|
||||||
TruncateToken,
|
TruncateToken,
|
||||||
decode_truncate_token,
|
decode_truncate_token,
|
||||||
@@ -66,14 +65,13 @@ async def get_truncated(token: str, *, max_chars: int = DEFAULT_TOOL_RESULT_MAX_
|
|||||||
fresh token, so chunks chain indefinitely. File chunks resume in the same
|
fresh token, so chunks chain indefinitely. File chunks resume in the same
|
||||||
view as the original read: numbered reads continue as numbered lines
|
view as the original read: numbered reads continue as numbered lines
|
||||||
(editable via ``edit_lineno``), plain reads as raw text with a 1-based
|
(editable via ``edit_lineno``), plain reads as raw text with a 1-based
|
||||||
``L{begin}-{end}`` range. Memory chunks (generic tool output) are served
|
``L{begin}-{end}`` range. Memory chunks (generic tool output, and ``fetch``
|
||||||
from a short-lived in-memory store that is forgotten when the agent exits.
|
bodies) are served from a short-lived in-memory store that is forgotten when
|
||||||
|
the agent exits, so resuming a fetch never repeats the request.
|
||||||
"""
|
"""
|
||||||
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":
|
|
||||||
return await fetch.handler(parsed.location, offset=parsed.offset, max_chars=parsed.limit)
|
|
||||||
if parsed.kind == "memory":
|
if parsed.kind == "memory":
|
||||||
return _memory_chunk(parsed, max_chars)
|
return _memory_chunk(parsed, max_chars)
|
||||||
text, err = read_raw_text(parsed.location)
|
text, err = read_raw_text(parsed.location)
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ from plyngent.tools.net.policy import (
|
|||||||
parse_fetch_url,
|
parse_fetch_url,
|
||||||
soft_confirm_reason,
|
soft_confirm_reason,
|
||||||
)
|
)
|
||||||
from plyngent.tools.truncate_token import truncate_with_token
|
from plyngent.tools.truncate_token import truncate_to_memory
|
||||||
|
|
||||||
|
|
||||||
def format_fetch_result(
|
def format_fetch_result(
|
||||||
@@ -37,16 +37,10 @@ def format_fetch_result(
|
|||||||
offset: int = 0,
|
offset: int = 0,
|
||||||
) -> str:
|
) -> str:
|
||||||
start = max(0, offset)
|
start = max(0, offset)
|
||||||
body = body_text[start:]
|
# The body is already in hand, so the remainder is kept in memory rather than
|
||||||
body, token = truncate_with_token(
|
# re-requested: resuming a fetch must not cost another round trip (latency)
|
||||||
body,
|
# nor re-hit the remote (rate limits / risk control).
|
||||||
max_chars,
|
body, token = truncate_to_memory(body_text[start:], max_chars)
|
||||||
kind="http",
|
|
||||||
location=final_url,
|
|
||||||
offset=start,
|
|
||||||
limit=max_chars,
|
|
||||||
total_len=len(body_text),
|
|
||||||
)
|
|
||||||
char_truncated = token is not None
|
char_truncated = token is not None
|
||||||
warn_line = "; ".join(warnings) if warnings else ""
|
warn_line = "; ".join(warnings) if warnings else ""
|
||||||
parts = [
|
parts = [
|
||||||
|
|||||||
@@ -12,11 +12,13 @@ re-requesting the whole source or raising limits. Tokens chain: each truncated
|
|||||||
chunk carries a fresh token, so ``get_truncated`` keeps resuming until the
|
chunk carries a fresh token, so ``get_truncated`` keeps resuming until the
|
||||||
source is exhausted (then no marker is emitted).
|
source is exhausted (then no marker is emitted).
|
||||||
|
|
||||||
Resumable sources (``file`` / ``http``) are re-read by ``location`` +
|
A ``file`` remainder is re-read from disk by ``location`` + ``offset``, so it
|
||||||
``offset``. Arbitrary tool output (run_argv stdout, todo renders, …) has no
|
survives for as long as the file does. Everything else — arbitrary tool output
|
||||||
resumable source, so its remainder lives in a short-lived in-memory store
|
(run_argv stdout, todo renders, …) and ``fetch`` bodies — has no source worth
|
||||||
(``kind="memory"``) that is forgotten when the agent process exits — the model
|
re-running (a re-fetch costs a round trip and re-hits the remote), so its
|
||||||
must re-run the tool after that.
|
remainder lives in a short-lived in-memory store (``kind="memory"``) that is
|
||||||
|
forgotten when the agent process exits; the model must re-run the tool after
|
||||||
|
that.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -27,18 +29,18 @@ from typing import Literal
|
|||||||
|
|
||||||
import msgspec
|
import msgspec
|
||||||
|
|
||||||
type TruncateKind = Literal["file", "http", "memory"]
|
type TruncateKind = Literal["file", "memory"]
|
||||||
|
|
||||||
|
|
||||||
class TruncateToken(msgspec.Struct, frozen=True):
|
class TruncateToken(msgspec.Struct, frozen=True):
|
||||||
"""Opaque cursor into a truncated source.
|
"""Opaque cursor into a truncated source.
|
||||||
|
|
||||||
``offset``/``limit`` are in characters (http bodies, memory remainders) or
|
``offset``/``limit`` are in characters (memory remainders) or lines (files
|
||||||
lines (files is line-based via ``read_file``); ``get_truncated`` interprets
|
is line-based via ``read_file``); ``get_truncated`` interprets by ``kind``.
|
||||||
by ``kind``. For ``memory``, ``location`` is a key into the in-memory
|
For ``memory``, ``location`` is a key into the in-memory remainder store and
|
||||||
remainder store and ``offset``/``limit`` index characters of that stored
|
``offset``/``limit`` index characters of that stored remainder. ``numbered``
|
||||||
remainder. ``numbered`` marks file tokens from ``read_file(with_lineno)``:
|
marks file tokens from ``read_file(with_lineno)``: the continuation resumes
|
||||||
the continuation resumes the numbered view instead of raw text.
|
the numbered view instead of raw text.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
kind: TruncateKind
|
kind: TruncateKind
|
||||||
@@ -134,17 +136,17 @@ def truncate_with_token(
|
|||||||
return text + truncation_marker(max_chars, omitted, token), token
|
return text + truncation_marker(max_chars, omitted, token), token
|
||||||
|
|
||||||
|
|
||||||
def truncate_generic(text: str, max_chars: int) -> str:
|
def truncate_to_memory(text: str, max_chars: int) -> tuple[str, TruncateToken | None]:
|
||||||
"""Cap arbitrary tool output; embed a memory truncate token (chainable).
|
"""Cap ``text`` (a whole source) with a memory token over a stored copy.
|
||||||
|
|
||||||
The full remainder is kept in the in-memory store so ``get_truncated`` can
|
Nothing is stored when the text fits, so a caller that only needs the
|
||||||
resume it. Returns ``text`` unchanged (no marker) when it fits within
|
bounded text does not pay for a remainder. Returns ``(bounded_text, token)``
|
||||||
``max_chars``.
|
where ``token`` is ``None`` when the source is exhausted.
|
||||||
"""
|
"""
|
||||||
if max_chars < 1 or len(text) <= max_chars:
|
if max_chars < 1 or len(text) <= max_chars:
|
||||||
return text
|
return text, None
|
||||||
key = store_remainder(text)
|
key = store_remainder(text)
|
||||||
bounded, _ = truncate_with_token(
|
return truncate_with_token(
|
||||||
text,
|
text,
|
||||||
max_chars,
|
max_chars,
|
||||||
kind="memory",
|
kind="memory",
|
||||||
@@ -153,4 +155,14 @@ def truncate_generic(text: str, max_chars: int) -> str:
|
|||||||
limit=max_chars,
|
limit=max_chars,
|
||||||
total_len=len(text),
|
total_len=len(text),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def truncate_generic(text: str, max_chars: int) -> str:
|
||||||
|
"""Cap arbitrary tool output; embed a memory truncate token (chainable).
|
||||||
|
|
||||||
|
The full remainder is kept in the in-memory store so ``get_truncated`` can
|
||||||
|
resume it. Returns ``text`` unchanged (no marker) when it fits within
|
||||||
|
``max_chars``.
|
||||||
|
"""
|
||||||
|
bounded, _ = truncate_to_memory(text, max_chars)
|
||||||
return bounded
|
return bounded
|
||||||
|
|||||||
@@ -269,13 +269,10 @@ async def test_fetch_char_truncation_embeds_token(workspace: Path, http_server:
|
|||||||
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"
|
# The body is kept in memory, so resuming never repeats the request.
|
||||||
|
assert token.kind == "memory"
|
||||||
assert token.limit == 1000
|
assert token.limit == 1000
|
||||||
assert 0 < token.offset < 1000
|
assert 0 < token.offset < 1000
|
||||||
# Continuation resumes exactly where the chunk stopped (no gap).
|
|
||||||
out2 = await call_async(fetch, f"{http_server}/big", max_chars=1000, offset=token.offset)
|
|
||||||
body2 = out2.split("--- body ---", 1)[-1].strip()
|
|
||||||
assert body2.startswith("x")
|
|
||||||
|
|
||||||
|
|
||||||
async def test_fetch_offset_skips_body(workspace: Path, http_server: str) -> None:
|
async def test_fetch_offset_skips_body(workspace: Path, http_server: str) -> None:
|
||||||
|
|||||||
@@ -84,18 +84,40 @@ async def test_get_truncated_file_missing(workspace: Path) -> None:
|
|||||||
assert out.startswith("error:")
|
assert out.startswith("error:")
|
||||||
|
|
||||||
|
|
||||||
async def test_get_truncated_http_passthrough(monkeypatch) -> None:
|
async def test_get_truncated_serves_a_fetch_body_without_refetching(monkeypatch) -> None:
|
||||||
calls: list[tuple[str, int, int]] = []
|
"""A fetch body's token reads the stored remainder; the request is not repeated."""
|
||||||
|
from plyngent.tools.net.fetch import format_fetch_result
|
||||||
|
|
||||||
async def fake_fetch(url: str, *, offset: int, max_chars: int) -> str:
|
fetch_module = importlib.import_module("plyngent.tools.net.fetch")
|
||||||
calls.append((url, offset, max_chars))
|
|
||||||
return f"body-from-{offset}"
|
|
||||||
|
|
||||||
monkeypatch.setattr(get_truncated_module.fetch, "handler", fake_fetch)
|
async def _must_not_refetch(*_args: object, **_kwargs: object) -> str:
|
||||||
token = encode_truncate_token(TruncateToken(kind="http", location="https://example.com/x", offset=500, limit=400))
|
msg = "resuming a fetch body must not re-run the request"
|
||||||
|
raise AssertionError(msg)
|
||||||
|
|
||||||
|
monkeypatch.setattr(fetch_module, "http_fetch", _must_not_refetch)
|
||||||
|
first = format_fetch_result(
|
||||||
|
status=200,
|
||||||
|
final_url="https://example.com/x",
|
||||||
|
content_type="text/plain",
|
||||||
|
body_text="x" * 600,
|
||||||
|
body_bytes=600,
|
||||||
|
truncated=False,
|
||||||
|
redirects=0,
|
||||||
|
method="GET",
|
||||||
|
security="public",
|
||||||
|
warnings=[],
|
||||||
|
body_kind="text",
|
||||||
|
max_chars=200,
|
||||||
|
)
|
||||||
|
token = _extract_token(first)
|
||||||
|
assert token is not None
|
||||||
|
parsed = decode_truncate_token(token)
|
||||||
|
assert parsed is not None
|
||||||
|
assert parsed.kind == "memory"
|
||||||
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"
|
content, _, _ = out.partition("\n[Truncated")
|
||||||
assert calls == [("https://example.com/x", 500, 400)]
|
assert content == "x" * len(content)
|
||||||
|
assert 0 < len(content) <= 200
|
||||||
|
|
||||||
|
|
||||||
async def test_get_truncated_memory_chains() -> None:
|
async def test_get_truncated_memory_chains() -> None:
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ def test_decode_garbage_returns_none() -> None:
|
|||||||
def test_decode_rejects_invalid_fields() -> None:
|
def test_decode_rejects_invalid_fields() -> None:
|
||||||
bad = TruncateToken(kind="file", location="x", offset=-1, limit=10)
|
bad = TruncateToken(kind="file", location="x", offset=-1, limit=10)
|
||||||
assert decode_truncate_token(encode_truncate_token(bad)) is None
|
assert decode_truncate_token(encode_truncate_token(bad)) is None
|
||||||
bad_limit = TruncateToken(kind="http", location="u", offset=0, limit=0)
|
bad_limit = TruncateToken(kind="file", location="u", offset=0, limit=0)
|
||||||
assert decode_truncate_token(encode_truncate_token(bad_limit)) is None
|
assert decode_truncate_token(encode_truncate_token(bad_limit)) is None
|
||||||
|
|
||||||
|
|
||||||
@@ -95,7 +95,7 @@ def test_truncate_with_token_end_reached_no_token() -> None:
|
|||||||
out, token = truncate_with_token(
|
out, token = truncate_with_token(
|
||||||
text,
|
text,
|
||||||
40,
|
40,
|
||||||
kind="http",
|
kind="file",
|
||||||
location="https://example.com/x",
|
location="https://example.com/x",
|
||||||
offset=0,
|
offset=0,
|
||||||
limit=40,
|
limit=40,
|
||||||
@@ -111,14 +111,14 @@ def test_truncate_with_token_nonzero_start() -> None:
|
|||||||
out, token = truncate_with_token(
|
out, token = truncate_with_token(
|
||||||
text,
|
text,
|
||||||
200,
|
200,
|
||||||
kind="http",
|
kind="file",
|
||||||
location="https://example.com/x",
|
location="https://example.com/x",
|
||||||
offset=150,
|
offset=150,
|
||||||
limit=200,
|
limit=200,
|
||||||
total_len=700,
|
total_len=700,
|
||||||
)
|
)
|
||||||
assert token is not None
|
assert token is not None
|
||||||
assert token.kind == "http"
|
assert token.kind == "file"
|
||||||
content, _, _ = out.partition("\n[Truncated")
|
content, _, _ = out.partition("\n[Truncated")
|
||||||
assert token.offset == 150 + len(content)
|
assert token.offset == 150 + len(content)
|
||||||
assert "[Truncated (200 chars max.; 350 omitted). truncate_token=" in out
|
assert "[Truncated (200 chars max.; 350 omitted). truncate_token=" in out
|
||||||
|
|||||||
Reference in New Issue
Block a user