diff --git a/src/plyngent/cli/display.py b/src/plyngent/cli/display.py index 98c6c5a..f097122 100644 --- a/src/plyngent/cli/display.py +++ b/src/plyngent/cli/display.py @@ -808,6 +808,66 @@ def _resume_detail(_args_json: str, result: str) -> str: return _range_detail(result) or _pretty_segments(("(done)", "green")) +def _doc_read_prefix(args_json: str) -> str: + document_id = _json_str_arg(args_json, "document_id") or "?" + return _pretty_prefix(f"* Doc Read '{document_id}' ") + + +def _doc_read_detail(_args_json: str, result: str) -> str: + """``read_document`` outcome: the line range read, or an error.""" + if result.startswith("error:"): + return _pretty_segments((f"({result})", "red")) + return _range_detail(result) or _pretty_segments(("(done)", "green")) + + +def _doc_search_prefix(args_json: str) -> str: + pattern = _json_str_arg(args_json, "pattern") or "?" + document_id = _json_str_arg(args_json, "document_id") or "?" + return _pretty_prefix(f"* Doc Search '{_preview(pattern, _ASK_SUBJECT_PREVIEW)}' in '{document_id}' ") + + +def _doc_search_detail(_args_json: str, result: str) -> str: + """``search_document`` outcome: the match count, none, or an error.""" + if result.startswith("error:"): + return _pretty_segments((f"({result})", "red")) + if result == "(no matches)": + return _pretty_segments(("(no matches)", "dim")) + matches = [line for line in result.splitlines() if line and not line.startswith("...[truncated")] + unit = "match" if len(matches) == 1 else "matches" + return _pretty_segments((f"({len(matches)} {unit})", None)) + + +def _doc_info_prefix(args_json: str) -> str: + document_id = _json_str_arg(args_json, "document_id") or "?" + return _pretty_prefix(f"* Doc Info '{document_id}' ") + + +def _doc_info_detail(_args_json: str, result: str) -> str: + """``document_info`` outcome: the described size, or an error.""" + if result.startswith("error:"): + return _pretty_segments((f"({result})", "red")) + fields = _status_fields(result) + lines = fields.get("lines", "") + chars = fields.get("chars", "") + text = f"{lines} lines, {chars} chars" if lines and chars else "done" + return _pretty_segments((f"({text})", None)) + + +def _doc_list_prefix(_args_json: str) -> str: + return _pretty_prefix("* Doc List ") + + +def _doc_list_detail(_args_json: str, result: str) -> str: + """``list_documents`` outcome: the document count, none, or an error.""" + if result.startswith("error:"): + return _pretty_segments((f"({result})", "red")) + if result == "(no documents)": + return _pretty_segments(("(no documents)", "dim")) + count = len(result.splitlines()) + unit = "document" if count == 1 else "documents" + return _pretty_segments((f"({count} {unit})", None)) + + def _vcs_kind_detail(_args_json: str, result: str) -> str: """``vcs_kind`` outcome: the detected kind, or an error.""" if result.startswith("error:"): @@ -952,6 +1012,10 @@ _PRETTY_LINES: dict[str, PrettyLine] = { "glob_paths": PrettyLine(prefix=_glob_prefix, detail=_glob_detail), "regex_files": PrettyLine(prefix=_regex_prefix, detail=_regex_detail), "get_truncated": PrettyLine(prefix=_resume_prefix, detail=_resume_detail), + "read_document": PrettyLine(prefix=_doc_read_prefix, detail=_doc_read_detail), + "search_document": PrettyLine(prefix=_doc_search_prefix, detail=_doc_search_detail), + "document_info": PrettyLine(prefix=_doc_info_prefix, detail=_doc_info_detail), + "list_documents": PrettyLine(prefix=_doc_list_prefix, detail=_doc_list_detail), "run_argv": PrettyLine(prefix=_run_argv_prefix, detail=_run_argv_detail), "run_argv_batch": PrettyLine(prefix=_run_argv_batch_prefix, detail=_run_argv_batch_detail), "fetch": PrettyLine(prefix=_fetch_prefix, detail=_fetch_detail), diff --git a/src/plyngent/cli/state.py b/src/plyngent/cli/state.py index e19e68b..eb9eec8 100644 --- a/src/plyngent/cli/state.py +++ b/src/plyngent/cli/state.py @@ -228,6 +228,18 @@ class ReplState: if hasattr(self, "agent") and self.agent.tools is not None: self.agent.tools.set_session_state(self.session_state) self.agent.tools.set_instance_state(self.instance_state) + self._bind_document_store() + + def _bind_document_store(self) -> None: + """Expose this session's tool results to the document tools. + + The store resolves a ``document_id`` against the agent's live history, + which reads whole (uncapped) tool results and re-reads whatever a resume + loaded from the durable store. + """ + from plyngent.tools.documents import DocumentStore + + self.instance_state.extras["documents"] = DocumentStore(lambda: self.agent.messages) async def persist_todo_stack(self) -> None: """Write the in-memory todo stack to the active session row.""" diff --git a/src/plyngent/tools/documents.py b/src/plyngent/tools/documents.py new file mode 100644 index 0000000..99361f1 --- /dev/null +++ b/src/plyngent/tools/documents.py @@ -0,0 +1,118 @@ +"""Documents: addressable views over tool-result messages. + +A tool result is stored whole in its own ``ToolChatMessage``. When the agent's +request-prep pass caps a long result, the marker it leaves names a +``document_id`` (the message's ``tool_call_id``); the model reads the rest with +the document tools. The body stays the message itself, so a document survives a +process restart with the session that holds it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from msgspec import UNSET + +from plyngent.lmproto.openai_compatible.model import AssistantChatMessage, ToolChatMessage + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from plyngent.lmproto.openai_compatible.model import AnyChatMessage + + +@dataclass(frozen=True, slots=True) +class Document: + """One tool-result message exposed as a document.""" + + document_id: str + text: str + tool_name: str | None + index: int + + @property + def chars(self) -> int: + return len(self.text) + + @property + def lines(self) -> int: + return len(self.text.splitlines()) + + +def document_marker(document_id: str, shown: int, total: int) -> str: + """The request-prep cap hint naming the document that holds the remainder.""" + return ( + f"\n[Truncated: {shown} of {total} chars shown. document_id={document_id}. " + f"Continue with read_document(document_id, offset=...) or " + f"search_document(document_id, pattern=...).]" + ) + + +def _tool_names(messages: Sequence[AnyChatMessage]) -> dict[str, str]: + """Map each tool-call id to the tool name the assistant asked for.""" + names: dict[str, str] = {} + for message in messages: + if not isinstance(message, AssistantChatMessage): + continue + tool_calls = message.tool_calls + if tool_calls is UNSET or not tool_calls: + continue + for call in tool_calls: + function = getattr(call, "function", None) + name = getattr(function, "name", None) if function is not None else None + call_id = getattr(call, "id", None) + if isinstance(call_id, str) and isinstance(name, str): + names[call_id] = name + return names + + +class DocumentStore: + """Resolve document ids to the tool-result messages that carry their body. + + *resolver* returns the live message history (the agent's own list), so the + store reflects the session as it grows and reads whatever a resume reloaded + from the durable store. + """ + + _resolver: Callable[[], Sequence[AnyChatMessage]] + + def __init__(self, resolver: Callable[[], Sequence[AnyChatMessage]]) -> None: + self._resolver = resolver + + def _documents(self) -> list[Document]: + messages = self._resolver() + names = _tool_names(messages) + documents: list[Document] = [] + for message in messages: + if not isinstance(message, ToolChatMessage): + continue + documents.append( + Document( + document_id=message.tool_call_id, + text=message.content, + tool_name=names.get(message.tool_call_id), + index=len(documents), + ) + ) + return documents + + def get(self, document_id: str) -> Document | None: + for document in self._documents(): + if document.document_id == document_id: + return document + return None + + def list(self) -> list[Document]: + return self._documents() + + +def get_document_store() -> DocumentStore | None: + """The host-bound document store for the current tool context, if any.""" + from plyngent.tools.context import get_instance + + instance = get_instance() + if instance is None: + return None + store = instance.extras.get("documents") + return store if isinstance(store, DocumentStore) else None diff --git a/src/plyngent/tools/file/__init__.py b/src/plyngent/tools/file/__init__.py index 3ef640d..54cb697 100644 --- a/src/plyngent/tools/file/__init__.py +++ b/src/plyngent/tools/file/__init__.py @@ -1,5 +1,9 @@ from plyngent.tools.temp_workspace import new_temporary_workspace as new_temporary_workspace +from .document import document_info as document_info +from .document import list_documents as list_documents +from .document import read_document as read_document +from .document import search_document as search_document from .edit_lineno import edit_lineno as edit_lineno from .edit_replace import edit_replace as edit_replace from .fs_ops import copy_path as copy_path @@ -29,4 +33,8 @@ FILE_TOOLS = [ delete_path, new_temporary_workspace, get_truncated, + read_document, + search_document, + document_info, + list_documents, ] diff --git a/src/plyngent/tools/file/document.py b/src/plyngent/tools/file/document.py new file mode 100644 index 0000000..68b556c --- /dev/null +++ b/src/plyngent/tools/file/document.py @@ -0,0 +1,217 @@ +"""Document tools: read and search the tool results the request-prep pass capped. + +A long tool result is stored whole in its tool message and capped on the wire +with a ``document_id`` marker. These tools resolve that id back to the message +body, so reading and searching a capped result need no re-run of the tool. +""" + +from __future__ import annotations + +import re + +from plyngent.agent import ToolTag, tool +from plyngent.agent.budget import DEFAULT_TOOL_RESULT_MAX_CHARS +from plyngent.tools.documents import Document, get_document_store +from plyngent.tools.file.read import format_with_lineno + +DEFAULT_MAX_MATCHES = 100 +DEFAULT_LINE_PREVIEW = 200 + +_MISSING_STORE = "error: documents are in scope only while a chat session holds tool results" +_UNKNOWN = "error: unknown document {document_id!r} (see list_documents)" + + +def _resolve(document_id: str) -> Document | str: + """The document for *document_id*, or a model-facing error string.""" + store = get_document_store() + if store is None: + return _MISSING_STORE + document = store.get(document_id) + if document is None: + return _UNKNOWN.format(document_id=document_id) + return document + + +def _cap_body(body: str, max_chars: int) -> tuple[str, bool]: + if max_chars < 1 or len(body) <= max_chars: + return body, False + return body[:max_chars], True + + +def _preview(line: str) -> str: + return line if len(line) <= DEFAULT_LINE_PREVIEW else line[:DEFAULT_LINE_PREVIEW] + "…" + + +def _compile_pattern(pattern: str, *, case_insensitive: bool) -> re.Pattern[str] | str: + if not pattern: + return "error: pattern must not be empty" + flags = re.MULTILINE + if case_insensitive: + flags |= re.IGNORECASE + try: + return re.compile(pattern, flags) + except re.error as exc: + return f"error: invalid regex: {exc}" + + +def _collect_hits(lines: list[str], regex: re.Pattern[str], max_matches: int) -> list[str]: + hits: list[str] = [] + for number, line in enumerate(lines, start=1): + if regex.search(line) is None: + continue + hits.append(f"L{number}: {_preview(line)}") + if len(hits) >= max_matches: + break + return hits + + +def _search_with_context( + lines: list[str], + regex: re.Pattern[str], + max_matches: int, + context_lines: int, +) -> str: + matches = [number for number, line in enumerate(lines, start=1) if regex.search(line) is not None] + if not matches: + return "(no matches)" + truncated = len(matches) > max_matches + ranges: list[tuple[int, int]] = [] + for number in matches[:max_matches]: + low = max(1, number - context_lines) + high = min(len(lines), number + context_lines) + if ranges and low <= ranges[-1][1] + 1: + ranges[-1] = (ranges[-1][0], max(ranges[-1][1], high)) + else: + ranges.append((low, high)) + blocks = [ + "\n".join(f"L{number}: {_preview(lines[number - 1])}" for number in range(low, high + 1)) + for low, high in ranges + ] + body = "\n--\n".join(blocks) + if truncated: + body += f"\n...[truncated at {max_matches} matches]" + return body + + +def _search(text: str, regex: re.Pattern[str], max_matches: int, context_lines: int) -> str: + lines = text.splitlines() + if context_lines > 0: + return _search_with_context(lines, regex, max_matches, context_lines) + hits = _collect_hits(lines, regex, max_matches) + if not hits: + return "(no matches)" + body = "\n".join(hits) + if len(hits) >= max_matches: + body += f"\n...[truncated at {max_matches} matches]" + return body + + +@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.READ_ONLY) +async def read_document( + document_id: str, + *, + offset: int = 0, + limit: int | None = None, + with_lineno: bool = False, + max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS, +) -> str: + """Read a capped tool result as a document, one window at a time. + + A long tool result is stored whole in its tool message and capped on the wire + with a ``document_id`` marker; pass that id here to read past the cap. + ``offset`` is a 0-based line start (0 = first line), ``limit`` a line count + (None = the rest), and ``with_lineno`` prefixes each line with its 1-based + number (`` N|…``). When the window is cut by ``max_chars`` the reply names + the next ``offset`` to send. ``list_documents`` lists the session's ids. + """ + document = _resolve(document_id) + if isinstance(document, str): + return document + if offset < 0: + return "error: offset must be >= 0" + lines = document.text.splitlines(keepends=True) + if offset >= len(lines): + return "" + end = len(lines) if limit is None else min(len(lines), offset + limit) + slice_lines = lines[offset:end] + body = format_with_lineno(slice_lines, start_lineno=offset + 1) if with_lineno else "".join(slice_lines) + body, cut = _cap_body(body, max_chars) + shown_lines = max(1, body.count("\n")) if cut else len(slice_lines) + next_offset = offset + shown_lines + if next_offset < len(lines): + separator = "" if body.endswith("\n") else "\n" + body = f"{body}{separator}[... more: next offset={next_offset}]" + if with_lineno: + return body + return f"L{offset + 1}-{next_offset}\n{body}" + + +@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.READ_ONLY) +async def search_document( + document_id: str, + pattern: str, + *, + case_insensitive: bool = False, + max_matches: int = DEFAULT_MAX_MATCHES, + context_lines: int = 0, +) -> str: + """Search a document's lines with a Python ``re`` regex, per document. + + Like ``regex_files`` but over one tool result. ``pattern`` is a Python regex + matched per line (``re.MULTILINE``), not a glob. Returns ``L{line}: {content}`` + hits, long lines clipped with ``…``, capped at ``max_matches`` (a + ``...[truncated at N matches]`` marker follows a cut). ``context_lines`` adds + that many lines around each hit, blocks joined with ``--``. ``list_documents`` + lists the session's ids. + """ + if max_matches < 1: + return "error: max_matches must be >= 1" + if context_lines < 0: + return "error: context_lines must be >= 0" + document = _resolve(document_id) + if isinstance(document, str): + return document + regex = _compile_pattern(pattern, case_insensitive=case_insensitive) + if isinstance(regex, str): + return regex + return _search(document.text, regex, max_matches, context_lines) + + +@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.READ_ONLY) +async def document_info(document_id: str) -> str: + """Describe one document: source tool, size, and first/last line previews.""" + document = _resolve(document_id) + if isinstance(document, str): + return document + lines = document.text.splitlines() + first = _preview(lines[0]) if lines else "" + last = _preview(lines[-1]) if lines else "" + return "\n".join( + [ + f"document_id={document.document_id}", + f"tool={document.tool_name or '(unknown)'}", + f"chars={document.chars}", + f"lines={document.lines}", + f"first={first}", + f"last={last}", + ] + ) + + +@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.READ_ONLY) +async def list_documents() -> str: + """List the session's documents (capped tool results) with their ids and sizes.""" + store = get_document_store() + if store is None: + return _MISSING_STORE + documents = store.list() + if not documents: + return "(no documents)" + rows: list[str] = [] + for document in documents: + first_line = document.text.splitlines()[0] if document.text else "" + rows.append( + f"{document.index}. document_id={document.document_id} tool={document.tool_name or '(unknown)'} " + f"chars={document.chars} lines={document.lines} | {_preview(first_line)}" + ) + return "\n".join(rows) diff --git a/src/plyngent/tools/file/read.py b/src/plyngent/tools/file/read.py index b84f969..7fe8717 100644 --- a/src/plyngent/tools/file/read.py +++ b/src/plyngent/tools/file/read.py @@ -11,7 +11,7 @@ from plyngent.tools.workspace import AccessMode, resolve_path _LINENO_WIDTH = 6 -def _format_with_lineno(lines: list[str], *, start_lineno: int) -> str: +def format_with_lineno(lines: list[str], *, start_lineno: int) -> str: """Prefix each line with a 1-based absolute line number (``edit_lineno`` style).""" out: list[str] = [] for index, line in enumerate(lines): @@ -94,7 +94,7 @@ async def read_file( if start >= len(lines): return "" slice_lines = lines[start:end] - body = _format_with_lineno(slice_lines, start_lineno=start + 1) if with_lineno else "".join(slice_lines) + body = format_with_lineno(slice_lines, start_lineno=start + 1) if with_lineno else "".join(slice_lines) char_start = len("".join(lines[:start])) if max_chars >= 1: body, token = truncate_with_token( diff --git a/tests/test_tools/test_catalog.py b/tests/test_tools/test_catalog.py index 9463783..236d999 100644 --- a/tests/test_tools/test_catalog.py +++ b/tests/test_tools/test_catalog.py @@ -61,6 +61,10 @@ def test_read_only_classification() -> None: "read_file", "view_image", "get_truncated", + "read_document", + "search_document", + "document_info", + "list_documents", "glob_paths", "regex_files", "listdir", diff --git a/tests/test_tools/test_document.py b/tests/test_tools/test_document.py new file mode 100644 index 0000000..8160cce --- /dev/null +++ b/tests/test_tools/test_document.py @@ -0,0 +1,158 @@ +"""Document tools: read/search/list the tool results the request-prep pass capped.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import msgspec + +from plyngent.lmproto.openai_compatible.model import ( + AnyChatMessage, + AssistantChatMessage, + AssistantFunctionTool, + AssistantFunctionToolCall, + ToolChatMessage, + UserChatMessage, +) +from plyngent.tools.context import InstanceState, bind_instance +from plyngent.tools.documents import DocumentStore +from plyngent.tools.file import document_info, list_documents, read_document, search_document +from tests.test_tools.helpers import call_async + +if TYPE_CHECKING: + from collections.abc import Generator, Sequence + + +def _history(text: str, *, tool_call_id: str = "call-1", tool_name: str = "some_tool") -> list[AnyChatMessage]: + return [ + UserChatMessage(content="go"), + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id=tool_call_id, + function=AssistantFunctionTool(name=tool_name, arguments="{}"), + ) + ], + ), + ToolChatMessage(content=text, tool_call_id=tool_call_id), + ] + + +@contextmanager +def bound_document( + text: str, + *, + tool_call_id: str = "call-1", + tool_name: str = "some_tool", +) -> Generator[DocumentStore]: + """Bind an instance whose document store reads one tool result.""" + messages = _history(text, tool_call_id=tool_call_id, tool_name=tool_name) + instance = InstanceState() + store = DocumentStore(lambda: messages) + instance.extras["documents"] = store + with bind_instance(instance): + yield store + + +_TEN_LINES = "".join(f"line {i}\n" for i in range(1, 11)) + + +async def test_read_document_without_store(workspace: object) -> None: + del workspace + out = await call_async(read_document, "call-1") + assert out.startswith("error:") + + +async def test_read_document_unknown_id() -> None: + with bound_document(_TEN_LINES): + out = await call_async(read_document, "nope") + assert out.startswith("error:") + assert "unknown document" in out + + +async def test_read_document_window_and_range() -> None: + with bound_document(_TEN_LINES): + window = await call_async(read_document, "call-1", offset=0, limit=3) + whole = await call_async(read_document, "call-1") + mid = await call_async(read_document, "call-1", offset=2, limit=2) + assert window == "L1-3\nline 1\nline 2\nline 3\n[... more: next offset=3]" + assert whole == f"L1-10\n{_TEN_LINES}" + assert "next offset" not in whole + assert mid.startswith("L3-4\nline 3\nline 4\n") + + +async def test_read_document_numbered() -> None: + with bound_document(_TEN_LINES): + out = await call_async(read_document, "call-1", offset=1, limit=2, with_lineno=True) + assert out.startswith(" 2|line 2\n 3|line 3\n") + + +async def test_read_document_cut_names_next_offset() -> None: + text = "".join(f"a very long line number {i}\n" for i in range(1, 21)) + with bound_document(text): + out = await call_async(read_document, "call-1", max_chars=40) + assert "next offset=" in out + offset = int(out.rsplit("next offset=", 1)[1].split("]", 1)[0]) + nxt = await call_async(read_document, "call-1", offset=offset) + assert nxt.startswith(f"L{offset + 1}-") + + +async def test_read_document_offset_past_end() -> None: + with bound_document(_TEN_LINES): + out = await call_async(read_document, "call-1", offset=99) + assert out == "" + + +async def test_search_document_hits_and_miss() -> None: + with bound_document(_TEN_LINES): + hits = await call_async(search_document, "call-1", r"line \d+") + miss = await call_async(search_document, "call-1", "absent") + assert hits.splitlines()[0] == "L1: line 1" + assert "L10: line 10" in hits + assert miss == "(no matches)" + + +async def test_search_document_context_and_cap() -> None: + with bound_document(_TEN_LINES): + block = await call_async(search_document, "call-1", "line 5", context_lines=1) + capped = await call_async(search_document, "call-1", r"line \d+", max_matches=2) + assert block == "L4: line 4\nL5: line 5\nL6: line 6" + assert "...[truncated at 2 matches]" in capped + + +async def test_search_document_invalid_regex() -> None: + with bound_document(_TEN_LINES): + out = await call_async(search_document, "call-1", "(") + assert out.startswith("error: invalid regex") + + +async def test_document_info_and_list() -> None: + with bound_document(_TEN_LINES, tool_name="run_argv"): + info = await call_async(document_info, "call-1") + listing = await call_async(list_documents) + assert "tool=run_argv" in info + assert "lines=10" in info + assert "first=line 1" in info + assert "last=line 10" in info + assert "document_id=call-1 tool=run_argv" in listing + assert "chars=" in listing + + +async def test_list_documents_without_store(workspace: object) -> None: + del workspace + assert (await call_async(list_documents)).startswith("error:") + + +async def test_document_survives_serialization_roundtrip() -> None: + """A resume reloads messages from JSON; the store still resolves the id.""" + messages: Sequence[AnyChatMessage] = _history(_TEN_LINES) + reloaded: list[AnyChatMessage] = [ + msgspec.convert(msgspec.to_builtins(message), type=AnyChatMessage) for message in messages + ] + instance = InstanceState() + instance.extras["documents"] = DocumentStore(lambda: reloaded) + with bind_instance(instance): + out = await call_async(read_document, "call-1", offset=0, limit=1) + assert out.startswith("L1-1\nline 1\n")