From f21529ab2b6555b17659ae43393b1fd7401435a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Sun, 23 Aug 2026 20:41:51 +0200 Subject: [PATCH 1/3] hindsight: let a run choose the retain extraction mode and chunk size MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ingest cost on a large split is dominated by fact extraction, and extraction cost is `corpus_chars / retain_chunk_size` LLM calls. BEAM-10M is 10 conversations totalling ~468M characters, so at the server-side default chunk size of 3000 it is **~156,000 extraction calls for a single run** — which is what makes that split impractical to ingest rather than merely slow. The provider had no way to influence either. `_bank_kwargs` sent only `enable_observations` and, for BEAM, a retain mission, so every bank ran the server default (`concise`, i.e. full LLM extraction) at the default chunk size. Two env vars now pass through to `create_bank`, which has accepted both all along: - `AMB_HINDSIGHT_EXTRACTION_MODE=chunks` skips the LLM entirely and stores each chunk as its own unit. That makes a run of this size tractable, but it is **not the same measurement**: `_BEAM_RETAIN_MISSION` is an extraction prompt, so chunks mode ignores it and stores raw text with no fact extraction and no entities. A chunks-mode score is not comparable to an extracted one. - `AMB_HINDSIGHT_CHUNK_SIZE` trades the same axis more gently — doubling it roughly halves the call count while keeping extraction, at coarser granularity. Both are UNSET by default, so every existing result stays exactly the run it was. This is deliberately a new capability rather than a change of default: which mode BEAM should be scored under is a benchmark-semantics decision, and this only makes the choice expressible. `scripts/test_bank_kwargs.py` checks the default sends neither control and keeps the BEAM mission, that each control arrives when set (with the chunk size coerced to int), and that non-BEAM datasets get the controls without the BEAM mission. --- scripts/test_bank_kwargs.py | 61 ++++++++++++++++++++++++++++ src/memory_bench/memory/hindsight.py | 22 ++++++++++ 2 files changed, 83 insertions(+) create mode 100644 scripts/test_bank_kwargs.py diff --git a/scripts/test_bank_kwargs.py b/scripts/test_bank_kwargs.py new file mode 100644 index 0000000..8c443fe --- /dev/null +++ b/scripts/test_bank_kwargs.py @@ -0,0 +1,61 @@ +"""Check that the retain-side ingest controls reach bank creation. + +Ingest cost on a large split is dominated by fact extraction, and extraction cost is +``corpus_chars / retain_chunk_size`` LLM calls — for BEAM-10M (~468M characters) that is ~156,000 +calls per run at the server-side default chunk size of 3000. These two env vars are the levers, so +they are worth a check that they actually arrive rather than being silently dropped. + + uv run python scripts/test_bank_kwargs.py +""" +import os + +from memory_bench.memory.hindsight import _HindsightBase + + +def _kwargs(dataset: str | None, **env) -> dict: + prev = {k: os.environ.get(k) for k in env} + os.environ.update({k: v for k, v in env.items() if v is not None}) + for k, v in env.items(): + if v is None: + os.environ.pop(k, None) + try: + p = _HindsightBase.__new__(_HindsightBase) + p._dataset = dataset + return p._bank_kwargs() + finally: + for k, v in prev.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + +def main() -> None: + # Default: neither control is sent, so an existing run is byte-for-byte the run it always was. + k = _kwargs("beam", AMB_HINDSIGHT_EXTRACTION_MODE=None, AMB_HINDSIGHT_CHUNK_SIZE=None) + assert "retain_extraction_mode" not in k, k + assert "retain_chunk_size" not in k, k + assert "retain_mission" in k, "BEAM must still get its extraction mission by default" + print("default -> unchanged, mission present ok") + + # chunks mode skips the LLM entirely. It is a DIFFERENT measurement — the mission above is an + # extraction prompt and chunks mode ignores it — so this is opt-in, never a default. + k = _kwargs("beam", AMB_HINDSIGHT_EXTRACTION_MODE="chunks", AMB_HINDSIGHT_CHUNK_SIZE=None) + assert k["retain_extraction_mode"] == "chunks", k + print("mode=chunks -> retain_extraction_mode=chunks ok") + + # Chunk size is the gentler lever: it keeps extraction and halves the call count per doubling. + k = _kwargs("beam", AMB_HINDSIGHT_EXTRACTION_MODE=None, AMB_HINDSIGHT_CHUNK_SIZE="12000") + assert k["retain_chunk_size"] == 12000 and isinstance(k["retain_chunk_size"], int), k + print("chunk_size=12000 -> retain_chunk_size=12000 (int) ok") + + # Non-BEAM datasets get the controls too, just no BEAM mission. + k = _kwargs("locomo", AMB_HINDSIGHT_EXTRACTION_MODE="chunks", AMB_HINDSIGHT_CHUNK_SIZE=None) + assert k["retain_extraction_mode"] == "chunks" and "retain_mission" not in k, k + print("non-beam -> controls apply, no beam mission ok") + + print("\nall ok") + + +if __name__ == "__main__": + main() diff --git a/src/memory_bench/memory/hindsight.py b/src/memory_bench/memory/hindsight.py index 1e010c3..d6eceab 100644 --- a/src/memory_bench/memory/hindsight.py +++ b/src/memory_bench/memory/hindsight.py @@ -125,6 +125,28 @@ def _bank_kwargs(self, bank_id: str | None = None) -> dict: kwargs: dict = dict(enable_observations=False) if self._dataset == "beam": kwargs["retain_mission"] = self._BEAM_RETAIN_MISSION + + # Retain-side ingest controls, both unset by default so every existing result stays + # comparable. They exist because ingest cost is dominated by fact extraction, and extraction + # cost is `corpus_chars / retain_chunk_size` LLM calls: + # + # BEAM-10M is 10 conversations totalling ~468M characters. At the server-side default + # chunk size of 3000 chars that is ~156,000 extraction calls for a single run. + # + # `AMB_HINDSIGHT_EXTRACTION_MODE=chunks` skips the LLM entirely and stores each chunk as its + # own unit. That makes a run of that size tractable, but it is NOT the same measurement: + # `_BEAM_RETAIN_MISSION` above is an extraction prompt, so chunks mode ignores it and stores + # raw text with no fact extraction and no entities. Use it deliberately, and do not compare + # a chunks-mode score against an extracted one. + # + # `AMB_HINDSIGHT_CHUNK_SIZE` trades the same axis more gently: doubling it roughly halves the + # call count while keeping extraction, at coarser granularity. + mode = os.environ.get("AMB_HINDSIGHT_EXTRACTION_MODE") + if mode: + kwargs["retain_extraction_mode"] = mode + chunk_size = os.environ.get("AMB_HINDSIGHT_CHUNK_SIZE") + if chunk_size: + kwargs["retain_chunk_size"] = int(chunk_size) return kwargs def _create_bank(self, bank_id: str, force_reset: bool = True) -> None: From 6c296682f119e5f1729042ab9deb52ea979ca94f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 25 Aug 2026 10:07:55 +0200 Subject: [PATCH 2/3] beam: chunk BEAM-10M into bounded documents like every other split `load_documents` caps documents at `_MAX_DOC_CHARS = 100_000` and splits sessions to stay under it. BEAM-10M never reached that code. It nests two levels deeper than the other splits: chat[i]["plan-N"][batch]["turns"] -> list of turn GROUPS, each a list of turn dicts Nothing there is a `list` at the top level, so `sessions` came out empty and the conversation fell through to the "unusual structure" branch, which emits it whole: 10m: 10 documents, median 47,280,119 chars <- 470x the cap 100k: 170 documents, median 92,351 chars That is 470x the limit this same function enforces on every other split, for exactly the reason the limit was introduced. The consequence was not a slow ingest but an impossible one. Retain cost is per-CALL, not per-byte -- measured against a live API, one item takes 1.29s and fifty take 1.39s -- and a backend that serializes retains per document gets no parallelism at all when a 10-conversation split is only 10 documents. Ingest ran at ~1,300 chars/s, ~100 hours for the split, so the harness hit its 300s-per-operation timeout, gave up, and scored a corpus that was 0.27% loaded. That is what the published 10m result (`ingested_docs: 1`, `accuracy: 0.0`) is. `_sessions_from_plans` flattens each batch's turn groups into one session, so the existing chunk loop applies unchanged: 10m: 10 -> 5,265 documents, median 97,409 chars 100k: 170 documents (unchanged) 1m: 1,830 documents (unchanged) Only the split that was taking the fallback moves; every other split loads byte-identically, so existing results stay the runs they were. Total content is preserved to within 10,510 chars of 468,288,866 (0.002%), the difference being per-session formatting separators. This is necessary but not sufficient: at ~1,300 chars/s a 5-document batch is still ~375s, over the 300s `_await_operation` timeout, so the ingest rate has to come down too before a 10m score means anything. Claude-Session: https://claude.ai/code/session_01CJvQzszqZck74S5oC8dEAZ --- src/memory_bench/dataset/beam.py | 54 +++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/src/memory_bench/dataset/beam.py b/src/memory_bench/dataset/beam.py index 4f3dfd8..d336134 100644 --- a/src/memory_bench/dataset/beam.py +++ b/src/memory_bench/dataset/beam.py @@ -245,6 +245,54 @@ def categories(self, split: str) -> list[str] | None: def category_type(self, split: str, category: str) -> str: return "query" + @staticmethod + def _sessions_from_plans(chat: list) -> list[list[dict]]: + """Sessions for the BEAM-10M shape, which nests two levels deeper than the other splits. + + Most splits store `chat` as a list of sessions, each a list of turn dicts, and the loader + chunks those into bounded documents. BEAM-10M instead stores:: + + chat[i]["plan-N"][batch]["turns"] -> list of turn GROUPS, each a list of turn dicts + + Nothing there is a `list` at the top level, so it fell through to the "unusual structure" + branch and was emitted as ONE document per conversation -- a median of 47,280,119 chars, + 470x the `_MAX_DOC_CHARS` this same loader enforces on every other split, and for the same + reason that limit exists. + + The consequence was not a slow ingest but an impossible one. Retain cost is per-CALL rather + than per-byte (measured against a live API: 1 item 1.29s, 50 items 1.39s), and a backend + that serializes retains per document gets no parallelism when a 10-conversation split is + only 10 documents. Ingest ran at ~1,300 chars/s -- ~100 hours for the split -- so the + harness hit its 300s-per-operation timeout and scored a corpus 0.27% loaded. + + Flattening each batch's turn groups into one session restores the normal path: many + bounded documents, chunked by the same code as every other split. + """ + sessions: list[list[dict]] = [] + for element in chat: + if not isinstance(element, dict): + continue + # plan-1, plan-2, ... -- sorted so document ids are stable across runs. A plan can be + # null (plan-10 routinely is), which is why the list check is not an assertion. + for _plan, batches in sorted(element.items()): + if not isinstance(batches, list): + continue + for batch in batches: + if not isinstance(batch, dict): + continue + groups = batch.get("turns") + if not isinstance(groups, list): + continue + turns = [ + t + for group in groups + for t in (group if isinstance(group, list) else [group]) + if isinstance(t, dict) and "role" in t + ] + if turns: + sessions.append(turns) + return sessions + def load_documents( self, split: str, @@ -274,6 +322,8 @@ def load_documents( # Max ~100k chars per document to keep PostgreSQL happy. _MAX_DOC_CHARS = 100_000 sessions = [s for s in chat if isinstance(s, list)] + if not sessions: + sessions = self._sessions_from_plans(chat) if sessions: doc_idx = 0 for s_idx, session in enumerate(sessions): @@ -311,7 +361,9 @@ def load_documents( doc_idx += 1 chunk_start = chunk_end else: - # Fallback for unusual structures (e.g., BEAM-10M flat turns) + # Neither sessions nor plan batches -- keep the conversation whole rather than drop + # it. Reaching here means a structure this loader does not model, so there is + # nothing to chunk on. documents.append(Document( id=conv_id, content=self._format_chat(chat), From 767314dfc281e3f9acae7ebb966819d4e6888950 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 25 Aug 2026 11:19:18 +0200 Subject: [PATCH 3/3] beam: scale the async-operation timeout to the corpus instead of a fixed 300s Reverts the client-side document chunking from the previous commit. Splitting the conversation in the harness changes the document shape the benchmark measures, and it was aimed at the wrong problem: nothing actually stops a 47MB document. Measured in-cluster against a chunks-mode bank with consolidation off: async submit 100,000 chars 0.1s async submit 1,000,000 chars 0.3s async submit 5,000,000 chars 1.0s async submit 20,000,000 chars 3.9s The submit returns an operation id immediately and scales fine -- a 47MB document submits in about 9s. The work behind it then runs at ~31,000 chars/s, so a BEAM-10M conversation (median 47,280,119 chars) needs ~25 MINUTES to complete, and the whole 468,288,866-char split about 4.2 hours. `_await_operation` waited a fixed 300s. So it gave up five minutes in, logged a warning, and let the harness query a corpus that was still loading. That is what the published 10m result is: `ingested_docs: 1`, `accuracy: 0.0` -- not a failed run, a run that scored an empty bank. Two changes: - The timeout defaults to 7200s and is overridable with `AMB_OPERATION_TIMEOUT_S`, which covers the largest single document in any current split with room to spare. - Abandoning an operation is now an ERROR that says the score is invalid, not a warning followed by "continuing anyway". A run that silently scores a partially loaded corpus is worse than one that fails, because the number looks real. Claude-Session: https://claude.ai/code/session_01CJvQzszqZck74S5oC8dEAZ --- src/memory_bench/dataset/beam.py | 54 +--------------------------- src/memory_bench/memory/hindsight.py | 29 ++++++++++++--- 2 files changed, 25 insertions(+), 58 deletions(-) diff --git a/src/memory_bench/dataset/beam.py b/src/memory_bench/dataset/beam.py index d336134..4f3dfd8 100644 --- a/src/memory_bench/dataset/beam.py +++ b/src/memory_bench/dataset/beam.py @@ -245,54 +245,6 @@ def categories(self, split: str) -> list[str] | None: def category_type(self, split: str, category: str) -> str: return "query" - @staticmethod - def _sessions_from_plans(chat: list) -> list[list[dict]]: - """Sessions for the BEAM-10M shape, which nests two levels deeper than the other splits. - - Most splits store `chat` as a list of sessions, each a list of turn dicts, and the loader - chunks those into bounded documents. BEAM-10M instead stores:: - - chat[i]["plan-N"][batch]["turns"] -> list of turn GROUPS, each a list of turn dicts - - Nothing there is a `list` at the top level, so it fell through to the "unusual structure" - branch and was emitted as ONE document per conversation -- a median of 47,280,119 chars, - 470x the `_MAX_DOC_CHARS` this same loader enforces on every other split, and for the same - reason that limit exists. - - The consequence was not a slow ingest but an impossible one. Retain cost is per-CALL rather - than per-byte (measured against a live API: 1 item 1.29s, 50 items 1.39s), and a backend - that serializes retains per document gets no parallelism when a 10-conversation split is - only 10 documents. Ingest ran at ~1,300 chars/s -- ~100 hours for the split -- so the - harness hit its 300s-per-operation timeout and scored a corpus 0.27% loaded. - - Flattening each batch's turn groups into one session restores the normal path: many - bounded documents, chunked by the same code as every other split. - """ - sessions: list[list[dict]] = [] - for element in chat: - if not isinstance(element, dict): - continue - # plan-1, plan-2, ... -- sorted so document ids are stable across runs. A plan can be - # null (plan-10 routinely is), which is why the list check is not an assertion. - for _plan, batches in sorted(element.items()): - if not isinstance(batches, list): - continue - for batch in batches: - if not isinstance(batch, dict): - continue - groups = batch.get("turns") - if not isinstance(groups, list): - continue - turns = [ - t - for group in groups - for t in (group if isinstance(group, list) else [group]) - if isinstance(t, dict) and "role" in t - ] - if turns: - sessions.append(turns) - return sessions - def load_documents( self, split: str, @@ -322,8 +274,6 @@ def load_documents( # Max ~100k chars per document to keep PostgreSQL happy. _MAX_DOC_CHARS = 100_000 sessions = [s for s in chat if isinstance(s, list)] - if not sessions: - sessions = self._sessions_from_plans(chat) if sessions: doc_idx = 0 for s_idx, session in enumerate(sessions): @@ -361,9 +311,7 @@ def load_documents( doc_idx += 1 chunk_start = chunk_end else: - # Neither sessions nor plan batches -- keep the conversation whole rather than drop - # it. Reaching here means a structure this loader does not model, so there is - # nothing to chunk on. + # Fallback for unusual structures (e.g., BEAM-10M flat turns) documents.append(Document( id=conv_id, content=self._format_chat(chat), diff --git a/src/memory_bench/memory/hindsight.py b/src/memory_bench/memory/hindsight.py index d6eceab..0503fa1 100644 --- a/src/memory_bench/memory/hindsight.py +++ b/src/memory_bench/memory/hindsight.py @@ -158,8 +158,26 @@ def _create_bank(self, bank_id: str, force_reset: bool = True) -> None: pass self._client.create_bank(bank_id=bank_id, name=f"Benchmark Bank ({bank_id})", **kwargs) - async def _await_operation(self, client, bank_id: str, operation_id: str, max_wait_s: int = 300) -> None: - """Poll until an async retain operation completes (5-minute timeout).""" + async def _await_operation(self, client, bank_id: str, operation_id: str, max_wait_s: int | None = None) -> None: + """Poll until an async retain operation completes. + + The timeout has to scale with the corpus, not sit at a constant. A retain is submitted + asynchronously and returns an operation id immediately -- measured in-cluster, submitting a + 20,000,000-char document takes 3.9s -- but the work behind it runs at roughly 31,000 chars/s, + so a BEAM-10M conversation (median 47,280,119 chars) needs ~25 MINUTES to finish. + + At the old fixed 300s this method gave up five minutes in, logged a warning, and let the + harness query a corpus that was still loading. That is what produced the published 10m + result: `ingested_docs: 1`, `accuracy: 0.0`. The run did not fail, it scored an empty bank. + + `AMB_OPERATION_TIMEOUT_S` overrides it; the default is two hours, which covers the largest + single document in any current split with room to spare. + """ + if max_wait_s is None: + try: + max_wait_s = int(os.environ.get("AMB_OPERATION_TIMEOUT_S", "7200")) + except ValueError: + max_wait_s = 7200 from hindsight_client_api.api.operations_api import OperationsApi ops_api = OperationsApi(client._api_client) waited = 0 @@ -179,9 +197,10 @@ async def _await_operation(self, client, bank_id: str, operation_id: str, max_wa waited += 1 if waited >= max_wait_s: import logging - logging.getLogger(__name__).warning( - f"_await_operation timed out after {max_wait_s}s for bank={bank_id} op={operation_id} " - f"last_status={last_status!r}; continuing anyway." + logging.getLogger(__name__).error( + f"_await_operation GAVE UP after {max_wait_s}s for bank={bank_id} op={operation_id} " + f"last_status={last_status!r}. Ingestion is INCOMPLETE and any score from this run " + f"measures a partially loaded corpus, not the system. Raise AMB_OPERATION_TIMEOUT_S." ) # ── Bank creation (async) ─────────────────────────────────────────────────