diff --git a/backend/src/agents/builtin_tools/document_read_tool.py b/backend/src/agents/builtin_tools/document_read_tool.py new file mode 100644 index 00000000..e99ab8b9 --- /dev/null +++ b/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", +] diff --git a/backend/src/agents/main_agent/core/tool_result_offload.py b/backend/src/agents/main_agent/core/tool_result_offload.py index bab3b21a..17ab9523 100644 --- a/backend/src/agents/main_agent/core/tool_result_offload.py +++ b/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 diff --git a/backend/src/agents/main_agent/session/document_context.py b/backend/src/agents/main_agent/session/document_context.py new file mode 100644 index 00000000..82e240c3 --- /dev/null +++ b/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:", " 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"] diff --git a/backend/src/agents/main_agent/session/document_rehydration.py b/backend/src/agents/main_agent/session/document_rehydration.py new file mode 100644 index 00000000..5991fcc3 --- /dev/null +++ b/backend/src/agents/main_agent/session/document_rehydration.py @@ -0,0 +1,229 @@ +"""Rehydrate stripped documents on restore — digest plus a live handle, not a +contentless placeholder (docs/specs/document-context-offload.md §4E, PR-3). + +Restore must drop inline document bytes from history: Bedrock rejects any +request where two document blocks share a name across the conversation, and +a re-attached file collides with the copy already in history. Until PR-3 the +replacement was ``[Document placeholder: name=…, format=…, original_size=…]`` +— zero content — so a returning user's document was silently gone (the 14% +same-file re-upload rate). This module replaces the bytes with the document's +``DocumentDigest`` rendered as a ```` block that names the +``upload_id`` the ``document_read`` tool (PR-1) reads pages back from. + +How a block finds its file: the session's upload rows (``SessionIndex``, one +query per restore that has a document to rehydrate, none otherwise) are +matched on the sanitized filename — the same sanitizer ``PromptBuilder`` used +when the block was built, allowing for its ``_2`` / ``_3`` duplicate suffix — +and the byte size. A block with no matching row (a direct base64 attachment +that never went through the upload flow, or a deleted file) keeps today's +placeholder, so nothing here can be worse than before. + +Digests missing from the row (uploads that predate PR-2, agent-written files) +are built lazily from the bytes that are *right there* in the restored +message — the outline only, no model call on the restore path — and +persisted so the next restore renders the same bytes. The upload-path build +(PR-2) may later overwrite an outline-only digest with one that has an +abstract; that changes the rendered block once, which is one prefix +re-write, accepted and recorded (``document_rehydrated`` carries the digest +tokens, so the change is visible). + +Everything is synchronous: it runs inside ``TurnBasedSessionManager.initialize`` +under the Strands agent constructor. It never raises — every failure falls +back to the placeholder for that block. +""" + +from __future__ import annotations + +import copy +import logging +import os +import re +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set + +from agents.main_agent.multimodal.file_sanitizer import FileSanitizer +from agents.main_agent.session.compaction_policy import CHARS_PER_TOKEN + +logger = logging.getLogger(__name__) + +_DUPLICATE_SUFFIX = re.compile(r"_(\d+)$") + + +def document_rehydration_enabled() -> bool: + """Default ON with a kill switch (house style): ``DOCUMENT_REHYDRATE_ENABLED=false`` + restores today's placeholder behavior without touching PR-1/PR-2.""" + return os.environ.get("DOCUMENT_REHYDRATE_ENABLED", "").strip().lower() != "false" + + +def placeholder_text(name: str, fmt: str, size: int) -> str: + """The pre-PR-3 stand-in, kept byte-identical for blocks that cannot be + matched to an upload (restore output must be stable across restores).""" + return f"[Document placeholder: name={name}, format={fmt}, original_size={size} bytes]" + + +@dataclass +class RehydrationResult: + messages: List[Dict[str, Any]] + rehydrated: int = 0 + stripped: int = 0 + digest_tokens: int = 0 + stripped_tokens: int = 0 + lazy_digests: int = 0 + upload_ids: List[str] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Matching +# --------------------------------------------------------------------------- + + +def _base_name(block_name: str) -> str: + """``"policy_2"`` → ``"policy"`` (PromptBuilder's duplicate suffix).""" + return _DUPLICATE_SUFFIX.sub("", block_name) + + +def match_document( + name: str, + fmt: str, + size: int, + candidates: List[Any], + used: Set[str], +) -> Optional[Any]: + """The upload row a document block came from, or ``None``. + + Preference order: same sanitized filename **and** same byte size; then + same filename; then same format and size (a renamed re-upload). Rows + already claimed by another block in this restore are skipped, so two + copies of one file map to two rows, newest first. + """ + sanitize = FileSanitizer.sanitize_filename + wanted = {name, _base_name(name)} + fresh = [c for c in candidates if c.upload_id not in used] + by_name = [c for c in fresh if sanitize(c.filename) in wanted] + for c in by_name: + if c.size_bytes == size: + return c + if by_name: + return by_name[0] + for c in fresh: + if c.size_bytes == size and (c.file_format or "") == fmt: + return c + return None + + +# --------------------------------------------------------------------------- +# Data access (sync — see module docstring) +# --------------------------------------------------------------------------- + + +def load_session_documents(session_id: str, user_id: Optional[str]) -> List[Any]: + """This user's READY document-class uploads in the session, newest first.""" + from apis.shared.files.document_read import is_document_class + from apis.shared.files.models import FileStatus + from apis.shared.files.repository import get_file_upload_repository + + rows = get_file_upload_repository().list_session_files_sync(session_id, status=FileStatus.READY) + return [ + meta for meta in rows + if (not user_id or meta.user_id == user_id) and is_document_class(meta.mime_type, meta.filename) + ] + + +def digest_for(meta: Any, raw: bytes) -> Optional[Any]: + """The row's ready ``DocumentDigest``, or one built now from ``raw`` + (outline only) and persisted best-effort. ``None`` if neither works.""" + from apis.shared.files.document_digest import DocumentDigest, estimate_tokens, extract_outline, render_digest + from apis.shared.files.repository import get_file_upload_repository + + existing = DocumentDigest.from_item(getattr(meta, "digest", None)) + if existing is not None and existing.status == "ready": + return existing + + fmt = meta.file_format + if not fmt: + return None + digest = extract_outline(fmt, raw) + digest.tokens = estimate_tokens(render_digest(digest, filename=meta.filename, upload_id=meta.upload_id)) + try: + get_file_upload_repository().update_file_digest_sync(meta.user_id, meta.upload_id, digest.to_item()) + except Exception: # noqa: BLE001 - the digest still serves this restore + logger.debug("Lazy digest not persisted for upload %s", meta.upload_id, exc_info=True) + return digest + + +# --------------------------------------------------------------------------- +# The pass +# --------------------------------------------------------------------------- + + +def rehydrate_documents( + messages: List[Dict[str, Any]], + *, + session_id: Optional[str], + user_id: Optional[str], +) -> RehydrationResult: + """Replace every inline document block with its digest block (matched to + an upload row) or, failing that, the placeholder. Never raises.""" + from apis.shared.files.document_digest import render_digest + + out = copy.deepcopy(messages) + result = RehydrationResult(messages=out) + enabled = document_rehydration_enabled() and bool(session_id) + candidates: Optional[List[Any]] = None + used: Set[str] = set() + + for msg in out: + content = msg.get("content", []) if isinstance(msg, dict) else None + if not isinstance(content, list): + continue + for idx, block in enumerate(content): + if not isinstance(block, dict) or "document" not in block: + continue + doc = block["document"] if isinstance(block["document"], dict) else {} + source = doc.get("source") or {} + if not isinstance(source, dict) or "bytes" not in source: + continue + raw = source.get("bytes", b"") + size = len(raw) if isinstance(raw, (bytes, bytearray)) else 0 + name = str(doc.get("name", "unknown")) + fmt = str(doc.get("format", "unknown")) + + replacement: Optional[str] = None + if enabled: + try: + if candidates is None: + candidates = load_session_documents(session_id, user_id) + meta = match_document(name, fmt, size, candidates, used) + if meta is not None: + digest = digest_for(meta, bytes(raw) if isinstance(raw, (bytes, bytearray)) else b"") + if digest is not None: + was_lazy = not (isinstance(getattr(meta, "digest", None), dict) and meta.digest.get("status") == "ready") + replacement = render_digest(digest, filename=meta.filename, upload_id=meta.upload_id) + used.add(meta.upload_id) + result.rehydrated += 1 + result.digest_tokens += len(replacement) // CHARS_PER_TOKEN + result.lazy_digests += 1 if was_lazy else 0 + result.upload_ids.append(meta.upload_id) + except Exception: # noqa: BLE001 - fall back to the placeholder for this block + logger.warning("Document rehydration failed for one block; using placeholder", exc_info=True) + if candidates is None: + candidates = [] # do not retry the lookup for every block + + if replacement is None: + replacement = placeholder_text(name, fmt, size) + result.stripped += 1 + result.stripped_tokens += size // CHARS_PER_TOKEN + content[idx] = {"text": replacement} + + return result + + +__all__ = [ + "RehydrationResult", + "digest_for", + "document_rehydration_enabled", + "load_session_documents", + "match_document", + "placeholder_text", + "rehydrate_documents", +] diff --git a/backend/src/agents/main_agent/session/hooks/context_ledger.py b/backend/src/agents/main_agent/session/hooks/context_ledger.py index 263d0bf9..85f6d276 100644 --- a/backend/src/agents/main_agent/session/hooks/context_ledger.py +++ b/backend/src/agents/main_agent/session/hooks/context_ledger.py @@ -15,6 +15,11 @@ and whatever a future scheduling policy records: forced cut, floor unreachable). Each carries the summary's token size at that moment, which is what proves a summary cap shrank summaries without another scan. +- ``documentReads`` — how many ``document_read`` retrievals the call + requested and how many pages / bytes they pulled back into context + (docs/specs/document-context-offload.md §6.1). With the per-row document + context fields the coordinator adds, this is what says whether a turn + answered from the full document, from a digest, or from retrieved pages. Per-turn, per-model-call, held on the agent wrapper exactly like the tool census: the stream coordinator reads ``ledger_for_call(idx)`` at turn end and @@ -30,7 +35,13 @@ import logging from typing import Any, Dict, List, Optional -from strands.hooks import BeforeInvocationEvent, BeforeModelCallEvent, HookProvider, HookRegistry +from strands.hooks import ( + AfterToolCallEvent, + BeforeInvocationEvent, + BeforeModelCallEvent, + HookProvider, + HookRegistry, +) from apis.shared.feature_flags import cost_diagnostics_enabled @@ -40,6 +51,40 @@ #: grow a cost row without bound. _MAX_EVENTS_PER_CALL = 8 +#: The document-retrieval tool whose results the ledger tallies +#: (``documentReads``): calls, pages and bytes the model pulled back into +#: context. Read from the tool result's own ``json`` block — numbers only. +DOCUMENT_READ_TOOL_NAME = "document_read" + + +def _document_read_counts(event: Any) -> Optional[Dict[str, int]]: + """``{"calls": 1, "pages": n, "bytes": b}`` for a successful + ``document_read`` result, ``None`` for anything else. Only the result's + numeric metadata is read; the document block itself is never inspected.""" + tool_use = getattr(event, "tool_use", None) or {} + if not isinstance(tool_use, dict) or tool_use.get("name") != DOCUMENT_READ_TOOL_NAME: + return None + result = getattr(event, "result", None) + if not isinstance(result, dict) or result.get("status") != "success": + return None + pages = 0 + size = 0 + for block in result.get("content") or []: + if not isinstance(block, dict): + continue + payload = block.get("json") + if isinstance(payload, dict): + try: + pages += int(payload.get("pages_returned") or 0) + except (TypeError, ValueError): + pass + document = block.get("document") + if isinstance(document, dict): + raw = (document.get("source") or {}).get("bytes") + if isinstance(raw, (bytes, bytearray)): + size += len(raw) + return {"calls": 1, "pages": pages, "bytes": size} + def _removed_message_count(agent: Any) -> Optional[int]: manager = getattr(agent, "conversation_manager", None) @@ -81,6 +126,25 @@ def __init__(self) -> None: def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: registry.add_callback(BeforeInvocationEvent, self._on_turn_start) registry.add_callback(BeforeModelCallEvent, self._on_before_model_call) + registry.add_callback(AfterToolCallEvent, self._on_after_tool_call) + + def _on_after_tool_call(self, event: AfterToolCallEvent) -> None: + """Tally ``document_read`` retrievals against the model call that + requested them (same cycle attribution as the tool census: a tool + running during cycle N was requested by call N). The pages come back + into context for call N+1; the row of call N says the model asked.""" + if not cost_diagnostics_enabled(): + return + try: + counts = _document_read_counts(event) + if counts is None: + return + entry = self._ledger.setdefault(self._cycle, {}) + reads = entry.setdefault("documentReads", {"calls": 0, "pages": 0, "bytes": 0}) + for key, value in counts.items(): + reads[key] = reads.get(key, 0) + value + except Exception as e: # noqa: BLE001 - a ledger must never break a turn + logger.debug("Context ledger skipped a tool result: %s", e) def ledger_for_call(self, call_index: int) -> Optional[Dict[str, Any]]: """The ledger entry for model call ``call_index`` (0-based), or diff --git a/backend/src/agents/main_agent/session/turn_based_session_manager.py b/backend/src/agents/main_agent/session/turn_based_session_manager.py index 7520bf3e..3da89c60 100644 --- a/backend/src/agents/main_agent/session/turn_based_session_manager.py +++ b/backend/src/agents/main_agent/session/turn_based_session_manager.py @@ -51,7 +51,18 @@ #: Compaction decisions the per-call ledger will record. Reserved kinds exist #: so a scheduling policy can report them without a schema change. -COMPACTION_EVENT_KINDS = frozenset({"applied", "checkpoint", "forced", "floor_unreachable"}) +#: Compaction-ledger event kinds. The ``document_*`` kinds are the document +#: lifecycle (docs/specs/document-context-offload.md §6.1): ``document_stripped`` +#: — restore replaced inline documents with contentless placeholders (the +#: defect PR-3 fixes; recorded from PR-1 so its reach is measured before the +#: fix lands); ``document_rehydrated`` — restore replaced them with a digest +#: plus a live ``document_read`` handle (PR-3); ``document_offload`` — the +#: post-turn trigger swapped an inline document for its digest, with the +#: cache gap at the moment it fired (PR-4). +COMPACTION_EVENT_KINDS = frozenset({ + "applied", "checkpoint", "forced", "floor_unreachable", + "document_stripped", "document_rehydrated", "document_offload", +}) _MAX_PENDING_COMPACTION_EVENTS = 8 @@ -1446,7 +1457,8 @@ def _sanitize_restored_content_blocks(self, messages: List[Dict]) -> List[Dict]: return sanitized def _strip_document_bytes(self, messages: List[Dict]) -> List[Dict]: - """Replace document content blocks' inline bytes with a text placeholder. + """Replace document content blocks' inline bytes with their digest — or, + when the block cannot be matched to an upload, a text placeholder. Called unconditionally on every session restore — independent of whether compaction is enabled. Document blocks with ``source.bytes`` must never @@ -1454,42 +1466,45 @@ def _strip_document_bytes(self, messages: List[Dict]) -> List[Dict]: document blocks share the same sanitized name across the conversation (ValidationException: "Messages can't contain duplicate document names"). + Since PR-3 of the offload spec the replacement is a ```` + block carrying the document's outline, abstract and ``upload_id``, so the + model still knows what the document says and can pull any page back with + ``document_read`` (see ``document_rehydration.py``). Blocks with no + upload row keep the contentless placeholder, exactly as before. + Images are handled the same way inside ``_truncate_tool_contents``, but that method is gated on compaction being enabled. This one is not. - The ``[Attached files: …]`` text marker already present in the user - message preserves the reference for the model without re-sending bytes. + Two content-free ledger events record what happened on the next cost + row: ``document_rehydrated`` (documents and their digest tokens) and + ``document_stripped`` (documents that fell back to the placeholder and + the tokens they were). The second going to zero is PR-3's gate. """ - stripped_messages = copy.deepcopy(messages) - strip_count = 0 - - for msg in stripped_messages: - content = msg.get("content", []) - if not isinstance(content, list): - continue - for block_idx, block in enumerate(content): - if not isinstance(block, dict) or "document" not in block: - continue - doc_data = block["document"] - source = doc_data.get("source", {}) - # Only replace blocks that carry inline bytes — s3Location - # blocks (enhancement #401) have no bytes to strip and are - # safe to leave as-is since they don't accumulate in history. - if "bytes" not in source: - continue - doc_name = doc_data.get("name", "unknown") - doc_format = doc_data.get("format", "unknown") - original_bytes = source.get("bytes", b"") - original_size = len(original_bytes) if isinstance(original_bytes, bytes) else 0 - content[block_idx] = { - "text": f"[Document placeholder: name={doc_name}, format={doc_format}, original_size={original_size} bytes]" - } - strip_count += 1 + from agents.main_agent.session.document_rehydration import rehydrate_documents - if strip_count > 0: - logger.debug(f"Stripped inline bytes from {strip_count} document block(s) in history") - - return stripped_messages + result = rehydrate_documents( + messages, + session_id=getattr(self.config, "session_id", None), + user_id=self.user_id, + ) + if result.rehydrated: + logger.info( + "Rehydrated %d document block(s) as digests (%d built lazily) in restored history", + result.rehydrated, result.lazy_digests, + ) + self.record_compaction_event( + "document_rehydrated", + documents=result.rehydrated, + documentTokens=result.digest_tokens, + ) + if result.stripped: + logger.debug(f"Stripped inline bytes from {result.stripped} unmatched document block(s) in history") + self.record_compaction_event( + "document_stripped", + documents=result.stripped, + documentTokens=result.stripped_tokens, + ) + return result.messages # ========================================================================= # Tool-pairing / role-alternation repair (restore-time safety net) diff --git a/backend/src/agents/main_agent/streaming/stream_coordinator.py b/backend/src/agents/main_agent/streaming/stream_coordinator.py index e807c1c9..f7ddd68b 100644 --- a/backend/src/agents/main_agent/streaming/stream_coordinator.py +++ b/backend/src/agents/main_agent/streaming/stream_coordinator.py @@ -3111,10 +3111,30 @@ async def _store_message_metadata( events = context_ledger.get("compactionEvents") if events: metadata_kwargs["compactionEvents"] = events + # document_read retrievals this call requested (calls / + # pages / bytes) — numbers read off the tool's own + # metadata, never the document. + reads = context_ledger.get("documentReads") + if reads: + metadata_kwargs["documentReads"] = reads if strands_agent is not None and cost_diagnostics_enabled(): prefix_tokens = get_prefix_token_split(strands_agent) if prefix_tokens: metadata_kwargs["prefixTokens"] = prefix_tokens + # The attachment footprint of the live context: inline + # documents (count, estimated tokens, format mix), digest + # stand-ins, and retrieved page slices. Flat fields so a + # query can split rows by hasDocuments / documentDigests + # without reading the conversation + # (docs/specs/document-context-offload.md §6.1). + try: + from agents.main_agent.session.document_context import summarize_document_context + + footprint = summarize_document_context(getattr(strands_agent, "messages", None)) + if footprint: + metadata_kwargs.update(footprint) + except Exception as doc_err: # noqa: BLE001 - never block the cost row + logger.debug(f"Skipping document context summary: {doc_err}") message_metadata = MessageMetadata(**metadata_kwargs) diff --git a/backend/src/apis/app_api/admin/costs/models.py b/backend/src/apis/app_api/admin/costs/models.py index 1754da8a..6df653ec 100644 --- a/backend/src/apis/app_api/admin/costs/models.py +++ b/backend/src/apis/app_api/admin/costs/models.py @@ -131,6 +131,24 @@ class CompactionEvent(BaseModel): retained_messages: Optional[int] = Field(None, alias="retainedMessages") truncated_tool_results: Optional[int] = Field(None, alias="truncatedToolResults") input_tokens: Optional[int] = Field(None, alias="inputTokens") + # Document lifecycle kinds (`document_stripped` / `document_rehydrated` / + # `document_offload`): how many documents the event touched, their + # estimated token weight, and — for an offload — the prompt-cache gap at + # the moment it fired (an offload while the cache is live is the + # regression the trigger must never produce). + documents: Optional[int] = None + document_tokens: Optional[int] = Field(None, alias="documentTokens") + cache_gap_seconds: Optional[int] = Field(None, alias="cacheGapSeconds") + + +class DocumentReads(BaseModel): + """``document_read`` retrievals one model call requested: calls, pages + returned as native document blocks, and their byte size.""" + model_config = ConfigDict(populate_by_name=True) + + calls: int = 0 + pages: int = 0 + bytes: int = 0 class SessionCallRow(BaseModel): @@ -181,6 +199,21 @@ class SessionCallRow(BaseModel): # the prefix changed before this call. window_trimmed: Optional[int] = Field(None, alias="windowTrimmed") compaction_events: Optional[List[CompactionEvent]] = Field(None, alias="compactionEvents") + # Document context at this call (optional; absent on rows written before + # it shipped or with diagnostics off). `hasDocuments` + `documentDigests` + # classify the call: full document inline, digest only, or neither. + # `documentTokens` is the compaction estimator's heuristic (bytes/4, flat + # per image), comparable across rows; `documentMime` is keyed by Bedrock's + # format enum plus `image` — never a filename. + has_documents: Optional[bool] = Field(None, alias="hasDocuments") + document_count: Optional[int] = Field(None, alias="documentCount") + document_tokens: Optional[int] = Field(None, alias="documentTokens") + document_digests: Optional[int] = Field(None, alias="documentDigests") + documents_attached: Optional[int] = Field(None, alias="documentsAttached") + document_slices: Optional[int] = Field(None, alias="documentSlices") + document_slice_tokens: Optional[int] = Field(None, alias="documentSliceTokens") + document_mime: Optional[Dict[str, int]] = Field(None, alias="documentMime") + document_reads: Optional[DocumentReads] = Field(None, alias="documentReads") class SessionCostAnatomy(BaseModel): @@ -384,6 +417,10 @@ class AttachmentProfile(BaseModel): count: int = 0 total_bytes: int = Field(0, alias="totalBytes") by_mime: Dict[str, int] = Field(default_factory=dict, alias="byMime") + # DocumentDigest coverage: uploads with a ready digest and the rendered + # token estimate they would cost in context (offload spec §4A). + digested: int = 0 + digest_tokens: int = Field(0, alias="digestTokens") class ContextTrajectoryPoint(BaseModel): @@ -433,6 +470,7 @@ class DataCoverage(BaseModel): prefix_tokens: bool = Field(False, alias="prefixTokens") window_trim: bool = Field(False, alias="windowTrim") compaction_events: bool = Field(False, alias="compactionEvents") + documents: bool = False class SessionProfile(BaseModel): @@ -475,3 +513,12 @@ class SessionProfile(BaseModel): ) # The summary's token size at the most recent compaction decision. last_summary_tokens: Optional[int] = Field(None, alias="lastSummaryTokens") + # Document lifecycle across the session's calls: how many calls ran with + # the full document inline vs. a digest only (the digest-vs-full turn + # shares), the largest estimated document footprint seen, and what + # document_read pulled back in total. + full_document_calls: int = Field(0, alias="fullDocumentCalls") + digest_only_calls: int = Field(0, alias="digestOnlyCalls") + peak_document_tokens: Optional[int] = Field(None, alias="peakDocumentTokens") + document_read_calls: int = Field(0, alias="documentReadCalls") + document_read_pages: int = Field(0, alias="documentReadPages") diff --git a/backend/src/apis/app_api/admin/costs/service.py b/backend/src/apis/app_api/admin/costs/service.py index d355306c..8b2efef4 100644 --- a/backend/src/apis/app_api/admin/costs/service.py +++ b/backend/src/apis/app_api/admin/costs/service.py @@ -21,6 +21,7 @@ ) from .models import ( CompactionEvent, + DocumentReads, PrefixTokens, AttachmentProfile, ContextTrajectoryPoint, @@ -99,15 +100,76 @@ class _CallLedger: removed: Optional[int] = None trimmed: Optional[int] = None events: List[CompactionEvent] = _field(default_factory=list) + #: The row's document context fields, decoded (``None`` = not tracked). + documents: Optional[Dict[str, Any]] = None + document_reads: Optional[DocumentReads] = None + + +_DOCUMENT_INT_FIELDS = ( + "documentCount", "documentTokens", "documentDigests", "documentsAttached", + "documentSlices", "documentSliceTokens", +) + + +def _call_documents(record: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """The row's document context (``hasDocuments`` and the counts), or + ``None`` when the row predates the fields. Ints coerced, the format map + kept as ``{format: count}``.""" + if "hasDocuments" not in record: + return None + out: Dict[str, Any] = {"hasDocuments": bool(record.get("hasDocuments"))} + for key in _DOCUMENT_INT_FIELDS: + value = _as_int(record.get(key)) + if value is not None: + out[key] = value + mime = record.get("documentMime") + if isinstance(mime, dict): + out["documentMime"] = { + str(k): (_as_int(v) or 0) for k, v in mime.items() + } + return out + + +def _call_document_reads(record: Dict[str, Any]) -> Optional[DocumentReads]: + raw = record.get("documentReads") + if not isinstance(raw, dict): + return None + return DocumentReads( + calls=_as_int(raw.get("calls")) or 0, + pages=_as_int(raw.get("pages")) or 0, + bytes=_as_int(raw.get("bytes")) or 0, + ) + + +def _document_row_fields(ledger: "_CallLedger") -> Dict[str, Any]: + """``SessionCallRow`` kwargs for the row's document context — empty when + the row predates the fields, so they render as null ("not tracked").""" + fields: Dict[str, Any] = {} + docs = ledger.documents + if docs is not None: + fields["has_documents"] = docs.get("hasDocuments") + fields["document_count"] = docs.get("documentCount") + fields["document_tokens"] = docs.get("documentTokens") + fields["document_digests"] = docs.get("documentDigests") + fields["documents_attached"] = docs.get("documentsAttached") + fields["document_slices"] = docs.get("documentSlices") + fields["document_slice_tokens"] = docs.get("documentSliceTokens") + fields["document_mime"] = docs.get("documentMime") + if ledger.document_reads is not None: + fields["document_reads"] = ledger.document_reads + return fields def _call_ledger(record: Dict[str, Any], previous_removed: Optional[int]) -> _CallLedger: """Decode a cost row's ``prefixTokens`` / ``windowRemovedMessages`` / - ``compactionEvents`` and derive ``trimmed`` (messages removed since the - previous ledger-bearing row). Absent fields stay ``None`` — "not tracked", - never 0 — and malformed ones are ignored rather than raised. + ``compactionEvents`` / document context and derive ``trimmed`` (messages + removed since the previous ledger-bearing row). Absent fields stay + ``None`` — "not tracked", never 0 — and malformed ones are ignored rather + than raised. """ ledger = _CallLedger() + ledger.documents = _call_documents(record) + ledger.document_reads = _call_document_reads(record) raw_prefix = record.get("prefixTokens") if isinstance(raw_prefix, dict): try: @@ -732,6 +794,7 @@ async def get_session_cost_anatomy(self, session_id: str) -> SessionCostAnatomy: window_removed_messages=ledger.removed, window_trimmed=ledger.trimmed, compaction_events=ledger.events or None, + **_document_row_fields(ledger), )) cache_traffic = total_cache_read + total_cache_write @@ -957,13 +1020,20 @@ async def _attachment_profile(self, session_id: str) -> AttachmentProfile: return AttachmentProfile() by_mime: Counter = Counter() total_bytes = 0 + digested = digest_tokens = 0 for item in stats: by_mime[item.get("mimeType") or "unknown"] += 1 total_bytes += _as_int(item.get("sizeBytes")) or 0 + digest = item.get("digest") + if isinstance(digest, dict) and digest.get("status") == "ready": + digested += 1 + digest_tokens += _as_int(digest.get("tokens")) or 0 return AttachmentProfile( count=len(stats), total_bytes=total_bytes, by_mime=dict(by_mime), + digested=digested, + digest_tokens=digest_tokens, ) async def get_session_profile(self, session_id: str) -> Optional[SessionProfile]: @@ -1005,6 +1075,10 @@ async def get_session_profile(self, session_id: str) -> Optional[SessionProfile] window_trim_calls = 0 compaction_event_counts: Counter = Counter() last_summary_tokens: Optional[int] = None + any_documents = False + full_document_calls = digest_only_calls = 0 + document_read_calls = document_read_pages = 0 + peak_document_tokens: Optional[int] = None for index, record in enumerate(records): usage = record.get("tokenUsage") or {} @@ -1065,6 +1139,19 @@ async def get_session_profile(self, session_id: str) -> Optional[SessionProfile] compaction_event_counts[event.kind] += 1 if event.summary_tokens is not None: last_summary_tokens = event.summary_tokens + if ledger.documents is not None: + any_documents = True + if ledger.documents.get("hasDocuments"): + full_document_calls += 1 + elif (ledger.documents.get("documentDigests") or 0) > 0: + digest_only_calls += 1 + doc_tokens = ledger.documents.get("documentTokens") + if doc_tokens is not None: + peak_document_tokens = max(peak_document_tokens or 0, doc_tokens) + if ledger.document_reads is not None: + any_documents = True + document_read_calls += ledger.document_reads.calls + document_read_pages += ledger.document_reads.pages trajectory.append(ContextTrajectoryPoint( call_index=index, timestamp=record.get("timestamp", ""), @@ -1142,12 +1229,29 @@ async def get_session_profile(self, session_id: str) -> Optional[SessionProfile] compaction_events=( any_compaction_events or row.get("compactionAppliedCount") is not None ), + documents=any_documents or row.get("fullDocumentCalls") is not None, ), prefix_tokens=prefix_tokens, window_trim_calls=window_trim_calls, window_removed_messages=last_removed, compaction_event_counts=dict(compaction_event_counts), last_summary_tokens=last_summary_tokens, + # Rows are authoritative when present; the session rollups cover + # calls whose rows have expired (365-day TTL) or a session read + # without its rows. + full_document_calls=( + full_document_calls if any_documents else (_as_int(row.get("fullDocumentCalls")) or 0) + ), + digest_only_calls=( + digest_only_calls if any_documents else (_as_int(row.get("digestOnlyCalls")) or 0) + ), + peak_document_tokens=peak_document_tokens, + document_read_calls=( + document_read_calls if any_documents else (_as_int(row.get("documentReadCalls")) or 0) + ), + document_read_pages=( + document_read_pages if any_documents else (_as_int(row.get("documentReadPages")) or 0) + ), ) async def get_dashboard( diff --git a/backend/src/apis/app_api/files/service.py b/backend/src/apis/app_api/files/service.py index 2c4fc027..ae9890fe 100644 --- a/backend/src/apis/app_api/files/service.py +++ b/backend/src/apis/app_api/files/service.py @@ -9,7 +9,7 @@ import logging import uuid from datetime import datetime, timedelta, timezone -from typing import List, Optional +from typing import List, Optional, Set import boto3 from botocore.config import Config @@ -128,6 +128,11 @@ class FileUploadService: # Stored alongside the original so cleanup happens with the file. THUMBNAIL_KEY_NAME = "_thumb.png" + # Strong refs to fire-and-forget digest builds. The event loop only holds + # weak references to tasks, so a bare `create_task` can be collected + # mid-run (the trap web_sources/deletion_service.py documents). + _digest_tasks: Set["asyncio.Task[None]"] = set() + def __init__( self, repository: Optional[FileUploadRepository] = None, @@ -363,6 +368,13 @@ async def complete_upload( # Increment quota await self.repository.increment_quota(user_id, file_meta.size_bytes) + # DocumentDigest, off the request path (offload spec §4A / PR-2): + # documents get an outline + abstract persisted on their row so a + # later turn can carry the digest instead of the bytes. Non-documents + # (spreadsheets, decks, images) and the kill switch skip it. The + # response does not wait for it. + self._schedule_digest(file_meta) + logger.info("Completed file upload") return CompleteUploadResponse( @@ -754,6 +766,56 @@ async def _render_and_store_thumbnail( f"({len(png_bytes)} bytes)" ) + # ========================================================================= + # DocumentDigest (docs/specs/document-context-offload.md §4A, PR-2) + # ========================================================================= + + def _schedule_digest(self, file_meta: FileMetadata) -> Optional["asyncio.Task[None]"]: + """Queue a digest build for a document upload; ``None`` when skipped. + + Skipped for non-documents, when ``DOCUMENT_DIGEST_ENABLED=false``, or + when no event loop is running (a synchronous caller). Never raises. + """ + try: + from apis.shared.files.document_digest import document_digest_enabled + from apis.shared.files.document_read import is_document_class + + if not document_digest_enabled(): + return None + if not is_document_class(file_meta.mime_type, file_meta.filename): + return None + task = asyncio.get_running_loop().create_task(self._generate_and_store_digest(file_meta)) + except Exception: # noqa: BLE001 - scheduling must never fail the upload + logger.warning("Document digest not scheduled", exc_info=True) + return None + self._digest_tasks.add(task) + task.add_done_callback(self._digest_tasks.discard) + return task + + async def _generate_and_store_digest(self, file_meta: FileMetadata) -> None: + """Read the original, build the digest, persist it. Never raises.""" + from apis.shared.files.document_digest import build_digest, record_digest + + try: + response = await asyncio.to_thread( + self._s3_client.get_object, Bucket=self.bucket_name, Key=file_meta.s3_key + ) + raw = await asyncio.to_thread(response["Body"].read) + digest = await build_digest( + raw=raw, + mime_type=file_meta.mime_type, + filename=file_meta.filename, + upload_id=file_meta.upload_id, + ) + record_digest(digest) + stored = await self.repository.update_file_digest( + file_meta.user_id, file_meta.upload_id, digest.to_item() + ) + if stored is None: + logger.info("Document digest discarded: file row no longer exists") + except Exception: # noqa: BLE001 - a digest is never worth an error + logger.warning("Document digest build failed", exc_info=True) + def _delete_thumbnail_object(self, file_meta: FileMetadata) -> None: """ Best-effort delete of the thumbnail sibling. diff --git a/backend/src/apis/inference_api/chat/routes.py b/backend/src/apis/inference_api/chat/routes.py index bf83762f..9198cded 100644 --- a/backend/src/apis/inference_api/chat/routes.py +++ b/backend/src/apis/inference_api/chat/routes.py @@ -10,6 +10,7 @@ import asyncio import json import logging +from collections import OrderedDict from typing import AsyncGenerator, Optional, Union from fastapi import APIRouter, Depends, HTTPException, status @@ -751,6 +752,98 @@ def _build_memory_tools(agent_memory, user_id: str, user_email: str) -> list: return tools +# ============================================================ +# Document Read Tool Injection (docs/specs/document-context-offload.md §4B) +# ============================================================ + +#: Sessions known to carry a readable document. The gate is one DynamoDB +#: query per turn otherwise; a positive answer is memoized because it is +#: monotonic in practice (an upload stays unless the user deletes it, and a +#: stale tool on a session whose files were deleted just returns "not found"). +#: Negative answers are never memoized — the next turn may be the upload. +_DOCUMENT_SESSIONS: "OrderedDict[str, bool]" = OrderedDict() +_DOCUMENT_SESSIONS_MAX = 10_000 + + +def _remember_document_session(session_id: str) -> None: + _DOCUMENT_SESSIONS[session_id] = True + _DOCUMENT_SESSIONS.move_to_end(session_id) + while len(_DOCUMENT_SESSIONS) > _DOCUMENT_SESSIONS_MAX: + _DOCUMENT_SESSIONS.popitem(last=False) + + +async def _session_has_documents( + session_id: str, + user_id: str, + turn_upload_ids: list | None = None, +) -> bool: + """Whether ``document_read`` should exist on this turn. + + True when this turn attaches uploads (their metadata rows already exist, + so no query is needed), when the session was seen carrying a document + earlier in this process, or when the session's upload rows include at + least one readable document (PDF, Word, text, markdown, HTML — not + spreadsheets, decks or images, which have other paths). Fail-closed on + error: a turn without the tool is today's behavior, never a broken turn. + """ + if turn_upload_ids: + _remember_document_session(session_id) + return True + if _DOCUMENT_SESSIONS.get(session_id): + return True + try: + from apis.shared.files.document_read import session_has_documents + + present = await session_has_documents(user_id, session_id) + except Exception: # noqa: BLE001 - the gate must never fail a turn + logger.warning("document_read gate lookup failed; tool not injected this turn", exc_info=True) + return False + if present: + _remember_document_session(session_id) + return present + + +async def _document_tools_gate( + session_id: str, + user_id: str, + turn_upload_ids: list | None = None, +) -> bool: + """The single answer to "does this turn carry ``document_read``" — the + builder and the resume path's cache key both read it, so the two can + never disagree (a disagreement orphans a paused agent).""" + from apis.shared.feature_flags import document_read_enabled + + if not document_read_enabled(): + return False + if not session_id or not user_id: + return False + return await _session_has_documents(session_id, user_id, turn_upload_ids) + + +async def _build_document_tools( + session_id: str, + user_id: str, + turn_upload_ids: list | None = None, +) -> list: + """Context-bound ``document_read`` for a session that has a readable attachment. + + **Not gated on ``enabled_tools``** — the governing capability is the user's + own attachment, exactly as the Memory-Space tools are governed by an + Agent's binding. Its id stays out of ``INJECTED_TOOL_IDS``. Kill switch: + ``DOCUMENT_READ_ENABLED=false``. The gate's answer also feeds the agent + cache key (``has_document_tools``), so an agent cached before the first + upload is never served without the tool afterwards. + """ + if not await _document_tools_gate(session_id, user_id, turn_upload_ids): + return [] + + from agents.builtin_tools.document_read_tool import make_document_read_tool + + tools = [make_document_read_tool(session_id, user_id)] + logger.info("Created document_read tool (session has a readable document)") + return tools + + # ============================================================ # Attachment Partitioning (#206) # ============================================================ @@ -2402,6 +2495,12 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g mantle_region=snapshot.mantle_region, agent_type=snapshot.agent_type, is_resume=True, + # The original turn's key carried whether the session had a + # readable document; the gate is monotonic, so re-asking it + # rebuilds the same key (an orphaned paused agent otherwise). + has_document_tools=await _document_tools_gate( + input_data.session_id, user_id + ), # Resume must rebuild the SAME cache key the original turn used, # or the paused agent is orphaned. New snapshots carry the # original turn's exact effective set in enabled_skills, so @@ -2566,6 +2665,18 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g ) extra_tools = extra_tools + memory_tools + # document_read for any session that carries a readable attachment + # (this turn's uploads count). Gated on session state, not the + # picker; its presence goes into the cache key below rather than + # vetoing the cache, so an attachment session that could keep a + # warm agent still does. + document_tools = await _build_document_tools( + session_id=input_data.session_id, + user_id=user_id, + turn_upload_ids=input_data.file_upload_ids, + ) + extra_tools = extra_tools + document_tools + # Can this turn's agent be cached despite carrying injected tools? # Only when every builder that fired closes over values the cache # key already carries (session, user, enabled_tools). Derived from @@ -2594,6 +2705,7 @@ async def invocations(request: InvocationRequest, current_user: User = Depends(g is_resume=False, accessible_skill_ids=effective_skill_ids, extra_tools_key_described=extra_tools_key_described, + has_document_tools=bool(document_tools), ) # Resume requests must target interrupts that the cached agent diff --git a/backend/src/apis/inference_api/chat/service.py b/backend/src/apis/inference_api/chat/service.py index 93f6279d..b9f7edd5 100644 --- a/backend/src/apis/inference_api/chat/service.py +++ b/backend/src/apis/inference_api/chat/service.py @@ -61,10 +61,19 @@ def _create_cache_key( freshness_hash: str, agent_type: Optional[str], skills_hash: str = "", + document_tools: bool = False, ) -> Tuple: """ Create a cache key for agent instances. + `document_tools` is whether the turn built the session-state-gated + ``document_read`` tool (the session has a readable attachment). It is not + in `enabled_tools`, so without this element an agent cached before the + first upload would be served — without the tool — to every turn after it. + The gate is monotonic in practice (files stay once uploaded), so the key + flips at most once per session, on the attach turn, when restored history + carries no document yet to lose. + `freshness_hash` is a short digest of the enabled tools' current `updated_at` values (see `freshness.get_freshness_hash`). When an admin edits a tool's config, the hash changes and the cache misses, @@ -95,6 +104,7 @@ def _create_cache_key( provider or "bedrock", freshness_hash, agent_type or "chat", + bool(document_tools), skills_hash, ) @@ -240,6 +250,7 @@ async def get_agent( accessible_skill_ids: Optional[List[str]] = None, extra_tools_key_described: bool = False, cache_write: bool = True, + has_document_tools: bool = False, ) -> BaseAgent: """ Get or create agent instance with current configuration for session @@ -314,6 +325,7 @@ async def get_agent( freshness_hash=freshness_hash, agent_type=agent_type, skills_hash=skills_hash, + document_tools=has_document_tools, ) # Whether this turn's injected tools (if any) let it use the cache at all. diff --git a/backend/src/apis/shared/feature_flags.py b/backend/src/apis/shared/feature_flags.py index 31555cb9..d2e78aef 100644 --- a/backend/src/apis/shared/feature_flags.py +++ b/backend/src/apis/shared/feature_flags.py @@ -92,6 +92,20 @@ def workspace_tools_enabled() -> bool: return os.environ.get("WORKSPACE_TOOLS_ENABLED", "").strip().lower() != "false" +def document_read_enabled() -> bool: + """Whether the ``document_read`` agent tool is injected for sessions that + carry a readable attachment (``docs/specs/document-context-offload.md`` + §4B). **Default ON with a kill switch** (house style): unset or empty + resolves to enabled; only the literal ``"false"`` disables. + + This is the *only* control on the tool. It is deliberately not gated on + RBAC or the tool picker: the governing capability is the user's own + attachment, and the ``workspace_files`` catalog key is granted to no prod + role, so an RBAC gate would ship the recovery path dark. + """ + return os.environ.get("DOCUMENT_READ_ENABLED", "").strip().lower() != "false" + + def agents_enabled() -> bool: """Whether the Agent Designer surface is enabled for this environment. diff --git a/backend/src/apis/shared/files/document_digest.py b/backend/src/apis/shared/files/document_digest.py new file mode 100644 index 00000000..ef1aef0e --- /dev/null +++ b/backend/src/apis/shared/files/document_digest.py @@ -0,0 +1,436 @@ +"""DocumentDigest — a bounded, once-per-upload description of a document. + +The offload design (``docs/specs/document-context-offload.md`` §4A) keeps a +document inline on the turn that introduces it and hands later turns a +*digest* plus a ``document_read`` handle instead. This module builds that +digest, off the model path, when an upload completes (PR-2). Nothing in the +chat path reads it yet; PR-3 renders it in place of the restore placeholder. + +Two parts, deliberately split: + +* **Outline** — deterministic, no model: page / paragraph / line count, a + heading outline with the unit each heading starts on, table / figure + mentions, and a text sample. PDF text comes from ``pypdfium2`` (already a + dependency), DOCX from the stdlib extractor in ``document_read``, the text + family from the bytes. A document with no detectable headings gets an + outline sampled from page first-lines so the model still has anchors. +* **Abstract** — 3–5 sentences from a cheap text model over the sample plus + the outline (Nova Micro, the same model the tool-batch summaries and the + compaction summary use; ``DOCUMENT_DIGEST_MODEL_ID`` overrides). Fail-open: + a model error leaves the digest without an abstract, never without an + outline. + +The rendered form (``render_digest``) is what enters context later. It is +hard-capped at ``DOCUMENT_DIGEST_MAX_TOKENS`` (default 1,500, chars/4): the +outline is trimmed from the end first, then the abstract. Every digest +records its rendered token estimate so the cap is a stored fact per file. + +Content policy: the digest carries model prose about the user's document +(``abstract``) and heading text (``sections``). Both are denylisted in +``content_policy.CONTENT_BEARING``; the numeric fields (``status``, +``tokens``, ``count``, ``format``) are projectable. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import re +import time +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from pydantic import BaseModel, Field + +from .document_read import _open_pdf, _pdf_page_text, docx_paragraphs, document_format_for + +logger = logging.getLogger(__name__) + +DIGEST_VERSION = 1 + +#: Same heuristic the compaction estimator uses (``compaction_policy.CHARS_PER_TOKEN``); +#: duplicated here because ``apis.shared`` must not import the agent layer. +_CHARS_PER_TOKEN = 4 + +#: Rendered budget. The spec's ceiling: a digest must be strictly smaller than +#: any document worth offloading (the trigger's floor is 5,000 tokens). +DOCUMENT_DIGEST_MAX_TOKENS = int(os.environ.get("DOCUMENT_DIGEST_MAX_TOKENS", 1_500)) +#: Outline entries kept (evenly sampled when the document has more headings). +DOCUMENT_DIGEST_MAX_SECTIONS = int(os.environ.get("DOCUMENT_DIGEST_MAX_SECTIONS", 40)) +#: Characters of document text handed to the abstract model. +DOCUMENT_DIGEST_SAMPLE_CHARS = int(os.environ.get("DOCUMENT_DIGEST_SAMPLE_CHARS", 12_000)) +#: Cheap text model for the abstract. Nova Micro is what the tool-batch +#: summaries and the compaction summary already run on; it is text-only, +#: which is fine because the outline extractor hands it text. +DOCUMENT_DIGEST_MODEL_ID = os.environ.get("DOCUMENT_DIGEST_MODEL_ID", "").strip() or "us.amazon.nova-micro-v1:0" +_ABSTRACT_MAX_OUTPUT_TOKENS = 320 +_HEADING_MAX_CHARS = 90 +_SECTION_TITLE_CHARS = 80 +_SAMPLE_PER_PAGE_CHARS = 600 + + +def document_digest_enabled() -> bool: + """Default ON with a kill switch (house style): only the literal "false" disables.""" + return os.environ.get("DOCUMENT_DIGEST_ENABLED", "").strip().lower() != "false" + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- + + +class DigestSection(BaseModel): + """One outline entry: the unit (page / paragraph / line, 1-indexed) a + heading starts on and the heading text, bounded.""" + + start: int + title: str + + +class DocumentDigest(BaseModel): + """What is persisted on ``FileMetadata.digest`` (as a plain dict).""" + + version: int = DIGEST_VERSION + status: str = "ready" # ready | failed + format: Optional[str] = None + unit: str = "page" # page | paragraph | line + count: int = 0 + sections: List[DigestSection] = Field(default_factory=list) + tables: int = 0 + figures: int = 0 + chars: int = 0 + abstract: Optional[str] = None + tokens: int = 0 # rendered-form estimate (chars/4) + model_id: Optional[str] = None + extractor_ms: int = 0 + error: Optional[str] = None # exception class name only, never a message + generated_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + + def to_item(self) -> Dict[str, Any]: + return self.model_dump(exclude_none=True) + + @classmethod + def from_item(cls, item: Any) -> Optional["DocumentDigest"]: + if not isinstance(item, dict): + return None + try: + return cls(**item) + except Exception: # noqa: BLE001 - a malformed row reads as "no digest" + return None + + +# --------------------------------------------------------------------------- +# Outline extraction (deterministic) +# --------------------------------------------------------------------------- + +_NUMBERED = re.compile( + r"^(?:\d+(?:\.\d+)*[.)]?\s+\S|[IVXLC]{1,6}[.)]\s+\S|(?:article|section|chapter|part|appendix|schedule|exhibit|title)\s+[\w-]+)", + re.IGNORECASE, +) +_TABLE = re.compile(r"\btable\s+\d+|\btable\b\s*[:\-]", re.IGNORECASE) +_FIGURE = re.compile(r"\b(?:figure|fig\.|chart|diagram)\s+\d+", re.IGNORECASE) +_HTML_TAG = re.compile(r"<[^>]+>") +_MD_HEADING = re.compile(r"^\s{0,3}#{1,6}\s+(.+?)\s*#*\s*$") + + +def looks_like_heading(line: str, *, standalone: bool = False) -> bool: + """Cheap, format-agnostic heading test: short, no sentence punctuation at + the end, and either numbered, ALL CAPS, or title-cased with few words. + + ``standalone`` is for units that are already a whole paragraph (DOCX, + text lines): a one-word capitalized paragraph ("Risks") is a heading + there, while a one-word line inside a PDF page is usually noise. + """ + text = line.strip() + if not (3 <= len(text) <= _HEADING_MAX_CHARS) or text.endswith((".", ",", ";", ":")): + return False + if _NUMBERED.match(text): + return True + letters = [c for c in text if c.isalpha()] + if len(letters) >= 2 and all(c.isupper() for c in letters): + return True + words = [w for w in re.split(r"\s+", text) if w] + if len(words) <= 10: + capitalized = sum(1 for w in words if w[0].isupper() or not w[0].isalpha()) + return capitalized / len(words) >= 0.7 and len(words) >= (1 if standalone else 2) + return False + + +def _strip_html(text: str) -> str: + return _HTML_TAG.sub(" ", text) + + +def _units_for(fmt: str, raw: bytes) -> Tuple[str, List[str]]: + """``(unit name, unit texts)`` — one string per page / paragraph / line.""" + if fmt == "pdf": + pdf = _open_pdf(raw) + try: + return "page", [_pdf_page_text(pdf, i) for i in range(len(pdf))] + finally: + pdf.close() + if fmt == "docx": + return "paragraph", docx_paragraphs(raw) + text = raw.decode("utf-8", errors="replace") + if fmt == "html": + text = _strip_html(text) + return "line", text.splitlines() + + +def _sample_sections(sections: List[DigestSection], limit: int) -> List[DigestSection]: + if len(sections) <= limit: + return sections + step = len(sections) / limit + return [sections[int(i * step)] for i in range(limit)] + + +def extract_outline(fmt: str, raw: bytes) -> DocumentDigest: + """The deterministic half of the digest (no abstract, no tokens yet).""" + unit, texts = _units_for(fmt, raw) + sections: List[DigestSection] = [] + tables = figures = chars = 0 + for index, text in enumerate(texts, start=1): + chars += len(text) + tables += len(_TABLE.findall(text)) + figures += len(_FIGURE.findall(text)) + lines = text.splitlines() if unit == "page" else [text] + for line in lines: + md = _MD_HEADING.match(line) if fmt == "md" else None + title = md.group(1) if md else ( + line.strip() if looks_like_heading(line, standalone=unit != "page") else None + ) + if title: + sections.append(DigestSection(start=index, title=title[:_SECTION_TITLE_CHARS])) + if unit != "page": + break + if not sections and unit == "page" and texts: + # No headings (a scan, a form): anchor the outline on page first-lines. + for index, text in enumerate(texts, start=1): + first = next((ln.strip() for ln in text.splitlines() if ln.strip()), "") + if first: + sections.append(DigestSection(start=index, title=first[:_SECTION_TITLE_CHARS])) + return DocumentDigest( + format=fmt, + unit=unit, + count=len(texts), + sections=_sample_sections(sections, DOCUMENT_DIGEST_MAX_SECTIONS), + tables=tables, + figures=figures, + chars=chars, + ) + + +def text_sample(fmt: str, raw: bytes, limit: int = DOCUMENT_DIGEST_SAMPLE_CHARS) -> str: + """Up to ``limit`` characters spread across the document (not just its + head), so the abstract sees the middle and the end of a long file.""" + _, texts = _units_for(fmt, raw) + if not texts: + return "" + if fmt == "pdf": + per_page = max(_SAMPLE_PER_PAGE_CHARS, limit // max(1, len(texts))) + pieces = [t.strip()[:per_page] for t in texts if t.strip()] + return "\n\n".join(pieces)[:limit] + return "\n".join(t for t in texts if t.strip())[:limit] + + +# --------------------------------------------------------------------------- +# Abstract (cheap model, fail-open) +# --------------------------------------------------------------------------- + +_ABSTRACT_SYSTEM_PROMPT = ( + "You write short, factual abstracts of documents for an assistant that will " + "later retrieve specific pages on demand. Write 3 to 5 plain sentences: what " + "the document is, who it is for or from, its main parts, and any notable " + "numbers, dates or decisions. No preamble, no bullet points, no quotes." +) + + +def _abstract_prompt(outline: DocumentDigest, sample: str) -> str: + lines = [f"Document type: {outline.format}, {outline.count} {outline.unit}s."] + if outline.sections: + lines.append("Outline:") + lines.extend(f"- ({outline.unit} {s.start}) {s.title}" for s in outline.sections[:25]) + lines.append("Text sample:") + lines.append(sample) + return "\n".join(lines) + + +async def generate_abstract(outline: DocumentDigest, sample: str, model_id: str = DOCUMENT_DIGEST_MODEL_ID) -> Optional[str]: + """3–5 sentences from the cheap model, or ``None`` (never raises).""" + if not sample.strip(): + return None + try: + import boto3 + except ImportError: # pragma: no cover + return None + try: + client = boto3.client("bedrock-runtime", region_name=os.environ.get("AWS_REGION", "us-west-2")) + response = await asyncio.to_thread( + client.converse, + modelId=model_id, + messages=[{"role": "user", "content": [{"text": _abstract_prompt(outline, sample)}]}], + system=[{"text": _ABSTRACT_SYSTEM_PROMPT}], + inferenceConfig={"temperature": 0.2, "maxTokens": _ABSTRACT_MAX_OUTPUT_TOKENS, "topP": 0.9}, + ) + if response.get("stopReason") == "max_tokens": + logger.debug("Document abstract hit the token ceiling; discarding") + return None + text = response["output"]["message"]["content"][0]["text"].strip() + return re.sub(r"\s+", " ", text) or None + except Exception: # noqa: BLE001 - an abstract is never worth a failed upload + logger.debug("Document abstract generation skipped", exc_info=True) + return None + + +# --------------------------------------------------------------------------- +# Rendering (what PR-3 puts in context) and the budget +# --------------------------------------------------------------------------- + + +def _xml_escape(text: str) -> str: + return text.replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """) + + +def estimate_tokens(text: str) -> int: + return len(text) // _CHARS_PER_TOKEN + + +def render_digest( + digest: DocumentDigest, + *, + filename: str, + upload_id: str, + budget_tokens: int = DOCUMENT_DIGEST_MAX_TOKENS, +) -> str: + """The ```` block, trimmed to ``budget_tokens``. + + Sections are dropped from the end first (the model can always ask + ``document_read`` for more), then the abstract is truncated. The opening + tag always fits: it is the ``document_read`` handle. + """ + header = ( + f'{_xml_escape(digest.abstract)}" if digest.abstract else "" + section_lines = [ + f'
{_xml_escape(s.title)}
' for s in digest.sections + ] + + def _build(abstract: str, sections: Sequence[str]) -> str: + parts = [header] + if abstract: + parts.append(abstract) + parts.extend(sections) + parts.append(footer) + return "\n".join(parts) + + budget_chars = max(0, budget_tokens) * _CHARS_PER_TOKEN + kept = list(section_lines) + text = _build(abstract_line, kept) + while len(text) > budget_chars and kept: + kept.pop() + text = _build(abstract_line, kept) + if len(text) > budget_chars and abstract_line: + room = budget_chars - len(_build("", kept)) - len(" ") - 1 + if room > 20 and digest.abstract: + trimmed = _xml_escape(digest.abstract[: max(0, room - 1)].rstrip()) + "…" + text = _build(f" {trimmed}", kept) + else: + text = _build("", kept) + return text + + +# --------------------------------------------------------------------------- +# Build +# --------------------------------------------------------------------------- + + +async def build_digest( + *, + raw: bytes, + mime_type: str, + filename: str, + upload_id: str, + model_id: str = DOCUMENT_DIGEST_MODEL_ID, + with_abstract: bool = True, +) -> DocumentDigest: + """Outline + abstract + rendered token estimate. Never raises: an + extraction failure yields ``status="failed"`` with the exception class + name; an abstract failure yields a digest without one.""" + started = time.monotonic() + fmt = document_format_for(mime_type, filename) + if fmt is None: + return DocumentDigest(status="failed", error="NotADocument", extractor_ms=0) + try: + outline = await asyncio.to_thread(extract_outline, fmt, raw) + sample = await asyncio.to_thread(text_sample, fmt, raw) if with_abstract else "" + except Exception as e: # noqa: BLE001 + logger.warning("Document digest extraction failed (%s)", type(e).__name__, exc_info=True) + return DocumentDigest( + status="failed", format=fmt, error=type(e).__name__, + extractor_ms=int((time.monotonic() - started) * 1000), + ) + if with_abstract: + outline.abstract = await generate_abstract(outline, sample, model_id=model_id) + outline.model_id = model_id if outline.abstract else None + outline.extractor_ms = int((time.monotonic() - started) * 1000) + outline.tokens = estimate_tokens(render_digest(outline, filename=filename, upload_id=upload_id)) + return outline + + +def record_digest(digest: DocumentDigest) -> None: + """One content-free EMF record per digest build. Never raises.""" + logger.info( + "document_digest: status=%s format=%s %s=%d sections=%d tokens=%d abstract=%s ms=%d", + digest.status, digest.format, digest.unit, digest.count, len(digest.sections), + digest.tokens, bool(digest.abstract), digest.extractor_ms, + ) + 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={ + "DocumentDigestGenerated": 1, + "DocumentDigestTokens": int(digest.tokens), + "DocumentDigestMs": int(digest.extractor_ms), + }, + properties={ + "format": digest.format, + "outcome": ( + "failed" if digest.status != "ready" + else ("ready" if digest.abstract else "no_abstract") + ), + }, + units={"DocumentDigestTokens": "Count", "DocumentDigestMs": "Milliseconds"}, + ) + except Exception as e: # noqa: BLE001 + logger.debug("document_digest EMF skipped: %s", e) + + +__all__ = [ + "DIGEST_VERSION", + "DOCUMENT_DIGEST_MAX_SECTIONS", + "DOCUMENT_DIGEST_MAX_TOKENS", + "DOCUMENT_DIGEST_MODEL_ID", + "DigestSection", + "DocumentDigest", + "build_digest", + "document_digest_enabled", + "extract_outline", + "generate_abstract", + "looks_like_heading", + "record_digest", + "render_digest", + "text_sample", +] diff --git a/backend/src/apis/shared/files/document_read.py b/backend/src/apis/shared/files/document_read.py new file mode 100644 index 00000000..fc90e6dc --- /dev/null +++ b/backend/src/apis/shared/files/document_read.py @@ -0,0 +1,576 @@ +"""Document read service — page-range and pattern retrieval over uploaded documents. + +Backs the ``document_read`` agent tool (``agents/builtin_tools/document_read_tool.py``; +design in ``docs/specs/document-context-offload.md`` §4B). The tool is the +recovery path for a document that is no longer inline in the model's context: +history restore strips inline document bytes (Bedrock rejects duplicate +document names across a conversation), and the offload trigger will later +replace a document with a digest on purpose. Either way the model needs a way +to pull *part* of a document back at native fidelity. + +Three retrieval modes, modelled on Strands' ``retrieve_offloaded_content``: + +* **pages** — a PDF page range, re-assembled server-side into a new PDF that + contains only those pages and returned as a native ``document`` block. The + model sees the pages as it would the original (text layer plus page + images), not as flattened text. Hard-capped by ``max_pages`` so one call can + never re-inject a whole document. +* **pattern** — a case-insensitive regex over the document's text layer. + Returns the matching lines with their page numbers (PDF) or line / + paragraph numbers (text, DOCX), so the model can then ask for the pages + that matter. Bounded match count. +* **text** — bounded text for the text-family formats (txt / md / html via the + workspace read path, DOCX via a stdlib extractor), with ``offset`` + continuation. + +Plus a **list** mode (no ``upload_id``) enumerating the session's readable +documents, and an **index** mode (a PDF with neither range nor pattern) that +returns the page count and a short per-page snippet index. + +Every read goes through the DynamoDB user-files table (ownership by key shape, +``PK = USER#{userId}``) and reuses the workspace service's S3 client. No +operation accepts a model-supplied S3 key. Nothing here emits metrics or +touches the agent layer — the tool does that. +""" + +from __future__ import annotations + +import asyncio +import io +import logging +import os +import re +import uuid +import zipfile +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Sequence, Tuple +from xml.etree import ElementTree + +from .models import FileMetadata, FileStatus, get_file_format, is_presentation_file, is_tabular_file +from .repository import get_file_upload_repository +from .workspace import ( + WORKSPACE_READ_MAX_BYTES, + WorkspaceError, + WorkspaceFileNotFoundError, + WorkspaceValidationError, + _get_owned_ready_file, + _require_identity, + _s3, + is_text_mime, + read_workspace_file, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Bounds (env-backed; every one caps a per-turn payload) +# --------------------------------------------------------------------------- + +#: Default page-range size when the model does not say. 8 pages of a dense PDF +#: is ~10–20k tokens — a bounded, deliberate spend, not a whole document. +DOCUMENT_READ_MAX_PAGES = int(os.environ.get("DOCUMENT_READ_MAX_PAGES", 8)) +#: Absolute ceiling on one call's page count, whatever ``max_pages`` the model +#: asks for. The tool's contract is "one call cannot re-inject the document". +DOCUMENT_READ_HARD_MAX_PAGES = int(os.environ.get("DOCUMENT_READ_HARD_MAX_PAGES", 20)) +#: Pattern mode returns at most this many matching lines. +DOCUMENT_READ_MAX_MATCHES = int(os.environ.get("DOCUMENT_READ_MAX_MATCHES", 40)) +#: Index mode lists a snippet for at most this many pages. +DOCUMENT_READ_INDEX_PAGES = int(os.environ.get("DOCUMENT_READ_INDEX_PAGES", 40)) +#: Characters kept per matching / snippet line. +_LINE_CHARS = 200 +_SNIPPET_CHARS = 90 +#: Lines of context on either side of a pattern match. +_CONTEXT_LINES = 1 +#: Documents larger than this are not searched or sliced (the upload cap is +#: 4 MB, so this only guards a misconfigured bucket). +_MAX_DOCUMENT_BYTES = int(os.environ.get("DOCUMENT_READ_MAX_SOURCE_BYTES", 8 * 1024 * 1024)) + +#: Bedrock document formats the tool can read. Tabular and presentation files +#: have their own tools; images have no page/text structure to retrieve. +DOCUMENT_CLASS_FORMATS = frozenset({"pdf", "docx", "txt", "html", "md"}) +_TEXT_FORMATS = frozenset({"txt", "html", "md"}) + +_DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" +_W_NS = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}" + + +class DocumentReadError(WorkspaceError): + """A document_read failure surfaced conversationally by the tool.""" + + +# --------------------------------------------------------------------------- +# Classification +# --------------------------------------------------------------------------- + + +def document_format_for(mime_type: str, filename: str = "") -> Optional[str]: + """The Bedrock document format the tool can read this file as, or ``None``. + + ``None`` for tabular files (spreadsheet tools), presentations (PowerPoint + tools), images, and anything the upload allowlist does not know. + """ + if is_tabular_file(filename or "", mime_type or "") or is_presentation_file(filename or "", mime_type or ""): + return None + fmt = get_file_format((mime_type or "").lower().split(";")[0].strip()) + if fmt is None and is_text_mime(mime_type or ""): + fmt = "txt" + return fmt if fmt in DOCUMENT_CLASS_FORMATS else None + + +def is_document_class(mime_type: str, filename: str = "") -> bool: + """True when ``document_read`` has something to retrieve from this file.""" + return document_format_for(mime_type, filename) is not None + + +# --------------------------------------------------------------------------- +# Result shape +# --------------------------------------------------------------------------- + + +@dataclass +class DocumentReadResult: + """What one read produced. ``payload`` is JSON for the tool result; the + optional ``document_block`` is a native Bedrock document block appended + after it. The counters are what the tool records (content-free).""" + + mode: str + payload: Dict[str, Any] + document_block: Optional[Dict[str, Any]] = None + pages_returned: int = 0 + bytes_returned: int = 0 + format: Optional[str] = None + extra_blocks: List[Dict[str, Any]] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Listing +# --------------------------------------------------------------------------- + + +async def list_session_documents(user_id: str, session_id: str) -> Dict[str, Any]: + """The session's READY documents the tool can read, newest first.""" + _require_identity(user_id, session_id) + files = await get_file_upload_repository().list_session_files(session_id, status=FileStatus.READY) + documents: List[Dict[str, Any]] = [] + for meta in files: + if meta.user_id != user_id: + continue + fmt = document_format_for(meta.mime_type, meta.filename) + if fmt is None: + continue + documents.append( + { + "upload_id": meta.upload_id, + "filename": meta.filename, + "format": fmt, + "size_bytes": meta.size_bytes, + "modes": _modes_for(fmt), + } + ) + return {"documents": documents, "count": len(documents)} + + +def _modes_for(fmt: str) -> List[str]: + if fmt == "pdf": + return ["page_range", "pattern"] + return ["pattern", "text"] + + +async def session_has_documents(user_id: str, session_id: str) -> bool: + """Whether the session has at least one READY document the tool can read. + + The tool-injection gate. One ``SessionIndex`` query; callers memoize the + positive answer because it is monotonic for practical purposes (a file, + once uploaded, stays unless the user deletes it). + """ + _require_identity(user_id, session_id) + files = await get_file_upload_repository().list_session_files(session_id, status=FileStatus.READY) + return any(meta.user_id == user_id and is_document_class(meta.mime_type, meta.filename) for meta in files) + + +# --------------------------------------------------------------------------- +# Reading +# --------------------------------------------------------------------------- + + +def parse_page_range(value: Any) -> Optional[Tuple[int, int]]: + """``"4-7"`` / ``"4"`` / ``"4..7"`` / ``"4:7"`` / ``{"start": 4, "end": 7}`` + → ``(4, 7)`` (1-indexed, inclusive). ``None`` for empty input. Raises + ``WorkspaceValidationError`` on anything else.""" + if value is None: + return None + if isinstance(value, dict): + start, end = value.get("start"), value.get("end", value.get("start")) + elif isinstance(value, (list, tuple)) and len(value) in (1, 2): + start, end = value[0], value[-1] + elif isinstance(value, int): + start = end = value + else: + text = str(value).strip() + if not text: + return None + match = re.fullmatch(r"\s*(\d+)\s*(?:(?:-|\.\.|:|–|to)\s*(\d+))?\s*", text) + if not match: + raise WorkspaceValidationError( + f"Unrecognized page_range '{value}'. Use 'start-end' with 1-indexed page numbers, e.g. '4-7'." + ) + start, end = match.group(1), match.group(2) or match.group(1) + try: + start_i, end_i = int(start), int(end) + except (TypeError, ValueError): + raise WorkspaceValidationError(f"Unrecognized page_range '{value}'.") from None + if start_i < 1 or end_i < start_i: + raise WorkspaceValidationError( + f"Invalid page_range '{value}': pages are 1-indexed and end must not precede start." + ) + return start_i, end_i + + +def clamp_max_pages(max_pages: Any) -> int: + try: + requested = int(max_pages) if max_pages is not None else DOCUMENT_READ_MAX_PAGES + except (TypeError, ValueError): + requested = DOCUMENT_READ_MAX_PAGES + return max(1, min(requested, DOCUMENT_READ_HARD_MAX_PAGES)) + + +async def read_document( + user_id: str, + session_id: str, + upload_id: str, + *, + page_range: Any = None, + pattern: Optional[str] = None, + max_pages: Any = None, + offset: int = 0, +) -> DocumentReadResult: + """Read part of one document. See the module docstring for the modes.""" + _require_identity(user_id, session_id) + if not upload_id: + raise WorkspaceValidationError("upload_id is required (call with no arguments to list documents)") + if offset < 0: + raise WorkspaceValidationError("offset must be >= 0") + + meta = await _get_owned_ready_file(user_id, upload_id) + fmt = document_format_for(meta.mime_type, meta.filename) + if fmt is None: + raise DocumentReadError( + f"'{meta.filename}' is not a readable document. Spreadsheets go through " + "analyze_spreadsheet, presentations through read_powerpoint_presentation, " + "and images are only available inline." + ) + if meta.size_bytes > _MAX_DOCUMENT_BYTES: + raise DocumentReadError(f"'{meta.filename}' is too large to read ({meta.size_bytes} bytes).") + + pages = parse_page_range(page_range) + pattern_text = (pattern or "").strip() or None + limit = clamp_max_pages(max_pages) + + base = { + "upload_id": meta.upload_id, + "filename": meta.filename, + "format": fmt, + "size_bytes": meta.size_bytes, + } + + if fmt == "pdf": + raw = await _fetch_bytes(meta) + if pattern_text: + return await asyncio.to_thread(_pdf_pattern, raw, pattern_text, base) + if pages: + return await asyncio.to_thread(_pdf_pages, raw, pages, limit, base, meta.filename) + return await asyncio.to_thread(_pdf_index, raw, base) + + if pages: + raise WorkspaceValidationError( + f"page_range applies to PDFs only; '{meta.filename}' is {fmt}. Use pattern or offset instead." + ) + + if fmt == "docx": + raw = await _fetch_bytes(meta) + return await asyncio.to_thread(_docx_read, raw, pattern_text, offset, base) + + # txt / md / html + if pattern_text: + raw = await _fetch_bytes(meta) + return await asyncio.to_thread(_text_pattern, raw.decode("utf-8", errors="replace"), pattern_text, base, "line") + result = await read_workspace_file(user_id, upload_id, offset=offset) + payload = { + **base, + "mode": "text", + "content": result.get("content", ""), + "offset": result.get("offset", offset), + "truncated": bool(result.get("truncated")), + "next_offset": result.get("next_offset"), + } + return DocumentReadResult( + mode="text", payload=payload, bytes_returned=len(payload["content"].encode("utf-8")), format=fmt + ) + + +# --------------------------------------------------------------------------- +# S3 +# --------------------------------------------------------------------------- + + +def _get_object_bytes(bucket: str, key: str) -> bytes: + return _s3().get_object(Bucket=bucket, Key=key)["Body"].read() + + +async def _fetch_bytes(meta: FileMetadata) -> bytes: + raw = await asyncio.to_thread(_get_object_bytes, meta.s3_bucket, meta.s3_key) + if len(raw) > _MAX_DOCUMENT_BYTES: + raise DocumentReadError(f"'{meta.filename}' is too large to read ({len(raw)} bytes).") + return raw + + +# --------------------------------------------------------------------------- +# PDF (pypdfium2 — already a dependency for attachment thumbnails) +# --------------------------------------------------------------------------- + + +def _open_pdf(raw: bytes): + import pypdfium2 as pdfium # lazy: native lib, not needed on non-PDF paths + + try: + return pdfium.PdfDocument(io.BytesIO(raw)) + except Exception as e: # noqa: BLE001 + raise DocumentReadError(f"Could not open the PDF: {e}") from e + + +def _pdf_page_text(pdf, index: int) -> str: + page = pdf[index] + try: + textpage = page.get_textpage() + try: + return textpage.get_text_bounded() or "" + finally: + textpage.close() + finally: + page.close() + + +def _pdf_pages(raw: bytes, pages: Tuple[int, int], limit: int, base: Dict[str, Any], filename: str) -> DocumentReadResult: + import pypdfium2 as pdfium + + start, end = pages + pdf = _open_pdf(raw) + try: + count = len(pdf) + if start > count: + raise WorkspaceValidationError(f"page_range starts at {start} but the document has {count} pages.") + end = min(end, count) + truncated = False + if end - start + 1 > limit: + end = start + limit - 1 + truncated = True + sub = pdfium.PdfDocument.new() + try: + sub.import_pages(pdf, pages=list(range(start - 1, end))) + buffer = io.BytesIO() + sub.save(buffer) + finally: + sub.close() + finally: + pdf.close() + + sliced = buffer.getvalue() + returned = end - start + 1 + name = _unique_document_name(filename, f"p{start}-{end}") + payload = { + **base, + "mode": "pages", + "page_count": count, + "pages": {"start": start, "end": end}, + "pages_returned": returned, + "truncated_to_max_pages": truncated, + "next_start": end + 1 if (truncated and end < count) else None, + "page_numbering": ( + f"The attached document contains original pages {start}-{end} only; " + f"its page k is original page {start} + k - 1." + ), + } + block = {"document": {"format": "pdf", "name": name, "source": {"bytes": sliced}}} + return DocumentReadResult( + mode="pages", payload=payload, document_block=block, + pages_returned=returned, bytes_returned=len(sliced), format="pdf", + ) + + +def _pdf_pattern(raw: bytes, pattern: str, base: Dict[str, Any]) -> DocumentReadResult: + regex = _compile(pattern) + pdf = _open_pdf(raw) + try: + count = len(pdf) + matches: List[Dict[str, Any]] = [] + pages_matched: List[int] = [] + for index in range(count): + if len(matches) >= DOCUMENT_READ_MAX_MATCHES: + break + lines = _pdf_page_text(pdf, index).splitlines() + page_no = index + 1 + hit = False + for entry in _grep_lines(lines, regex, DOCUMENT_READ_MAX_MATCHES - len(matches)): + entry["page"] = page_no + matches.append(entry) + hit = True + if hit: + pages_matched.append(page_no) + finally: + pdf.close() + payload = { + **base, + "mode": "pattern", + "pattern": pattern, + "page_count": count, + "pages_matched": pages_matched, + "match_count": len(matches), + "truncated": len(matches) >= DOCUMENT_READ_MAX_MATCHES, + "matches": matches, + "hint": "Call again with page_range to read the matching pages at full fidelity.", + } + return DocumentReadResult(mode="pattern", payload=payload, pages_returned=0, format="pdf") + + +def _pdf_index(raw: bytes, base: Dict[str, Any]) -> DocumentReadResult: + pdf = _open_pdf(raw) + try: + count = len(pdf) + index: List[Dict[str, Any]] = [] + for i in range(min(count, DOCUMENT_READ_INDEX_PAGES)): + snippet = "" + for line in _pdf_page_text(pdf, i).splitlines(): + line = line.strip() + if line: + snippet = line[:_SNIPPET_CHARS] + break + index.append({"page": i + 1, "snippet": snippet}) + finally: + pdf.close() + payload = { + **base, + "mode": "index", + "page_count": count, + "page_index": index, + "index_truncated": count > DOCUMENT_READ_INDEX_PAGES, + "hint": ( + f"Read specific pages with page_range (max {DOCUMENT_READ_MAX_PAGES} per call, " + f"hard cap {DOCUMENT_READ_HARD_MAX_PAGES}) or locate them first with pattern." + ), + } + return DocumentReadResult(mode="index", payload=payload, format="pdf") + + +# --------------------------------------------------------------------------- +# DOCX (stdlib: zipfile + ElementTree over word/document.xml) +# --------------------------------------------------------------------------- + + +def docx_paragraphs(raw: bytes) -> List[str]: + """Paragraph text of a .docx, in document order. Tables contribute their + cell paragraphs; runs are joined; tabs and line breaks become whitespace. + Good enough for pattern search and bounded reading, not a renderer.""" + try: + with zipfile.ZipFile(io.BytesIO(raw)) as archive: + xml = archive.read("word/document.xml") + except (zipfile.BadZipFile, KeyError) as e: + raise DocumentReadError(f"Could not open the Word document: {e}") from e + try: + root = ElementTree.fromstring(xml) + except ElementTree.ParseError as e: + raise DocumentReadError(f"Could not parse the Word document: {e}") from e + paragraphs: List[str] = [] + for para in root.iter(f"{_W_NS}p"): + parts: List[str] = [] + for node in para.iter(): + if node.tag == f"{_W_NS}t": + parts.append(node.text or "") + elif node.tag in (f"{_W_NS}tab",): + parts.append("\t") + elif node.tag in (f"{_W_NS}br", f"{_W_NS}cr"): + parts.append(" ") + text = "".join(parts).strip() + if text: + paragraphs.append(text) + return paragraphs + + +def _docx_read(raw: bytes, pattern: Optional[str], offset: int, base: Dict[str, Any]) -> DocumentReadResult: + paragraphs = docx_paragraphs(raw) + if pattern: + result = _text_pattern("\n".join(paragraphs), pattern, base, "paragraph") + result.format = "docx" + return result + text = "\n\n".join(paragraphs) + encoded = text.encode("utf-8") + if offset and offset >= len(encoded): + raise WorkspaceValidationError(f"offset {offset} is beyond the end of the document text ({len(encoded)} bytes)") + chunk = encoded[offset: offset + WORKSPACE_READ_MAX_BYTES] + end = offset + len(chunk) + truncated = end < len(encoded) + payload = { + **base, + "mode": "text", + "paragraph_count": len(paragraphs), + "content": chunk.decode("utf-8", errors="replace"), + "offset": offset, + "truncated": truncated, + "next_offset": end if truncated else None, + } + return DocumentReadResult(mode="text", payload=payload, bytes_returned=len(chunk), format="docx") + + +# --------------------------------------------------------------------------- +# Text helpers +# --------------------------------------------------------------------------- + + +def _compile(pattern: str) -> "re.Pattern[str]": + try: + return re.compile(pattern, re.IGNORECASE) + except re.error: + return re.compile(re.escape(pattern), re.IGNORECASE) + + +def _grep_lines(lines: Sequence[str], regex: "re.Pattern[str]", budget: int) -> List[Dict[str, Any]]: + """Matching lines with ``_CONTEXT_LINES`` of context, 1-indexed.""" + out: List[Dict[str, Any]] = [] + for i, line in enumerate(lines): + if len(out) >= budget: + break + if not regex.search(line): + continue + before = [ln.strip()[:_LINE_CHARS] for ln in lines[max(0, i - _CONTEXT_LINES): i]] + after = [ln.strip()[:_LINE_CHARS] for ln in lines[i + 1: i + 1 + _CONTEXT_LINES]] + out.append({"line": i + 1, "text": line.strip()[:_LINE_CHARS], "before": before, "after": after}) + return out + + +def _text_pattern(text: str, pattern: str, base: Dict[str, Any], unit: str) -> DocumentReadResult: + regex = _compile(pattern) + lines = text.splitlines() + matches = _grep_lines(lines, regex, DOCUMENT_READ_MAX_MATCHES) + payload = { + **base, + "mode": "pattern", + "pattern": pattern, + "unit": unit, + f"{unit}_count": len(lines), + "match_count": len(matches), + "truncated": len(matches) >= DOCUMENT_READ_MAX_MATCHES, + "matches": matches, + } + return DocumentReadResult(mode="pattern", payload=payload, format=base.get("format")) + + +_NAME_UNSAFE = re.compile(r"[^a-zA-Z0-9\s\-\(\)\[\]]") + + +def _unique_document_name(filename: str, suffix: str) -> str: + """A Bedrock-safe document name that cannot collide with any other block in + the conversation. Bedrock allows alphanumerics, whitespace, hyphens, + parentheses and square brackets, and rejects duplicate names across the + whole message history — so every slice gets a fresh random tail.""" + stem = filename.rsplit(".", 1)[0] if "." in filename else filename + stem = re.sub(r"\s+", " ", _NAME_UNSAFE.sub(" ", stem)).strip()[:40] or "document" + return f"{stem} {suffix} {uuid.uuid4().hex[:6]}" diff --git a/backend/src/apis/shared/files/models.py b/backend/src/apis/shared/files/models.py index 18cf7934..52c35329 100644 --- a/backend/src/apis/shared/files/models.py +++ b/backend/src/apis/shared/files/models.py @@ -204,6 +204,14 @@ class FileMetadata(BaseModel): # never part of an access decision. source: str = Field(default="upload", description="Origin of the file") + # DocumentDigest (docs/specs/document-context-offload.md §4A), built once + # when a document upload completes and stored as a plain map — see + # ``apis.shared.files.document_digest``. ``None`` = never generated (files + # uploaded before PR-2, non-documents, or the flag off). Carries model + # prose (``abstract``) and heading text (``sections``): content-bearing, + # denylisted in the admin projections. + digest: Optional[dict] = Field(None, description="DocumentDigest map, when generated") + # Timestamps created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) @@ -230,7 +238,7 @@ def to_dynamo_item(self) -> dict: if ttl_value is None: ttl_value = int(self.created_at.timestamp()) + (365 * 24 * 60 * 60) - return { + item = { "PK": f"USER#{self.user_id}", "SK": f"FILE#{self.upload_id}", "GSI1PK": f"CONV#{self.session_id}", @@ -250,6 +258,9 @@ def to_dynamo_item(self) -> dict: "updatedAt": to_iso(self.updated_at), "ttl": ttl_value, } + if self.digest: + item["digest"] = self.digest + return item @classmethod def from_dynamo_item(cls, item: dict) -> "FileMetadata": @@ -271,6 +282,7 @@ def from_dynamo_item(cls, item: dict) -> "FileMetadata": created_at=from_iso(created_at) if created_at else datetime.now(timezone.utc), updated_at=from_iso(updated_at) if updated_at else datetime.now(timezone.utc), ttl=item.get("ttl"), + digest=item.get("digest") if isinstance(item.get("digest"), dict) else None, ) diff --git a/backend/src/apis/shared/files/repository.py b/backend/src/apis/shared/files/repository.py index cce1a1dc..3d3ba3f1 100644 --- a/backend/src/apis/shared/files/repository.py +++ b/backend/src/apis/shared/files/repository.py @@ -139,6 +139,40 @@ async def update_file_status( logger.error(f"Error updating file status {upload_id}: {e}") raise + async def update_file_digest( + self, user_id: str, upload_id: str, digest: dict + ) -> Optional[FileMetadata]: + """Store a ``DocumentDigest`` map on the file row (``SET digest``). + + Idempotent last-write-wins; ``None`` when the row no longer exists + (the file was deleted while the digest was being built). + """ + return self.update_file_digest_sync(user_id, upload_id, digest) + + def update_file_digest_sync( + self, user_id: str, upload_id: str, digest: dict + ) -> Optional[FileMetadata]: + """Synchronous body of :meth:`update_file_digest` — the restore path + persists a lazily built digest from synchronous code (see + :meth:`list_session_files_sync`).""" + try: + response = self._table.update_item( + Key={"PK": f"USER#{user_id}", "SK": f"FILE#{upload_id}"}, + UpdateExpression="SET digest = :digest, updatedAt = :now", + ExpressionAttributeValues={ + ":digest": self._convert_floats_to_decimals(dict(digest)), + ":now": utc_now_iso(), + }, + ConditionExpression="attribute_exists(PK)", + ReturnValues="ALL_NEW", + ) + return FileMetadata.from_dynamo_item(response["Attributes"]) + except ClientError as e: + if e.response["Error"]["Code"] == "ConditionalCheckFailedException": + return None + logger.error(f"Error updating file digest {upload_id}: {e}") + raise + async def delete_file(self, user_id: str, upload_id: str) -> Optional[FileMetadata]: """ Delete a file metadata record. @@ -257,6 +291,19 @@ async def list_session_files( Returns: List of FileMetadata """ + return self.list_session_files_sync(session_id, status) + + def list_session_files_sync( + self, session_id: str, status: Optional[FileStatus] = None + ) -> List[FileMetadata]: + """Synchronous body of :meth:`list_session_files`. + + The session manager's restore path (``TurnBasedSessionManager.initialize``) + runs synchronously inside the Strands agent constructor, under a + running event loop it cannot re-enter, and needs the session's upload + rows to rehydrate stripped documents. boto3 is synchronous anyway; the + async method is a thin wrapper over this. + """ try: query_params = { "IndexName": "SessionIndex", diff --git a/backend/src/apis/shared/observability/content_policy.py b/backend/src/apis/shared/observability/content_policy.py index 5867f121..d5f5213b 100644 --- a/backend/src/apis/shared/observability/content_policy.py +++ b/backend/src/apis/shared/observability/content_policy.py @@ -64,6 +64,8 @@ "filename", # user-chosen "s3Key", # embeds the filename "s3Uri", + "digest.abstract", # model-generated abstract of the document + "digest.sections", # heading text lifted from the document }) #: The one path a content-free reader may *request* but must never *return*. @@ -132,6 +134,12 @@ def is_content_bearing(path: str) -> bool: "compactionAppliedCount", "compactionForcedCount", "compactionFloorUnreachableCount", + # Document lifecycle rollups (per-call document fields summed; see + # `apis.shared.sessions.metadata.DOCUMENT_ROLLUP_ATTRS`) + "fullDocumentCalls", + "digestOnlyCalls", + "documentReadCalls", + "documentReadPages", ) #: C# rows for the cost anatomy and the session profile's trajectory. @@ -158,6 +166,19 @@ def is_content_bearing(path: str) -> bool: "prefixTokens", "windowRemovedMessages", "compactionEvents", + # Document context, optional: the attachment footprint of the live + # context at this call (counts, estimated tokens, a format→count map keyed + # by Bedrock's format enum) and the document_read retrievals the call + # requested. Numbers and enum keys only — never a filename or a byte. + "hasDocuments", + "documentCount", + "documentTokens", + "documentDigests", + "documentsAttached", + "documentSlices", + "documentSliceTokens", + "documentMime", + "documentReads", ) #: FILE# rows for the session profile's attachment summary. @@ -169,6 +190,12 @@ def is_content_bearing(path: str) -> bool: "source", "status", "createdAt", + # DocumentDigest coverage (numbers and the format enum only; the + # abstract and section titles are denylisted above). + "digest.status", + "digest.format", + "digest.count", + "digest.tokens", ) ALL_PROJECTIONS: Dict[str, Tuple[str, ...]] = { diff --git a/backend/src/apis/shared/sessions/metadata.py b/backend/src/apis/shared/sessions/metadata.py index f58a0486..63714f9a 100644 --- a/backend/src/apis/shared/sessions/metadata.py +++ b/backend/src/apis/shared/sessions/metadata.py @@ -1750,6 +1750,13 @@ async def _bump_session_aggregates( for kind, attr in _COMPACTION_EVENT_COUNTERS.items(): update_parts_add.append(f"{attr} :{attr}") values[f":{attr}"] = _compaction_event_count(message_metadata, kind) + # Document lifecycle rollups (docs/specs/document-context-offload.md + # §6.1): how many calls ran with the full document inline vs. a + # digest only, and how much document_read pulled back. Written + # as 0 while the diagnostics are on, like the counters above. + for attr, value in _document_rollups(message_metadata).items(): + update_parts_add.append(f"{attr} :{attr}") + values[f":{attr}"] = value update_expression = ( "ADD " + ", ".join(update_parts_add) + " SET " + ", ".join(update_parts_set) @@ -1782,6 +1789,40 @@ async def _bump_session_aggregates( } +#: Session-row counters derived from a call's document fields. ``fullDocumentCalls`` +#: and ``digestOnlyCalls`` are the digest-vs-full turn shares; the two +#: ``documentRead*`` counters sum the call's ``documentReads`` ledger entry. +DOCUMENT_ROLLUP_ATTRS = ("fullDocumentCalls", "digestOnlyCalls", "documentReadCalls", "documentReadPages") + + +def _document_rollups(message_metadata: Any) -> Dict[str, int]: + """``{attr: delta}`` for every ``DOCUMENT_ROLLUP_ATTRS`` entry, from the + call's ``hasDocuments`` / ``documentDigests`` / ``documentReads`` extras. + Absent or malformed fields count as zero — the bump must never fail.""" + extra = getattr(message_metadata, "model_extra", None) + extra = extra if isinstance(extra, dict) else {} + has_documents = bool(extra.get("hasDocuments")) + try: + digests = int(extra.get("documentDigests") or 0) + except (TypeError, ValueError): + digests = 0 + reads = extra.get("documentReads") + reads = reads if isinstance(reads, dict) else {} + + def _int(value: Any) -> int: + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + return { + "fullDocumentCalls": 1 if has_documents else 0, + "digestOnlyCalls": 1 if (digests > 0 and not has_documents) else 0, + "documentReadCalls": _int(reads.get("calls")), + "documentReadPages": _int(reads.get("pages")), + } + + def _compaction_event_count(message_metadata: Any, kind: str) -> int: """How many events of ``kind`` the call's ``compactionEvents`` extra carries. diff --git a/backend/src/apis/shared/tools/injected.py b/backend/src/apis/shared/tools/injected.py index 2d7b0067..6b36051c 100644 --- a/backend/src/apis/shared/tools/injected.py +++ b/backend/src/apis/shared/tools/injected.py @@ -42,9 +42,19 @@ # Session workspace files. A single toggle that provisions list/read/write. WORKSPACE_TOOL_IDS = frozenset({"workspace_files"}) -# Every id owned by a per-request factory. Memory-Space tools are deliberately -# absent: they are gated on an Agent's memory binding rather than on -# ``enabled_tools``, so they never reach the filter. +# Document retrieval (bound to session/user). Listed for the record, and +# deliberately NOT folded into ``INJECTED_TOOL_IDS`` below: ``document_read`` +# is gated on the session having a readable attachment, not on +# ``enabled_tools`` — there is no catalog entry, no RBAC grant and no picker +# toggle (docs/specs/document-context-offload.md §4B; the ``workspace_files`` +# key it would otherwise hang off is granted to no prod role). Like the +# Memory-Space tools it never reaches ``ToolFilter``. +DOCUMENT_TOOL_IDS = frozenset({"document_read"}) + +# Every id owned by a per-request factory. Memory-Space and document tools are +# deliberately absent: they are gated on an Agent's memory binding / the +# session's attachments rather than on ``enabled_tools``, so they never reach +# the filter. INJECTED_TOOL_IDS = frozenset( SPREADSHEET_TOOL_IDS | ARTIFACT_TOOL_IDS diff --git a/backend/tests/agents/builtin_tools/test_document_read_tool.py b/backend/tests/agents/builtin_tools/test_document_read_tool.py new file mode 100644 index 00000000..97bcc007 --- /dev/null +++ b/backend/tests/agents/builtin_tools/test_document_read_tool.py @@ -0,0 +1,239 @@ +"""``document_read`` tool factory, its session-state gate, and the two places +its presence has to be accounted for: the agent cache key and the tool-result +offloader's exemption list. + +The gate is the load-bearing decision of docs/specs/document-context-offload.md +§4B: the tool exists for any session with a readable attachment, whatever the +user's RBAC grants or picker state, and its id never enters +``INJECTED_TOOL_IDS``. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from agents.builtin_tools.document_read_tool import ( + DOCUMENT_READ_TOOL_NAME, + make_document_read_tool, + record_document_read, +) +from apis.shared.files.document_read import DocumentReadResult +from apis.shared.files.workspace import WorkspaceFileNotFoundError, WorkspaceStorageNotConfiguredError +from apis.shared.tools.injected import DOCUMENT_TOOL_IDS, INJECTED_TOOL_IDS + +TOOL_MODULE = "agents.builtin_tools.document_read_tool" +ROUTES = "apis.inference_api.chat.routes" + + +async def _call(tool, *args, **kwargs): + fn = getattr(tool, "__wrapped__", None) or tool + return await fn(*args, **kwargs) + + +# --------------------------------------------------------------------------- +# Tool factory +# --------------------------------------------------------------------------- + + +class TestTool: + def test_identity_is_mandatory_and_the_name_is_stable(self): + tool = make_document_read_tool("s1", "u1") + assert tool.tool_name == DOCUMENT_READ_TOOL_NAME == "document_read" + with pytest.raises(ValueError): + make_document_read_tool("", "u1") + with pytest.raises(ValueError): + make_document_read_tool("s1", "") + + def test_id_is_recorded_but_never_an_injected_catalog_id(self): + assert DOCUMENT_TOOL_IDS == {"document_read"} + assert not (DOCUMENT_TOOL_IDS & INJECTED_TOOL_IDS) + + @pytest.mark.asyncio + async def test_no_upload_id_lists_the_sessions_documents(self, monkeypatch): + listing = AsyncMock(return_value={"documents": [{"upload_id": "up-1"}], "count": 1}) + monkeypatch.setattr(f"{TOOL_MODULE}.list_session_documents", listing) + result = await _call(make_document_read_tool("s1", "u1")) + listing.assert_awaited_once_with("u1", "s1") + assert result["status"] == "success" + payload = result["content"][0]["json"] + assert payload["count"] == 1 and "hint" in payload + + @pytest.mark.asyncio + async def test_page_read_appends_the_native_block_after_the_metadata(self, monkeypatch): + block = {"document": {"format": "pdf", "name": "policy p4-7 abc123", "source": {"bytes": b"%PDF"}}} + read = AsyncMock(return_value=DocumentReadResult( + mode="pages", payload={"pages_returned": 4}, document_block=block, pages_returned=4, bytes_returned=4, format="pdf", + )) + monkeypatch.setattr(f"{TOOL_MODULE}.read_document", read) + result = await _call(make_document_read_tool("s1", "u1"), upload_id="up-1", page_range="4-7", max_pages=6) + read.assert_awaited_once_with("u1", "s1", "up-1", page_range="4-7", pattern=None, max_pages=6, offset=0) + assert result["content"] == [{"json": {"pages_returned": 4}}, block] + + @pytest.mark.asyncio + async def test_empty_strings_mean_absent_arguments(self, monkeypatch): + read = AsyncMock(return_value=DocumentReadResult(mode="index", payload={})) + monkeypatch.setattr(f"{TOOL_MODULE}.read_document", read) + await _call(make_document_read_tool("s1", "u1"), upload_id="up-1", page_range="", pattern="") + assert read.await_args.kwargs["page_range"] is None and read.await_args.kwargs["pattern"] is None + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "exc,fragment", + [ + (WorkspaceStorageNotConfiguredError("no bucket"), "not configured"), + (WorkspaceFileNotFoundError("No file with id 'up-9'"), "up-9"), + (RuntimeError("s3 down"), "s3 down"), + ], + ) + async def test_failures_surface_as_error_results_not_exceptions(self, monkeypatch, exc, fragment): + monkeypatch.setattr(f"{TOOL_MODULE}.read_document", AsyncMock(side_effect=exc)) + result = await _call(make_document_read_tool("s1", "u1"), upload_id="up-9") + assert result["status"] == "error" + assert fragment in result["content"][0]["text"] + + def test_metric_is_content_free_and_in_the_compaction_namespace(self, monkeypatch): + calls = [] + monkeypatch.setattr( + "apis.shared.observability.emf.emit_emf_metrics", + lambda ns, metrics, properties=None, units=None: calls.append((ns, metrics, properties)), + ) + monkeypatch.setattr("apis.shared.observability.prompt_cache.prompt_cache_observability_enabled", lambda: True) + record_document_read(DocumentReadResult( + mode="pages", payload={"filename": "secret.pdf"}, pages_returned=4, bytes_returned=900, format="pdf", + )) + assert calls == [( + "AgentCoreStack/Compaction", + {"DocumentRead": 1, "DocumentReadPages": 4, "DocumentReadBytes": 900}, + {"mode": "pages", "format": "pdf"}, + )] + assert "secret" not in repr(calls) + + +# --------------------------------------------------------------------------- +# Routes gate +# --------------------------------------------------------------------------- + + +@pytest.fixture +def clear_memo(): + from apis.inference_api.chat import routes + + routes._DOCUMENT_SESSIONS.clear() + yield + routes._DOCUMENT_SESSIONS.clear() + + +class TestRoutesGate: + @pytest.mark.asyncio + async def test_this_turns_uploads_build_the_tool_without_a_query(self, monkeypatch, clear_memo): + from apis.inference_api.chat.routes import _build_document_tools + + lookup = AsyncMock(return_value=False) + monkeypatch.setattr("apis.shared.files.document_read.session_has_documents", lookup) + tools = await _build_document_tools("s1", "u1", turn_upload_ids=["up-1"]) + assert [t.tool_name for t in tools] == ["document_read"] + lookup.assert_not_awaited() + + @pytest.mark.asyncio + async def test_a_session_with_a_document_gets_the_tool_regardless_of_enabled_tools(self, monkeypatch, clear_memo): + from apis.inference_api.chat.routes import _build_document_tools + + lookup = AsyncMock(return_value=True) + monkeypatch.setattr("apis.shared.files.document_read.session_has_documents", lookup) + assert len(await _build_document_tools("s1", "u1")) == 1 + lookup.assert_awaited_once_with("u1", "s1") + # Memoized: the second turn does not query again. + assert len(await _build_document_tools("s1", "u1")) == 1 + assert lookup.await_count == 1 + + @pytest.mark.asyncio + async def test_no_document_means_no_tool_and_no_memo(self, monkeypatch, clear_memo): + from apis.inference_api.chat.routes import _DOCUMENT_SESSIONS, _build_document_tools + + lookup = AsyncMock(return_value=False) + monkeypatch.setattr("apis.shared.files.document_read.session_has_documents", lookup) + assert await _build_document_tools("s1", "u1") == [] + assert await _build_document_tools("s1", "u1") == [] + assert lookup.await_count == 2 and "s1" not in _DOCUMENT_SESSIONS + + @pytest.mark.asyncio + async def test_lookup_failure_fails_closed(self, monkeypatch, clear_memo): + from apis.inference_api.chat.routes import _build_document_tools + + monkeypatch.setattr("apis.shared.files.document_read.session_has_documents", AsyncMock(side_effect=RuntimeError("ddb"))) + assert await _build_document_tools("s1", "u1") == [] + + @pytest.mark.asyncio + async def test_kill_switch_and_missing_identity(self, monkeypatch, clear_memo): + from apis.inference_api.chat.routes import _build_document_tools, _document_tools_gate + + monkeypatch.setenv("DOCUMENT_READ_ENABLED", "false") + assert await _build_document_tools("s1", "u1", turn_upload_ids=["up-1"]) == [] + assert await _document_tools_gate("s1", "u1", turn_upload_ids=["up-1"]) is False + monkeypatch.setenv("DOCUMENT_READ_ENABLED", "") + assert await _document_tools_gate("s1", "u1", turn_upload_ids=["up-1"]) is True + assert await _document_tools_gate("", "u1", turn_upload_ids=["up-1"]) is False + + def test_memo_is_bounded(self, monkeypatch, clear_memo): + from apis.inference_api.chat import routes + + monkeypatch.setattr(routes, "_DOCUMENT_SESSIONS_MAX", 3) + for i in range(5): + routes._remember_document_session(f"s{i}") + assert list(routes._DOCUMENT_SESSIONS) == ["s2", "s3", "s4"] + + +# --------------------------------------------------------------------------- +# Agent cache key +# --------------------------------------------------------------------------- + + +def test_cache_key_carries_the_document_tool_bit_without_moving_the_skills_hash(): + from apis.inference_api.chat import service + + base = dict( + session_id="s", user_id="u", enabled_tools=["t"], model_id="m", inference_params={}, + system_prompt=None, caching_enabled=False, provider="bedrock", freshness_hash="f", agent_type="chat", + ) + without = service._create_cache_key(**base, skills_hash="k") + with_docs = service._create_cache_key(**base, skills_hash="k", document_tools=True) + assert without != with_docs + assert without[-1] == with_docs[-1] == "k" + assert without[-2] is False and with_docs[-2] is True + + +# --------------------------------------------------------------------------- +# Offloader exemption +# --------------------------------------------------------------------------- + + +class TestOffloaderExemption: + def test_document_read_results_are_never_offloaded(self): + from agents.main_agent.core.tool_result_offload import OFFLOAD_EXEMPT_TOOLS, _exempt_tool, _should_offload + + assert "document_read" in OFFLOAD_EXEMPT_TOOLS + assert _exempt_tool(SimpleNamespace(tool_use={"name": "document_read"})) is True + assert _exempt_tool(SimpleNamespace(tool_use={"name": "gmail_search"})) is False + assert _exempt_tool(SimpleNamespace(tool_use=None)) is False + assert _should_offload("document_read", 90_000) is False + assert _should_offload("gmail_search", 90_000) is True + + @pytest.mark.asyncio + async def test_mixin_returns_before_counting_tokens_for_an_exempt_tool(self): + from agents.main_agent.core.tool_result_offload import _OffloaderMixin + + class _Base: + async def _handle_tool_result(self, event): # pragma: no cover - must not run + raise AssertionError("base offloader reached for an exempt tool") + + class _Sut(_OffloaderMixin, _Base): + _max_result_tokens = 10 + + big = "x" * 4_000 + result = {"toolUseId": "t1", "status": "success", "content": [{"text": big}]} + event = SimpleNamespace(result=result, tool_use={"toolUseId": "t1", "name": "document_read"}) + await _Sut()._handle_tool_result(event) + assert event.result is result diff --git a/backend/tests/agents/main_agent/session/test_document_context.py b/backend/tests/agents/main_agent/session/test_document_context.py new file mode 100644 index 00000000..ff660fdc --- /dev/null +++ b/backend/tests/agents/main_agent/session/test_document_context.py @@ -0,0 +1,145 @@ +"""The document-lifecycle analytics (docs/specs/document-context-offload.md §6.1): + +- ``summarize_document_context`` — the attachment footprint of the live + context, persisted per cost row (content-free by construction). +- the ``document_stripped`` compaction-ledger event ``_strip_document_bytes`` + records, so the restore defect is measured before PR-3 fixes it. +- the ``documentReads`` ledger entry the ``ContextLedgerHook`` tallies from + ``document_read`` results. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from agents.main_agent.session.compaction_policy import CHARS_PER_TOKEN, IMAGE_TOKEN_ESTIMATE +from agents.main_agent.session.document_context import summarize_document_context +from agents.main_agent.session.hooks.context_ledger import ContextLedgerHook + + +@pytest.fixture(autouse=True) +def diagnostics_enabled(monkeypatch): + monkeypatch.delenv("COST_DIAGNOSTICS_ENABLED", raising=False) + + +def _doc(name="a", fmt="pdf", size=4000): + return {"document": {"format": fmt, "name": name, "source": {"bytes": b"x" * size}}} + + +def _image(size=100): + return {"image": {"format": "png", "source": {"bytes": b"i" * size}}} + + +class TestSummary: + def test_empty_or_malformed_is_none(self): + assert summarize_document_context(None) is None + assert summarize_document_context([]) is None + + def test_counts_tokens_formats_digests_and_the_last_prompts_attachments(self): + messages = [ + {"role": "user", "content": [{"text": "q1"}, _doc("a", "pdf", 4000), _doc("b", "docx", 800), _image()]}, + {"role": "assistant", "content": [{"text": "a1"}]}, + {"role": "user", "content": [{"text": "q2"}, {"text": "[Document placeholder: name=c, format=pdf, original_size=9 bytes]"}]}, + {"role": "assistant", "content": [{"toolUse": {"toolUseId": "t", "name": "document_read", "input": {}}}]}, + {"role": "user", "content": [{"toolResult": {"toolUseId": "t", "status": "success", "content": [ + {"json": {"pages_returned": 2}}, _doc("slice", "pdf", 2000)]}}]}, + {"role": "user", "content": [{"text": " "}, {"text": "q3"}]}, + ] + assert summarize_document_context(messages) == { + "hasDocuments": True, + "documentCount": 3, + "documentTokens": 4000 // CHARS_PER_TOKEN + 800 // CHARS_PER_TOKEN + IMAGE_TOKEN_ESTIMATE, + "documentDigests": 2, + "documentsAttached": 0, # the last prompt (q3) attached nothing + "documentSlices": 1, + "documentSliceTokens": 2000 // CHARS_PER_TOKEN, + "documentMime": {"pdf": 1, "docx": 1, "image": 1}, + } + + def test_a_digest_only_context_reads_as_no_documents(self): + messages = [ + {"role": "user", "content": [{"text": "[Document placeholder: name=a, format=pdf, original_size=1 bytes]"}, {"text": "q"}]}, + {"role": "assistant", "content": [{"text": "a"}]}, + ] + summary = summarize_document_context(messages) + assert summary["hasDocuments"] is False + assert summary["documentCount"] == 0 and summary["documentDigests"] == 1 + + def test_attach_turn_counts_its_own_attachments(self): + messages = [{"role": "user", "content": [{"text": "q"}, _doc(), _doc("b")]}] + assert summarize_document_context(messages)["documentsAttached"] == 2 + + def test_never_emits_names_or_bytes(self): + summary = summarize_document_context([{"role": "user", "content": [_doc("Secret Contract")]}]) + assert "Secret" not in repr(summary) and b"x" not in repr(summary).encode() + + +class TestStripEvent: + """With no upload rows to match, restore falls back to the placeholder and + records ``document_stripped`` — the pre-PR-3 behavior, kept as the floor. + The rehydrated path is pinned in ``test_document_rehydration.py``.""" + + @pytest.fixture(autouse=True) + def no_upload_rows(self, monkeypatch): + monkeypatch.setattr( + "agents.main_agent.session.document_rehydration.load_session_documents", lambda *_: [] + ) + + def test_strip_records_one_content_free_event(self, make_session_manager): + manager = make_session_manager() + messages = [ + {"role": "user", "content": [{"text": "q"}, _doc("a", "pdf", 8000), _doc("b", "txt", 400)]}, + {"role": "assistant", "content": [{"text": "a"}]}, + ] + stripped = manager._strip_document_bytes(messages) + assert all("document" not in b for b in stripped[0]["content"]) + assert manager.drain_compaction_events() == [ + {"kind": "document_stripped", "documents": 2, "documentTokens": 8400 // CHARS_PER_TOKEN} + ] + + def test_nothing_to_strip_records_nothing(self, make_session_manager): + manager = make_session_manager() + manager._strip_document_bytes([{"role": "user", "content": [{"text": "q"}]}]) + assert manager.drain_compaction_events() == [] + + def test_kill_switch_still_strips_but_records_nothing(self, make_session_manager, monkeypatch): + monkeypatch.setenv("COST_DIAGNOSTICS_ENABLED", "false") + manager = make_session_manager() + stripped = manager._strip_document_bytes([{"role": "user", "content": [_doc()]}]) + assert "text" in stripped[0]["content"][0] + assert manager.drain_compaction_events() == [] + + +def _tool_event(name="document_read", status="success", pages=3, size=500): + content = [{"json": {"pages_returned": pages}}] + if size: + content.append(_doc("slice", size=size)) + return MagicMock(tool_use={"toolUseId": "t", "name": name}, result={"toolUseId": "t", "status": status, "content": content}) + + +class TestLedgerDocumentReads: + def test_reads_are_attributed_to_the_requesting_call(self): + hook = ContextLedgerHook() + agent = MagicMock() + agent.conversation_manager.removed_message_count = None + agent._session_manager = None + hook._on_turn_start(MagicMock()) + hook._on_before_model_call(MagicMock(agent=agent)) # call 0 + hook._on_after_tool_call(_tool_event(pages=3, size=500)) + hook._on_after_tool_call(_tool_event(pages=2, size=300)) + hook._on_before_model_call(MagicMock(agent=agent)) # call 1 + hook._on_after_tool_call(_tool_event(name="calculator")) + hook._on_after_tool_call(_tool_event(status="error")) + + assert hook.ledger_for_call(0) == {"documentReads": {"calls": 2, "pages": 5, "bytes": 800}} + assert hook.ledger_for_call(1) is None + + def test_kill_switch_records_nothing(self, monkeypatch): + monkeypatch.setenv("COST_DIAGNOSTICS_ENABLED", "false") + hook = ContextLedgerHook() + hook._on_turn_start(MagicMock()) + hook._on_before_model_call(MagicMock()) + hook._on_after_tool_call(_tool_event()) + assert hook.ledger_for_call(0) is None diff --git a/backend/tests/agents/main_agent/session/test_document_rehydration.py b/backend/tests/agents/main_agent/session/test_document_rehydration.py new file mode 100644 index 00000000..9f28c92d --- /dev/null +++ b/backend/tests/agents/main_agent/session/test_document_rehydration.py @@ -0,0 +1,185 @@ +"""Restore rehydrates stripped documents as digests (offload spec §4E, PR-3). + +Pins: block→upload matching (sanitized name incl. PromptBuilder's `_2` +suffix, size tiebreak, no double-claiming), the rendered digest replacing the +bytes with the `upload_id` handle inside, lazy outline-only digests for rows +without one (persisted, no model call), the placeholder fallback for +unmatched blocks, the never-raise contract, the kill switch, and the two +ledger events the session manager records. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from agents.main_agent.session import document_rehydration as dr +from apis.shared.files.models import FileMetadata, FileStatus +from tests.shared.test_document_read import build_pdf + +MODULE = "agents.main_agent.session.document_rehydration" + + +@pytest.fixture(autouse=True) +def diagnostics_enabled(monkeypatch): + monkeypatch.delenv("COST_DIAGNOSTICS_ENABLED", raising=False) + monkeypatch.delenv("DOCUMENT_REHYDRATE_ENABLED", raising=False) + + +def _meta(upload_id="up-1", filename="BBR Policy.pdf", size=1000, mime="application/pdf", digest=None, user="u1"): + return FileMetadata( + upload_id=upload_id, user_id=user, session_id="s1", filename=filename, mime_type=mime, + size_bytes=size, s3_key="k", s3_bucket="b", status=FileStatus.READY, digest=digest, + ) + + +def _doc(name="BBR Policy_pdf", fmt="pdf", raw=b"x" * 1000): + """PromptBuilder sanitizes the WHOLE filename, dot included: `BBR Policy.pdf` → `BBR Policy_pdf`.""" + return {"document": {"format": fmt, "name": name, "source": {"bytes": raw}}} + + +READY_DIGEST = { + "version": 1, "status": "ready", "format": "pdf", "unit": "page", "count": 47, + "sections": [{"start": 1, "title": "Declarations"}, {"start": 4, "title": "Insuring Agreements"}], + "abstract": "A cyber liability policy form.", "tokens": 80, +} + + +class TestMatching: + def test_prefers_name_and_size_then_name_then_format_and_size(self): + a = _meta("a", "BBR Policy.pdf", 1000) + b = _meta("b", "BBR Policy.pdf", 2000) + c = _meta("c", "Other.pdf", 1000) + assert dr.match_document("BBR Policy_pdf", "pdf", 2000, [a, b, c], set()) is b + assert dr.match_document("BBR Policy_pdf", "pdf", 999, [a, b, c], set()) is a + assert dr.match_document("Renamed", "pdf", 1000, [b, c], set()) is c + assert dr.match_document("Renamed", "docx", 1000, [b, c], set()) is None + + def test_duplicate_suffix_and_claimed_rows(self): + first = _meta("a", "memo.pdf", 10) + second = _meta("b", "memo.pdf", 10) + used: set = set() + assert dr.match_document("memo_pdf", "pdf", 10, [first, second], used) is first + used.add("a") + assert dr.match_document("memo_pdf_2", "pdf", 10, [first, second], used) is second + used.add("b") + assert dr.match_document("memo_pdf_3", "pdf", 10, [first, second], used) is None + + +class TestRehydrate: + def test_replaces_bytes_with_the_rendered_digest_and_handle(self, monkeypatch): + meta = _meta(digest=READY_DIGEST) + monkeypatch.setattr(f"{MODULE}.load_session_documents", lambda *_: [meta]) + messages = [ + {"role": "user", "content": [{"text": "q\n\n[Attached files: BBR Policy.pdf]"}, _doc()]}, + {"role": "assistant", "content": [{"text": "a"}]}, + ] + result = dr.rehydrate_documents(messages, session_id="s1", user_id="u1") + + block = result.messages[0]["content"][1] + assert set(block) == {"text"} + assert block["text"].startswith('') + assert "A cyber liability policy form." in block["text"] + assert '
Insuring Agreements
' in block["text"] + assert (result.rehydrated, result.stripped, result.lazy_digests) == (1, 0, 0) + assert result.digest_tokens > 0 and result.upload_ids == ["up-1"] + # The originals are untouched (deep copy) and no bytes survive. + assert "document" in messages[0]["content"][1] + assert b"x" * 1000 not in repr(result.messages).encode() + + def test_lazy_digest_is_built_from_the_restored_bytes_and_persisted(self, monkeypatch): + meta = _meta(digest=None, size=5) + monkeypatch.setattr(f"{MODULE}.load_session_documents", lambda *_: [meta]) + repo = MagicMock() + monkeypatch.setattr("apis.shared.files.repository.get_file_upload_repository", lambda: repo) + raw = build_pdf(["1. Intro\nbody", "2. Terms"]) + messages = [{"role": "user", "content": [_doc(raw=raw)]}] + + result = dr.rehydrate_documents(messages, session_id="s1", user_id="u1") + + text = result.messages[0]["content"][0]["text"] + assert 'upload_id="up-1"' in text and '
2. Terms
' in text + assert "" not in text # outline only — no model call on the restore path + assert result.lazy_digests == 1 + repo.update_file_digest_sync.assert_called_once() + user, upload, item = repo.update_file_digest_sync.call_args.args + assert (user, upload, item["status"], item["count"]) == ("u1", "up-1", "ready", 2) + assert item["tokens"] > 0 + + def test_lazy_digest_persist_failure_still_serves_this_restore(self, monkeypatch): + meta = _meta(digest=None) + monkeypatch.setattr(f"{MODULE}.load_session_documents", lambda *_: [meta]) + repo = MagicMock() + repo.update_file_digest_sync.side_effect = RuntimeError("ddb") + monkeypatch.setattr("apis.shared.files.repository.get_file_upload_repository", lambda: repo) + result = dr.rehydrate_documents([{"role": "user", "content": [_doc(raw=build_pdf(["Hello"]))]}], session_id="s1", user_id="u1") + assert result.rehydrated == 1 and " 0 + assert events[1] == {"kind": "document_stripped", "documents": 1, "documentTokens": 100} + + def test_restore_output_is_stable_across_restores(self, make_session_manager, monkeypatch): + rows = [_meta("up-1", "BBR Policy.pdf", 1000, digest=READY_DIGEST)] + monkeypatch.setattr(f"{MODULE}.load_session_documents", lambda *_: rows) + manager = make_session_manager() + messages = [{"role": "user", "content": [_doc()]}] + first = manager._strip_document_bytes(messages) + second = manager._strip_document_bytes(messages) + assert first == second diff --git a/backend/tests/apis/app_api/admin/costs/test_document_fields.py b/backend/tests/apis/app_api/admin/costs/test_document_fields.py new file mode 100644 index 00000000..9dde52bb --- /dev/null +++ b/backend/tests/apis/app_api/admin/costs/test_document_fields.py @@ -0,0 +1,138 @@ +"""Document-context fields on the admin cost surfaces +(docs/specs/document-context-offload.md §6.1). + +Rows carry the attachment footprint at each call; the anatomy projects it +per call, the profile rolls it up (digest-vs-full turn shares, peak document +tokens, document_read totals) and reports coverage. The session-row rollups +the write path ADDs are pinned too, so the two never disagree on a definition. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from apis.app_api.admin.costs.models import CompactionEvent +from apis.app_api.admin.costs.service import AdminCostService, _call_ledger +from apis.shared.sessions.metadata import DOCUMENT_ROLLUP_ATTRS, _document_rollups + + +def _record(i, **extra): + rec = { + "timestamp": f"2026-09-16T00:00:{i:02d}Z", + "messageId": i, + "tokenUsage": {"inputTokens": 100, "outputTokens": 5, "cacheReadInputTokens": 0, "cacheWriteInputTokens": 0}, + "modelInfo": {"modelId": "m1"}, + "cost": {"total": 0.01}, + "cacheStatus": "hit", + } + rec.update(extra) + return rec + + +FULL = dict(hasDocuments=True, documentCount=2, documentTokens=12_000, documentDigests=0, + documentsAttached=2, documentSlices=0, documentSliceTokens=0, documentMime={"pdf": 2}) +DIGEST = dict(hasDocuments=False, documentCount=0, documentTokens=0, documentDigests=2, + documentsAttached=0, documentSlices=1, documentSliceTokens=900, documentMime={}, + documentReads={"calls": 1, "pages": 4, "bytes": 3600}) + + +class TestCallLedger: + def test_absent_fields_stay_untracked(self): + ledger = _call_ledger(_record(0), None) + assert ledger.documents is None and ledger.document_reads is None + + def test_fields_decode_with_decimal_coercion(self): + from decimal import Decimal + + rec = _record(0, hasDocuments=True, documentCount=Decimal("1"), documentTokens=Decimal("500"), + documentMime={"pdf": Decimal("1")}, documentReads={"calls": Decimal("2"), "pages": 3}) + ledger = _call_ledger(rec, None) + assert ledger.documents == {"hasDocuments": True, "documentCount": 1, "documentTokens": 500, "documentMime": {"pdf": 1}} + assert (ledger.document_reads.calls, ledger.document_reads.pages, ledger.document_reads.bytes) == (2, 3, 0) + + def test_compaction_event_accepts_the_document_kinds(self): + event = CompactionEvent(kind="document_offload", documents=1, documentTokens=9000, cacheGapSeconds=420) + assert (event.documents, event.document_tokens, event.cache_gap_seconds) == (1, 9000, 420) + + +def _service(records, row=None, files=None): + service = AdminCostService.__new__(AdminCostService) + service.storage = AsyncMock() + service.storage.get_session_cost_records = AsyncMock(return_value=records) + service.storage.get_session_diagnostic_row = AsyncMock(return_value=row) + service.storage.get_user_cost_summary = AsyncMock(return_value=None) + service._file_repository = AsyncMock() + service._file_repository.list_session_file_stats = AsyncMock(return_value=files or []) + return service + + +class TestAnatomy: + @pytest.mark.asyncio + async def test_rows_project_the_document_fields_or_null(self): + anatomy = await _service([_record(0, **FULL), _record(1, **DIGEST), _record(2)]).get_session_cost_anatomy("s1") + full, digest, bare = anatomy.calls + assert full.has_documents is True and full.document_tokens == 12_000 and full.document_mime == {"pdf": 2} + assert full.document_reads is None + assert digest.has_documents is False and digest.document_digests == 2 + assert digest.document_reads.pages == 4 and digest.document_slice_tokens == 900 + assert bare.has_documents is None and bare.document_reads is None + wire = anatomy.model_dump(by_alias=True)["calls"][1] + assert wire["documentReads"] == {"calls": 1, "pages": 4, "bytes": 3600} + + +def _row(**overrides): + row = { + "sessionId": "s1", "userId": "u1", "status": "active", "messageCount": 4, "totalCost": 0.5, + "lastContextTokens": 20_000, "contextWindow": 200_000, "totalCacheReadTokens": 1, "totalCacheWriteTokens": 1, + "preferences": {"lastModel": "m1", "enabledTools": []}, "compaction": {}, + } + row.update(overrides) + return row + + +class TestProfile: + @pytest.mark.asyncio + async def test_rollups_from_rows(self): + records = [_record(0, **FULL), _record(1, **{**FULL, "documentTokens": 15_000}), _record(2, **DIGEST), _record(3)] + profile = await _service(records, row=_row()).get_session_profile("s1") + assert profile.data_coverage.documents is True + assert (profile.full_document_calls, profile.digest_only_calls) == (2, 1) + assert profile.peak_document_tokens == 15_000 + assert (profile.document_read_calls, profile.document_read_pages) == (1, 4) + + @pytest.mark.asyncio + async def test_session_row_rollups_cover_rows_that_expired(self): + row = _row(fullDocumentCalls=7, digestOnlyCalls=2, documentReadCalls=3, documentReadPages=11) + profile = await _service([_record(0)], row=row).get_session_profile("s1") + assert profile.data_coverage.documents is True + assert (profile.full_document_calls, profile.digest_only_calls) == (7, 2) + assert (profile.document_read_calls, profile.document_read_pages) == (3, 11) + assert profile.peak_document_tokens is None + + @pytest.mark.asyncio + async def test_untracked_reads_as_not_tracked(self): + profile = await _service([_record(0)], row=_row()).get_session_profile("s1") + assert profile.data_coverage.documents is False + assert profile.full_document_calls == 0 and profile.peak_document_tokens is None + + +class TestWriteSideRollups: + def test_attrs_are_the_ones_the_profile_reads(self): + assert DOCUMENT_ROLLUP_ATTRS == ("fullDocumentCalls", "digestOnlyCalls", "documentReadCalls", "documentReadPages") + + @pytest.mark.parametrize( + "extra,expected", + [ + ({}, (0, 0, 0, 0)), + (FULL, (1, 0, 0, 0)), + (DIGEST, (0, 1, 1, 4)), + ({"hasDocuments": True, "documentDigests": 3}, (1, 0, 0, 0)), # inline wins over digests + ({"documentReads": {"calls": "x", "pages": None}}, (0, 0, 0, 0)), + ], + ) + def test_deltas_from_the_calls_extras(self, extra, expected): + rollups = _document_rollups(SimpleNamespace(model_extra=dict(extra))) + assert tuple(rollups[a] for a in DOCUMENT_ROLLUP_ATTRS) == expected diff --git a/backend/tests/apis/app_api/admin/costs/test_drilldown_routes.py b/backend/tests/apis/app_api/admin/costs/test_drilldown_routes.py index 0d5d9f56..4eeda6a4 100644 --- a/backend/tests/apis/app_api/admin/costs/test_drilldown_routes.py +++ b/backend/tests/apis/app_api/admin/costs/test_drilldown_routes.py @@ -105,6 +105,7 @@ def test_session_profile_returns_200(): assert body["dataCoverage"] == { "toolCensus": False, "compactionCount": False, "fingerprints": False, "cost": False, "prefixTokens": False, "windowTrim": False, "compactionEvents": False, + "documents": False, } service.get_session_profile.assert_awaited_once_with("s1") diff --git a/backend/tests/apis/app_api/test_file_digest_scheduling.py b/backend/tests/apis/app_api/test_file_digest_scheduling.py new file mode 100644 index 00000000..3916f999 --- /dev/null +++ b/backend/tests/apis/app_api/test_file_digest_scheduling.py @@ -0,0 +1,119 @@ +"""DocumentDigest scheduling on upload completion (files/service.py). + +``complete_upload`` queues a digest build for document uploads only, off the +request path, with a strong task reference; the build reads S3, calls the +builder and persists the map; nothing on this path can fail the upload or +raise out of the task. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from apis.app_api.files.service import FileUploadService +from apis.shared.files.document_digest import DocumentDigest +from apis.shared.files.models import FileMetadata, FileStatus + +DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + + +def _meta(mime="application/pdf", filename="policy.pdf", status=FileStatus.PENDING): + return FileMetadata( + upload_id="up-1", user_id="u1", session_id="s1", filename=filename, mime_type=mime, + size_bytes=10, s3_key="k", s3_bucket="b", status=status, + ) + + +def _service(meta, body=b"%PDF"): + repository = MagicMock() + repository.get_file = AsyncMock(return_value=meta) + repository.update_file_status = AsyncMock() + repository.increment_quota = AsyncMock() + repository.update_file_digest = AsyncMock(return_value=meta) + s3 = MagicMock() + s3.get_object.return_value = {"Body": MagicMock(read=MagicMock(return_value=body))} + return FileUploadService(repository=repository, s3_client=s3, bucket_name="b") + + +@pytest.fixture(autouse=True) +def _flag(monkeypatch): + monkeypatch.delenv("DOCUMENT_DIGEST_ENABLED", raising=False) + + +async def _drain(service): + tasks = list(service._digest_tasks) + if tasks: + await asyncio.gather(*tasks) + return tasks + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mime,filename,scheduled", + [ + ("application/pdf", "a.pdf", True), + (DOCX, "a.docx", True), + ("text/markdown", "a.md", True), + ("text/csv", "a.csv", False), + ("image/png", "a.png", False), + ("application/vnd.openxmlformats-officedocument.presentationml.presentation", "a.pptx", False), + ], +) +async def test_complete_upload_schedules_a_digest_for_documents_only(monkeypatch, mime, filename, scheduled): + built = DocumentDigest(status="ready", format="pdf", count=1, tokens=50) + build = AsyncMock(return_value=built) + monkeypatch.setattr("apis.shared.files.document_digest.build_digest", build) + service = _service(_meta(mime, filename)) + + response = await service.complete_upload("u1", "up-1") + assert response.status == "ready" + tasks = await _drain(service) + + assert bool(tasks) is scheduled + if scheduled: + build.assert_awaited_once() + assert build.await_args.kwargs["raw"] == b"%PDF" + service.repository.update_file_digest.assert_awaited_once_with("u1", "up-1", built.to_item()) + else: + service.repository.update_file_digest.assert_not_awaited() + assert not service._digest_tasks # done-callback dropped the strong ref + + +@pytest.mark.asyncio +async def test_kill_switch_skips_scheduling(monkeypatch): + monkeypatch.setenv("DOCUMENT_DIGEST_ENABLED", "false") + service = _service(_meta()) + await service.complete_upload("u1", "up-1") + assert service._schedule_digest(_meta()) is None and not service._digest_tasks + + +@pytest.mark.asyncio +async def test_build_failures_never_escape_the_task(monkeypatch): + monkeypatch.setattr("apis.shared.files.document_digest.build_digest", AsyncMock(side_effect=RuntimeError("boom"))) + service = _service(_meta()) + await service.complete_upload("u1", "up-1") + await _drain(service) # gather would re-raise if the task did + service.repository.update_file_digest.assert_not_awaited() + + service.repository.get_file = AsyncMock(return_value=_meta()) + service._s3_client.get_object.side_effect = RuntimeError("s3 down") + monkeypatch.setattr("apis.shared.files.document_digest.build_digest", AsyncMock()) + await service.complete_upload("u1", "up-1") + await _drain(service) + + +@pytest.mark.asyncio +async def test_a_deleted_row_is_tolerated(monkeypatch): + monkeypatch.setattr("apis.shared.files.document_digest.build_digest", AsyncMock(return_value=DocumentDigest())) + service = _service(_meta()) + service.repository.update_file_digest = AsyncMock(return_value=None) + await service.complete_upload("u1", "up-1") + await _drain(service) + + +def test_scheduling_without_a_running_loop_is_a_noop(): + service = _service(_meta()) + assert service._schedule_digest(_meta()) is None diff --git a/backend/tests/shared/test_document_digest.py b/backend/tests/shared/test_document_digest.py new file mode 100644 index 00000000..4299e8f3 --- /dev/null +++ b/backend/tests/shared/test_document_digest.py @@ -0,0 +1,196 @@ +"""DocumentDigest (apis/shared/files/document_digest.py). + +The deterministic outline (headings, counts, table/figure mentions, sampled +sections), the rendered block and its hard token budget — including the +spec's gate, a 200-page PDF digest under 1,500 tokens — the fail-open +abstract, and ``build_digest``'s never-raise contract. Fixtures are the same +hand-built PDFs and DOCX archives the document_read tests use. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from apis.shared.files import document_digest as dd +from tests.shared.test_document_read import build_docx, build_pdf + +MODULE = "apis.shared.files.document_digest" +DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + + +class TestHeadings: + @pytest.mark.parametrize( + "line,expected", + [ + ("1. Introduction", True), + ("2.3 Coverage Limits", True), + ("ARTICLE IV", True), + ("Section 12 Termination", True), + ("DEFINITIONS", True), + ("Named Insured And Limits", True), + ("This policy covers the named insured for losses.", False), + ("the quick brown fox jumps over the lazy dog again", False), + ("ok", False), + ("Total:", False), + ], + ) + def test_heuristic(self, line, expected): + assert dd.looks_like_heading(line) is expected + + +class TestOutline: + def test_pdf_outline_carries_page_numbers_and_structure_counts(self): + pages = ["1. Declarations\nNamed insured Acme", "body text only here.", "2. Insuring Agreements\nsee Table 1 and Figure 2"] + outline = dd.extract_outline("pdf", build_pdf(pages)) + assert (outline.unit, outline.count) == ("page", 3) + assert [(s.start, s.title) for s in outline.sections] == [(1, "1. Declarations"), (3, "2. Insuring Agreements")] + assert (outline.tables, outline.figures) == (1, 1) + assert outline.chars > 0 and outline.abstract is None + + def test_pdf_without_headings_anchors_on_first_lines(self): + outline = dd.extract_outline("pdf", build_pdf(["scanned text one.", "", "scanned text three."])) + assert [(s.start, s.title) for s in outline.sections] == [(1, "scanned text one."), (3, "scanned text three.")] + + def test_docx_outline_uses_paragraph_numbers(self): + raw = build_docx(["Executive Summary", "The budget grows by 4 million this year.", "Risks", "Some risk text."]) + outline = dd.extract_outline("docx", raw) + assert (outline.unit, outline.count) == ("paragraph", 4) + assert [(s.start, s.title) for s in outline.sections] == [(1, "Executive Summary"), (3, "Risks")] + + def test_markdown_headings_win_and_html_tags_are_stripped(self): + md = dd.extract_outline("md", b"# Title\nbody\n## Second part\nmore") + assert [(s.start, s.title) for s in md.sections] == [(1, "Title"), (3, "Second part")] + html = dd.extract_outline("html", b"

OVERVIEW

\n

lower case body sentence.

") + assert [s.title for s in html.sections] == ["OVERVIEW"] + + def test_sections_are_sampled_evenly_past_the_cap(self, monkeypatch): + monkeypatch.setattr(dd, "DOCUMENT_DIGEST_MAX_SECTIONS", 5) + outline = dd.extract_outline("pdf", build_pdf([f"{i}. Heading {i}" for i in range(1, 21)])) + assert len(outline.sections) == 5 + assert [s.start for s in outline.sections] == [1, 5, 9, 13, 17] + + def test_text_sample_spreads_across_pages(self): + sample = dd.text_sample("pdf", build_pdf(["first page words", "middle page words", "last page words"]), limit=200) + assert "first" in sample and "middle" in sample and "last" in sample + assert len(dd.text_sample("md", b"x" * 50_000, limit=100)) == 100 + + +class TestRender: + def _digest(self, sections=3, abstract="An abstract."): + return dd.DocumentDigest( + format="pdf", unit="page", count=47, tables=2, figures=0, abstract=abstract, + sections=[dd.DigestSection(start=i + 1, title=f"Section {i + 1}") for i in range(sections)], + ) + + def test_block_shape(self): + text = dd.render_digest(self._digest(), filename='Policy <"A&B">.pdf', upload_id="up-1") + assert text.startswith('') + assert " An abstract." in text + assert '
Section 3
' in text + assert text.endswith("") + + def test_budget_drops_sections_first_then_trims_the_abstract(self): + digest = self._digest(sections=40, abstract="word " * 400) + full = dd.render_digest(digest, filename="a.pdf", upload_id="u") + assert dd.estimate_tokens(full) > 300 + tight = dd.render_digest(digest, filename="a.pdf", upload_id="u", budget_tokens=300) + assert dd.estimate_tokens(tight) <= 300 + assert '
' not in tight or "…" in tight + tiny = dd.render_digest(digest, filename="a.pdf", upload_id="u", budget_tokens=40) + assert dd.estimate_tokens(tiny) <= 40 + 4 # header always fits, may exceed a tiny budget by a line + assert tiny.startswith(" 0 and digest.extractor_ms >= 0 + item = digest.to_item() + assert item["version"] == dd.DIGEST_VERSION and "error" not in item + assert dd.DocumentDigest.from_item(item).sections[1].title == "2. Terms" + + @pytest.mark.asyncio + async def test_abstract_failure_keeps_the_outline(self, monkeypatch): + _bedrock(monkeypatch, fail=True) + digest = await dd.build_digest(raw=build_pdf(["1. Intro"]), mime_type="application/pdf", filename="p.pdf", upload_id="u1") + assert digest.status == "ready" and digest.abstract is None and digest.model_id is None + assert digest.sections[0].title == "1. Intro" + + @pytest.mark.asyncio + async def test_extraction_failure_is_a_failed_digest_with_a_class_name_only(self): + digest = await dd.build_digest(raw=b"not a pdf", mime_type="application/pdf", filename="p.pdf", upload_id="u1", with_abstract=False) + assert digest.status == "failed" and digest.error == "DocumentReadError" + assert "not a pdf" not in digest.to_item().values().__repr__() + + @pytest.mark.asyncio + async def test_non_documents_are_refused(self): + digest = await dd.build_digest(raw=b"a,b", mime_type="text/csv", filename="d.csv", upload_id="u1") + assert (digest.status, digest.error) == ("failed", "NotADocument") + + def test_malformed_rows_read_as_no_digest(self): + assert dd.DocumentDigest.from_item("nope") is None + assert dd.DocumentDigest.from_item({"count": "many"}) is None + + def test_metric_is_content_free(self, monkeypatch): + calls = [] + monkeypatch.setattr("apis.shared.observability.emf.emit_emf_metrics", lambda ns, metrics, properties=None, units=None: calls.append((ns, metrics, properties))) + monkeypatch.setattr("apis.shared.observability.prompt_cache.prompt_cache_observability_enabled", lambda: True) + dd.record_digest(dd.DocumentDigest(format="pdf", tokens=900, extractor_ms=1200, abstract="SECRET")) + assert calls == [("AgentCoreStack/Compaction", {"DocumentDigestGenerated": 1, "DocumentDigestTokens": 900, "DocumentDigestMs": 1200}, {"format": "pdf", "outcome": "ready"})] + assert "SECRET" not in repr(calls) + + def test_kill_switch(self, monkeypatch): + monkeypatch.setenv("DOCUMENT_DIGEST_ENABLED", "false") + assert dd.document_digest_enabled() is False + monkeypatch.setenv("DOCUMENT_DIGEST_ENABLED", "") + assert dd.document_digest_enabled() is True diff --git a/backend/tests/shared/test_document_read.py b/backend/tests/shared/test_document_read.py new file mode 100644 index 00000000..c8fe0342 --- /dev/null +++ b/backend/tests/shared/test_document_read.py @@ -0,0 +1,356 @@ +"""Document read service (apis/shared/files/document_read.py). + +The retrieval primitive behind ``document_read``: PDF page ranges re-assembled +as native document blocks (bounded by ``max_pages``), pattern search over the +text layer with page numbers, bounded text for DOCX and text-family files, the +session listing, and the tool-injection gate. Fixtures are hand-built PDFs +and DOCX archives — never user files — rendered through the same pypdfium2 +the service uses, so page identity under re-assembly is exercised for real. +""" + +from __future__ import annotations + +import io +import zipfile +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from apis.shared.files import document_read as dr +from apis.shared.files.models import FileMetadata, FileStatus +from apis.shared.files.workspace import WorkspaceFileNotFoundError, WorkspaceValidationError + +MODULE = "apis.shared.files.document_read" + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def build_pdf(page_texts): + """A minimal multi-page PDF with one Helvetica text line per page.""" + objs = [] + + def add(body: bytes) -> int: + objs.append(body) + return len(objs) + + font = add(b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>") + contents = [] + for text in page_texts: + esc = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + stream = f"BT /F1 12 Tf 72 720 Td ({esc}) Tj ET".encode() + contents.append(add(b"<< /Length %d >>\nstream\n" % len(stream) + stream + b"\nendstream")) + pages_id = len(objs) + len(page_texts) + 1 + page_ids = [ + add( + f"<< /Type /Page /Parent {pages_id} 0 R /MediaBox [0 0 612 792] " + f"/Resources << /Font << /F1 {font} 0 R >> >> /Contents {c} 0 R >>".encode() + ) + for c in contents + ] + kids = " ".join(f"{p} 0 R" for p in page_ids) + assert add(f"<< /Type /Pages /Kids [{kids}] /Count {len(page_ids)} >>".encode()) == pages_id + catalog = add(f"<< /Type /Catalog /Pages {pages_id} 0 R >>".encode()) + out = io.BytesIO() + out.write(b"%PDF-1.4\n") + offsets = [] + for i, body in enumerate(objs, 1): + offsets.append(out.tell()) + out.write(f"{i} 0 obj\n".encode() + body + b"\nendobj\n") + xref = out.tell() + out.write(f"xref\n0 {len(objs) + 1}\n".encode() + b"0000000000 65535 f \n") + for o in offsets: + out.write(f"{o:010d} 00000 n \n".encode()) + out.write( + f"trailer\n<< /Size {len(objs) + 1} /Root {catalog} 0 R >>\nstartxref\n{xref}\n%%EOF\n".encode() + ) + return out.getvalue() + + +def build_docx(paragraphs): + ns = 'xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"' + body = "".join(f"{p}" for p in paragraphs) + xml = f'{body}' + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as z: + z.writestr("[Content_Types].xml", "") + z.writestr("word/document.xml", xml) + return buf.getvalue() + + +PAGES = [f"Page {i} alpha beta" for i in range(1, 11)] + ["The termination clause survives", "Page 12 end"] + + +def _meta(upload_id="up-1", filename="policy.pdf", mime="application/pdf", size=1000, user="u1", session="s1"): + return FileMetadata( + upload_id=upload_id, user_id=user, session_id=session, filename=filename, mime_type=mime, + size_bytes=size, s3_key=f"user-files/{user}/{session}/{upload_id}/{filename}", s3_bucket="b", + status=FileStatus.READY, + ) + + +@pytest.fixture +def stored(monkeypatch): + """Patch the metadata lookup and the S3 fetch; returns a setter.""" + state = {} + + async def _owned(user_id, upload_id): + meta = state.get("meta") + if meta is None or meta.user_id != user_id or meta.upload_id != upload_id: + raise WorkspaceFileNotFoundError(f"No file with id '{upload_id}'") + return meta + + monkeypatch.setattr(f"{MODULE}._get_owned_ready_file", _owned) + monkeypatch.setattr(f"{MODULE}._get_object_bytes", lambda bucket, key: state["raw"]) + + def _set(meta, raw): + state["meta"], state["raw"] = meta, raw + + return _set + + +# --------------------------------------------------------------------------- +# Classification and argument parsing +# --------------------------------------------------------------------------- + + +class TestClassification: + @pytest.mark.parametrize( + "mime,filename,expected", + [ + ("application/pdf", "a.pdf", "pdf"), + ("application/vnd.openxmlformats-officedocument.wordprocessingml.document", "a.docx", "docx"), + ("text/plain", "a.txt", "txt"), + ("text/markdown", "a.md", "md"), + ("text/html", "a.html", "html"), + ("text/csv", "a.csv", None), # spreadsheet tools + ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "a.xlsx", None), + ("application/vnd.openxmlformats-officedocument.presentationml.presentation", "a.pptx", None), + ("image/png", "a.png", None), + ], + ) + def test_document_class(self, mime, filename, expected): + assert dr.document_format_for(mime, filename) == expected + assert dr.is_document_class(mime, filename) is (expected is not None) + + @pytest.mark.parametrize( + "value,expected", + [("4-7", (4, 7)), ("4", (4, 4)), (" 4 .. 7 ", (4, 7)), ("4:7", (4, 7)), ({"start": 2, "end": 3}, (2, 3)), + ([5, 6], (5, 6)), (3, (3, 3)), ("", None), (None, None)], + ) + def test_parse_page_range(self, value, expected): + assert dr.parse_page_range(value) == expected + + @pytest.mark.parametrize("value", ["0-3", "7-4", "abc", "1-2-3"]) + def test_bad_page_range_is_a_validation_error(self, value): + with pytest.raises(WorkspaceValidationError): + dr.parse_page_range(value) + + def test_max_pages_is_clamped_to_the_hard_cap(self): + assert dr.clamp_max_pages(None) == dr.DOCUMENT_READ_MAX_PAGES + assert dr.clamp_max_pages(0) == 1 + assert dr.clamp_max_pages(10_000) == dr.DOCUMENT_READ_HARD_MAX_PAGES + assert dr.clamp_max_pages("x") == dr.DOCUMENT_READ_MAX_PAGES + + +# --------------------------------------------------------------------------- +# PDF +# --------------------------------------------------------------------------- + + +class TestPdf: + @pytest.mark.asyncio + async def test_page_range_returns_only_those_pages_as_a_native_block(self, stored): + stored(_meta(), build_pdf(PAGES)) + result = await dr.read_document("u1", "s1", "up-1", page_range="4-7") + + assert result.mode == "pages" + assert result.pages_returned == 4 + assert result.payload["pages"] == {"start": 4, "end": 7} + assert result.payload["page_count"] == 12 + assert result.payload["truncated_to_max_pages"] is False + block = result.document_block + assert block["document"]["format"] == "pdf" + assert block["document"]["name"].startswith("policy p4-7 ") + # The slice really is pages 4–7 of the original, in order. + import pypdfium2 as pdfium + + sub = pdfium.PdfDocument(io.BytesIO(block["document"]["source"]["bytes"])) + assert len(sub) == 4 + assert dr._pdf_page_text(sub, 0) == "Page 4 alpha beta" + assert dr._pdf_page_text(sub, 3) == "Page 7 alpha beta" + sub.close() + assert "page k is original page 4 + k - 1" in result.payload["page_numbering"] + + @pytest.mark.asyncio + async def test_max_pages_caps_a_range_and_points_at_the_continuation(self, stored): + stored(_meta(), build_pdf(PAGES)) + result = await dr.read_document("u1", "s1", "up-1", page_range="1-12", max_pages=3) + assert result.pages_returned == 3 + assert result.payload["pages"] == {"start": 1, "end": 3} + assert result.payload["truncated_to_max_pages"] is True + assert result.payload["next_start"] == 4 + + @pytest.mark.asyncio + async def test_hard_cap_holds_whatever_the_model_asks(self, stored): + stored(_meta(), build_pdf([f"p{i}" for i in range(1, 41)])) + result = await dr.read_document("u1", "s1", "up-1", page_range="1-40", max_pages=999) + assert result.pages_returned == dr.DOCUMENT_READ_HARD_MAX_PAGES + + @pytest.mark.asyncio + async def test_range_past_the_end_clips_or_rejects(self, stored): + stored(_meta(), build_pdf(PAGES)) + result = await dr.read_document("u1", "s1", "up-1", page_range="11-30") + assert result.payload["pages"] == {"start": 11, "end": 12} + with pytest.raises(WorkspaceValidationError): + await dr.read_document("u1", "s1", "up-1", page_range="13-14") + + @pytest.mark.asyncio + async def test_pattern_returns_page_numbers_and_context(self, stored): + stored(_meta(), build_pdf(PAGES)) + result = await dr.read_document("u1", "s1", "up-1", pattern="TERMINATION clause") + assert result.mode == "pattern" + assert result.document_block is None + assert result.payload["pages_matched"] == [11] + assert result.payload["matches"][0]["page"] == 11 + assert "termination" in result.payload["matches"][0]["text"].lower() + + @pytest.mark.asyncio + async def test_invalid_regex_falls_back_to_a_literal_search(self, stored): + stored(_meta(), build_pdf(["cost (usd", "other"])) + result = await dr.read_document("u1", "s1", "up-1", pattern="cost (usd") + assert result.payload["pages_matched"] == [1] + + @pytest.mark.asyncio + async def test_no_range_no_pattern_returns_a_bounded_page_index(self, stored, monkeypatch): + monkeypatch.setattr(dr, "DOCUMENT_READ_INDEX_PAGES", 5) + stored(_meta(), build_pdf(PAGES)) + result = await dr.read_document("u1", "s1", "up-1") + assert result.mode == "index" + assert result.payload["page_count"] == 12 + assert len(result.payload["page_index"]) == 5 + assert result.payload["page_index"][0] == {"page": 1, "snippet": "Page 1 alpha beta"} + assert result.payload["index_truncated"] is True + + @pytest.mark.asyncio + async def test_corrupt_pdf_is_a_conversational_error(self, stored): + stored(_meta(), b"not a pdf") + with pytest.raises(dr.DocumentReadError): + await dr.read_document("u1", "s1", "up-1", page_range="1") + + def test_slice_names_are_bedrock_safe_and_never_collide(self): + a = dr._unique_document_name("BBR 5.0 Policy_Form (final).pdf", "p4-7") + b = dr._unique_document_name("BBR 5.0 Policy_Form (final).pdf", "p4-7") + assert a != b + assert a.startswith("BBR 5 0 Policy Form (final) p4-7 ") + import re + + assert re.fullmatch(r"[a-zA-Z0-9\s\-\(\)\[\]]+", a) and " " not in a + + +# --------------------------------------------------------------------------- +# DOCX and text +# --------------------------------------------------------------------------- + + +class TestDocxAndText: + DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + + @pytest.mark.asyncio + async def test_docx_text_is_bounded_with_offset_continuation(self, stored, monkeypatch): + monkeypatch.setattr(dr, "WORKSPACE_READ_MAX_BYTES", 20) + stored(_meta(filename="memo.docx", mime=self.DOCX_MIME), build_docx(["Alpha paragraph", "Beta paragraph"])) + first = await dr.read_document("u1", "s1", "up-1") + assert first.mode == "text" and first.payload["paragraph_count"] == 2 + assert first.payload["truncated"] is True and first.payload["next_offset"] == 20 + second = await dr.read_document("u1", "s1", "up-1", offset=first.payload["next_offset"]) + assert first.payload["content"] + second.payload["content"] == "Alpha paragraph\n\nBeta paragraph" + assert second.payload["truncated"] is False + + @pytest.mark.asyncio + async def test_docx_pattern_reports_paragraph_numbers(self, stored): + stored(_meta(filename="memo.docx", mime=self.DOCX_MIME), build_docx(["Intro", "Budget is 4 million", "Close"])) + result = await dr.read_document("u1", "s1", "up-1", pattern=r"\d+ million") + assert result.payload["unit"] == "paragraph" + assert result.payload["matches"] == [ + {"line": 2, "text": "Budget is 4 million", "before": ["Intro"], "after": ["Close"]} + ] + assert result.format == "docx" + + @pytest.mark.asyncio + async def test_page_range_on_a_non_pdf_is_rejected(self, stored): + stored(_meta(filename="memo.docx", mime=self.DOCX_MIME), build_docx(["x"])) + with pytest.raises(WorkspaceValidationError): + await dr.read_document("u1", "s1", "up-1", page_range="1-2") + + @pytest.mark.asyncio + async def test_text_family_delegates_to_the_workspace_read(self, stored, monkeypatch): + stored(_meta(filename="notes.md", mime="text/markdown"), b"# heading\nbody line\n") + delegate = AsyncMock(return_value={"content": "# heading", "offset": 0, "truncated": True, "next_offset": 9}) + monkeypatch.setattr(f"{MODULE}.read_workspace_file", delegate) + result = await dr.read_document("u1", "s1", "up-1") + delegate.assert_awaited_once_with("u1", "up-1", offset=0) + assert result.payload["content"] == "# heading" and result.payload["next_offset"] == 9 + + @pytest.mark.asyncio + async def test_text_family_pattern_greps_lines(self, stored): + stored(_meta(filename="notes.txt", mime="text/plain"), b"one\ntwo Fish\nthree\n") + result = await dr.read_document("u1", "s1", "up-1", pattern="fish") + assert result.payload["matches"][0]["line"] == 2 + + @pytest.mark.asyncio + async def test_non_document_files_are_refused(self, stored): + stored(_meta(filename="data.csv", mime="text/csv"), b"a,b\n") + with pytest.raises(dr.DocumentReadError): + await dr.read_document("u1", "s1", "up-1") + + @pytest.mark.asyncio + async def test_missing_upload_id_and_negative_offset_are_validation_errors(self): + with pytest.raises(WorkspaceValidationError): + await dr.read_document("u1", "s1", "") + with pytest.raises(WorkspaceValidationError): + await dr.read_document("u1", "s1", "up-1", offset=-1) + + +# --------------------------------------------------------------------------- +# Listing and the injection gate +# --------------------------------------------------------------------------- + + +class TestListingAndGate: + def _repo(self, monkeypatch, files): + repo = MagicMock() + repo.list_session_files = AsyncMock(return_value=files) + monkeypatch.setattr(f"{MODULE}.get_file_upload_repository", lambda: repo) + return repo + + @pytest.mark.asyncio + async def test_listing_keeps_only_this_users_readable_documents(self, monkeypatch): + repo = self._repo(monkeypatch, [ + _meta("up-1", "policy.pdf"), + _meta("up-2", "data.xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"), + _meta("up-3", "other.pdf", user="someone-else"), + _meta("up-4", "notes.md", "text/markdown"), + ]) + listing = await dr.list_session_documents("u1", "s1") + repo.list_session_files.assert_awaited_once_with("s1", status=FileStatus.READY) + assert [d["upload_id"] for d in listing["documents"]] == ["up-1", "up-4"] + assert listing["documents"][0]["modes"] == ["page_range", "pattern"] + assert listing["documents"][1]["modes"] == ["pattern", "text"] + assert listing["count"] == 2 + + @pytest.mark.asyncio + async def test_gate_is_true_only_with_a_readable_document(self, monkeypatch): + self._repo(monkeypatch, [_meta("up-2", "data.csv", "text/csv")]) + assert await dr.session_has_documents("u1", "s1") is False + self._repo(monkeypatch, [_meta("up-2", "data.csv", "text/csv"), _meta("up-1", "a.pdf")]) + assert await dr.session_has_documents("u1", "s1") is True + + @pytest.mark.asyncio + async def test_identity_is_mandatory(self): + with pytest.raises(Exception): + await dr.list_session_documents("", "s1") + with pytest.raises(Exception): + await dr.session_has_documents("u1", "") diff --git a/backend/tests/shared/test_files.py b/backend/tests/shared/test_files.py index 1a5dddc7..341b710a 100644 --- a/backend/tests/shared/test_files.py +++ b/backend/tests/shared/test_files.py @@ -44,6 +44,19 @@ async def test_update_status(self, file_repository): assert updated is not None assert updated.status == "ready" + @pytest.mark.asyncio + async def test_update_digest_round_trips_and_needs_the_row(self, file_repository): + await file_repository.create_file(_make_file()) + digest = {"version": 1, "status": "ready", "format": "pdf", "count": 3, + "sections": [{"start": 1, "title": "Intro"}], "tokens": 120, "abstract": "x"} + updated = await file_repository.update_file_digest("u1", "f1", digest) + assert updated is not None and updated.digest["count"] == 3 + assert (await file_repository.get_file("u1", "f1")).digest["sections"] == [{"start": 1, "title": "Intro"}] + assert await file_repository.update_file_digest("u1", "missing", digest) is None + # Rows without a digest read as None, never as an empty map. + await file_repository.create_file(_make_file("f2")) + assert (await file_repository.get_file("u1", "f2")).digest is None + @pytest.mark.asyncio async def test_delete_file(self, file_repository): await file_repository.create_file(_make_file()) diff --git a/docs/specs/document-context-offload.md b/docs/specs/document-context-offload.md index 1dc6cf65..5716878d 100644 --- a/docs/specs/document-context-offload.md +++ b/docs/specs/document-context-offload.md @@ -1,8 +1,16 @@ # Document context offload — bound the cost of conversations with attachments -**Status:** Draft (no branch yet) +**Status:** PR-1 built — `feature/document-offload-pr1`, off `develop` @ +`c28ccdbd` (after the compaction stack #1125 → #1128 → #1129 → #1131 → +#1132 merged). Revised 2026-09-16: analytics pulled forward from PR-7 into +PR-1 (§5, §6.1), decision log in §8. **Motivating measurement:** prod scan 2026-08-03, all 18,942 `C#` cost rows -joined to `boisestateai-v2-user-file-uploads` on `GSI1PK = CONV#{sessionId}` +joined to `boisestateai-v2-user-file-uploads` on `GSI1PK = CONV#{sessionId}`; +re-confirmed by the 2026-09-15 prod cost audit — of 95 September sessions +that peaked over 100k tokens, 25 (26%) had an attachment in their last three +turns, including 1–3-turn sessions at 280k–530k that are a single huge +upload. A compaction cut cannot touch those (thresholds spec §3.6 keeps +attachments out of scope on purpose); a digest plus a page-range read can. **Related:** [[project-prod-cache-write-premium]] (the compaction root cause this depends on), `docs/specs/session-workspace-tools.md` (the retrieval primitive this extends), `docs/specs/tool-search-token-bloat-strategy.md` (same tenet, @@ -383,15 +391,23 @@ Re-sequenced around defect 4: the loss fires on turn 2 for ~80% of attachment sessions, so the recovery path and the digest are the urgent half, and the turn-level offload — the cost work — comes after. +**Revised 2026-09-16.** The analytics that were PR-7 now ship *in* PR-1, so +the cost work is measured from the first day the recovery path exists and the +ship / widen-pinning / abandon decision in the evaluation spec (§4.2 there) +is decidable from stored rows rather than a one-off scan. The field design is +§6.1 below; the kaizen finding that the compaction work has no quality signal +joined to its cost data is addressed there too — as far as the platform +currently allows (see the outcome-signal row). + | PR | Scope | Gate | |----|-------|------| -| 1 | `document_read` (page-range + pattern), native `document` block reassembly, **gated on the session having an attachment**; ids kept out of `INJECTED_TOOL_IDS` | 47-page PDF: `page_range={4,7}` returns 4 pages, ≤6k tok; `max_pages` cap holds; tool present for a `student`-role session with no grants | -| 2 | `DocumentDigest` model + Haiku extractor + persist on `FileMetadata`, generated at upload; **not yet used in context** | digest ≤1,500 tok for a 200-page PDF; extractor p95 < 8s; no chat-path change | -| 3 | `_strip_document_bytes` → digest + live `document_read` handle instead of the placeholder (restore path only) | a session with `analyze_spreadsheet` enabled answers a document question correctly on **turn 2**; re-upload rate starts falling | -| 4 | Offload trigger in `update_after_turn` + pinning, behind `DOCUMENT_OFFLOAD_ENABLED` | slice-assignment guard test passes; offload fires at most once per document; never fires while cache is live | +| 1 — **built** | `document_read` (page-range + pattern + bounded text; native `document` block reassembly), **gated on the session having a readable attachment**, id kept out of `INJECTED_TOOL_IDS`, presence carried in the agent-cache key; **analytics**: per-call document context on `C#` rows, the `documentReads` ledger entry, the `document_stripped` ledger event, session-row rollups, anatomy + profile surfaces, EMF; `document_read` results exempt from the tool-result offloader | 12-page PDF: `page_range="4-7"` returns 4 pages as one native block whose page 1 is original page 4; `max_pages` and the hard cap hold; tool built for a session with no grants and no picker toggle; content-policy test walks the new fields; full backend suite green | +| 2 — **built** (`feature/document-offload-pr2`, stacked on PR-1) | `DocumentDigest` (`apis/shared/files/document_digest.py`): a deterministic outline (headings with page / paragraph / line anchors, table and figure mentions, counts) plus a 3–5-sentence abstract from the cheap text model (Nova Micro, `DOCUMENT_DIGEST_MODEL_ID`), built off the request path when a document upload completes and persisted as `FileMetadata.digest`; `render_digest` produces the `` block under a hard 1,500-token budget (sections dropped first, then the abstract) and stores the estimate as `digest.tokens`; `digest.abstract` / `digest.sections` denylisted, coverage on the attachment profile; **not yet used in context** | 200-page PDF digest ≤1,500 tokens (tested); extraction and abstract each fail open; kill switch `DOCUMENT_DIGEST_ENABLED`; no chat-path change; p95 latency to be read from `DocumentDigestMs` in dev | +| 3 — **built** (`feature/document-offload-pr3`, stacked on PR-2) | `_strip_document_bytes` → `session/document_rehydration.py`: each inline document block is matched to its upload row (sanitized filename incl. PromptBuilder's `_2` suffix, byte-size tiebreak, newest first, no double-claiming) and replaced by the rendered `` block; rows without a digest get an **outline-only digest built from the bytes already in the restored message** (no S3, no model call) and persisted; unmatched blocks keep the placeholder byte-for-byte; ledger records `document_rehydrated` and `document_stripped` separately; kill switch `DOCUMENT_REHYDRATE_ENABLED` | restore output is stable across restores (tested); a matched block carries the abstract, outline and handle; unmatched blocks are unchanged from today; `document_stripped` per session-day should fall to the unmatched residue (direct base64 attachments, deleted files) — read it from the ledger; re-upload rate (byte-identical) starts falling | +| 4 | Offload trigger in `update_after_turn` + pinning, behind `DOCUMENT_OFFLOAD_ENABLED` as a percentage rollout keyed on `hash(session_id)`; records `document_offload` with `cacheGapSeconds`; old `document_read` slices in history treated like documents | slice-assignment guard test passes; offload fires at most once per document; **zero** `document_offload` events with `cacheGapSeconds` under the TTL | | 5 | Fix the cache-live guard on the *existing* truncation deferral (same predicate as PR-4) | truncation events while cache live: 72% → ~0 | | 6 | Backend guard on a turn's aggregate inline attachment bytes (~7.5 MB), mirroring the SPA's `MAX_FILES_PER_MESSAGE` | oversized turn degrades to the `oversized_inline` guidance path, never to `SessionException` | -| 7 | `hasDocuments` / `documentTokens` on `MessageMetadata`; admin cost anatomy shows document share | the two-table join this spec required becomes a dashboard column | +| 7 | **Outcome signal**: a content-free thumbs up/down on assistant messages, persisted keyed on `(sessionId, messageId)` so it joins the `C#` row's `hasDocuments` / `documentDigests` / `documentReads` in one key; fleet-level document-share column on the admin dashboard | down-thumb rate reported by turn class (full / digest-only / retrieved) with n per class | **PRs 1–3 are the correctness fix and should ship together as a unit.** None carries cost-regression risk: the digest is strictly smaller than the document, @@ -400,7 +416,9 @@ must precede PR-3 — a digest that points at a tool nobody has is no better tha today's placeholder. PR-4 is the cost work and is independently revertible. PR-6 is unrelated to both -and can go whenever. +and can go whenever. PR-7 exists because **no feedback surface exists today** +(`MessageMetadata` carries only a `# feedback: …` placeholder comment); the +rows PR-1 writes make the join a one-key lookup the moment one does. **Out of scope, but surfaced by this work:** the `extra_tools` agent-cache bypass makes ~76% of *all* sessions rebuild their Agent every turn. Fixing that @@ -431,6 +449,90 @@ quality regression**, measured as: digest is too good to be true and the model is answering without the source; if it is >3/turn the digest is too thin +### 6.1 Analytics field design (PR-1, built) + +**Rule: content-free by construction.** Every field below is an id, a count, +a byte size, a token estimate or an enum key. No document text, title or +filename is ever persisted in a metric, a cost row or an EMF record — the +`document_read` tool result carries filenames (that is conversation content, +where the model needs them); the ledger reads only the result's numbers. The +content-policy test walks the new projections and response models. + +**Alignment.** Same mechanism as the compaction ledger (#1130): per-call +facts ride the `C#` cost row as flat extra fields next to `prefixTokens` / +`windowRemovedMessages` / `compactionEvents`; lifecycle events use +`record_compaction_event` and the `compactionEvents` list; session rollups are +`ADD`ed on the `S#` row beside `compactionAppliedCount`; EMF goes to +`AgentCoreStack/Compaction`; the anatomy (`GET /admin/costs/sessions/{id}/calls`) +and the profile read them. One namespace, one ledger, one anatomy view. + +**Per model call (`C#` row).** Written by the stream coordinator from +`agent.messages` at turn end (`session/document_context.py`); gated by +`COST_DIAGNOSTICS_ENABLED` like the rest of the ledger, so an absent field +reads "not tracked", never 0. + +| field | meaning | +|---|---| +| `hasDocuments` | ≥1 inline attachment block (document or image bytes) is in the live context | +| `documentCount` | inline attachment blocks on user prompts | +| `documentTokens` | their estimated weight — the compaction estimator's bytes/4 per document, flat per image. Heuristic, comparable across rows; the measured total it is a share of is `tokenUsage` (input + cacheRead + cacheWrite) minus `prefixTokens.system + prefixTokens.tools` | +| `documentDigests` | digest / placeholder stand-ins in context (`[Document placeholder:` today, ` 0` and not `hasDocuments`), *retrieved* +(`documentReads.pages > 0`), or *none*. This is the digest-vs-full share. + +**Lifecycle events** (`compactionEvents[].kind`, numbers only): + +| kind | when | fields | +|---|---|---| +| `document_stripped` | restore replaced inline documents with contentless placeholders (the defect; recorded from PR-1 so its reach is measured before and after PR-3) | `documents`, `documentTokens` | +| `document_rehydrated` | PR-3: restore replaced them with a digest + live handle | `documents`, `documentTokens` | +| `document_offload` | PR-4: the post-turn trigger swapped a document for its digest | `documents`, `documentTokens`, **`cacheGapSeconds`** at the moment it fired | + +**Session rollups (`S#` row, `ADD`):** `fullDocumentCalls`, `digestOnlyCalls`, +`documentReadCalls`, `documentReadPages` — so the profile can answer without +the rows once they expire, and so a fleet query needs one table. + +**EMF (`AgentCoreStack/Compaction`):** `DocumentRead` / `DocumentReadPages` / +`DocumentReadBytes` with properties `mode` (`list` / `index` / `pages` / +`pattern` / `text`) and `format`. PR-4 adds `DocumentOffloaded` / +`DocumentOffloadedTokens` beside `ToolResultOffloaded`. + +**Admin surfaces:** anatomy rows carry every per-call field (`documentReads` +as a map); the profile carries `fullDocumentCalls`, `digestOnlyCalls`, +`peakDocumentTokens`, `documentReadCalls`, `documentReadPages` and +`dataCoverage.documents`; the SPA anatomy page shows a per-row `doc` / +`digest` / `+Np` badge, a Documents line in the expanded row, and the +consumption summary under the Attachments card. + +**What is decidable from the rows alone** (the point of pulling this forward): + +- *Document share of the prefix*, per call: `documentTokens / (context − + prefixTokens.system − prefixTokens.tools)`. Summed over cold rows + (`cacheStatus ∉ {hit}`), `cacheWriteInputTokens × share` is the document + part of every re-write — the recoverable envelope the evaluation spec's §4.1 + says was unmeasured, and the number that turns "−15%" into a derived target. +- *Was the offload ever wrong*: `document_offload` events whose + `cacheGapSeconds` is under the cache TTL. Target zero; any other value is + the compaction PR-3 rule being broken. +- *Did the recovery path reach users*: `document_stripped` per session-day + before PR-3, `document_rehydrated` after; `documentReadCalls` per attachment + session against the evaluation's health band (≈0.3–3; ≈0 means the digest is + answering unaided, >3/turn means it is too thin). +- *Digest-vs-full shares and the B/C arms*: `fullDocumentCalls : + digestOnlyCalls` per session, split by the PR-4 rollout bucket. +- *The outcome signal* — **not buildable today.** There is no feedback + surface (`MessageMetadata` has a `# feedback:` placeholder only). The rows + are keyed so that a thumbs row on `(sessionId, messageId)` joins the turn + class in one lookup; PR-7 adds that surface. Until then the quality gate + stays the evaluation spec's offline harness, and this remains the kaizen + gap it has been for the compaction work too. + ### Quality gate — this must not ship on cost numbers alone **Full evaluation design: `docs/specs/document-offload-evaluation.md`** (three @@ -498,6 +600,128 @@ visual fidelity before any digest comparison is scored. --- +## 8. Decision log — PR-1 (2026-09-16) + +1. **Analytics ship first, not last.** PR-7's fields moved into PR-1 (§6.1) so + the cost work is measured from day one and the evaluation's stopping rule + is decidable from stored rows. The `document_stripped` event is recorded + *before* the fix that removes it, so PR-3's effect is a before/after on one + counter rather than a scan. +2. **Gate = session state; presence lives in the agent-cache key, not a + cache veto.** `document_read` is built when the session has a READY + document-class upload (PDF, DOCX, TXT, MD, HTML — not tabular, decks or + images, which have other paths) or this turn attaches one. Positive + answers are memoized per process (monotonic in practice). Vetoing the + agent cache instead — the Memory-Space pattern — would have made every + attachment session rebuild its agent each turn and hit the strip on turn + 2 for the ~19% that keep a warm agent today: a quality regression PR-1 + alone would introduce. A `document_tools` element in `_create_cache_key` + flips at most once per session, on the attach turn, when restored history + has no document to lose; the resume path recomputes the same gate so a + paused agent's key is reproduced. +3. **Id recorded, never injected-filtered.** `DOCUMENT_TOOL_IDS` exists for + the record and is deliberately not in `INJECTED_TOOL_IDS`; the only control + is `DOCUMENT_READ_ENABLED` (default on, `=false` kills). +4. **`document_read` results are exempt from the tool-result offloader** + (`OFFLOAD_EXEMPT_TOOLS` in `core/tool_result_offload.py`, plus the + plugin's own `should_offload`). Offloading the slice the model just asked + for would undo the read. The tool's `max_pages` (default 8, hard cap 20) + is the bound instead. +5. **Slices persist in history.** A retrieved page range is a tool-result + document block with a unique Bedrock-safe name (`" p4-7 <6 hex>"`), + so it survives restore (the strip only touches top-level blocks) and never + collides. They are counted on the row (`documentSlices`); PR-4 should age + them like documents rather than let them re-write forever. +6. **Page identity.** The slice is a new PDF numbered 1..k; the payload's + `page_numbering` note and the tool description say "page k is original + page start+k−1 — cite original page numbers". The evaluation's + citation-page-identity family is the test of whether models follow it. +7. **No new dependency.** PDF slicing and text extraction use `pypdfium2` + (already pinned for thumbnails); DOCX text comes from a stdlib + `zipfile` + `ElementTree` extractor over `word/document.xml`. `page_range` + is PDF-only; DOCX and text use `pattern` or bounded `offset` reads + (the workspace read bound, 48 KB). +8. **Token numbers are heuristics.** Bedrock reports no per-block usage; + `documentTokens` is bytes/4 (documents) and a flat estimate per image, + the compaction estimator's own numbers. They are comparable across rows and + against the measured messages partition, which is all the decisions in + §6.1 need. +9. **Metrics share the compaction namespace** (`AgentCoreStack/Compaction`) + on purpose: the cut record, the tool-result offload record and the + document read are three answers to one question — what bounded this + session's prefix — and belong on one dashboard. +10. **`docs/specs/document-conversations-cost.md` does not exist in the + repo.** The numbers it would hold are §1 here and the validation report; + nothing in this spec depends on it. +11. **Validation nits carried in:** `INJECTED_TOOL_IDS` is in + `apis/shared/tools/injected.py` (line drifted); `WORKSPACE_READ_MAX_BYTES` + lives in `apis/shared/files/workspace.py`. Citations remain out of PR-1 + (§6 probe note). + +**PR-2 (2026-09-16)** + +12. **Outline is deterministic; only the abstract uses a model.** Page / + paragraph / line counts, headings (numbered, ALL CAPS, or title-cased + short lines; markdown `#`), table / figure mentions and a text sample + come from `pypdfium2` and the PR-1 DOCX extractor. The abstract is 3–5 + sentences from Nova Micro over the sample plus outline — the same cheap + text model the tool-batch summaries and the compaction summary run on, + and the extraction step never needs a frontier model. The §4A sketch said + Haiku over the document; a text model over extracted text is cheaper, + text-only is enough for an abstract, and the visual channel is preserved + by `document_read`, not by the digest. +13. **Built at `complete_upload` only, fire-and-forget.** The SPA upload flow + is where documents enter; agent-written files (`workspace_write`, Word / + Excel / PowerPoint tools) register `FileMetadata` directly and get no + digest. Uploads that predate PR-2 have none either. **PR-3 must generate + lazily on the restore path when `digest` is absent** (same builder, + `with_abstract` optional under a latency budget) rather than assume it. +14. **The digest is content-bearing.** `digest.abstract` and + `digest.sections` are denylisted; `digest.status` / `format` / `count` / + `tokens` are projected so the attachment profile can report coverage and + the per-file token cost the 1,500 ceiling is enforced against. +15. **Budget enforcement is in the renderer**, not the extractor: sections + are dropped from the end, then the abstract is truncated, and the opening + tag (the `document_read` handle) always survives. `digest.tokens` records + the rendered estimate, so "digest ≤1,500 tokens" is a stored fact per + file rather than a claim. + +**PR-3 (2026-09-16)** + +16. **Matching is by sanitized filename, not by id.** A document block + carries no upload id — only the name `PromptBuilder` gave it, which is + `FileSanitizer.sanitize_filename(filename)` (the extension's dot becomes + an underscore: `BBR Policy.pdf` → `BBR Policy_pdf`, duplicates get + `_2`/`_3`) — so restore matches blocks to the session's upload rows on + that name with the byte size as tiebreak, newest row first, each row + claimed once. Adding an upload id to the block itself would be cleaner + but changes the attach-turn bytes (a prefix change for every attachment + turn); left for a later PR that touches the prompt builder anyway. +17. **Lazy digests are outline-only and come from the restored bytes.** The + bytes are in `agent.messages` at the moment of the strip (§4E), so no S3 + read is needed, and the restore path is synchronous inside the agent + constructor, so no model call is made there. The digest is persisted so + the next restore renders identical bytes; the upload-path build (PR-2) + may later overwrite it with one that has an abstract, which changes the + rendered block once — one prefix re-write, visible as a jump in + `document_rehydrated.documentTokens`. +18. **Restore output is byte-stable** given the same upload rows: the + renderer is deterministic over `FileMetadata.filename`, `upload_id` and + the persisted digest, and unmatched blocks keep the pre-PR-3 placeholder + verbatim. Tested. The cache contract in CLAUDE.md is therefore honoured + on the restore path exactly as before. +19. **Never worse than before.** Every failure — lookup, matching, digest + build, persistence — falls back to the placeholder for that block and + is recorded as `document_stripped`; the lookup is attempted once per + restore, and only when the history holds a document block. +20. **The synchronous constraint is real.** `TurnBasedSessionManager.initialize` + runs under the Strands agent constructor with an event loop already + running, so it cannot await; the repository gained `_sync` bodies for the + session query and the digest write (boto3 was synchronous underneath all + along) rather than a thread-with-its-own-loop. + +--- + ## Sources - [Effective context engineering for AI agents — Anthropic](https://www.anthropic.com/engineering/effective-context-engineering-for-ai-agents) diff --git a/frontend/ai.client/src/app/admin/costs/models/admin-cost.models.ts b/frontend/ai.client/src/app/admin/costs/models/admin-cost.models.ts index 0f500c9e..753ebed1 100644 --- a/frontend/ai.client/src/app/admin/costs/models/admin-cost.models.ts +++ b/frontend/ai.client/src/app/admin/costs/models/admin-cost.models.ts @@ -138,6 +138,21 @@ export interface CompactionEvent { retainedMessages?: number | null; truncatedToolResults?: number | null; inputTokens?: number | null; + /** + * Document lifecycle kinds (`document_stripped` / `document_rehydrated` / + * `document_offload`): documents touched, their estimated token weight, + * and — for an offload — the prompt-cache gap when it fired. + */ + documents?: number | null; + documentTokens?: number | null; + cacheGapSeconds?: number | null; +} + +/** `document_read` retrievals one model call requested. */ +export interface DocumentReads { + calls: number; + pages: number; + bytes: number; } /** One model call within a session's cost anatomy. */ @@ -181,6 +196,22 @@ export interface SessionCallRow { /** Messages trimmed since the previous ledger-bearing call; > 0 means the prefix changed before this call. */ windowTrimmed?: number | null; compactionEvents?: CompactionEvent[] | null; + /** + * Document context at this call (absent on rows written before it shipped). + * `hasDocuments` + `documentDigests` classify the call: full document inline, + * digest only, or neither. `documentTokens` is a heuristic (bytes/4, flat per + * image), comparable across rows; `documentMime` is keyed by Bedrock's format + * enum plus `image` — never a filename. + */ + hasDocuments?: boolean | null; + documentCount?: number | null; + documentTokens?: number | null; + documentDigests?: number | null; + documentsAttached?: number | null; + documentSlices?: number | null; + documentSliceTokens?: number | null; + documentMime?: Record | null; + documentReads?: DocumentReads | null; } /** Per-call cost anatomy for one session (admin cache-miss forensics). */ @@ -338,6 +369,9 @@ export interface AttachmentProfile { count: number; totalBytes: number; byMime: Record; + /** Uploads with a ready DocumentDigest, and the rendered tokens they would cost in context. */ + digested?: number; + digestTokens?: number; } /** One model call's context occupancy (input + cacheRead + cacheWrite). */ @@ -377,6 +411,7 @@ export interface DataCoverage { prefixTokens?: boolean; windowTrim?: boolean; compactionEvents?: boolean; + documents?: boolean; } /** The content-free diagnostic profile of one conversation. */ @@ -405,6 +440,16 @@ export interface SessionProfile { compactionEventCounts?: Record; /** The summary's token size at the most recent compaction decision. */ lastSummaryTokens?: number | null; + /** + * Document lifecycle across the session's calls: calls that ran with the + * full document inline vs. a digest only, the largest estimated document + * footprint seen, and what `document_read` pulled back in total. + */ + fullDocumentCalls?: number; + digestOnlyCalls?: number; + peakDocumentTokens?: number | null; + documentReadCalls?: number; + documentReadPages?: number; } // ========== API Request Options ========== diff --git a/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.spec.ts b/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.spec.ts index 78572a1f..cefeb693 100644 --- a/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.spec.ts +++ b/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.spec.ts @@ -265,6 +265,51 @@ describe('SessionCostAnatomyPage', () => { page.compactionEventTitle({ kind: 'floor_unreachable', checkpoint: 12, summaryTokens: 900, retainedMessages: 30 }), ).toBe('floor unreachable · checkpoint 12 · summary 900 · 30 messages retained'); expect(page.compactionEventTitle({ kind: 'checkpoint' })).toBe('checkpoint'); + expect( + page.compactionEventTitle({ kind: 'document_offload', documents: 1, documentTokens: 9_000, cacheGapSeconds: 420 }), + ).toBe('document offload · 1 document · ~9.0K tokens · cache gap 7m 0s'); + }); + + it('summarises how the documents were consumed, or falls back when untracked', async () => { + const profile = vi.fn().mockReturnValue( + of({ + ...MOCK_PROFILE, + dataCoverage: { ...MOCK_PROFILE.dataCoverage, documents: true }, + fullDocumentCalls: 2, + digestOnlyCalls: 1, + documentReadCalls: 1, + documentReadPages: 4, + peakDocumentTokens: 12_000, + }), + ); + const page = setup(vi.fn().mockReturnValue(of(MOCK_ANATOMY)), profile).componentInstance; + await vi.waitFor(() => expect(page.profileResource.hasValue()).toBe(true)); + expect(page.documentsLine()).toBe('2 full · 1 digest-only · read 4 pages in 1 calls · peak ~12.0K'); + + const untracked = setup(vi.fn().mockReturnValue(of(MOCK_ANATOMY))).componentInstance; + await vi.waitFor(() => expect(untracked.profileResource.hasValue()).toBe(true)); + expect(untracked.documentsLine()).toBe(''); + }); + + it('badges a call by what the model had of the documents', () => { + const page = setup(vi.fn().mockReturnValue(of(MOCK_ANATOMY))).componentInstance; + const base = MOCK_ANATOMY.calls[0]; + expect(page.documentBadge({ ...base })).toBe(''); + expect(page.documentBadge({ ...base, hasDocuments: true, documentCount: 1 })).toBe('doc'); + expect(page.documentBadge({ ...base, hasDocuments: false, documentDigests: 2 })).toBe('digest'); + expect(page.documentBadge({ ...base, hasDocuments: false, documentReads: { calls: 1, pages: 4, bytes: 9 } })).toBe('+4p'); + expect(page.documentDetail({ ...base })).toBe(''); + expect( + page.documentDetail({ + ...base, + hasDocuments: true, + documentCount: 2, + documentTokens: 12_000, + documentMime: { pdf: 2 }, + documentsAttached: 2, + documentReads: { calls: 1, pages: 4, bytes: 9 }, + }), + ).toBe('2 inline ~12.0K (pdf×2) · 2 attached this turn · document_read ×1 → 4 pages'); }); it('survives a missing profile without touching the anatomy', async () => { diff --git a/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.ts b/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.ts index 63e37c00..c7e1f9cb 100644 --- a/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.ts +++ b/frontend/ai.client/src/app/admin/costs/pages/session-cost-anatomy.page.ts @@ -23,6 +23,7 @@ import { SpinnerComponent } from '../../../components/spinner/spinner.component' import { ContextTrajectoryChartComponent } from '../components/context-trajectory-chart.component'; import { CacheStatus, DiagnosisSeverity, SessionDiagnosis, CompactionEvent, + SessionCallRow, } from '../models'; import { AnatomyRow, @@ -141,7 +142,13 @@ import {

Attachments

{{ profile.attachments.count }}

- @if (profile.attachments.count > 0) { + @if (documentsLine(); as documents) { + +

{{ documents }}

+ } @else if (profile.attachments.count > 0) {

{{ bytes(profile.attachments.totalBytes) }}

}
@@ -517,6 +524,16 @@ import { >{{ event.kind }} } + @if (documentBadge(row.call); as badge) { + + {{ badge }} + } {{ formatGap(row.call.cacheGapSeconds) }} @@ -626,6 +643,14 @@ import { } +
+
+ Documents +
+
+ {{ documentDetail(row.call) || '—' }} +
+
@for (key of fingerprintKeys; track key) {
@@ -787,13 +812,59 @@ export class SessionCostAnatomyPage { }); compactionEventTitle(event: CompactionEvent): string { - const parts = [event.kind.replace('_', ' ')]; + const parts = [event.kind.replace(/_/g, ' ')]; if (event.checkpoint != null) parts.push(`checkpoint ${event.checkpoint}`); if (event.summaryTokens != null) parts.push(`summary ${this.formatTokens(event.summaryTokens)}`); if (event.summarizedTurns != null) parts.push(`${event.summarizedTurns} turns summarized`); if (event.retainedMessages != null) parts.push(`${event.retainedMessages} messages retained`); if (event.truncatedToolResults) parts.push(`${event.truncatedToolResults} tool results truncated`); + if (event.documents != null) parts.push(`${event.documents} document${event.documents === 1 ? '' : 's'}`); + if (event.documentTokens != null) parts.push(`~${this.formatTokens(event.documentTokens)} tokens`); + if (event.cacheGapSeconds != null) parts.push(`cache gap ${this.formatGap(event.cacheGapSeconds)}`); + return parts.join(' · '); + } + + /** + * "2 full · 1 digest-only · read 4 pages · peak ~12K" — how the session's + * documents were consumed, call by call. Empty when the rows predate the + * document fields, so the card falls back to the upload byte total. + */ + readonly documentsLine = computed(() => { + if (!this.profileResource.hasValue()) return ''; + const p = this.profileResource.value(); + if (!p.dataCoverage.documents) return ''; + const parts: string[] = []; + if (p.fullDocumentCalls) parts.push(`${p.fullDocumentCalls} full`); + if (p.digestOnlyCalls) parts.push(`${p.digestOnlyCalls} digest-only`); + if (p.documentReadCalls) parts.push(`read ${p.documentReadPages ?? 0} pages in ${p.documentReadCalls} calls`); + if (p.peakDocumentTokens != null) parts.push(`peak ~${this.formatTokens(p.peakDocumentTokens)}`); return parts.join(' · '); + }); + + /** Short row badge: `doc`, `digest`, or `+Np` for pages retrieved this call. */ + documentBadge(call: SessionCallRow): string { + const reads = call.documentReads; + if (reads && reads.pages > 0) return `+${reads.pages}p`; + if (call.hasDocuments) return 'doc'; + if (call.documentDigests) return 'digest'; + return ''; + } + + /** The expanded-row line for the call's document context. */ + documentDetail(call: SessionCallRow): string { + if (call.hasDocuments == null && !call.documentReads) return ''; + const parts: string[] = []; + if (call.hasDocuments) { + const mime = Object.entries(call.documentMime ?? {}) + .map(([fmt, n]) => `${fmt}×${n}`) + .join(' '); + parts.push(`${call.documentCount ?? 0} inline ~${this.formatTokens(call.documentTokens ?? 0)}${mime ? ` (${mime})` : ''}`); + } + if (call.documentDigests) parts.push(`${call.documentDigests} digest${call.documentDigests === 1 ? '' : 's'}`); + if (call.documentsAttached) parts.push(`${call.documentsAttached} attached this turn`); + if (call.documentSlices) parts.push(`${call.documentSlices} retrieved slice${call.documentSlices === 1 ? '' : 's'} ~${this.formatTokens(call.documentSliceTokens ?? 0)}`); + if (call.documentReads?.calls) parts.push(`document_read ×${call.documentReads.calls} → ${call.documentReads.pages} pages`); + return parts.length ? parts.join(' · ') : 'no documents in context'; } readonly unexplainedMisses = computed(() => {