feat(documents): document_read tool + document-context analytics (offload PR-1)

backend/src/agents/builtin_tools/document_read_tool.py

@@ -0,0 +1,189 @@
+"""``document_read`` — page-range and pattern retrieval over the session's documents.
+
+The recovery half of ``docs/specs/document-context-offload.md`` (§4B). A
+document the user attached lives in the model's context only until history
+is restored (restore strips inline bytes; see
+``TurnBasedSessionManager._strip_document_bytes``) or, later, until the
+offload trigger swaps it for a digest. This tool is how the model gets the
+part it needs back — a PDF page range as a native ``document`` block (full
+fidelity: text layer plus page images), a regex over the text layer to find
+the right pages, or bounded text for text-family documents.
+
+Design notes
+------------
+* **Gated on session state, not on the tool picker.** The tool is built for
+  any session that has a readable document (``_build_document_tools`` in
+  ``apis/inference_api/chat/routes.py``), whatever the user's RBAC grants.
+  The ``workspace_files`` catalog key is granted to no prod role, so gating
+  there would ship the recovery path dark. Its id therefore stays OUT of
+  ``INJECTED_TOOL_IDS`` — same reasoning as the Memory-Space tools: the
+  governing capability is the user's own attachment.
+* Identity is captured by closure (``make_document_read_tool``), never read
+  from ``invocation_state``.
+* One call can never re-inject a whole document: page ranges are capped by
+  ``max_pages`` (default 8, hard cap 20), pattern results by match count,
+  text by the workspace read bound.
+* **Content-free record per call** in ``AgentCoreStack/Compaction``
+  (``DocumentRead``, ``DocumentReadPages``, ``DocumentReadBytes``; properties
+  ``mode`` / ``format``) — the same namespace as the compaction cut and
+  tool-result offload records, so one dashboard explains what bounded a
+  session's prefix. Nothing from the document is ever emitted.
+* Tool results carrying a native ``document`` block are exempt from the
+  tool-result offloader (``core/tool_result_offload.py``): offloading the
+  slice the model just asked for would defeat the read.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, Optional
+
+from strands import tool
+
+from apis.shared.files.document_read import (
+    DOCUMENT_READ_HARD_MAX_PAGES,
+    DOCUMENT_READ_MAX_PAGES,
+    DocumentReadResult,
+    list_session_documents,
+    read_document,
+)
+from apis.shared.files.workspace import WorkspaceError, WorkspaceStorageNotConfiguredError
+
+logger = logging.getLogger(__name__)
+
+#: The tool's name — also the census / ledger key and the offloader exemption.
+DOCUMENT_READ_TOOL_NAME = "document_read"
+
+_NO_STORAGE_MESSAGE = (
+    "❌ Document storage is not configured (S3_USER_FILES_BUCKET_NAME is not set on the runtime)."
+)
+
+
+def _error(text: str) -> Dict[str, Any]:
+    return {"content": [{"text": text}], "status": "error"}
+
+
+def _success(result: DocumentReadResult) -> Dict[str, Any]:
+    content = [{"json": result.payload}]
+    if result.document_block is not None:
+        content.append(result.document_block)
+    content.extend(result.extra_blocks)
+    return {"content": content, "status": "success"}
+
+
+def record_document_read(result: DocumentReadResult) -> None:
+    """One content-free EMF record per successful read. Never raises."""
+    logger.info(
+        "document_read: mode=%s format=%s pages=%d bytes=%d",
+        result.mode, result.format, result.pages_returned, result.bytes_returned,
+    )
+    try:
+        from apis.shared.observability.prompt_cache import prompt_cache_observability_enabled
+        from apis.shared.observability.emf import emit_emf_metrics
+
+        if not prompt_cache_observability_enabled():
+            return
+        emit_emf_metrics(
+            "AgentCoreStack/Compaction",
+            metrics={
+                "DocumentRead": 1,
+                "DocumentReadPages": int(result.pages_returned),
+                "DocumentReadBytes": int(result.bytes_returned),
+            },
+            properties={"mode": result.mode, "format": result.format},
+            units={"DocumentReadPages": "Count", "DocumentReadBytes": "Bytes"},
+        )
+    except Exception as e:  # noqa: BLE001
+        logger.debug("document_read EMF skipped: %s", e)
+
+
+def make_document_read_tool(session_id: str, user_id: str):
+    """Create a ``document_read`` tool bound to the given identity."""
+    if not session_id or not user_id:
+        raise ValueError("document_read requires a session_id and a user_id")
+
+    @tool(name=DOCUMENT_READ_TOOL_NAME)
+    async def document_read(
+        upload_id: str = "",
+        page_range: str = "",
+        pattern: str = "",
+        max_pages: int = DOCUMENT_READ_MAX_PAGES,
+        offset: int = 0,
+    ) -> Any:
+        """Read part of a document the user attached to this conversation.
+
+        Use this when a document you need is no longer inline (its content
+        was replaced by a placeholder or digest after the conversation was
+        restored) or when you need specific pages of a long document.
+
+        Modes, by argument:
+        - No arguments: list this conversation's readable documents with
+          their upload_id, format and size.
+        - upload_id + page_range (PDF only, e.g. "4-7"): returns those pages
+          as an attached document at full fidelity (text and page images),
+          at most max_pages per call. The attachment's own page k is
+          original page start+k-1 — cite original page numbers.
+        - upload_id + pattern: case-insensitive regex over the text; returns
+          matching lines with page numbers (PDF) or line/paragraph numbers
+          (text, Word), so you can then read the right pages. Use this first
+          on a long document instead of paging through it.
+        - upload_id alone: for a PDF, the page count and a per-page snippet
+          index; for text and Word documents, the text itself (bounded —
+          continue with offset = next_offset when truncated).
+
+        Prefer pattern then a narrow page_range over reading many pages.
+        Spreadsheets are read with analyze_spreadsheet and presentations
+        with read_powerpoint_presentation, not here.
+
+        Args:
+            upload_id: The document's id (from the listing, the attachment
+                note, or an earlier call). Empty to list documents.
+            page_range: 1-indexed inclusive page range for PDFs, "start-end"
+                or a single page "5".
+            pattern: Regular expression (case-insensitive) to search for.
+            max_pages: Cap on pages returned by page_range (default 8, hard
+                cap 20).
+            offset: Byte offset to continue a truncated text read.
+
+        Returns:
+            JSON metadata, plus the requested pages as an attached document
+            for page_range reads.
+        """
+        try:
+            if not upload_id:
+                listing = await list_session_documents(user_id, session_id)
+                listing["hint"] = (
+                    "Call again with upload_id plus page_range (PDF), pattern, or neither."
+                )
+                result = DocumentReadResult(mode="list", payload=listing)
+            else:
+                result = await read_document(
+                    user_id,
+                    session_id,
+                    upload_id,
+                    page_range=page_range or None,
+                    pattern=pattern or None,
+                    max_pages=max_pages,
+                    offset=offset,
+                )
+        except WorkspaceStorageNotConfiguredError:
+            return _error(_NO_STORAGE_MESSAGE)
+        except WorkspaceError as exc:
+            return _error(f"❌ {exc}")
+        except Exception as exc:  # noqa: BLE001 - surface conversationally, never raise through the loop
+            logger.error("document_read error: %s", exc, exc_info=True)
+            return _error(f"❌ Failed to read document '{upload_id}': {exc}")
+
+        record_document_read(result)
+        return _success(result)
+
+    return document_read
+
+
+__all__ = [
+    "DOCUMENT_READ_HARD_MAX_PAGES",
+    "DOCUMENT_READ_MAX_PAGES",
+    "DOCUMENT_READ_TOOL_NAME",
+    "make_document_read_tool",
+    "record_document_read",
+]

backend/src/agents/main_agent/core/tool_result_offload.py

@@ -57,6 +57,13 @@
 # half the gate leaves that margin twice over.
 PREFILTER_RATIO = 0.5
 
+# Tools whose results are never offloaded. ``document_read`` returns the page
+# slice the model just asked for as a native document block; offloading it
+# to S3 and handing back a text preview would undo the read and cost a second
+# round trip (docs/specs/document-context-offload.md §4B). Its own
+# ``max_pages`` cap is the bound.
+OFFLOAD_EXEMPT_TOOLS = frozenset({"document_read"})
+
 
 def tool_result_offload_enabled() -> bool:
     """Default ON with a kill switch (house style): only the literal "false" disables."""
@@ -77,6 +84,21 @@ def _import_plugin():
     return ContextOffloader
 
 
+def _exempt_tool(event: Any) -> bool:
+    """True when the event's tool is in ``OFFLOAD_EXEMPT_TOOLS``."""
+    try:
+        name = (getattr(event, "tool_use", None) or {}).get("name")
+    except Exception:  # noqa: BLE001
+        return False
+    return name in OFFLOAD_EXEMPT_TOOLS
+
+
+def _should_offload(tool_name: str, _token_count: int) -> bool:
+    """The plugin's own ``should_offload`` hook — defense in depth behind the
+    mixin's early return, for the path where the base class is reached."""
+    return tool_name not in OFFLOAD_EXEMPT_TOOLS
+
+
 class _OffloaderMixin:
     """Behavior layered on the vended plugin; kept separate so it can be tested
     against a stub base class without importing the real one."""
@@ -89,6 +111,8 @@ async def _handle_tool_result(self, event: Any) -> None:  # type: ignore[overrid
             content = result.get("content")
             if not isinstance(content, list):
                 return
+            if _exempt_tool(event):
+                return
             # Cheap pre-filter: skip the CountTokens round trip for results
             # that cannot be anywhere near the gate.
             estimate = estimate_message_tokens({"role": "user", "content": [{"toolResult": result}]})
@@ -179,6 +203,7 @@ def build_tool_result_offloader(
             max_result_tokens=max_tokens,
             preview_tokens=preview,
             include_retrieval_tool=True,
+            should_offload=_should_offload,
             evict_after_cycles=None,
         )
     except Exception:  # noqa: BLE001

backend/src/agents/main_agent/session/document_context.py

@@ -0,0 +1,142 @@
+"""Content-free summary of the attachments in a conversation's live context.
+
+The cost rows could say what a call cost and (since the context ledger) what
+compaction did to the history before it, but not *how much of the context
+was documents* — the one quantity ``docs/specs/document-context-offload.md``
+needs to decide whether its cost work is worth shipping (evaluation spec §4.1:
+the recoverable envelope is the document share of every cold re-write, and
+that share was unmeasured). This module measures it from ``agent.messages``
+at turn end, and the stream coordinator persists the numbers on each of the
+turn's ``C#`` rows next to ``prefixTokens`` and ``compactionEvents``.
+
+Everything here is a count, a byte size or a token estimate. No filename,
+title, text or document byte ever leaves this function — the MIME map is
+keyed by Bedrock's ``document.format`` enum (``pdf`` / ``docx`` / …) plus
+``image``.
+
+Token numbers are the compaction estimator's heuristics (bytes/4 for a
+document, a flat figure per image), not Bedrock counts: Bedrock reports no
+per-block usage. They are comparable across rows and against
+``contextBreakdown.messages``, which is the measured total they are a share of.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from agents.main_agent.session.compaction_policy import CHARS_PER_TOKEN, IMAGE_TOKEN_ESTIMATE
+
+#: A restore-time stand-in for a document that is no longer inline — today the
+#: contentless placeholder ``_strip_document_bytes`` writes; from PR-3 the
+#: digest block. Both count as "the model has a reference, not the bytes".
+DIGEST_MARKERS = ("[Document placeholder:", "<document-digest")
+
+
+def _inline_bytes(block: Dict[str, Any], key: str) -> Optional[int]:
+    """Byte length of an inline ``document`` / ``image`` block, or ``None``
+    when the block is not one (or carries no bytes)."""
+    payload = block.get(key)
+    if not isinstance(payload, dict):
+        return None
+    raw = (payload.get("source") or {}).get("bytes")
+    if isinstance(raw, (bytes, bytearray)):
+        return len(raw)
+    return None
+
+
+def _is_digest(block: Dict[str, Any]) -> bool:
+    text = block.get("text")
+    return isinstance(text, str) and text.lstrip().startswith(DIGEST_MARKERS)
+
+
+def _is_turn_prompt(message: Dict[str, Any]) -> bool:
+    """The user message that carries the turn's prompt (never a tool-result
+    message, which is also ``role: user``)."""
+    if message.get("role") != "user":
+        return False
+    content = message.get("content")
+    if not isinstance(content, list):
+        return isinstance(content, str)
+    return not any(isinstance(b, dict) and "toolResult" in b for b in content)
+
+
+def summarize_document_context(messages: Optional[List[Dict[str, Any]]]) -> Optional[Dict[str, Any]]:
+    """The attachment footprint of ``messages`` (the live conversation).
+
+    Returns ``None`` for an empty or malformed list, otherwise a dict with:
+
+    - ``hasDocuments`` — at least one inline attachment (document or image
+      bytes) is in context.
+    - ``documentCount`` / ``documentTokens`` — inline attachment blocks on
+      user prompts and their estimated token weight.
+    - ``documentDigests`` — digest / placeholder stand-ins in context (a turn
+      with digests and no inline documents answered from the digest).
+    - ``documentsAttached`` — attachment blocks on the most recent prompt
+      (an attach turn vs. a follow-up).
+    - ``documentSlices`` / ``documentSliceTokens`` — document blocks that
+      ``document_read`` returned inside tool results and that still live in
+      the history (the re-injected pages, bounded by the tool's cap).
+    - ``documentMime`` — ``{format: count}`` over the inline attachments,
+      ``image`` for image blocks.
+    """
+    if not isinstance(messages, list) or not messages:
+        return None
+
+    count = tokens = digests = slices = slice_tokens = 0
+    mime: Dict[str, int] = {}
+    last_prompt_attachments = 0
+
+    for message in messages:
+        if not isinstance(message, dict):
+            continue
+        content = message.get("content")
+        if not isinstance(content, list):
+            continue
+        prompt = _is_turn_prompt(message)
+        attached_here = 0
+        for block in content:
+            if not isinstance(block, dict):
+                continue
+            size = _inline_bytes(block, "document")
+            if size is not None:
+                count += 1
+                attached_here += 1
+                tokens += size // CHARS_PER_TOKEN
+                fmt = str((block.get("document") or {}).get("format") or "unknown")
+                mime[fmt] = mime.get(fmt, 0) + 1
+                continue
+            size = _inline_bytes(block, "image")
+            if size is not None:
+                count += 1
+                attached_here += 1
+                tokens += IMAGE_TOKEN_ESTIMATE
+                mime["image"] = mime.get("image", 0) + 1
+                continue
+            if _is_digest(block):
+                digests += 1
+                continue
+            tool_result = block.get("toolResult")
+            if isinstance(tool_result, dict):
+                for inner in tool_result.get("content") or []:
+                    if not isinstance(inner, dict):
+                        continue
+                    inner_size = _inline_bytes(inner, "document")
+                    if inner_size is not None:
+                        slices += 1
+                        slice_tokens += inner_size // CHARS_PER_TOKEN
+        if prompt:
+            last_prompt_attachments = attached_here
+
+    return {
+        "hasDocuments": count > 0,
+        "documentCount": count,
+        "documentTokens": tokens,
+        "documentDigests": digests,
+        "documentsAttached": last_prompt_attachments,
+        "documentSlices": slices,
+        "documentSliceTokens": slice_tokens,
+        "documentMime": mime,
+    }
+
+
+__all__ = ["DIGEST_MARKERS", "summarize_document_context"]