diff --git a/CHANGELOG.md b/CHANGELOG.md index ea1dcf51f7..d74c2ac245 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## Unreleased - Control Channel: `inspect ctl sample list` no longer recomputes sample summaries on every request, so polling an eval buffering many large samples (e.g. a retry's carried transcripts) can no longer stall the eval process. +- Control Channel: paged event reads served from the realtime sample buffer now load only the message/call pool entries and attachments the page references, instead of the sample's full pools and every attachment body. - Logging: Building a sample summary no longer serializes large structured metadata values just to exclude them, avoiding stalls when sample metadata embeds large data. ## 0.3.248 (17 July 2026) diff --git a/src/inspect_ai/event/_pool.py b/src/inspect_ai/event/_pool.py index 5f2ce60c1e..272f9d899b 100644 --- a/src/inspect_ai/event/_pool.py +++ b/src/inspect_ai/event/_pool.py @@ -26,7 +26,7 @@ import dataclasses import json from collections.abc import Iterable, Mapping, Sequence -from typing import Final, TypeVar, cast +from typing import Final, Literal, NamedTuple, TypeVar, cast from pydantic import BaseModel, JsonValue from pydantic_core import to_jsonable_python @@ -41,8 +41,8 @@ def materialize_pooled_events( events: Iterable[object], - message_pool: list[ChatMessage], - call_pool: list[JsonValue], + message_pool: Sequence[ChatMessage] | Mapping[int, ChatMessage], + call_pool: Sequence[JsonValue] | Mapping[int, JsonValue], ) -> list[Event]: materialized = validate_events(list(events)) materialized = resolve_model_event_inputs(materialized, message_pool) @@ -244,18 +244,153 @@ def _compress_refs(indices: list[int]) -> list[tuple[int, int]]: def _expand_refs( refs: list[tuple[int, int]], - pool: list[_T], + pool: Sequence[_T] | Mapping[int, _T], ) -> list[_T]: """Expand range-encoded refs against a pool. Each element is ``(start, end_exclusive)``: yields ``pool[start:end_exclusive]``. + Position-keyed mappings (page-scoped buffer reads carry only the positions + their events reference) are indexed per position, skipping absent positions + to mirror slice truncation of out-of-range refs. """ result: list[_T] = [] - for start, end_exclusive in refs: - result.extend(pool[start:end_exclusive]) + if isinstance(pool, Mapping): + for start, end_exclusive in refs: + result.extend(pool[i] for i in range(start, end_exclusive) if i in pool) + else: + for start, end_exclusive in refs: + result.extend(pool[start:end_exclusive]) return result +class PoolRefField(NamedTuple): + """Location of a range-encoded pool-ref field on a condensed event.""" + + pool: Literal["message", "call"] + """Which pool the refs index into.""" + + path: tuple[str, ...] + """Key path from the event root to the refs list, in raw JSON form.""" + + +POOL_REF_FIELDS: Final[tuple[PoolRefField, ...]] = ( + PoolRefField(pool="message", path=("input_refs",)), + PoolRefField(pool="call", path=("call", "call_refs")), +) +"""Registry of every event field that carries range-encoded pool refs. + +Everything that reads or rewrites pool refs must agree on which event +fields hold them, and this registry is their single source of truth: + +- the condense/resolve functions in this module, which write and read the + typed fields; +- :func:`collect_pool_ref_positions`, which the buffer's page-scoped reads + use to load only the pool entries a page's events reference; +- :func:`remap_pool_refs`, which export paths use to translate refs after + pool entries are assigned new positions in a destination store; +- ``test_pool_ref_registry_covers_all_ref_fields`` (in + ``tests/log/test_sample_history.py``), which introspects the event models + for ``*_refs`` fields and fails when one is missing here. + +A pool-ref field that isn't registered here would make page-scoped reads +silently drop the entries it references (:func:`_expand_refs` skips absent +positions), so any new ``*_refs`` field MUST be added here alongside its +condense/resolve support — the test enforces registration. +""" + + +class PoolRefPositions(NamedTuple): + message_positions: set[int] + call_positions: set[int] + + +def collect_pool_ref_positions( + events: Iterable[Mapping[str, JsonValue]], +) -> PoolRefPositions: + """Pool positions referenced by condensed events in raw JSON form. + + Walks :data:`POOL_REF_FIELDS` so callers that load partial pools (the + buffer's page-scoped reads) stay in sync with the condense/resolve + functions in this module. + """ + positions = PoolRefPositions(message_positions=set(), call_positions=set()) + pool_positions: dict[str, set[int]] = { + "message": positions.message_positions, + "call": positions.call_positions, + } + for event in events: + for field in POOL_REF_FIELDS: + value: object = event + for key in field.path: + value = value.get(key) if isinstance(value, Mapping) else None + _accumulate_ref_positions(value, pool_positions[field.pool]) + return positions + + +def _accumulate_ref_positions(refs: object, positions: set[int]) -> None: + """Accumulate positions covered by range-encoded ``(start, end)`` refs.""" + if not isinstance(refs, list): + return + for ref in refs: + if not isinstance(ref, (list, tuple)) or len(ref) != 2: + continue + start, end = ref + if isinstance(start, int) and isinstance(end, int): + positions.update(range(start, end)) + + +def remap_pool_refs( + event: Mapping[str, JsonValue], + message_pos_map: Mapping[int, int], + call_pos_map: Mapping[int, int], +) -> dict[str, JsonValue]: + """Rewrite a condensed event's pool refs through position maps. + + Used when a sample's pool entries are exported into another store and + assigned new positions there. Walks :data:`POOL_REF_FIELDS` so exporters + stay in sync with the condense/resolve functions in this module. + """ + pos_maps: dict[str, Mapping[int, int]] = { + "message": message_pos_map, + "call": call_pos_map, + } + remapped: dict[str, JsonValue] = dict(event) + for field in POOL_REF_FIELDS: + remapped = _remap_refs_at_path(remapped, field.path, pos_maps[field.pool]) + return remapped + + +def _remap_refs_at_path( + node: dict[str, JsonValue], path: tuple[str, ...], pos_map: Mapping[int, int] +) -> dict[str, JsonValue]: + """Return ``node`` with the refs at ``path`` remapped, copying dicts on the path.""" + key, rest = path[0], path[1:] + child = node.get(key) + if rest: + if not isinstance(child, dict): + return node + new_child = _remap_refs_at_path(child, rest, pos_map) + return node if new_child is child else {**node, key: new_child} + if not isinstance(child, list): + return node + return {**node, key: cast(JsonValue, _remap_refs(child, pos_map))} + + +def _remap_refs( + refs: Sequence[object], pos_map: Mapping[int, int] +) -> list[tuple[int, int]]: + """Translate range-encoded refs through a position map and re-compress.""" + indices: list[int] = [] + for ref in refs: + if not isinstance(ref, (list, tuple)) or len(ref) != 2: + continue + start, end = ref + if not isinstance(start, int) or not isinstance(end, int): + continue + indices.extend(pos_map[index] for index in range(start, end)) + return _compress_refs(indices) + + def condense_model_event_calls( events: Sequence[Event], next_index: int, @@ -338,7 +473,7 @@ def condense_model_event_calls( def resolve_model_event_calls( events: list[Event], - call_pool: list[JsonValue], + call_pool: Sequence[JsonValue] | Mapping[int, JsonValue], ) -> list[Event]: """Restore call.request messages from call_pool references.""" if not call_pool: @@ -364,7 +499,7 @@ def resolve_model_event_calls( def resolve_model_event_inputs( events: list[Event], - message_pool: list[ChatMessage], + message_pool: Sequence[ChatMessage] | Mapping[int, ChatMessage], ) -> list[Event]: """Resolve ModelEvent input_refs back to full input lists.""" if not message_pool: diff --git a/src/inspect_ai/log/_recorders/buffer/database.py b/src/inspect_ai/log/_recorders/buffer/database.py index 36d5632029..4a22b521ca 100644 --- a/src/inspect_ai/log/_recorders/buffer/database.py +++ b/src/inspect_ai/log/_recorders/buffer/database.py @@ -5,12 +5,18 @@ import sqlite3 import threading import time -from collections.abc import Sequence from contextlib import contextmanager from logging import getLogger from pathlib import Path from sqlite3 import Connection, OperationalError -from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Literal, cast +from typing import ( + TYPE_CHECKING, + Callable, + Iterable, + Iterator, + Literal, + TypeVar, +) import psutil from pydantic import BaseModel, JsonValue @@ -30,9 +36,10 @@ from inspect_ai.event._model import ModelEvent from inspect_ai.event._pool import ( _call_pool_json, - _compress_refs, _msg_pool_json, _msg_pool_jsonable, + collect_pool_ref_positions, + remap_pool_refs, ) from inspect_ai.event._pool_index import ( CallPoolIndex, @@ -45,6 +52,7 @@ from ..._condense import ( ATTACHMENT_PROTOCOL, WalkContext, + attachment_refs_from_value, attachments_content_fn, walk_chat_message, walk_events, @@ -670,9 +678,7 @@ def attachment_lookup(hash: str) -> str | None: for row in self._get_events(conn, id, epoch, latest_only=True): transcript_store.merge_condensed_event( row.event_id, - self._remap_pool_refs( - row.event, message_pos_map, call_pos_map - ), + remap_pool_refs(row.event, message_pos_map, call_pos_map), attachment_lookup, ) seed_count += 1 @@ -682,27 +688,6 @@ def attachment_lookup(hash: str) -> str | None: conn.rollback() raise - @staticmethod - def _remap_pool_refs( - event: JsonData, message_pos_map: dict[int, int], call_pos_map: dict[int, int] - ) -> JsonData: - """Rewrite a condensed event's pool refs after exporting its pool entries.""" - remapped = dict(event) - input_refs = remapped.get("input_refs") - if isinstance(input_refs, list): - remapped["input_refs"] = cast( - JsonValue, _remap_refs(input_refs, message_pos_map) - ) - call = remapped.get("call") - if isinstance(call, dict): - call_refs = call.get("call_refs") - if isinstance(call_refs, list): - remapped["call"] = { - **call, - "call_refs": cast(JsonValue, _remap_refs(call_refs, call_pos_map)), - } - return remapped - @contextmanager def open_sample_history_tail( self, @@ -711,14 +696,18 @@ def open_sample_history_tail( n: int, ) -> Iterator[SampleHistory]: if n <= 0: - yield SampleHistory([], [], [], {}) + yield SampleHistory([], {}, {}, {}, page_scoped=True) return with self._acquire_sample_read_lease(id, epoch): with self._get_connection() as conn: conn.execute("BEGIN") history = self._sample_history( - conn, id, epoch, self._get_events_tail(conn, id, epoch, n) + conn, + id, + epoch, + self._get_events_tail(conn, id, epoch, n), + page_scoped=True, ) conn.commit() yield history @@ -739,6 +728,7 @@ def open_sample_history_from( id, epoch, self._get_events_from(conn, id, epoch, start, limit), + page_scoped=True, ) conn.commit() yield history @@ -767,18 +757,107 @@ def _sample_history( id: str | int, epoch: int, events: Iterable[EventData], + *, + page_scoped: bool = False, ) -> SampleHistory: - message_pool = [ - json.loads(entry.data) for entry in self._get_message_pool(conn, id, epoch) - ] - call_pool = [ - json.loads(entry.data) for entry in self._get_call_pool(conn, id, epoch) - ] - attachments = { - entry.hash: entry.content - for entry in self._get_attachments(conn, id, epoch) - } - return SampleHistory(list(events), message_pool, call_pool, attachments) + """Assemble a ``SampleHistory`` for the given event rows. + + With ``page_scoped``, only the pool entries and attachments actually + referenced by ``events`` are loaded, keeping per-page read cost + proportional to the page rather than the whole sample. Whole-history + reads load everything: their events reference (nearly) all of it + anyway, and the .eval recorder needs the complete pools to serialize + ``events_data``. + """ + event_rows = list(events) + if page_scoped: + positions = collect_pool_ref_positions(row.event for row in event_rows) + message_pool = { + pos: json.loads(data) + for pos, data in self._get_pool_entries_at( + conn, "message_pool", id, epoch, positions.message_positions + ) + } + call_pool = { + pos: json.loads(data) + for pos, data in self._get_pool_entries_at( + conn, "call_pool", id, epoch, positions.call_positions + ) + } + hashes: set[str] = set() + for row in event_rows: + hashes.update(attachment_refs_from_value(row.event)) + for pool_value in (*message_pool.values(), *call_pool.values()): + hashes.update(attachment_refs_from_value(pool_value)) + attachments = self._get_sample_attachments_content(conn, id, epoch, hashes) + else: + message_pool = { + pos: json.loads(entry.data) + for pos, entry in enumerate(self._get_message_pool(conn, id, epoch)) + } + call_pool = { + pos: json.loads(entry.data) + for pos, entry in enumerate(self._get_call_pool(conn, id, epoch)) + } + attachments = { + entry.hash: entry.content + for entry in self._get_attachments(conn, id, epoch) + } + return SampleHistory( + event_rows, message_pool, call_pool, attachments, page_scoped=page_scoped + ) + + def _get_pool_entries_at( + self, + conn: Connection, + table: Literal["message_pool", "call_pool"], + id: str | int, + epoch: int, + positions: set[int], + ) -> Iterator[tuple[int, str]]: + """Pool entries at the given .eval positional indices. + + Position equals insertion rank within the sample (rows are inserted in + pool-position order — see ``_condense_model_event``), so rank is + recovered with ROW_NUMBER over the autoincrement id. + """ + if not positions: + return + query = f""" + WITH ordered AS ( + SELECT ROW_NUMBER() OVER (ORDER BY id) - 1 AS pos, data + FROM {table} + WHERE sample_id = ? AND sample_epoch = ? + ) + SELECT pos, data FROM ordered WHERE pos IN + """ + for chunk in _chunked(sorted(positions), _SQLITE_MAX_VARIABLES): + placeholders = ",".join("?" * len(chunk)) + for row in conn.execute( + f"{query} ({placeholders})", [str(id), epoch, *chunk] + ): + yield int(row["pos"]), str(row["data"]) + + def _get_sample_attachments_content( + self, + conn: Connection, + id: str | int, + epoch: int, + hashes: set[str], + ) -> dict[str, str]: + """The sample's attachment bodies for the given hashes (missing omitted).""" + attachments: dict[str, str] = {} + for chunk in _chunked(sorted(hashes), _SQLITE_MAX_VARIABLES): + placeholders = ",".join("?" * len(chunk)) + for row in conn.execute( + f""" + SELECT hash, content FROM attachments + WHERE sample_id = ? AND sample_epoch = ? AND hash IN ({placeholders}) + """, + [str(id), epoch, *chunk], + ): + attachments[str(row["hash"])] = str(row["content"]) + return attachments @contextmanager def _acquire_sample_read_lease( @@ -1731,19 +1810,15 @@ def maximum_ids( return event_id, attachment_id, message_pool_id, call_pool_id -def _remap_refs( - refs: Sequence[object], pos_map: dict[int, int] -) -> list[tuple[int, int]]: - """Translate pooled ref ranges after exporting pool entries into a new store.""" - indices: list[int] = [] - for ref in refs: - if not isinstance(ref, (list, tuple)) or len(ref) != 2: - continue - start, end = ref - if not isinstance(start, int) or not isinstance(end, int): - continue - indices.extend(pos_map[index] for index in range(start, end)) - return _compress_refs(indices) +# stay safely under SQLite's historical 999 bound-parameter limit +_SQLITE_MAX_VARIABLES = 500 + +_T = TypeVar("_T") + + +def _chunked(items: list[_T], size: int) -> Iterator[list[_T]]: + for i in range(0, len(items), size): + yield items[i : i + size] def cleanup_sample_buffer_databases(db_dir: Path | None = None) -> None: diff --git a/src/inspect_ai/log/_recorders/buffer/history.py b/src/inspect_ai/log/_recorders/buffer/history.py index c72f86791f..4b089f41c1 100644 --- a/src/inspect_ai/log/_recorders/buffer/history.py +++ b/src/inspect_ai/log/_recorders/buffer/history.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Iterator -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TypeAlias from pydantic import JsonValue, TypeAdapter @@ -23,20 +23,51 @@ @dataclass class SampleHistory: - """Latest logical sample history with .eval positional pool refs.""" + """Latest logical sample history with .eval positional pool refs. + + Pools are keyed by .eval positional index: full-history reads carry every + entry (contiguous positions from 0), while page-scoped reads carry only + the positions referenced by the page's events. + """ raw_event_rows: list[EventData] - message_pool: list[ChatMessage] - call_pool: list[JsonValue] + message_pool: dict[int, ChatMessage] + call_pool: dict[int, JsonValue] attachments: dict[str, str] - events_data: EventsData = field(init=False) + page_scoped: bool = False + """True when pools/attachments carry only the entries referenced by this + page's events (which may still form a contiguous prefix, so completeness + can't be inferred from the keys).""" def __post_init__(self) -> None: - self.message_pool = validate_chat_messages( - self.message_pool, context={"deserializing": True} + validated_messages = validate_chat_messages( + list(self.message_pool.values()), context={"deserializing": True} + ) + self.message_pool = dict( + zip(self.message_pool, validated_messages, strict=True) + ) + validated_calls = _json_value_list_adapter.validate_python( + list(self.call_pool.values()) + ) + self.call_pool = dict(zip(self.call_pool, validated_calls, strict=True)) + + @property + def events_data(self) -> EventsData: + """Dense pools for embedding in an ``EvalSample`` / .eval log. + + Positional refs index into the dense lists, so this requires a + full-history read (complete pools with contiguous positions from 0); + page-scoped histories raise rather than silently misalign refs. + """ + if self.page_scoped: + raise RuntimeError( + "events_data requires a full sample history; this history " + "is page-scoped and carries only referenced pool entries" + ) + return EventsData( + messages=[self.message_pool[i] for i in range(len(self.message_pool))], + calls=[self.call_pool[i] for i in range(len(self.call_pool))], ) - self.call_pool = _json_value_list_adapter.validate_python(self.call_pool) - self.events_data = EventsData(messages=self.message_pool, calls=self.call_pool) @property def event_count(self) -> int: diff --git a/src/inspect_ai/log/_recorders/buffer/transcript_history_provider.py b/src/inspect_ai/log/_recorders/buffer/transcript_history_provider.py index 4bf3a54bd2..4c36e8a6ad 100644 --- a/src/inspect_ai/log/_recorders/buffer/transcript_history_provider.py +++ b/src/inspect_ai/log/_recorders/buffer/transcript_history_provider.py @@ -5,6 +5,8 @@ from contextlib import AbstractContextManager, contextmanager from typing import TYPE_CHECKING, Protocol +from pydantic import JsonValue + from inspect_ai.event._base import BaseEvent from inspect_ai.event._event import Event from inspect_ai.event._pool import ( @@ -14,8 +16,8 @@ ) from inspect_ai.event._validate import validate_events from inspect_ai.log._condense import resolve_events_attachments -from inspect_ai.log._log import EventsData from inspect_ai.log._recorders.buffer.types import JsonData +from inspect_ai.model import ChatMessage if TYPE_CHECKING: from inspect_ai.log._transcript import TranscriptHistoryProvider @@ -24,8 +26,16 @@ class SampleHistoryLike(Protocol): - events_data: EventsData - attachments: dict[str, str] + # read-only properties: page-scoped histories carry sparse position-keyed + # pools, so consumers must resolve refs by position rather than slicing + @property + def message_pool(self) -> Mapping[int, ChatMessage]: ... + + @property + def call_pool(self) -> Mapping[int, JsonValue]: ... + + @property + def attachments(self) -> Mapping[str, str]: ... def iter_events(self) -> Iterator[Event | JsonData]: ... @@ -184,8 +194,8 @@ def _iter_events(self) -> Iterator[Event]: def _materialize_events(history: SampleHistoryLike) -> list[Event]: events = materialize_pooled_events( history.iter_events(), - history.events_data["messages"], - history.events_data["calls"], + history.message_pool, + history.call_pool, ) # Resolve content attachments (large text / images) back to their # underlying values so callers get usable events, not bare @@ -195,8 +205,8 @@ def _materialize_events(history: SampleHistoryLike) -> list[Event]: def _iter_materialized_events(history: SampleHistoryLike) -> Iterator[Event]: - message_pool = history.events_data["messages"] - call_pool = history.events_data["calls"] + message_pool = history.message_pool + call_pool = history.call_pool attachments = history.attachments for raw_event in history.iter_events(): event = ( diff --git a/src/inspect_ai/log/_recorders/buffer/types.py b/src/inspect_ai/log/_recorders/buffer/types.py index f7a89dd088..10c448a624 100644 --- a/src/inspect_ai/log/_recorders/buffer/types.py +++ b/src/inspect_ai/log/_recorders/buffer/types.py @@ -173,7 +173,13 @@ def open_sample_history_tail( epoch: int, n: int, ) -> AbstractContextManager["SampleHistory"]: - """Open a consistent snapshot of the last ``n`` sample events.""" + """Open a consistent snapshot of the last ``n`` sample events. + + The yielded history is page-scoped: its pools carry only the entries + referenced by the page's events (position-keyed, possibly sparse) and + ``SampleHistory.events_data`` raises ``RuntimeError``. Use + ``open_sample_history`` when dense full pools are required. + """ ... @abc.abstractmethod @@ -189,6 +195,11 @@ def open_sample_history_from( ``limit`` caps the number of events read (``None`` = through the end), so cursored page readers don't materialize the full remaining history to serve one page. + + The yielded history is page-scoped: its pools carry only the entries + referenced by the page's events (position-keyed, possibly sparse) and + ``SampleHistory.events_data`` raises ``RuntimeError``. Use + ``open_sample_history`` when dense full pools are required. """ ... diff --git a/src/inspect_ai/log/_recorders/streaming.py b/src/inspect_ai/log/_recorders/streaming.py index 07c87a607b..7f0e0be47e 100644 --- a/src/inspect_ai/log/_recorders/streaming.py +++ b/src/inspect_ai/log/_recorders/streaming.py @@ -41,10 +41,11 @@ def eval_retry_error_from_history( elif suffix: suffix.append(event) + events_data = history.events_data events = materialize_pooled_events( suffix, - history.events_data["messages"], - history.events_data["calls"], + events_data["messages"], + events_data["calls"], ) return EvalRetryError( diff --git a/tests/_eval/test_retry_error_events.py b/tests/_eval/test_retry_error_events.py index 2401c72f39..f47027b273 100644 --- a/tests/_eval/test_retry_error_events.py +++ b/tests/_eval/test_retry_error_events.py @@ -4,12 +4,13 @@ from types import SimpleNamespace import pytest +from pydantic import JsonValue from test_helpers.transcript import make_model_event from inspect_ai._eval.task.log import TaskLogger from inspect_ai._eval.task.run import _eval_retry_error, _sample_transcript_config from inspect_ai.event import Event, InfoEvent, ModelEvent -from inspect_ai.log import EvalError, EventsData, Transcript +from inspect_ai.log import EvalError, Transcript from inspect_ai.log._recorders.buffer.database import SampleBufferDatabase from inspect_ai.log._recorders.buffer.transcript_history_provider import ( BufferTranscriptHistoryProvider, @@ -17,7 +18,7 @@ from inspect_ai.log._recorders.streaming import eval_retry_error_from_history from inspect_ai.log._recorders.types import SampleEvent from inspect_ai.log._transcript import _transcript, init_transcript -from inspect_ai.model import ChatMessageUser, GenerateConfig, ModelOutput +from inspect_ai.model import ChatMessage, ChatMessageUser, GenerateConfig, ModelOutput def _model(uuid: str, content: str) -> ModelEvent: @@ -192,7 +193,8 @@ def test_buffer_provider_iter_events_streams_first_event_before_later_rows() -> class LazyHistory: attachments: dict[str, str] = {} - events_data: EventsData = {"messages": [], "calls": []} + message_pool: dict[int, ChatMessage] = {} + call_pool: dict[int, JsonValue] = {} def iter_events(self) -> Iterator[Event]: yield first diff --git a/tests/log/test_sample_history.py b/tests/log/test_sample_history.py index 49d7a40a6f..ffa995107b 100644 --- a/tests/log/test_sample_history.py +++ b/tests/log/test_sample_history.py @@ -1,14 +1,22 @@ import sqlite3 +from collections.abc import Iterator import pytest +from inspect_ai._util.hash import mm3_hash from inspect_ai._util.json import to_json_str_safe from inspect_ai.event import InfoEvent, ModelEvent from inspect_ai.log._log import EvalSample, EventsData from inspect_ai.log._recorders.buffer.database import SampleBufferDatabase from inspect_ai.log._recorders.streaming import materialize_streaming_sample from inspect_ai.log._recorders.types import SampleEvent -from inspect_ai.model import ChatMessageUser, GenerateConfig, ModelCall, ModelOutput +from inspect_ai.model import ( + ChatMessage, + ChatMessageUser, + GenerateConfig, + ModelCall, + ModelOutput, +) def _model( @@ -16,11 +24,14 @@ def _model( completion: str, pending: bool | None = None, call: ModelCall | None = None, + input: list[ChatMessage] | None = None, ) -> ModelEvent: return ModelEvent( uuid=uuid, model="mockllm/model", - input=[ChatMessageUser(id="input-message", content="question")], + input=input + if input is not None + else [ChatMessageUser(id="input-message", content="question")], tools=[], tool_choice="none", config=GenerateConfig(), @@ -426,6 +437,151 @@ def test_open_sample_history_from_honors_limit(tmp_path): assert [row.event["data"] for row in rows] == [4] +def test_page_scoped_history_loads_only_referenced_pool_entries(tmp_path): + db = SampleBufferDatabase(str(tmp_path / "test.eval"), db_dir=tmp_path) + msg_a = ChatMessageUser(id="msg-a", content="question one") + msg_b = ChatMessageUser(id="msg-b", content="question two") + db.log_events( + [ + SampleEvent( + id="sample", epoch=1, event=_model("event-1", "a", input=[msg_a]) + ), + SampleEvent( + id="sample", epoch=1, event=_model("event-2", "b", input=[msg_b]) + ), + ] + ) + + # the page's events reference only pool position 1 + with db.open_sample_history_from("sample", 1, 1) as history: + assert [row.event_id for row in history.raw_event_rows] == ["event-2"] + assert set(history.message_pool) == {1} + assert history.message_pool[1].content == "question two" + + with db.open_sample_history_tail("sample", 1, 1) as history: + assert [row.event_id for row in history.raw_event_rows] == ["event-2"] + assert set(history.message_pool) == {1} + + # full-history reads still carry the complete pool + with db.open_sample_history("sample", 1) as history: + assert set(history.message_pool) == {0, 1} + assert len(history.events_data["messages"]) == 2 + + +def test_page_scoped_history_rejects_events_data(tmp_path): + db = SampleBufferDatabase(str(tmp_path / "test.eval"), db_dir=tmp_path) + msg_a = ChatMessageUser(id="msg-a", content="question one") + msg_b = ChatMessageUser(id="msg-b", content="question two") + db.log_events( + [ + SampleEvent( + id="sample", epoch=1, event=_model("event-1", "a", input=[msg_a]) + ), + SampleEvent( + id="sample", epoch=1, event=_model("event-2", "b", input=[msg_b]) + ), + ] + ) + + with db.open_sample_history_from("sample", 1, 1) as history: + with pytest.raises(RuntimeError, match="page-scoped"): + _ = history.events_data + + # a page referencing a contiguous pool prefix (positions 0..k) is + # indistinguishable from a full pool by its keys, but must still raise + with db.open_sample_history_from("sample", 1, 0, 1) as history: + assert set(history.message_pool) == {0} + with pytest.raises(RuntimeError, match="page-scoped"): + _ = history.events_data + + +def test_page_scoped_history_loads_only_referenced_attachments(tmp_path): + db = SampleBufferDatabase(str(tmp_path / "test.eval"), db_dir=tmp_path) + big_first = "a" * 200 + big_second = "b" * 200 + db.log_events( + [ + SampleEvent( + id="sample", epoch=1, event=InfoEvent(uuid="info-1", data=big_first) + ), + SampleEvent( + id="sample", epoch=1, event=InfoEvent(uuid="info-2", data=big_second) + ), + ] + ) + + with db.open_sample_history_from("sample", 1, 1) as history: + assert set(history.attachments) == {mm3_hash(big_second)} + + with db.open_sample_history("sample", 1) as history: + assert set(history.attachments) == {mm3_hash(big_first), mm3_hash(big_second)} + + +def test_page_scoped_history_materializes_pool_entry_attachments(tmp_path): + """A page read resolves attachments referenced from its pool entries. + + Large message content is stored as an attachment *inside* the pooled + message, so page-scoped attachment collection must walk the loaded pool + entries, not just the event rows. + """ + from inspect_ai.log._recorders.buffer.transcript_history_provider import ( + BufferTranscriptHistoryProvider, + ) + + db = SampleBufferDatabase(str(tmp_path / "test.eval"), db_dir=tmp_path) + big_content = "b" * 200 + msg_a = ChatMessageUser(id="msg-a", content="short question") + msg_b = ChatMessageUser(id="msg-b", content=big_content) + db.log_events( + [ + SampleEvent( + id="sample", epoch=1, event=_model("event-1", "a", input=[msg_a]) + ), + SampleEvent( + id="sample", epoch=1, event=_model("event-2", "b", input=[msg_b]) + ), + ] + ) + + with db.open_sample_history_from("sample", 1, 1) as history: + assert mm3_hash(big_content) in history.attachments + + events = BufferTranscriptHistoryProvider(db, "sample", 1).events_from(1, 1) + assert len(events) == 1 + model_event = events[0] + assert isinstance(model_event, ModelEvent) + assert model_event.input_refs is None + assert model_event.input[0].content == big_content + + +def test_page_scoped_history_loads_referenced_call_pool_entries(tmp_path): + db = SampleBufferDatabase(str(tmp_path / "test.eval"), db_dir=tmp_path) + + def call(content: str) -> ModelCall: + return ModelCall( + request={"messages": [{"role": "user", "content": content}]}, + response={}, + ) + + db.log_events( + [ + SampleEvent( + id="sample", epoch=1, event=_model("event-1", "a", call=call("one")) + ), + SampleEvent( + id="sample", epoch=1, event=_model("event-2", "b", call=call("two")) + ), + ] + ) + + with db.open_sample_history_from("sample", 1, 1) as history: + assert set(history.call_pool) == {1} + assert history.call_pool[1] == {"role": "user", "content": "two"} + + with db.open_sample_history("sample", 1) as history: + assert set(history.call_pool) == {0, 1} + + def test_read_only_connection_does_not_recreate_deleted_db(tmp_path): """A read-only buffer reader cannot resurrect a deleted database. @@ -586,3 +742,52 @@ def test_provider_translates_store_failures_to_domain_error(tmp_path): racing._close_all_connections() with pytest.raises(TranscriptHistoryUnavailableError): provider.events_from(0, 10) + + +def test_pool_ref_registry_covers_all_ref_fields(): + """``POOL_REF_FIELDS`` must register every ``*_refs`` field on event models. + + Page-scoped buffer reads load only the pool positions collected via the + registry; a pool-ref field missing from it would make those reads silently + drop the entries it references (``_expand_refs`` skips absent positions). + This walks every pydantic model reachable from the ``Event`` union and + fails when the ``*_refs`` fields found don't match the registry exactly. + """ + from typing import get_args + + from pydantic import BaseModel + + from inspect_ai.event._event import Event + from inspect_ai.event._pool import POOL_REF_FIELDS + + def nested_models(annotation: object) -> Iterator[type[BaseModel]]: + if isinstance(annotation, type) and issubclass(annotation, BaseModel): + yield annotation + for arg in get_args(annotation): + yield from nested_models(arg) + + found: set[tuple[str, ...]] = set() + + def walk( + model: type[BaseModel], path: tuple[str, ...], seen: frozenset[type] + ) -> None: + if model in seen: + return + seen |= {model} + for name, field in model.model_fields.items(): + if name.endswith("_refs"): + found.add(path + (name,)) + for nested in nested_models(field.annotation): + walk(nested, path + (name,), seen) + + for event_type in get_args(Event): + walk(event_type, (), frozenset()) + + registered = {field.path for field in POOL_REF_FIELDS} + assert found == registered, ( + "Pool-ref fields on the event models don't match POOL_REF_FIELDS in " + "inspect_ai.event._pool. Register new *_refs fields there (with " + "condense/resolve support) so page-scoped buffer reads load their " + f"pool entries. Found on models: {sorted(found)}; " + f"registered: {sorted(registered)}" + )