SIGN IN SIGN UP

fix(retain): bound retain's memory by a budget instead of by the document (#3756) (#3763)

* fix(retain): bound retain's memory by a budget instead of by the document (#3756)

Retaining one large document held state proportional to the document rather
than to a working set. The reported peak was blamed on embeddings held as
list[float]; measured, that was not where it went. Two whole-document
operations dominated, both running before a single fact or embedding existed.

Peak Python bytes allocated, 45 MB body (tracemalloc, order-independent):

                          before     after
  sizing   count_tokens   384.7 MB     9.6 MB   (windowed; flat at any size)
  chunking split_text     200.8 MB     0.0 MB   (streamed; flat at any size)

The facts themselves were never the problem: the streaming pipeline already
bounds them to retain_chunk_batch_size chunks, ~1700 facts, ~21 MB.

Four changes, and a bound that is now explicit:

1. count_tokens_windowed() sizes a body a megabyte at a time instead of
   building one boxed int per token. Every caller compares against a batch
   budget or logs the number, so the one-token-per-window boundary error
   (45 tokens in 11.6M) is unobservable.

2. iter_chunks() streams chunks instead of materialising them, including a
   lazy re-implementation of the RecursiveCharacterTextSplitter configuration
   retain used. langchain leaves the runtime dependencies with it, kept only
   as the reference the differential test diffs against.

   Boundaries are content hashes for delta retain and chunk_ids by index, so
   test_chunking_streams.py pins the output against langchain directly rather
   than trusting the rewrite.

3. An oversized item's body is Memory Defense screened AND content-hashed
   once, not once per slice. That hash was the last piece of work scaling
   with (sub-batches x document size): ~0.9s repeated ~1,200 times for a
   45 MB body, about 18 minutes spent re-deriving one value.

4. ProcessedFact.embedding is array("f") rather than list[float] - 1,616
   bytes against 12,344 for 384 dims. float32 is what pgvector stores, so
   the rounding just happens one step earlier and the stored bytes match.

RetainMemoryBudget then turns the pipeline's bound from a count of chunks
into a ceiling in bytes (HINDSIGHT_API_RETAIN_MEMORY_BUDGET_MB, default 128).
A count is only a memory bound if chunks cost a predictable amount, and they
do not - a chunk carries however many facts the extractor found in it. The
producer reserves a chunk's estimated cost before queueing and the consumer
releases it once written, so over budget extraction waits for the write path
instead of growing.

What still costs a copy of the document is the document: the submitted string
and the sub-batch slices cut from it (45.7 MB for a 45 MB body). Nothing in
the front half is superlinear any more, and nothing but those two scales at
all.

Speed is unchanged or better - sizing 2.42s -> 2.46s, chunking 0.29s -> 0.16s
on a 45 MB body.

Measured by hindsight-dev/benchmarks/perf/retain_memory.py, which reports
tracemalloc rather than RSS: RSS cannot attribute an allocation to the code
that made it (arenas are mapped on first touch and reused silently), and
reading it that way is what made this issue's original diagnosis wrong twice.

* test(retain): pin the new chunker against the one it replaced, over a corpus (#3756)

The streaming chunker was diffed against langchain only for the plain-text
splitter, on hand-written strings. That leaves the paths retain actually takes
on real input — JSON conversations, JSONL logs, the structured-limit branch —
covered by nothing but the assumption that they were untouched.

This runs the whole pre-#3756 `chunk_text`, copied verbatim and still calling
langchain, against the live one over 21 document shapes x 7 chunk sizes, the
same again across three structured-chunk-size settings, and 200 seeded random
documents. It also asserts `iter_chunks` and `chunk_text` agree everywhere and
that every emitted chunk re-chunks to itself (#2301's invariant).

Why the copy is verbatim and must stay that way: the moment the reference
shares a code path with the live implementation the comparison proves nothing.
`test_the_legacy_reference_really_is_the_old_implementation` pins the reference
to langchain's own output so that cannot happen quietly.

Verified non-vacuous by mutation: shifting the packer's budget by one character
fails 41 of the 126 tests. (An exactly-budget piece re-packs to itself whichever
branch it takes, so `<` vs `<=` is an equivalent mutation rather than an
uncaught one — noted in the corpus helper so nobody re-derives it.)

* fix(retain): four regressions CI caught in the memory-budget work (#3756)

1. estimate_chunk_bytes() called len() on a null context and failed the retain
   it was sizing. ProcessedFact.context is annotated str, but a converted file
   upload retains without one and puts None there — every file-retain test
   returned HTTP 500. A budget heuristic must never be able to break the
   operation, so every string now goes through a None-tolerant helper.

2. RetainMemoryBudget accepted a non-int limit and hung. A MagicMock config
   makes every comparison inside it truthy, so the producer waited for room
   that could never be reported — test_consumer_failure_cancels_in_flight_
   extractions timed out at 300s instead of failing. It now rejects a limit
   that is not an int, so a mis-mocked config is a clear TypeError rather than
   a wedged retain, and the test models the field it needs.

3. Two async-batch tests stubbed chunk_text, which retain no longer calls;
   they patch iter_chunks now, which covers both forms.

4. One mapping test compared ProcessedFact.embedding against a float list.
   It is packed since this branch — compare the unpacked values, since what
   the assertion is about is which facts survived, not the container.
N
Nicolò Boschi committed
e65973b1d97f7853a68dde7597c5e6256a7f1e8e
Parent: 85bd23d
Committed by GitHub <noreply@github.com> on 8/24/2026, 1:03:07 PM