From 61bf428ba9c97dca33165a5f93a9fce0d520ae52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 3 Mar 2026 14:52:48 +0100 Subject: [PATCH] perf: fetch all recall chunks in a single query instead of batched while-loop (#475) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the multi-round-trip while-loop in step 5.5 of recall_async with a single WHERE chunk_id = ANY($1) query covering all candidate chunk IDs. Token-budget accounting happens in Python after the single fetch. Measured on a 97K-unit / 98M-link bank (budget=HIGH, include_chunks, include_entities): p50: 1.209s → 0.611s (−49%) mean: 1.534s → 0.772s (−50%) p95: 3.366s → 2.316s (−31%) Also update recall_perf.py benchmark to use Budget.HIGH, include_chunks, include_entities, and a realistic mixed fact_type distribution. --- .../hindsight_api/engine/memory_engine.py | 81 +++++++------------ hindsight-dev/benchmarks/perf/recall_perf.py | 12 ++- 2 files changed, 42 insertions(+), 51 deletions(-) diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 1dc7362a..44d6488e 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -2648,66 +2648,47 @@ class MemoryEngine(MemoryEngineInterface): seen_chunk_ids.add(chunk_id) if chunk_ids_ordered: - # Estimate batch size based on retain_chunk_size * 2 (rough estimate) - # Chunk sizes vary per document, so we fetch in batches until budget is exhausted - bank_config = await self._config_resolver.resolve_full_config(bank_id, request_context) - estimated_batch_size = max(1, (max_chunk_tokens // bank_config.retain_chunk_size) * 2) - chunks_dict = {} encoding = _get_tiktoken_encoding() - chunk_offset = 0 - # Fetch chunks in batches until we run out of budget or chunks - while chunk_offset < len(chunk_ids_ordered) and total_chunk_tokens < max_chunk_tokens: - # Get next batch of chunk IDs - batch_chunk_ids = chunk_ids_ordered[chunk_offset : chunk_offset + estimated_batch_size] - chunk_offset += estimated_batch_size + # Fetch all candidate chunks in a single query. Token-budget accounting + # happens in Python after the fetch — one round-trip is always faster + # than multiple batched round-trips when the candidate set is large. + async with acquire_with_retry(pool) as conn: + chunks_rows = await conn.fetch( + f""" + SELECT chunk_id, chunk_text, chunk_index + FROM {fq_table("chunks")} + WHERE chunk_id = ANY($1::text[]) + """, + chunk_ids_ordered, + ) - # Fetch chunk data from database - async with acquire_with_retry(pool) as conn: - chunks_rows = await conn.fetch( - f""" - SELECT chunk_id, chunk_text, chunk_index - FROM {fq_table("chunks")} - WHERE chunk_id = ANY($1::text[]) - """, - batch_chunk_ids, - ) + chunks_lookup = {row["chunk_id"]: row for row in chunks_rows} - # Create a lookup dict for fast access (preserves order from batch_chunk_ids) - chunks_lookup = {row["chunk_id"]: row for row in chunks_rows} + # Process chunks in relevance order, respecting token budget + for chunk_id in chunk_ids_ordered: + if chunk_id not in chunks_lookup: + continue - # Process chunks in order, respecting token budget - for chunk_id in batch_chunk_ids: - if chunk_id not in chunks_lookup: - continue + row = chunks_lookup[chunk_id] + chunk_text = row["chunk_text"] + chunk_tokens = len(encoding.encode(chunk_text)) - row = chunks_lookup[chunk_id] - chunk_text = row["chunk_text"] - chunk_tokens = len(encoding.encode(chunk_text)) - - # Check if adding this chunk would exceed the limit - if total_chunk_tokens + chunk_tokens > max_chunk_tokens: - # Truncate the chunk to fit within the remaining budget - remaining_tokens = max_chunk_tokens - total_chunk_tokens - if remaining_tokens > 0: - # Truncate to remaining tokens - truncated_text = encoding.decode(encoding.encode(chunk_text)[:remaining_tokens]) - chunks_dict[chunk_id] = ChunkInfo( - chunk_text=truncated_text, chunk_index=row["chunk_index"], truncated=True - ) - total_chunk_tokens = max_chunk_tokens - # Budget exhausted - stop fetching more batches - break - else: + if total_chunk_tokens + chunk_tokens > max_chunk_tokens: + remaining_tokens = max_chunk_tokens - total_chunk_tokens + if remaining_tokens > 0: + truncated_text = encoding.decode(encoding.encode(chunk_text)[:remaining_tokens]) chunks_dict[chunk_id] = ChunkInfo( - chunk_text=chunk_text, chunk_index=row["chunk_index"], truncated=False + chunk_text=truncated_text, chunk_index=row["chunk_index"], truncated=True ) - total_chunk_tokens += chunk_tokens - - # If we hit the budget limit in this batch, stop fetching more batches - if total_chunk_tokens >= max_chunk_tokens: + total_chunk_tokens = max_chunk_tokens break + else: + chunks_dict[chunk_id] = ChunkInfo( + chunk_text=chunk_text, chunk_index=row["chunk_index"], truncated=False + ) + total_chunk_tokens += chunk_tokens # Step 6: Token budget filtering step_start = time.time() diff --git a/hindsight-dev/benchmarks/perf/recall_perf.py b/hindsight-dev/benchmarks/perf/recall_perf.py index 711580c8..c3c5cad0 100644 --- a/hindsight-dev/benchmarks/perf/recall_perf.py +++ b/hindsight-dev/benchmarks/perf/recall_perf.py @@ -503,9 +503,14 @@ def _make_fact_callback() -> tuple[Callable[[list[dict], str], Any], list[int]]: """ call_counter = [0] + # Realistic fact_type distribution matching production observations: + # ~60% world, ~30% experience, ~10% mental_model + _FACT_TYPE_CYCLE = (["world"] * 6 + ["experience"] * 3 + ["mental_model"] * 1) * 10 # 100-element cycle + def callback(messages: list[dict], scope: str) -> Any: if scope == "retain_extract_facts": idx = call_counter[0] % len(FACT_TEMPLATES) + fact_type = _FACT_TYPE_CYCLE[call_counter[0] % len(_FACT_TYPE_CYCLE)] call_counter[0] += 1 template = FACT_TEMPLATES[idx] fact_text = _fill_template(template) @@ -521,7 +526,7 @@ def _make_fact_callback() -> tuple[Callable[[list[dict], str], Any], list[int]]: "where": "N/A", "who": "N/A", "why": "N/A", - "fact_type": "world", + "fact_type": fact_type, "entities": entities, } ] @@ -736,12 +741,17 @@ async def cmd_benchmark(bank_id: str, query: str, iterations: int, concurrency: all_phase_timings: dict[str, list[float]] = {} async def recall_one() -> float: + from hindsight_api.engine.memory_engine import Budget + t0 = time.perf_counter() result = await engine.recall_async( bank_id=bank_id, query=query, + budget=Budget.HIGH, max_tokens=4096, enable_trace=True, + include_chunks=True, + include_entities=True, request_context=request_context, _quiet=True, )