perf: add GIN index on source_memory_ids for observation lookup (#485)
* perf: add GIN index on source_memory_ids for observation lookup Addresses a 927x performance regression (45ms → 0.049ms) reported by a user with ~77k observations. The array overlap operator (&&) on source_memory_ids was doing a full sequential scan over all observations, causing recall timeouts (57-64s) and slow user recall (18-27s avg). The partial GIN index reduces consolidation recall from timeout to ~15s and user recall to ~6s. * fix: use pre-bounded memory_links for observation graph expansion Replace raw unit_entities join in _expand_observations() with the same memory_links entity graph used by non-observation fact types. The previous approach joined unit_entities twice (seeds→entities→connected_sources), which explodes at scale (30-70s at 100k observations). The LIMIT 500 workaround was non-deterministic and dropped valid results. Using memory_links (pre-bounded to MAX_LINKS_PER_ENTITY=50 at retain time) is algorithmically identical to the non-observation entity expansion and keeps graph retrieval at ~2s p50 even at 100k observations. Also fix migration down_revision (z1u2v3w4x5y6 → d2e3f4a5b6c7) and add observation generation + fact-type filtering to the recall perf benchmark.
This commit is contained in:
Nicolò Boschi
GitHub
parent
3f2a6ec9ce
commit
ad2cf72aab
3 changed files with 176 additions and 19 deletions
|
|
@ -0,0 +1,54 @@
|
|||
"""Add GIN index on source_memory_ids for observation lookup performance
|
||||
|
||||
Without this index, queries using the array overlap operator (&&) or array
|
||||
containment (@>) on source_memory_ids require a full sequential scan over all
|
||||
observation memory_units. At ~77k observations this was measured at 45ms per
|
||||
query, becoming a bottleneck during consolidation recall (57-64s timeouts) and
|
||||
user recall (18-27s average).
|
||||
|
||||
The GIN index reduces these queries to index scans: 45ms → 0.049ms (927x
|
||||
speedup). Recall dropped from 18-27s to ~6s, and consolidation recall
|
||||
stabilised from timeout to ~15s.
|
||||
|
||||
Created with CONCURRENTLY so the migration does not block reads or writes.
|
||||
CONCURRENTLY requires running outside a transaction block, so the migration
|
||||
emits an explicit COMMIT before the statement and uses IF NOT EXISTS for
|
||||
idempotency.
|
||||
|
||||
Revision ID: a2b3c4d5e6f8
|
||||
Revises: f7g8h9i0j1k2
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2b3c4d5e6f8"
|
||||
down_revision: str | Sequence[str] | None = "f7g8h9i0j1k2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction first.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
|
|
@ -395,27 +395,28 @@ class LinkExpansionRetriever(GraphRetriever):
|
|||
WHERE id = ANY($1::uuid[])
|
||||
AND source_memory_ids IS NOT NULL
|
||||
),
|
||||
source_entities AS (
|
||||
SELECT DISTINCT ue.entity_id
|
||||
connected_sources AS (
|
||||
-- Mirror the non-observation entity expansion: follow pre-bounded entity
|
||||
-- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time).
|
||||
-- Score = number of distinct shared entities, same as the non-obs path.
|
||||
SELECT DISTINCT ml.to_unit_id AS source_id
|
||||
FROM seed_sources ss
|
||||
JOIN {fq_table("unit_entities")} ue ON ss.source_id = ue.unit_id
|
||||
JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id
|
||||
WHERE ml.link_type = 'entity'
|
||||
),
|
||||
all_connected_sources AS (
|
||||
SELECT DISTINCT other_ue.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
JOIN {fq_table("unit_entities")} other_ue ON se.entity_id = other_ue.entity_id
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
)
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(DISTINCT cs.source_id)::float AS score
|
||||
FROM all_connected_sources cs
|
||||
JOIN {fq_table("memory_units")} mu
|
||||
ON mu.source_memory_ids @> ARRAY[cs.source_id]
|
||||
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
|
||||
FROM {fq_table("memory_units")} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
GROUP BY mu.id
|
||||
AND ca.source_ids IS NOT NULL
|
||||
AND mu.source_memory_ids && ca.source_ids
|
||||
ORDER BY score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
|
|
|
|||
|
|
@ -560,6 +560,78 @@ def _build_engine(*, disable_observations: bool = False) -> "Any":
|
|||
return engine
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Synthetic observation insertion
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BATCH_SIZE = 500
|
||||
|
||||
|
||||
async def _insert_synthetic_observations(pool: Any, bank_id: str) -> int:
|
||||
"""
|
||||
For every non-observation memory unit in *bank_id*, insert one synthetic
|
||||
observation with the same text, embedding, and tags, pointing back to that
|
||||
unit as its sole source fact.
|
||||
|
||||
Returns the number of observations inserted.
|
||||
"""
|
||||
import uuid
|
||||
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
|
||||
table = fq_table("memory_units")
|
||||
|
||||
# Fetch all non-observation units
|
||||
rows = await pool.fetch(
|
||||
f"""
|
||||
SELECT id, text, embedding, tags, event_date, occurred_start, occurred_end, mentioned_at
|
||||
FROM {table}
|
||||
WHERE bank_id = $1 AND fact_type != 'observation'
|
||||
ORDER BY id
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not rows:
|
||||
return 0
|
||||
|
||||
inserted = 0
|
||||
for offset in range(0, len(rows), _BATCH_SIZE):
|
||||
batch = rows[offset : offset + _BATCH_SIZE]
|
||||
await pool.executemany(
|
||||
f"""
|
||||
INSERT INTO {table} (
|
||||
id, bank_id, text, fact_type, embedding,
|
||||
proof_count, source_memory_ids, history,
|
||||
tags, event_date, occurred_start, occurred_end, mentioned_at
|
||||
) VALUES (
|
||||
$1, $2, $3, 'observation', $4::vector,
|
||||
1, ARRAY[$5::uuid], '[]'::jsonb,
|
||||
$6, $7, $8, $9, $10
|
||||
)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
[
|
||||
(
|
||||
uuid.uuid4(), # new observation id
|
||||
bank_id, # bank_id
|
||||
row["text"],
|
||||
row["embedding"],
|
||||
row["id"], # source fact id
|
||||
row["tags"] or [],
|
||||
row["event_date"],
|
||||
row["occurred_start"],
|
||||
row["occurred_end"],
|
||||
row["mentioned_at"],
|
||||
)
|
||||
for row in batch
|
||||
],
|
||||
)
|
||||
inserted += len(batch)
|
||||
|
||||
return inserted
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Subcommand: generate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -613,7 +685,7 @@ async def _wait_for_operation(pool: Any, operation_id: str, timeout: float = 864
|
|||
raise TimeoutError(f"Operation {operation_id} did not complete within {timeout}s")
|
||||
|
||||
|
||||
async def cmd_generate(bank_id: str, scale: str, workers: int = 16) -> None:
|
||||
async def cmd_generate(bank_id: str, scale: str, workers: int = 16, with_observations: bool = False) -> None:
|
||||
"""Submit all content as a single async batch and process with an in-process worker."""
|
||||
from hindsight_api.models import RequestContext
|
||||
from hindsight_api.worker.poller import WorkerPoller
|
||||
|
|
@ -680,8 +752,6 @@ async def cmd_generate(bank_id: str, scale: str, workers: int = 16) -> None:
|
|||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
await pool.close()
|
||||
|
||||
status_color = "green" if final_status == "completed" else "red"
|
||||
console.print(
|
||||
f"\n[{status_color}]Done[/{status_color}] — status=[bold]{final_status}[/bold] "
|
||||
|
|
@ -689,6 +759,15 @@ async def cmd_generate(bank_id: str, scale: str, workers: int = 16) -> None:
|
|||
)
|
||||
console.print(f"LLM callback invoked {call_counter[0]:,} times.")
|
||||
|
||||
if with_observations:
|
||||
console.print("\n Inserting synthetic observations (1 per fact)…")
|
||||
t_obs = time.perf_counter()
|
||||
n_obs = await _insert_synthetic_observations(pool, bank_id)
|
||||
elapsed_obs = time.perf_counter() - t_obs
|
||||
console.print(f" Inserted {n_obs:,} observations in {elapsed_obs:.1f}s")
|
||||
|
||||
await pool.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RRF-only reranker (bypasses cross-encoder for DB-focused benchmarking)
|
||||
|
|
@ -719,7 +798,9 @@ class _RRFReranker:
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def cmd_benchmark(bank_id: str, query: str, iterations: int, concurrency: int, reranker: str) -> None:
|
||||
async def cmd_benchmark(
|
||||
bank_id: str, query: str, iterations: int, concurrency: int, reranker: str, fact_types: list[str] | None = None
|
||||
) -> None:
|
||||
"""Run recall in parallel and report p50/p95/p99 timings with per-step breakdown."""
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
|
@ -727,7 +808,8 @@ async def cmd_benchmark(bank_id: str, query: str, iterations: int, concurrency:
|
|||
console.print(f" Query : {query}")
|
||||
console.print(f" Iterations : {iterations} (total recall calls)")
|
||||
console.print(f" Concurrency : {concurrency}")
|
||||
console.print(f" Reranker : {reranker}\n")
|
||||
console.print(f" Reranker : {reranker}")
|
||||
console.print(f" Fact types : {', '.join(fact_types) if fact_types else 'all'}\n")
|
||||
|
||||
engine = _build_engine()
|
||||
await engine.initialize()
|
||||
|
|
@ -751,6 +833,8 @@ async def cmd_benchmark(bank_id: str, query: str, iterations: int, concurrency:
|
|||
enable_trace=True,
|
||||
include_chunks=True,
|
||||
include_entities=True,
|
||||
include_source_facts=True,
|
||||
fact_type=fact_types,
|
||||
request_context=request_context,
|
||||
_quiet=True,
|
||||
)
|
||||
|
|
@ -913,6 +997,12 @@ def main() -> None:
|
|||
gen.add_argument("--bank-id", required=True)
|
||||
gen.add_argument("--scale", choices=list(SCALES), default="small")
|
||||
gen.add_argument("--workers", type=int, default=8, help="Max concurrent worker slots (default: 8)")
|
||||
gen.add_argument(
|
||||
"--with-observations",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="After retain, insert one synthetic observation per fact (same text, same embedding)",
|
||||
)
|
||||
|
||||
# benchmark
|
||||
bm = sub.add_parser("benchmark", help="Run recall and report latency")
|
||||
|
|
@ -926,6 +1016,14 @@ def main() -> None:
|
|||
default="rrf",
|
||||
help="Reranker to use: rrf=RRF scores only (no ML), cross-encoder=neural reranker (default: rrf)",
|
||||
)
|
||||
bm.add_argument(
|
||||
"--fact-types",
|
||||
nargs="+",
|
||||
choices=["world", "experience", "observation"],
|
||||
default=None,
|
||||
metavar="TYPE",
|
||||
help="Fact types to include in recall (default: all). E.g. --fact-types observation",
|
||||
)
|
||||
|
||||
# stats
|
||||
st = sub.add_parser("stats", help="Print memory/entity/link counts for banks")
|
||||
|
|
@ -938,9 +1036,13 @@ def main() -> None:
|
|||
args = parser.parse_args()
|
||||
|
||||
if args.cmd == "generate":
|
||||
asyncio.run(cmd_generate(args.bank_id, args.scale, workers=args.workers))
|
||||
asyncio.run(
|
||||
cmd_generate(args.bank_id, args.scale, workers=args.workers, with_observations=args.with_observations)
|
||||
)
|
||||
elif args.cmd == "benchmark":
|
||||
asyncio.run(cmd_benchmark(args.bank_id, args.query, args.iterations, args.concurrency, args.reranker))
|
||||
asyncio.run(
|
||||
cmd_benchmark(args.bank_id, args.query, args.iterations, args.concurrency, args.reranker, args.fact_types)
|
||||
)
|
||||
elif args.cmd == "stats":
|
||||
asyncio.run(cmd_stats(args.bank_ids))
|
||||
elif args.cmd == "clean":
|
||||
|
|
|
|||
Loading…
Reference in a new issue