holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

perf: fetch all recall chunks in a single query instead of batched while-loop (#475)

Browse source 
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.
This commit is contained in:
Nicolò Boschi 2026-03-03 14:52:48 +01:00 committed by GitHub
parent 73ef99e7b1
commit 61bf428ba9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 42 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
hindsight-api/hindsight_api/engine/memory_engine.py
View file

@ -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()

12
hindsight-dev/benchmarks/perf/recall_perf.py
View file

@ -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,
)