From ad2cf72aabd89bc85b56b6cb73490a290f062228 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Thu, 5 Mar 2026 10:15:40 +0100 Subject: [PATCH] perf: add GIN index on source_memory_ids for observation lookup (#485) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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. --- ...4d5e6f8_add_gin_index_source_memory_ids.py | 54 ++++++++ .../engine/search/link_expansion_retrieval.py | 25 ++-- hindsight-dev/benchmarks/perf/recall_perf.py | 116 ++++++++++++++++-- 3 files changed, 176 insertions(+), 19 deletions(-) create mode 100644 hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f8_add_gin_index_source_memory_ids.py diff --git a/hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f8_add_gin_index_source_memory_ids.py b/hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f8_add_gin_index_source_memory_ids.py new file mode 100644 index 00000000..e1aacc48 --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/a2b3c4d5e6f8_add_gin_index_source_memory_ids.py @@ -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") diff --git a/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py b/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py index 5720565f..a761546b 100644 --- a/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py +++ b/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py @@ -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 """, diff --git a/hindsight-dev/benchmarks/perf/recall_perf.py b/hindsight-dev/benchmarks/perf/recall_perf.py index e12d7e47..1b70f566 100644 --- a/hindsight-dev/benchmarks/perf/recall_perf.py +++ b/hindsight-dev/benchmarks/perf/recall_perf.py @@ -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":