perf: fetch all recall chunks in a single query instead of batched while-loop (#475)
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
GitHub
parent
73ef99e7b1
commit
61bf428ba9
2 changed files with 42 additions and 51 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in a new issue