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.

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")

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 (
+            connected_sources AS (
-                SELECT DISTINCT ue.entity_id
+                -- 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 (
+            connected_array AS (
-                SELECT DISTINCT other_ue.unit_id AS source_id
+                SELECT array_agg(source_id) AS source_ids FROM connected_sources
-                FROM source_entities se
-                JOIN {fq_table("unit_entities")} other_ue ON se.entity_id = other_ue.entity_id
             )
             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
+                (SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
-            FROM all_connected_sources cs
+            FROM {fq_table("memory_units")} mu, connected_array ca
-            JOIN {fq_table("memory_units")} mu
-                ON mu.source_memory_ids @> ARRAY[cs.source_id]
             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
             """,

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":