fix(performance): improve recall and retain performance on large banks (#469)

hindsight-api/hindsight_api/alembic/versions/b3c4d5e6f7g8_add_temporal_date_indexes.py

@@ -0,0 +1,68 @@
+"""Add partial indexes on memory_units temporal date fields for fast temporal retrieval
+
+Revision ID: b3c4d5e6f7g8
+Revises: c1a2b3d4e5f6
+Create Date: 2026-03-02
+
+The temporal retrieval entry-point query filters memory_units by occurred_start,
+occurred_end, and mentioned_at using OR conditions. Without dedicated indexes the
+planner falls back to a sequential scan of all bank rows after applying the
+(bank_id, fact_type) index, then re-checks each date field.
+
+These three partial indexes give the planner bitmap-index scan options for the
+three most common date predicates, dramatically reducing the row set before any
+embedding computation is required.
+
+All indexes are created CONCURRENTLY so the migration does not block writes on
+memory_units during production deployments. CONCURRENTLY requires running outside
+a transaction block; see migrations.py for how this is handled safely.
+"""
+
+from collections.abc import Sequence
+
+from alembic import context, op
+
+revision: str = "b3c4d5e6f7g8"
+down_revision: str | Sequence[str] | None = "c1a2b3d4e5f6"
+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()
+    # Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
+    op.execute("COMMIT")
+    op.execute(
+        f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
+        f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
+        f"WHERE occurred_start IS NOT NULL"
+    )
+    # Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
+    op.execute("COMMIT")
+    op.execute(
+        f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
+        f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
+        f"WHERE occurred_end IS NOT NULL"
+    )
+    # Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
+    op.execute("COMMIT")
+    op.execute(
+        f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
+        f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
+        f"WHERE mentioned_at 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_bank_mentioned_at")
+    op.execute("COMMIT")
+    op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
+    op.execute("COMMIT")
+    op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")

hindsight-api/hindsight_api/alembic/versions/c1a2b3d4e5f6_enable_pg_trgm_and_entities_trgm_index.py

@@ -0,0 +1,46 @@
+"""Enable pg_trgm extension and add GIN trigram index on entities.canonical_name
+
+Revision ID: c1a2b3d4e5f6
+Revises: b4c5d6e7f8a9
+Create Date: 2026-03-02
+
+Index is created CONCURRENTLY so the migration does not block writes on entities
+during production deployments. CONCURRENTLY requires running outside a transaction
+block; see migrations.py for how this is handled safely.
+"""
+
+from collections.abc import Sequence
+
+from alembic import context, op
+
+revision: str = "c1a2b3d4e5f6"
+down_revision: str | Sequence[str] | None = "b4c5d6e7f8a9"
+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:
+    # pg_trgm ships with every standard PostgreSQL installation as a contrib module.
+    # It enables fast similarity lookups via GIN indexes, used for entity name matching.
+    op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
+
+    schema = _get_schema_prefix()
+    # GIN index on canonical_name enables sub-millisecond trigram similarity queries
+    # (% operator, similarity()) instead of full-table scans across all bank entities.
+    op.execute("COMMIT")
+    op.execute(
+        f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
+        f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
+    )
+
+
+def downgrade() -> None:
+    schema = _get_schema_prefix()
+    op.execute("COMMIT")
+    op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
+    # Note: not dropping pg_trgm extension as other indexes may depend on it

hindsight-api/hindsight_api/alembic/versions/d2e3f4a5b6c7_add_memory_links_expansion_indexes.py

@@ -0,0 +1,83 @@
+"""Add covering and composite indexes to speed up link expansion graph retrieval.
+
+Two indexes target the two bottlenecks identified by EXPLAIN ANALYZE on a 17M-row
+memory_links table:
+
+1. idx_memory_links_to_type_weight  (to_unit_id, link_type, weight DESC)
+   The semantic incoming direction — finding facts that consider seeds as their
+   nearest neighbour — currently hits an expensive BitmapAnd of two separate
+   bitmap scans (to_unit_id bitmap ∩ link_type bitmap).  A composite index
+   on (to_unit_id, link_type) turns this into a single index scan and reduces
+   latency from ~36 ms to < 5 ms per query.
+
+2. idx_memory_links_entity_covering  (from_unit_id) INCLUDE (to_unit_id, entity_id)
+   WHERE link_type = 'entity'
+   The entity co-occurrence expansion uses COUNT(DISTINCT ml.entity_id) and
+   joins on ml.to_unit_id.  Without a covering index the planner must read
+   ~2 500 heap pages to fetch entity_id and to_unit_id after the bitmap index
+   scan, adding ~230 ms of random I/O.  INCLUDE adds those two columns to the
+   index leaf pages so the entire query can be served from the index (index-only
+   scan), eliminating the heap reads entirely.
+   Partial index (WHERE link_type = 'entity') keeps index size ~40 % smaller.
+
+Both indexes are created with CONCURRENTLY so the migration does not block
+concurrent reads or writes on memory_links.  CONCURRENTLY requires running
+outside a transaction block, so the migration emits an explicit COMMIT before
+each statement and uses IF NOT EXISTS for idempotency.
+
+Revision ID: d2e3f4a5b6c7
+Revises: b3c4d5e6f7g8
+Create Date: 2026-03-02
+"""
+
+from collections.abc import Sequence
+
+from alembic import context, op
+
+revision: str = "d2e3f4a5b6c7"
+down_revision: str | Sequence[str] | None = "b3c4d5e6f7g8"
+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, then issue each CONCURRENTLY
+    # statement in its own implicit autocommit transaction.
+    # IF NOT EXISTS makes each statement idempotent if the migration is retried.
+
+    # Index for the semantic *incoming* direction in link_expansion_retrieval.py.
+    # Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
+    # with a single composite index scan.
+    op.execute("COMMIT")
+    op.execute(
+        f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
+        f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
+    )
+
+    # Covering index for entity co-occurrence expansion.
+    # Enables an index-only scan: entity_id and to_unit_id are read from the
+    # index leaf pages instead of the heap, eliminating ~2 500 random heap-page
+    # reads per expansion query.
+    op.execute("COMMIT")
+    op.execute(
+        f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
+        f"ON {schema}memory_links(from_unit_id) "
+        f"INCLUDE (to_unit_id, entity_id) "
+        f"WHERE link_type = 'entity'"
+    )
+
+
+def downgrade() -> None:
+    schema = _get_schema_prefix()
+    op.execute("COMMIT")
+    op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
+    op.execute("COMMIT")
+    op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")

hindsight-api/hindsight_api/api/http.py

@@ -2385,148 +2385,32 @@ def _register_routes(app: FastAPI):
     ):
         """Get statistics about memory nodes and links for a memory bank."""
         try:
-            # Authenticate and set tenant schema
+            stats = await app.state.memory.get_bank_stats(bank_id, request_context=request_context)
-            await app.state.memory._authenticate_tenant(request_context)
+            nodes_by_type = stats["node_counts"]
-            if app.state.memory._operation_validator:
+            links_by_type = stats["link_counts"]
-                from hindsight_api.extensions import BankReadContext
+            links_by_fact_type = stats["link_counts_by_fact_type"]
-
+            links_breakdown: dict[str, dict[str, int]] = {}
-                ctx = BankReadContext(bank_id=bank_id, operation="get_bank_stats", request_context=request_context)
+            for row in stats["link_breakdown"]:
-                await app.state.memory._validate_operation(
+                ft = row["fact_type"]
-                    app.state.memory._operation_validator.validate_bank_read(ctx)
+                if ft not in links_breakdown:
-                )
+                    links_breakdown[ft] = {}
-            pool = await app.state.memory._get_pool()
+                links_breakdown[ft][row["link_type"]] = row["count"]
-            async with acquire_with_retry(pool) as conn:
+            ops = stats["operations"]
-                # Get node counts by fact_type
-                node_stats = await conn.fetch(
-                    f"""
-                    SELECT fact_type, COUNT(*) as count
-                    FROM {fq_table("memory_units")}
-                    WHERE bank_id = $1
-                    GROUP BY fact_type
-                    """,
-                    bank_id,
-                )
-
-                # Get link counts by link_type
-                link_stats = await conn.fetch(
-                    f"""
-                    SELECT ml.link_type, COUNT(*) as count
-                    FROM {fq_table("memory_links")} ml
-                    JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                    WHERE mu.bank_id = $1
-                    GROUP BY ml.link_type
-                    """,
-                    bank_id,
-                )
-
-                # Get link counts by fact_type (from nodes)
-                link_fact_type_stats = await conn.fetch(
-                    f"""
-                    SELECT mu.fact_type, COUNT(*) as count
-                    FROM {fq_table("memory_links")} ml
-                    JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                    WHERE mu.bank_id = $1
-                    GROUP BY mu.fact_type
-                    """,
-                    bank_id,
-                )
-
-                # Get link counts by fact_type AND link_type
-                link_breakdown_stats = await conn.fetch(
-                    f"""
-                    SELECT mu.fact_type, ml.link_type, COUNT(*) as count
-                    FROM {fq_table("memory_links")} ml
-                    JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                    WHERE mu.bank_id = $1
-                    GROUP BY mu.fact_type, ml.link_type
-                    """,
-                    bank_id,
-                )
-
-                # Get pending and failed operations counts
-                ops_stats = await conn.fetch(
-                    f"""
-                    SELECT status, COUNT(*) as count
-                    FROM {fq_table("async_operations")}
-                    WHERE bank_id = $1
-                    GROUP BY status
-                    """,
-                    bank_id,
-                )
-                ops_by_status = {row["status"]: row["count"] for row in ops_stats}
-                pending_operations = ops_by_status.get("pending", 0)
-                failed_operations = ops_by_status.get("failed", 0)
-
-                # Get document count
-                doc_count_result = await conn.fetchrow(
-                    f"""
-                    SELECT COUNT(*) as count
-                    FROM {fq_table("documents")}
-                    WHERE bank_id = $1
-                    """,
-                    bank_id,
-                )
-                total_documents = doc_count_result["count"] if doc_count_result else 0
-
-                # Get consolidation stats from memory-level tracking
-                consolidation_stats = await conn.fetchrow(
-                    f"""
-                    SELECT
-                        MAX(consolidated_at) as last_consolidated_at,
-                        COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending
-                    FROM {fq_table("memory_units")}
-                    WHERE bank_id = $1
-                    """,
-                    bank_id,
-                )
-                last_consolidated_at = consolidation_stats["last_consolidated_at"] if consolidation_stats else None
-                pending_consolidation = consolidation_stats["pending"] if consolidation_stats else 0
-
-                # Count total observations (consolidated knowledge)
-                observation_count_result = await conn.fetchrow(
-                    f"""
-                    SELECT COUNT(*) as count
-                    FROM {fq_table("memory_units")}
-                    WHERE bank_id = $1 AND fact_type = 'observation'
-                    """,
-                    bank_id,
-                )
-                total_observations = observation_count_result["count"] if observation_count_result else 0
-
-                # Format results
-                nodes_by_type = {row["fact_type"]: row["count"] for row in node_stats}
-                links_by_type = {row["link_type"]: row["count"] for row in link_stats}
-                links_by_fact_type = {row["fact_type"]: row["count"] for row in link_fact_type_stats}
-
-                # Build detailed breakdown: {fact_type: {link_type: count}}
-                links_breakdown = {}
-                for row in link_breakdown_stats:
-                    fact_type = row["fact_type"]
-                    link_type = row["link_type"]
-                    count = row["count"]
-                    if fact_type not in links_breakdown:
-                        links_breakdown[fact_type] = {}
-                    links_breakdown[fact_type][link_type] = count
-
-                total_nodes = sum(nodes_by_type.values())
-                total_links = sum(links_by_type.values())
-
             return BankStatsResponse(
                 bank_id=bank_id,
-                    total_nodes=total_nodes,
+                total_nodes=sum(nodes_by_type.values()),
-                    total_links=total_links,
+                total_links=sum(links_by_type.values()),
-                    total_documents=total_documents,
+                total_documents=stats["total_documents"],
                 nodes_by_fact_type=nodes_by_type,
                 links_by_link_type=links_by_type,
                 links_by_fact_type=links_by_fact_type,
                 links_breakdown=links_breakdown,
-                    pending_operations=pending_operations,
+                pending_operations=ops.get("pending", 0),
-                    failed_operations=failed_operations,
+                failed_operations=ops.get("failed", 0),
-                    last_consolidated_at=(last_consolidated_at.isoformat() if last_consolidated_at else None),
+                last_consolidated_at=stats["last_consolidated_at"],
-                    pending_consolidation=pending_consolidation,
+                pending_consolidation=stats["pending_consolidation"],
-                    total_observations=total_observations,
+                total_observations=stats["total_observations"],
             )
-
         except OperationValidationError as e:
             raise HTTPException(status_code=e.status_code, detail=e.reason)
         except (AuthenticationError, HTTPException):

hindsight-api/hindsight_api/config.py

@@ -260,6 +260,7 @@ ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE"
 ENV_RETAIN_MISSION = "HINDSIGHT_API_RETAIN_MISSION"
 ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
 ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
+ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
 ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
 ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
 
@@ -416,6 +417,7 @@ RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom")  # Allowed extraction
 DEFAULT_RETAIN_MISSION = None  # Declarative spec of what to retain (injected into any extraction mode)
 DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None  # Custom extraction guidelines (only used when mode="custom")
 DEFAULT_RETAIN_BATCH_TOKENS = 10_000  # ~40KB of text  # Max chars per sub-batch for async retain auto-splitting
+DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram"  # "full" or "trigram"
 DEFAULT_RETAIN_BATCH_ENABLED = False  # Use LLM Batch API for fact extraction (only when async=True)
 DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60  # Batch API polling interval in seconds
 
@@ -662,6 +664,7 @@ class HindsightConfig:
     retain_batch_tokens: int
     retain_batch_enabled: bool
     retain_batch_poll_interval_seconds: int
+    retain_entity_lookup: str  # "full" or "trigram"
 
     # File storage (static - server-level only)
     file_storage_type: str  # "native" (PostgreSQL) or "s3" (S3-compatible)
@@ -1084,6 +1087,7 @@ class HindsightConfig:
             retain_mission=os.getenv(ENV_RETAIN_MISSION) or DEFAULT_RETAIN_MISSION,
             retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS,
             retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
+            retain_entity_lookup=os.getenv(ENV_RETAIN_ENTITY_LOOKUP, DEFAULT_RETAIN_ENTITY_LOOKUP),
             retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
             == "true",
             retain_batch_poll_interval_seconds=int(

hindsight-api/hindsight_api/engine/db_utils.py

@@ -20,6 +20,7 @@ RETRYABLE_EXCEPTIONS = (
     asyncpg.exceptions.InterfaceError,
     asyncpg.exceptions.ConnectionDoesNotExistError,
     asyncpg.exceptions.TooManyConnectionsError,
+    asyncpg.exceptions.DeadlockDetectedError,
     OSError,
     ConnectionError,
     asyncio.TimeoutError,

hindsight-api/hindsight_api/engine/entity_resolver.py

@@ -5,6 +5,10 @@ Uses spaCy for entity extraction and implements resolution logic
 to disambiguate entities across memory units.
 """
 
+import asyncio
+import logging
+from collections import defaultdict
+from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from difflib import SequenceMatcher
 
@@ -14,6 +18,42 @@ from .db_utils import acquire_with_retry
 from .memory_engine import fq_table
 from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
 
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _EntityToCreate:
+    """An entity that needs to be inserted (no matching candidate found)."""
+
+    idx: int
+    name: str
+    event_date: datetime | None
+
+
+@dataclass
+class _EntityStat:
+    """Stat accumulation entry for a resolved entity (post-transaction update)."""
+
+    entity_id: str
+    event_date: datetime | None
+
+
+@dataclass
+class _EntityStatAgg:
+    """Aggregated stats used when flushing pending updates."""
+
+    count: int = 0
+    max_date: datetime | None = None
+
+
+@dataclass
+class _CooccurrencePair:
+    """A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
+
+    entity_id_1: str
+    entity_id_2: str
+
+
 # Load spaCy model (singleton)
 _nlp = None
 
@@ -23,14 +63,90 @@ class EntityResolver:
     Resolves entities to canonical IDs with disambiguation.
     """
 
-    def __init__(self, pool: asyncpg.Pool):
+    def __init__(self, pool: asyncpg.Pool, entity_lookup: str = "full"):
         """
         Initialize entity resolver.
 
         Args:
             pool: asyncpg connection pool
+            entity_lookup: Lookup strategy — "full" loads all bank entities then
+                matches in Python; "trigram" uses pg_trgm GIN index to fetch only
+                similar candidates per entity name (much faster for large banks).
         """
         self.pool = pool
+        self.entity_lookup = entity_lookup
+        # Keyed by asyncio task id so concurrent retain batches never mix their
+        # pending updates.  flush_pending_stats() pops only the calling task's items.
+        self._pending_stats: dict[int, list[_EntityStat]] = {}
+        self._pending_cooccurrences: dict[int, list[_CooccurrencePair]] = {}
+
+    def _task_key(self) -> int:
+        """Return a unique key for the current asyncio task (or 0 for non-task context)."""
+        task = asyncio.current_task()
+        return id(task) if task is not None else 0
+
+    async def flush_pending_stats(self) -> None:
+        """
+        Flush accumulated entity stats and co-occurrence counts for the current task.
+
+        Must be called AFTER the retain transaction commits.  Pops only the items
+        accumulated by the calling asyncio task so concurrent retain batches never
+        flush each other's uncommitted entity IDs.
+        """
+        if self.pool is None:
+            return
+
+        key = self._task_key()
+        stats = self._pending_stats.pop(key, [])
+        cooccurrences = self._pending_cooccurrences.pop(key, [])
+
+        if not stats and not cooccurrences:
+            return
+
+        async with acquire_with_retry(self.pool) as conn:
+            if stats:
+                # Aggregate: sum counts and find max date per entity_id.
+                agg: dict[str, _EntityStatAgg] = defaultdict(_EntityStatAgg)
+                for s in stats:
+                    entry = agg[s.entity_id]
+                    entry.count += 1
+                    if s.event_date is not None:
+                        entry.max_date = s.event_date if entry.max_date is None else max(entry.max_date, s.event_date)
+
+                # Sort by entity_id so all concurrent workers acquire row locks in
+                # the same order — prevents circular lock dependencies (deadlocks).
+                rows = sorted((eid, a.count, a.max_date) for eid, a in agg.items())
+                await conn.executemany(
+                    f"""
+                    UPDATE {fq_table("entities")} SET
+                        mention_count = mention_count + $2,
+                        last_seen     = GREATEST(last_seen, $3)
+                    WHERE id = $1::uuid
+                    """,
+                    rows,
+                )
+
+            if cooccurrences:
+                # Aggregate: count occurrences per (entity_id_1, entity_id_2) pair.
+                coo_agg: dict[tuple[str, str], int] = {}
+                for c in cooccurrences:
+                    pair = (c.entity_id_1, c.entity_id_2)
+                    coo_agg[pair] = coo_agg.get(pair, 0) + 1
+
+                now = datetime.now(UTC)
+                # Sort by (entity_id_1, entity_id_2) for consistent lock ordering.
+                await conn.executemany(
+                    f"""
+                    INSERT INTO {fq_table("entity_cooccurrences")}
+                        (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
+                    VALUES ($1, $2, $3, $4)
+                    ON CONFLICT (entity_id_1, entity_id_2)
+                    DO UPDATE SET
+                        cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + EXCLUDED.cooccurrence_count,
+                        last_cooccurred    = GREATEST({fq_table("entity_cooccurrences")}.last_cooccurred, EXCLUDED.last_cooccurred)
+                    """,
+                    sorted((e1, e2, count, now) for (e1, e2), count in coo_agg.items()),
+                )
 
     @staticmethod
     def _build_labels_lookup(entity_labels: list | None) -> set[str]:
@@ -85,6 +201,14 @@ class EntityResolver:
         unit_event_date,
         taxonomy_lookup: set[str] | None = None,
     ) -> list[str]:
+        if self.entity_lookup == "trigram":
+            return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
+        return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
+
+    async def _resolve_entities_batch_full(
+        self, conn, bank_id: str, entities_data: list[dict], unit_event_date
+    ) -> list[str]:
+        """Original strategy: load all bank entities then match in Python."""
         # Query ALL candidates for this bank
         all_entities = await conn.fetch(
             f"""
@@ -148,12 +272,103 @@ class EntityResolver:
                     matching.append((ent_id, canonical_name, metadata, last_seen, mention_count))
             all_candidates[entity_text] = matching
 
+        return await self._resolve_from_candidates(
+            conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
+        )
+
+    async def _resolve_entities_batch_trigram(
+        self, conn, bank_id: str, entities_data: list[dict], unit_event_date
+    ) -> list[str]:
+        """
+        Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
+
+        Instead of loading all bank entities (O(N)), uses a GIN trigram index to fetch
+        only the small set of candidates that are textually similar to each input name.
+        Reduces DB data transfer from 165K rows to ~5-20 rows per entity.
+        """
+        entity_texts = list(set(e["text"] for e in entities_data))
+
+        # Fetch candidates for all unique entity texts in a single batched query.
+        # The trigram % operator uses the GIN index; the substring conditions cover
+        # exact prefix/suffix matches that trigrams might miss at low similarity.
+        rows = await conn.fetch(
+            f"""
+            SELECT DISTINCT ON (e.id)
+                e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
+                q.query_text
+            FROM unnest($2::text[]) AS q(query_text)
+            JOIN {fq_table("entities")} e ON (
+                e.bank_id = $1
+                AND (
+                    e.canonical_name % q.query_text
+                    OR LOWER(e.canonical_name) LIKE '%' || LOWER(q.query_text) || '%'
+                    OR LOWER(q.query_text) LIKE '%' || LOWER(e.canonical_name) || '%'
+                )
+            )
+            """,
+            bank_id,
+            entity_texts,
+        )
+
+        # Group candidates by query_text
+        all_candidates: dict[str, list] = {t: [] for t in entity_texts}
+        candidate_ids: set = set()
+        for row in rows:
+            query_text = row["query_text"]
+            all_candidates[query_text].append(
+                (row["id"], row["canonical_name"], row["metadata"], row["last_seen"], row["mention_count"])
+            )
+            candidate_ids.add(row["id"])
+
+        # Fetch co-occurrences only for the candidate entities (not all bank entities)
+        cooccurrence_map: dict[str, set[str]] = {}
+        if candidate_ids:
+            candidate_id_list = list(candidate_ids)
+            cooc_rows = await conn.fetch(
+                f"""
+                SELECT ec.entity_id_1, ec.entity_id_2
+                FROM {fq_table("entity_cooccurrences")} ec
+                WHERE ec.entity_id_1 = ANY($1::uuid[])
+                   OR ec.entity_id_2 = ANY($1::uuid[])
+                """,
+                candidate_id_list,
+            )
+            # Build name lookup for co-occurrence mapping
+            id_to_name = {
+                row["id"]: row["canonical_name"].lower()
+                for cands in all_candidates.values()
+                for row in [{"id": c[0], "canonical_name": c[1]} for c in cands]
+            }
+            for row in cooc_rows:
+                eid1, eid2 = row["entity_id_1"], row["entity_id_2"]
+                if eid1 not in cooccurrence_map:
+                    cooccurrence_map[eid1] = set()
+                if eid2 not in cooccurrence_map:
+                    cooccurrence_map[eid2] = set()
+                if eid2 in id_to_name:
+                    cooccurrence_map[eid1].add(id_to_name[eid2])
+                if eid1 in id_to_name:
+                    cooccurrence_map[eid2].add(id_to_name[eid1])
+
+        return await self._resolve_from_candidates(
+            conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
+        )
+
+    async def _resolve_from_candidates(
+        self,
+        conn,
+        bank_id: str,
+        entities_data: list[dict],
+        unit_event_date,
+        all_candidates: dict[str, list],
+        cooccurrence_map: dict[str, set[str]],
+    ) -> list[str]:
+        """Shared scoring + upsert logic used by both lookup strategies."""
+
         # Resolve each entity using pre-fetched candidates
         entity_ids = [None] * len(entities_data)
-        entities_to_update = []  # (entity_id, event_date)
+        entities_to_update: list[_EntityStat] = []
-        entities_to_create = []  # (idx, entity_data, event_date)
+        entities_to_create: list[_EntityToCreate] = []
-
-        taxonomy_lookup = taxonomy_lookup or set()
 
         for idx, entity_data in enumerate(entities_data):
             entity_text = entity_data["text"]
@@ -161,16 +376,11 @@ class EntityResolver:
             # Use per-entity date if available, otherwise fall back to batch-level date
             entity_event_date = entity_data.get("event_date", unit_event_date)
 
-            # Taxonomy entities: skip fuzzy matching, use exact canonical name
-            if taxonomy_lookup and entity_text.lower() in taxonomy_lookup:
-                entities_to_create.append((idx, entity_data, entity_event_date))
-                continue
-
             candidates = all_candidates.get(entity_text, [])
 
             if not candidates:
                 # Will create new entity
-                entities_to_create.append((idx, entity_data, entity_event_date))
+                entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
                 continue
 
             # Score candidates
@@ -214,73 +424,83 @@ class EntityResolver:
 
             if best_score > threshold:
                 entity_ids[idx] = best_candidate
-                entities_to_update.append((best_candidate, entity_event_date))
+                entities_to_update.append(_EntityStat(entity_id=best_candidate, event_date=entity_event_date))
             else:
-                entities_to_create.append((idx, entity_data, entity_event_date))
+                entities_to_create.append(
-
+                    _EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
-        # Batch update existing entities
-        if entities_to_update:
-            await conn.executemany(
-                f"""
-                UPDATE {fq_table("entities")} SET
-                    mention_count = mention_count + 1,
-                    last_seen = $2
-                WHERE id = $1::uuid
-                """,
-                entities_to_update,
                 )
 
-        # Batch create new entities using COPY + INSERT for maximum speed
+        # Existing entities: IDs already known from the candidate SELECT above.
-        # This handles duplicates via ON CONFLICT and returns all IDs
+        # No in-transaction UPDATE — mention_count/last_seen are stats deferred to
+        # flush_pending_stats() which the orchestrator calls after the transaction.
+        pending: list[_EntityStat] = list(entities_to_update)
+
+        # New entities: INSERT with DO NOTHING to avoid row locks on concurrent races.
+        # ON CONFLICT DO NOTHING returns nothing for rows that conflicted; we handle
+        # that rare case with a fallback SELECT.
         if entities_to_create:
-            # Group entities by canonical name (lowercase) to handle duplicates within batch
+            # Group by lowercase name — deduplicate within the batch.
-            # For duplicates, we only insert once and reuse the ID, but track the count
+            @dataclass
-            unique_entities = {}  # lowercase_name -> (entity_data, event_date, [indices])
+            class _NameGroup:
-            for idx, entity_data, event_date in entities_to_create:
+                name: str
-                name_lower = entity_data["text"].lower()
+                event_date: datetime | None
-                if name_lower not in unique_entities:
+                indices: list[int] = field(default_factory=list)
-                    unique_entities[name_lower] = (entity_data, event_date, [idx])
-                else:
-                    # Same entity appears multiple times - add index to list
-                    unique_entities[name_lower][2].append(idx)
 
-            # Batch insert unique entities and get their IDs
+            groups: dict[str, _NameGroup] = {}
-            # Use a single query with unnest for speed
+            for e in entities_to_create:
-            entity_names = []
+                name_lower = e.name.lower()
-            entity_dates = []
+                if name_lower not in groups:
-            entity_counts = []  # Track how many times each entity appears in this batch
+                    groups[name_lower] = _NameGroup(name=e.name, event_date=e.event_date)
-            indices_map = []  # Maps result index -> list of original indices
+                groups[name_lower].indices.append(e.idx)
 
-            for name_lower, (entity_data, event_date, indices) in unique_entities.items():
+            # Sort by lowercase name for deterministic ordering.
-                entity_names.append(entity_data["text"])
+            sorted_groups = sorted(groups.items())
-                entity_dates.append(event_date)
+            entity_names = [g.name for _, g in sorted_groups]
-                entity_counts.append(len(indices))  # Count of occurrences in this batch
+            entity_dates = [g.event_date for _, g in sorted_groups]
-                indices_map.append(indices)
 
-            # Batch INSERT ... ON CONFLICT with RETURNING
+            # INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
-            # Uses the batch count for mention_count instead of always 1
+            inserted_rows = await conn.fetch(
-            rows = await conn.fetch(
                 f"""
                 INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
-                SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), cnt
+                SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 1
-                FROM unnest($2::text[], $3::timestamptz[], $4::int[]) AS t(name, event_date, cnt)
+                FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
                 ON CONFLICT (bank_id, LOWER(canonical_name))
-                DO UPDATE SET
+                DO NOTHING
-                    mention_count = {fq_table("entities")}.mention_count + EXCLUDED.mention_count,
+                RETURNING id, LOWER(canonical_name) AS name_lower
-                    last_seen = EXCLUDED.last_seen
-                RETURNING id
                 """,
                 bank_id,
                 entity_names,
                 entity_dates,
-                entity_counts,
             )
+            id_by_name: dict[str, str] = {row["name_lower"]: row["id"] for row in inserted_rows}
 
-            # Map returned IDs back to original indices
+            # Fallback SELECT for names that conflicted (another worker won the race).
-            for result_idx, row in enumerate(rows):
+            missing = [n for n, _ in sorted_groups if n not in id_by_name]
-                entity_id = row["id"]
+            if missing:
-                for original_idx in indices_map[result_idx]:
+                existing_rows = await conn.fetch(
+                    f"""
+                    SELECT id, LOWER(canonical_name) AS name_lower
+                    FROM {fq_table("entities")}
+                    WHERE bank_id = $1 AND LOWER(canonical_name) = ANY($2::text[])
+                    """,
+                    bank_id,
+                    missing,
+                )
+                for row in existing_rows:
+                    id_by_name[row["name_lower"]] = row["id"]
+
+            # Assign entity IDs back and queue for post-txn stats flush.
+            for name_lower, g in sorted_groups:
+                entity_id = id_by_name.get(name_lower)
+                if entity_id:
+                    for original_idx in g.indices:
                         entity_ids[original_idx] = entity_id
+                    pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
+
+        # Accumulate into the resolver's pending list; the orchestrator flushes
+        # these with await entity_resolver.flush_pending_stats() after the txn.
+        key = self._task_key()
+        self._pending_stats.setdefault(key, []).extend(pending)
 
         return entity_ids
 
@@ -566,19 +786,14 @@ class EntityResolver:
                         entity_id_1, entity_id_2 = entity_id_2, entity_id_1
                     cooccurrence_pairs.add((entity_id_1, entity_id_2))
 
-        # Batch update co-occurrences
+        # Accumulate co-occurrence pairs for post-transaction flush.
+        # The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
+        # row-level lock contention (ON CONFLICT DO UPDATE inside a long transaction
+        # serialises concurrent writers on popular entity pairs).
         if cooccurrence_pairs:
-            now = datetime.now(UTC)
+            key = self._task_key()
-            await conn.executemany(
+            self._pending_cooccurrences.setdefault(key, []).extend(
-                f"""
+                _CooccurrencePair(entity_id_1=e1, entity_id_2=e2) for e1, e2 in cooccurrence_pairs
-                INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
-                VALUES ($1, $2, $3, $4)
-                ON CONFLICT (entity_id_1, entity_id_2)
-                DO UPDATE SET
-                    cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
-                    last_cooccurred = EXCLUDED.last_cooccurred
-                """,
-                [(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs],
             )
 
     async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]:

hindsight-api/hindsight_api/engine/llm_wrapper.py

@@ -503,6 +503,14 @@ class LLMProvider:
 
             return result
 
+    def set_response_callback(self, fn: Any) -> None:
+        """Set a callback invoked on each call() instead of the fixed mock response."""
+        if self.provider == "mock":
+            from .providers.mock_llm import MockLLM
+
+            if isinstance(self._provider_impl, MockLLM):
+                self._provider_impl.set_response_callback(fn)
+
     def set_mock_response(self, response: Any) -> None:
         """Set the response to return from mock calls."""
         # Backward compatibility: Store in both wrapper and provider implementation

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -357,6 +357,7 @@ class MemoryEngine(MemoryEngineInterface):
         self._db_command_timeout = db_command_timeout if db_command_timeout is not None else config.db_command_timeout
         self._db_acquire_timeout = db_acquire_timeout if db_acquire_timeout is not None else config.db_acquire_timeout
         self._run_migrations = run_migrations
+        self._retain_entity_lookup = config.retain_entity_lookup
 
         # Initialize entity resolver (will be created in initialize())
         self.entity_resolver = None
@@ -1340,8 +1341,11 @@ class MemoryEngine(MemoryEngineInterface):
             timeout=self._db_acquire_timeout,  # Connection acquisition timeout (seconds)
         )
 
-        # Initialize entity resolver with pool
+        # Initialize entity resolver with pool and configured lookup strategy
-        self.entity_resolver = EntityResolver(self._pool)
+        self.entity_resolver = EntityResolver(
+            self._pool,
+            entity_lookup=self._retain_entity_lookup,
+        )
 
         # Initialize config resolver for hierarchical configuration
         from ..config_resolver import ConfigResolver
@@ -1485,110 +1489,6 @@ class MemoryEngine(MemoryEngineInterface):
         # Could check if day is significant (not 1st or 15th) and include it
         return f"{month_name} {year}"
 
-    async def _find_duplicate_facts_batch(
-        self,
-        conn,
-        bank_id: str,
-        texts: list[str],
-        embeddings: list[list[float]],
-        event_date: datetime,
-        time_window_hours: int = 24,
-        similarity_threshold: float = 0.95,
-    ) -> list[bool]:
-        """
-        Check which facts are duplicates using semantic similarity + temporal window.
-
-        For each new fact, checks if a semantically similar fact already exists
-        within the time window. Uses pgvector cosine similarity for efficiency.
-
-        Args:
-            conn: Database connection
-            bank_id: bank IDentifier
-            texts: List of fact texts to check
-            embeddings: Corresponding embeddings
-            event_date: Event date for temporal filtering
-            time_window_hours: Hours before/after event_date to search (default: 24)
-            similarity_threshold: Minimum cosine similarity to consider duplicate (default: 0.95)
-
-        Returns:
-            List of booleans - True if fact is a duplicate (should skip), False if new
-        """
-        if not texts:
-            return []
-
-        # Handle edge cases where event_date is at datetime boundaries
-        try:
-            time_lower = event_date - timedelta(hours=time_window_hours)
-        except OverflowError:
-            time_lower = datetime.min
-        try:
-            time_upper = event_date + timedelta(hours=time_window_hours)
-        except OverflowError:
-            time_upper = datetime.max
-
-        # Fetch ALL existing facts in time window ONCE (much faster than N queries)
-        import time as time_mod
-
-        fetch_start = time_mod.time()
-        existing_facts = await conn.fetch(
-            f"""
-            SELECT id, text, embedding
-            FROM {fq_table("memory_units")}
-            WHERE bank_id = $1
-              AND event_date BETWEEN $2 AND $3
-            """,
-            bank_id,
-            time_lower,
-            time_upper,
-        )
-
-        # If no existing facts, nothing is duplicate
-        if not existing_facts:
-            return [False] * len(texts)
-
-        # Compute similarities in Python (vectorized with numpy)
-        is_duplicate = []
-
-        # Convert existing embeddings to numpy for faster computation
-        embedding_arrays = []
-        for row in existing_facts:
-            raw_emb = row["embedding"]
-            # Handle different pgvector formats
-            if isinstance(raw_emb, str):
-                # Parse string format: "[1.0, 2.0, ...]"
-                import json
-
-                emb = np.array(json.loads(raw_emb), dtype=np.float32)
-            elif isinstance(raw_emb, (list, tuple)):
-                emb = np.array(raw_emb, dtype=np.float32)
-            else:
-                # Try direct conversion
-                emb = np.array(raw_emb, dtype=np.float32)
-            embedding_arrays.append(emb)
-
-        if not embedding_arrays:
-            existing_embeddings = np.array([])
-        elif len(embedding_arrays) == 1:
-            # Single embedding: reshape to (1, dim)
-            existing_embeddings = embedding_arrays[0].reshape(1, -1)
-        else:
-            # Multiple embeddings: vstack
-            existing_embeddings = np.vstack(embedding_arrays)
-
-        comp_start = time_mod.time()
-        for embedding in embeddings:
-            # Compute cosine similarity with all existing facts
-            emb_array = np.array(embedding)
-            # Cosine similarity = 1 - cosine distance
-            # For normalized vectors: cosine_sim = dot product
-            similarities = np.dot(existing_embeddings, emb_array)
-
-            # Check if any existing fact is too similar
-            max_similarity = np.max(similarities) if len(similarities) > 0 else 0
-            is_duplicate.append(max_similarity > similarity_threshold)
-
-        return is_duplicate
-
     def retain(
         self,
         bank_id: str,
@@ -1939,7 +1839,6 @@ class MemoryEngine(MemoryEngineInterface):
                     llm_config=self._retain_llm_config,
                     entity_resolver=self.entity_resolver,
                     format_date_fn=self._format_readable_date,
-                    duplicate_checker_fn=self._find_duplicate_facts_batch,
                     bank_id=bank_id,
                     contents_dicts=contents,
                     document_id=document_id,
@@ -2575,6 +2474,11 @@ class MemoryEngine(MemoryEngineInterface):
                         "temporal_count": len(temporal_results) if temporal_results else 0,
                     },
                 )
+                # Also expose each retrieval method as its own phase so
+                # benchmarks can pinpoint which sub-query drives latency.
+                for _method, _dur in aggregated_timings.items():
+                    if _dur > 0:
+                        tracer.add_phase_metric(f"retrieval_{_method}", _dur)
 
             # Step 3: Merge with RRF
             step_start = time.time()
@@ -5117,31 +5021,8 @@ class MemoryEngine(MemoryEngineInterface):
                 bank_id,
             )
 
-            # Get link counts by link_type
+            # Single query for all link stats — avoids triple join on memory_links (can be 21M+ rows).
-            link_stats = await conn.fetch(
+            # link_counts and link_counts_by_fact_type are derived in Python from the breakdown.
-                f"""
-                SELECT ml.link_type, COUNT(*) as count
-                FROM {fq_table("memory_links")} ml
-                JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                WHERE mu.bank_id = $1
-                GROUP BY ml.link_type
-                """,
-                bank_id,
-            )
-
-            # Get link counts by fact_type (from nodes)
-            link_fact_type_stats = await conn.fetch(
-                f"""
-                SELECT mu.fact_type, COUNT(*) as count
-                FROM {fq_table("memory_links")} ml
-                JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                WHERE mu.bank_id = $1
-                GROUP BY mu.fact_type
-                """,
-                bank_id,
-            )
-
-            # Get link counts by fact_type AND link_type
             link_breakdown_stats = await conn.fetch(
                 f"""
                 SELECT mu.fact_type, ml.link_type, COUNT(*) as count
@@ -5153,7 +5034,14 @@ class MemoryEngine(MemoryEngineInterface):
                 bank_id,
             )
 
-            # Get pending and failed operations counts
+            link_counts: dict[str, int] = {}
+            link_counts_by_fact_type: dict[str, int] = {}
+            for row in link_breakdown_stats:
+                link_counts[row["link_type"]] = link_counts.get(row["link_type"], 0) + row["count"]
+                link_counts_by_fact_type[row["fact_type"]] = (
+                    link_counts_by_fact_type.get(row["fact_type"], 0) + row["count"]
+                )
+
             ops_stats = await conn.fetch(
                 f"""
                 SELECT status, COUNT(*) as count
@@ -5163,17 +5051,39 @@ class MemoryEngine(MemoryEngineInterface):
                 """,
                 bank_id,
             )
+            doc_count_row = await conn.fetchrow(
+                f"SELECT COUNT(*) as count FROM {fq_table('documents')} WHERE bank_id = $1",
+                bank_id,
+            )
+            consolidation_row = await conn.fetchrow(
+                f"""
+                SELECT
+                    MAX(consolidated_at) as last_consolidated_at,
+                    COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending
+                FROM {fq_table("memory_units")}
+                WHERE bank_id = $1
+                """,
+                bank_id,
+            )
+
+            node_counts = {row["fact_type"]: row["count"] for row in node_stats}
+            ops_by_status = {row["status"]: row["count"] for row in ops_stats}
+            last_consolidated_at = consolidation_row["last_consolidated_at"] if consolidation_row else None
 
             return {
                 "bank_id": bank_id,
-                "node_counts": {row["fact_type"]: row["count"] for row in node_stats},
+                "node_counts": node_counts,
-                "link_counts": {row["link_type"]: row["count"] for row in link_stats},
+                "link_counts": link_counts,
-                "link_counts_by_fact_type": {row["fact_type"]: row["count"] for row in link_fact_type_stats},
+                "link_counts_by_fact_type": link_counts_by_fact_type,
                 "link_breakdown": [
                     {"fact_type": row["fact_type"], "link_type": row["link_type"], "count": row["count"]}
                     for row in link_breakdown_stats
                 ],
-                "operations": {row["status"]: row["count"] for row in ops_stats},
+                "operations": ops_by_status,
+                "total_documents": doc_count_row["count"] if doc_count_row else 0,
+                "last_consolidated_at": last_consolidated_at.isoformat() if last_consolidated_at else None,
+                "pending_consolidation": consolidation_row["pending"] if consolidation_row else 0,
+                "total_observations": node_counts.get("observation", 0),
             }
 
     async def get_entity(

hindsight-api/hindsight_api/engine/providers/mock_llm.py

@@ -6,6 +6,7 @@ without making actual API calls to external LLM services.
 """
 
 import logging
+from collections.abc import Callable
 from typing import Any
 
 from ..llm_interface import LLMInterface
@@ -66,6 +67,7 @@ class MockLLM(LLMInterface):
         self._mock_calls: list[dict] = []
         self._mock_response: Any = None
         self._mock_exception: Exception | None = None
+        self._response_callback: Callable[[list[dict], str], Any] | None = None
 
     async def verify_connection(self) -> None:
         """
@@ -147,7 +149,9 @@ class MockLLM(LLMInterface):
         )
 
         # Return mock response
-        if self._mock_response is not None:
+        if self._response_callback is not None:
+            result = self._response_callback(messages, scope)
+        elif self._mock_response is not None:
             result = self._mock_response
         elif response_format is not None:
             # Try to create a minimal valid instance of the response format
@@ -214,7 +218,15 @@ class MockLLM(LLMInterface):
 
         span_recorder = get_span_recorder()
 
-        if self._mock_response is not None:
+        if self._response_callback is not None:
+            cb_result = self._response_callback(messages, scope)
+            if isinstance(cb_result, LLMToolCallResult):
+                result = cb_result
+            else:
+                result = LLMToolCallResult(
+                    content=str(cb_result) if cb_result is not None else "mock response", finish_reason="stop"
+                )
+        elif self._mock_response is not None:
             if isinstance(self._mock_response, LLMToolCallResult):
                 result = self._mock_response
             elif isinstance(self._mock_response, list):
@@ -258,6 +270,16 @@ class MockLLM(LLMInterface):
         """Clean up resources (no-op for mock provider)."""
         pass
 
+    def set_response_callback(self, fn: Callable[[list[dict], str], Any]) -> None:
+        """
+        Set a callback invoked on each call() instead of _mock_response.
+
+        The callback receives (messages, scope) and returns the response.
+        Useful for returning different responses per call (e.g., cycling
+        through a corpus in a benchmark).
+        """
+        self._response_callback = fn
+
     def set_mock_response(self, response: Any) -> None:
         """
         Set the response to return from mock calls.

hindsight-api/hindsight_api/engine/query_analyzer.py

@@ -92,11 +92,17 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
         self._search_dates = None
 
     def load(self) -> None:
-        """Load dateparser (lazy import)."""
+        """Load dateparser and warm up internal data structures.
+
+        Triggers the real initialization cost (regex tables, timezone data) at
+        load time so the first actual recall doesn't pay the cold-start penalty.
+        """
         if self._search_dates is None:
             from dateparser.search import search_dates
 
             self._search_dates = search_dates
+            # Warm up: fire a dummy call to trigger lazy-loaded internal tables.
+            self._search_dates("today")
 
     def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
         """

hindsight-api/hindsight_api/engine/retain/__init__.py

@@ -5,7 +5,6 @@ This package contains modular components for the retain operation:
 - types: Type definitions for retain pipeline
 - fact_extraction: Extract facts from content
 - embedding_processing: Augment texts and generate embeddings
-- deduplication: Check for duplicate facts
 - entity_processing: Process and resolve entities
 - link_creation: Create temporal, semantic, entity, and causal links
 - chunk_storage: Handle chunk storage
@@ -14,7 +13,6 @@ This package contains modular components for the retain operation:
 
 from . import (
     chunk_storage,
-    deduplication,
     embedding_processing,
     entity_processing,
     fact_extraction,
@@ -35,7 +33,6 @@ __all__ = [
     # Modules
     "fact_extraction",
     "embedding_processing",
-    "deduplication",
     "entity_processing",
     "link_creation",
     "chunk_storage",

hindsight-api/hindsight_api/engine/retain/deduplication.py

@@ -1,85 +0,0 @@
-"""
-Deduplication logic for retain pipeline.
-
-Checks for duplicate facts using semantic similarity and temporal proximity.
-"""
-
-import logging
-from collections import defaultdict
-from datetime import UTC
-
-from .types import ProcessedFact
-
-logger = logging.getLogger(__name__)
-
-
-async def check_duplicates_batch(conn, bank_id: str, facts: list[ProcessedFact], duplicate_checker_fn) -> list[bool]:
-    """
-    Check which facts are duplicates using batched time-window queries.
-
-    Groups facts by 12-hour time buckets to efficiently check for duplicates
-    within a 24-hour window.
-
-    Args:
-        conn: Database connection
-        bank_id: Bank identifier
-        facts: List of ProcessedFact objects to check
-        duplicate_checker_fn: Async function(conn, bank_id, texts, embeddings, date, time_window_hours)
-                              that returns List[bool] indicating duplicates
-
-    Returns:
-        List of boolean flags (same length as facts) indicating if each fact is a duplicate
-    """
-    if not facts:
-        return []
-
-    # Group facts by event_date (rounded to 12-hour buckets) for efficient batching
-    time_buckets = defaultdict(list)
-    for idx, fact in enumerate(facts):
-        # Use occurred_start if available, otherwise use mentioned_at
-        # For deduplication purposes, we need a time reference
-        fact_date = fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at
-
-        # Defensive: if both are None (shouldn't happen), use now()
-        if fact_date is None:
-            from datetime import datetime
-
-            fact_date = datetime.now(UTC)
-
-        # Round to 12-hour bucket to group similar times
-        bucket_key = fact_date.replace(hour=(fact_date.hour // 12) * 12, minute=0, second=0, microsecond=0)
-        time_buckets[bucket_key].append((idx, fact))
-
-    # Process each bucket in batch
-    all_is_duplicate = [False] * len(facts)
-
-    for bucket_date, bucket_items in time_buckets.items():
-        indices = [item[0] for item in bucket_items]
-        texts = [item[1].fact_text for item in bucket_items]
-        embeddings = [item[1].embedding for item in bucket_items]
-
-        # Check duplicates for this time bucket
-        dup_flags = await duplicate_checker_fn(conn, bank_id, texts, embeddings, bucket_date, time_window_hours=24)
-
-        # Map results back to original indices
-        for idx, is_dup in zip(indices, dup_flags):
-            all_is_duplicate[idx] = is_dup
-
-    return all_is_duplicate
-
-
-def filter_duplicates(facts: list[ProcessedFact], is_duplicate_flags: list[bool]) -> list[ProcessedFact]:
-    """
-    Filter out duplicate facts based on duplicate flags.
-
-    Args:
-        facts: List of ProcessedFact objects
-        is_duplicate_flags: Boolean flags indicating which facts are duplicates
-
-    Returns:
-        List of non-duplicate facts
-    """
-    if len(facts) != len(is_duplicate_flags):
-        raise ValueError(f"Mismatch between facts ({len(facts)}) and flags ({len(is_duplicate_flags)})")
-
-    return [fact for fact, is_dup in zip(facts, is_duplicate_flags) if not is_dup]

hindsight-api/hindsight_api/engine/retain/embedding_utils.py

@@ -41,10 +41,9 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
         List of embeddings in same order as input texts
     """
     try:
-        # Run embeddings in thread pool to avoid blocking event loop
         loop = asyncio.get_event_loop()
         embeddings = await loop.run_in_executor(
-            None,  # Use default thread pool
+            None,
             embeddings_backend.encode,
             texts,
         )

hindsight-api/hindsight_api/engine/retain/link_utils.py

@@ -498,14 +498,13 @@ async def create_temporal_links_batch_per_fact(
             # Batch inserts to avoid timeout on large batches
             BATCH_SIZE = 1000
             for batch_start in range(0, len(links), BATCH_SIZE):
-                batch = links[batch_start : batch_start + BATCH_SIZE]
                 await conn.executemany(
                     f"""
                     INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
                     VALUES ($1, $2, $3, $4, $5)
                     ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
                     """,
-                    batch,
+                    links[batch_start : batch_start + BATCH_SIZE],
                 )
             _log(log_buffer, f"      [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
 
@@ -553,81 +552,45 @@ async def create_semantic_links_batch(
 
         import numpy as np
 
-        # Fetch ALL existing units with embeddings in ONE query
+        # Use pgvector ANN search (HNSW index) for each new unit instead of fetching
-        fetch_start = time_mod.time()
+        # all existing embeddings into Python. At large scale (100K+ units) the old
-        all_existing = await conn.fetch(
+        # approach would transfer 100K × 384 floats (~150 MB) per retain call; the
-            f"""
+        # ANN query completes in <5 ms and transfers only top_k rows.
-            SELECT id, embedding
+        ann_start = time_mod.time()
-            FROM {fq_table("memory_units")}
-            WHERE bank_id = $1
-              AND embedding IS NOT NULL
-              AND id::text != ALL($2)
-            """,
-            bank_id,
-            unit_ids,
-        )
-        _log(
-            log_buffer,
-            f"      [8.1] Fetch {len(all_existing)} existing embeddings (1 query): {time_mod.time() - fetch_start:.3f}s",
-        )
-
-        # Convert to numpy for vectorized similarity computation
-        compute_start = time_mod.time()
         all_links = []
 
-        if all_existing:
+        # Build UUID exclude list once for all ANN queries
-            # Convert existing embeddings to numpy array
+        import uuid as uuid_mod
-            existing_ids = [str(row["id"]) for row in all_existing]
-            # Stack embeddings as 2D array: (num_embeddings, embedding_dim)
-            embedding_arrays = []
-            for row in all_existing:
-                raw_emb = row["embedding"]
-                # Handle different pgvector formats
-                if isinstance(raw_emb, str):
-                    # Parse string format: "[1.0, 2.0, ...]"
-                    import json
 
-                    emb = np.array(json.loads(raw_emb), dtype=np.float32)
+        exclude_uuids = [uuid_mod.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids]
-                elif isinstance(raw_emb, (list, tuple)):
-                    emb = np.array(raw_emb, dtype=np.float32)
-                else:
-                    # Try direct conversion (works for numpy arrays, pgvector objects, etc.)
-                    emb = np.array(raw_emb, dtype=np.float32)
 
-                # Ensure it's 1D
-                if emb.ndim != 1:
-                    raise ValueError(f"Expected 1D embedding, got shape {emb.shape}")
-                embedding_arrays.append(emb)
-
-            if not embedding_arrays:
-                existing_embeddings = np.array([])
-            elif len(embedding_arrays) == 1:
-                # Single embedding: reshape to (1, dim)
-                existing_embeddings = embedding_arrays[0].reshape(1, -1)
-            else:
-                # Multiple embeddings: vstack
-                existing_embeddings = np.vstack(embedding_arrays)
-
-            # For each new unit, compute similarities with ALL existing units
         for unit_id, new_embedding in zip(unit_ids, embeddings):
-                new_emb_array = np.array(new_embedding)
+            emb_str = str(list(new_embedding) if not isinstance(new_embedding, list) else new_embedding)
+            rows = await conn.fetch(
+                f"""
+                SELECT id::text,
+                       1 - (embedding <=> $1::vector) AS similarity
+                FROM {fq_table("memory_units")}
+                WHERE bank_id = $2
+                  AND embedding IS NOT NULL
+                  AND id != ALL($3::uuid[])
+                ORDER BY embedding <=> $1::vector
+                LIMIT $4
+                """,
+                emb_str,
+                bank_id,
+                exclude_uuids,
+                top_k,
+            )
+            for row in rows:
+                sim = float(min(1.0, max(0.0, row["similarity"])))
+                if sim >= threshold:
+                    all_links.append((unit_id, str(row["id"]), "semantic", sim, None))
 
-                # Compute cosine similarities (dot product for normalized vectors)
+        _log(
-                similarities = np.dot(existing_embeddings, new_emb_array)
+            log_buffer,
-
+            f"      [8.1] ANN search for {len(unit_ids)} new units → {len(all_links)} candidate links: {time_mod.time() - ann_start:.3f}s",
-                # Find top-k above threshold
+        )
-                # Get indices of similarities above threshold
-                above_threshold = np.where(similarities >= threshold)[0]
-
-                if len(above_threshold) > 0:
-                    # Sort by similarity (descending) and take top-k
-                    sorted_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
-
-                    for idx in sorted_indices:
-                        similar_id = existing_ids[idx]
-                        # Clamp to [0, 1] to handle floating point precision issues
-                        similarity = float(min(1.0, max(0.0, similarities[idx])))
-                        all_links.append((unit_id, similar_id, "semantic", similarity, None))
 
         # Also compute similarities WITHIN the new batch (new units to each other)
         # Apply the same top_k limit per unit as we do for existing units
@@ -659,7 +622,7 @@ async def create_semantic_links_batch(
 
         _log(
             log_buffer,
-            f"      [8.2] Compute similarities & generate {len(all_links)} semantic links: {time_mod.time() - compute_start:.3f}s",
+            f"      [8.2] Within-batch similarities added {len(all_links)} total semantic links",
         )
 
         if all_links:
@@ -667,14 +630,13 @@ async def create_semantic_links_batch(
             # Batch inserts to avoid timeout on large batches
             BATCH_SIZE = 1000
             for batch_start in range(0, len(all_links), BATCH_SIZE):
-                batch = all_links[batch_start : batch_start + BATCH_SIZE]
                 await conn.executemany(
                     f"""
                     INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
                     VALUES ($1, $2, $3, $4, $5)
                     ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
                     """,
-                    batch,
+                    all_links[batch_start : batch_start + BATCH_SIZE],
                 )
             _log(
                 log_buffer, f"      [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s"
@@ -690,18 +652,18 @@ async def create_semantic_links_batch(
         raise
 
 
-async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: int = 50000):
+async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: int = 5000):
     """
-    Insert all entity links using COPY to temp table + INSERT for maximum speed.
+    Insert all entity links using COPY to temp table + chunked INSERT for reliability.
 
-    Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading,
+    Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading into a
-    then INSERT ... ON CONFLICT from temp table. This is the fastest
+    temp table, then INSERT ... ON CONFLICT in chunks of chunk_size. Chunking
-    method for bulk inserts with conflict handling.
+    prevents single-query timeouts on very large tables (100M+ rows).
 
     Args:
         conn: Database connection
         links: List of EntityLink objects
-        chunk_size: Number of rows per batch (default 50000)
+        chunk_size: Number of rows per INSERT chunk (default 5000)
     """
     if not links:
         return
@@ -710,10 +672,11 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
 
     total_start = time_mod.time()
 
-    # Create temp table for bulk loading
+    # Create temp table with serial for stable chunked access
     create_start = time_mod.time()
     await conn.execute("""
         CREATE TEMP TABLE IF NOT EXISTS _temp_entity_links (
+            _row_num SERIAL,
             from_unit_id uuid,
             to_unit_id uuid,
             link_type text,
@@ -730,9 +693,7 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
 
     # Convert EntityLink objects to tuples for COPY
     convert_start = time_mod.time()
-    records = []
+    records = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
-    for link in links:
-        records.append((link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id))
     logger.debug(f"      [9.3] Convert {len(records)} records: {time_mod.time() - convert_start:.3f}s")
 
     # Bulk load using COPY (fastest method)
@@ -744,15 +705,25 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
     )
     logger.debug(f"      [9.4] COPY {len(records)} records to temp table: {time_mod.time() - copy_start:.3f}s")
 
-    # Insert from temp table with ON CONFLICT (single query for all rows)
+    # Insert from temp table in chunks to avoid single-query timeouts on large tables
     insert_start = time_mod.time()
-    await conn.execute(f"""
+    total_rows = len(records)
+    chunks = 0
+    for chunk_start in range(0, total_rows, chunk_size):
+        chunk_end = chunk_start + chunk_size
+        await conn.execute(
+            f"""
             INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
             SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
             FROM _temp_entity_links
+            WHERE _row_num > $1 AND _row_num <= $2
             ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
-    """)
+            """,
-    logger.debug(f"      [9.5] INSERT from temp table: {time_mod.time() - insert_start:.3f}s")
+            chunk_start,
+            chunk_end,
+        )
+        chunks += 1
+    logger.debug(f"      [9.5] INSERT {total_rows} rows in {chunks} chunks: {time_mod.time() - insert_start:.3f}s")
     logger.debug(f"      [9.TOTAL] Entity links batch insert: {time_mod.time() - total_start:.3f}s")
 
 

hindsight-api/hindsight_api/engine/retain/orchestrator.py

@@ -55,7 +55,6 @@ def parse_datetime_flexible(value: Any) -> datetime:
 from ..response_models import TokenUsage
 from . import (
     chunk_storage,
-    deduplication,
     embedding_processing,
     entity_processing,
     fact_extraction,
@@ -73,7 +72,6 @@ async def retain_batch(
     llm_config,
     entity_resolver,
     format_date_fn,
-    duplicate_checker_fn,
     bank_id: str,
     contents_dicts: list[RetainContentDict],
     config,
@@ -94,7 +92,6 @@ async def retain_batch(
         llm_config: LLM configuration for fact extraction
         entity_resolver: Entity resolver for entity processing
         format_date_fn: Function to format datetime to readable string
-        duplicate_checker_fn: Function to check for duplicate facts
         bank_id: Bank identifier
         contents_dicts: List of content dictionaries
         config: Resolved HindsightConfig for this bank
@@ -165,8 +162,6 @@ async def retain_batch(
         docs_tracked = 0
         async with acquire_with_retry(pool) as conn:
             async with conn.transaction():
-                await fact_storage.ensure_bank_exists(conn, bank_id)
-
                 # Group contents by document_id (consistent with normal path)
                 contents_by_doc_early = defaultdict(list)
                 for idx, content_dict in enumerate(contents_dicts):
@@ -284,9 +279,6 @@ async def retain_batch(
     # Step 4: Database transaction
     async with acquire_with_retry(pool) as conn:
         async with conn.transaction():
-            # Ensure bank exists
-            await fact_storage.ensure_bank_exists(conn, bank_id)
-
             # Handle document tracking for all documents
             step_start = time.time()
             # Map None document_id to generated UUIDs
@@ -438,20 +430,7 @@ async def retain_batch(
                         actual_doc_id = document_id
                     processed_fact.document_id = actual_doc_id
 
-            # Deduplication
+            non_duplicate_facts = processed_facts
-            step_start = time.time()
-            is_duplicate_flags = await deduplication.check_duplicates_batch(
-                conn, bank_id, processed_facts, duplicate_checker_fn
-            )
-            log_buffer.append(
-                f"[4] Deduplication: {sum(is_duplicate_flags)} duplicates in {time.time() - step_start:.3f}s"
-            )
-
-            # Filter out duplicates
-            non_duplicate_facts = deduplication.filter_duplicates(processed_facts, is_duplicate_flags)
-
-            if not non_duplicate_facts:
-                return [[] for _ in contents], usage
 
             # Insert facts (document_id is now stored per-fact)
             step_start = time.time()
@@ -503,7 +482,11 @@ async def retain_batch(
             log_buffer.append(f"[10] Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
 
             # Map results back to original content items
-            result_unit_ids = _map_results_to_contents(contents, extracted_facts, is_duplicate_flags, unit_ids)
+            result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids)
+
+        # Flush entity stats (mention_count / last_seen) now that the transaction
+        # has committed.  Uses a fresh pool connection — no locks held.
+        await entity_resolver.flush_pending_stats()
 
         # Log final summary
         total_time = time.time() - start_time
@@ -521,28 +504,20 @@ async def retain_batch(
 def _map_results_to_contents(
     contents: list[RetainContent],
     extracted_facts: list[ExtractedFact],
-    is_duplicate_flags: list[bool],
     unit_ids: list[str],
 ) -> list[list[str]]:
-    """
+    """Map created unit IDs back to original content items."""
-    Map created unit IDs back to original content items.
+    facts_by_content: dict[int, list[int]] = {i: [] for i in range(len(contents))}
-
-    Accounts for duplicates when mapping back.
-    """
-    result_unit_ids = []
-    filtered_idx = 0
-
-    # Group facts by content_index
-    facts_by_content = {i: [] for i in range(len(contents))}
     for i, fact in enumerate(extracted_facts):
         facts_by_content[fact.content_index].append(i)
 
+    result_unit_ids = []
+    unit_idx = 0
     for content_index in range(len(contents)):
         content_unit_ids = []
-        for fact_idx in facts_by_content[content_index]:
+        for _ in facts_by_content[content_index]:
-            if not is_duplicate_flags[fact_idx]:
+            content_unit_ids.append(unit_ids[unit_idx])
-                content_unit_ids.append(unit_ids[filtered_idx])
+            unit_idx += 1
-                filtered_idx += 1
         result_unit_ids.append(content_unit_ids)
 
     return result_unit_ids

hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py

@@ -1,18 +1,28 @@
 """
 Link Expansion graph retrieval.
 
-A simple, fast graph retrieval that expands from seeds via:
+Expands from semantic/temporal seeds through three parallel, first-class signals
-1. Entity links: Find facts sharing entities with seeds (filtered by entity frequency)
+stored in memory_links:
-2. Causal links: Find facts causally linked to seeds (top-k by weight)
 
-Characteristics:
+1. Entity links  — precomputed co-occurrence graph (created at retain time, bounded to
-- 2-3 DB queries (seed finding + parallel entity/causal expansion)
+                   MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared
-- Sublinear: only touches connected facts via indexes
+                   entities between the seed set and each candidate.
-- No iteration, no propagation, no normalization
+2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
-- Target: <100ms
+                    similar existing facts at insert time, similarity >= 0.7). Checked
+                    in both directions since the graph is not symmetric. Score = weight.
+3. Causal links  — explicit causal chains (causes/caused_by/enables/prevents).
+                   Score = weight + 1.0 (boosted as highest-quality signal).
+
+All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
+at query time. Each expansion is a simple aggregation over a small result set.
+
+For non-observation fact types the three expansions are issued as a single CTE query
+(one roundtrip, one connection) with a `source` discriminator column so the Python
+merge step can apply per-signal score transformations.
 """
 
 import logging
+import math
 import time
 
 from ..db_utils import acquire_with_retry
@@ -65,27 +75,23 @@ class LinkExpansionRetriever(GraphRetriever):
     """
     Graph retrieval via direct link expansion from seeds.
 
-    Expands through entity co-occurrence and causal links in a single query.
+    Runs three expansions through precomputed memory_links: entity co-occurrence,
-    Fast and simple alternative to MPFP.
+    semantic kNN, and causal chains, all bounded at retain time.
+
+    For non-observation fact types the three expansions are issued as a single CTE
+    query (one roundtrip, one connection slot) with a `source` discriminator column.
+    The Python merge step applies per-signal score transformations.
     """
 
     def __init__(
         self,
-        max_entity_frequency: int = 500,
         causal_weight_threshold: float = 0.3,
-        causal_limit_per_seed: int = 10,
     ):
         """
-        Initialize link expansion retriever.
-
         Args:
-            max_entity_frequency: Skip entities appearing in more than this many facts
+            causal_weight_threshold: Minimum weight for causal links to follow.
-            causal_weight_threshold: Minimum weight for causal links
-            causal_limit_per_seed: Max causal links to follow per seed
         """
-        self.max_entity_frequency = max_entity_frequency
         self.causal_weight_threshold = causal_weight_threshold
-        self.causal_limit_per_seed = causal_limit_per_seed
 
     @property
     def name(self) -> str:
@@ -110,7 +116,7 @@ class LinkExpansionRetriever(GraphRetriever):
 
         Args:
             pool: Database connection pool
-            query_embedding_str: Query embedding (unused, kept for interface)
+            query_embedding_str: Query embedding as string
             bank_id: Memory bank ID
             fact_type: Fact type to filter
             budget: Maximum results to return
@@ -118,7 +124,7 @@ class LinkExpansionRetriever(GraphRetriever):
             semantic_seeds: Pre-computed semantic entry points
             temporal_seeds: Pre-computed temporal entry points
             adjacency: Unused, kept for interface compatibility
-            tags: Optional list of tags for visibility filtering (OR matching)
+            tags: Optional list of tags for visibility filtering
 
         Returns:
             Tuple of (results, timings)
@@ -126,8 +132,6 @@ class LinkExpansionRetriever(GraphRetriever):
         start_time = time.time()
         timings = MPFPTimings(fact_type=fact_type)
 
-        # Use single connection for all queries to reduce pool pressure
-        # (queries are fast ~50ms each, connection acquisition is the bottleneck)
         async with acquire_with_retry(pool) as conn:
             # Find seeds if not provided
             if semantic_seeds:
@@ -150,7 +154,6 @@ class LinkExpansionRetriever(GraphRetriever):
                     f"(tags={tags}, tags_match={tags_match})"
                 )
 
-            # Add temporal seeds if provided
             if temporal_seeds:
                 all_seeds.extend(temporal_seeds)
 
@@ -160,223 +163,61 @@ class LinkExpansionRetriever(GraphRetriever):
             seed_ids = list({s.id for s in all_seeds})
             timings.pattern_count = len(seed_ids)
 
-            # Run entity and causal expansion sequentially on same connection
             query_start = time.time()
 
-            # For observations, traverse through source_memory_ids to find entity connections.
-            # Observations don't have direct unit_entities - they inherit entities via their
-            # source world/experience facts.
-            #
-            # Path: observation → source_memory_ids → world fact → entities →
-            #       ALL world facts with those entities → their observations (excluding seeds)
             if fact_type == "observation":
-                # Debug: Check what source_memory_ids exist on seed observations
+                entity_rows, semantic_rows, causal_rows = await self._expand_observations(conn, seed_ids, budget)
-                debug_sources = await conn.fetch(
-                    f"""
-                    SELECT id, source_memory_ids
-                    FROM {fq_table("memory_units")}
-                    WHERE id = ANY($1::uuid[])
-                    """,
-                    seed_ids,
-                )
-                source_ids_found = []
-                for row in debug_sources:
-                    if row["source_memory_ids"]:
-                        source_ids_found.extend(row["source_memory_ids"])
-                logger.debug(
-                    f"[LinkExpansion] observation graph: {len(seed_ids)} seeds, "
-                    f"{len(source_ids_found)} source_memory_ids found"
-                )
-
-                entity_rows = await conn.fetch(
-                    f"""
-                    WITH seed_sources AS (
-                        -- Get source memory IDs from seed observations
-                        SELECT DISTINCT unnest(source_memory_ids) AS source_id
-                        FROM {fq_table("memory_units")}
-                        WHERE id = ANY($1::uuid[])
-                          AND source_memory_ids IS NOT NULL
-                    ),
-                    source_entities AS (
-                        -- Get entities from those source memories (filtered by frequency)
-                        SELECT DISTINCT ue.entity_id
-                        FROM seed_sources ss
-                        JOIN {fq_table("unit_entities")} ue ON ss.source_id = ue.unit_id
-                        JOIN {fq_table("entities")} e ON ue.entity_id = e.id
-                        WHERE e.mention_count < $2
-                    ),
-                    all_connected_sources AS (
-                        -- Find ALL world facts sharing those entities (don't exclude seed sources)
-                        -- The exclusion happens at the observation level, not the source level
-                        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
-                    )
-                    -- Find observations derived from connected source memories
-                    -- Only exclude the actual seed observations
-                    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]
-                    WHERE mu.fact_type = 'observation'
-                      AND mu.id != ALL($1::uuid[])
-                    GROUP BY mu.id
-                    ORDER BY score DESC
-                    LIMIT $3
-                    """,
-                    seed_ids,
-                    self.max_entity_frequency,
-                    budget,
-                )
-                logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
             else:
-                # For world/experience facts, use direct entity lookup
+                entity_rows, semantic_rows, causal_rows = await self._expand_combined(conn, seed_ids, fact_type, budget)
-                entity_rows = await conn.fetch(
-                    f"""
-                    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(*)::float AS score
-                    FROM {fq_table("unit_entities")} seed_ue
-                    JOIN {fq_table("entities")} e ON seed_ue.entity_id = e.id
-                    JOIN {fq_table("unit_entities")} other_ue ON seed_ue.entity_id = other_ue.entity_id
-                    JOIN {fq_table("memory_units")} mu ON other_ue.unit_id = mu.id
-                    WHERE seed_ue.unit_id = ANY($1::uuid[])
-                      AND e.mention_count < $2
-                      AND mu.id != ALL($1::uuid[])
-                      AND mu.fact_type = $3
-                    GROUP BY mu.id
-                    ORDER BY score DESC
-                    LIMIT $4
-                    """,
-                    seed_ids,
-                    self.max_entity_frequency,
-                    fact_type,
-                    budget,
-                )
-
-            causal_rows = await conn.fetch(
-                f"""
-                SELECT DISTINCT ON (mu.id)
-                    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,
-                    ml.weight + 1.0 AS score
-                FROM {fq_table("memory_links")} ml
-                JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
-                WHERE ml.from_unit_id = ANY($1::uuid[])
-                  AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
-                  AND ml.weight >= $2
-                  AND mu.fact_type = $3
-                ORDER BY mu.id, ml.weight DESC
-                LIMIT $4
-                """,
-                seed_ids,
-                self.causal_weight_threshold,
-                fact_type,
-                budget,
-            )
-
-            # Fallback: semantic/temporal/entity links from memory_links table
-            # These are secondary to entity links (via unit_entities) and causal links
-            # Weight is halved (0.5x) to prioritize primary link types
-            # Check both directions: seeds -> others AND others -> seeds
-            fallback_rows = await conn.fetch(
-                f"""
-                WITH outgoing AS (
-                    -- Links FROM seeds TO other facts
-                    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,
-                           ml.weight
-                    FROM {fq_table("memory_links")} ml
-                    JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
-                    WHERE ml.from_unit_id = ANY($1::uuid[])
-                      AND ml.link_type IN ('semantic', 'temporal', 'entity')
-                      AND ml.weight >= $2
-                      AND mu.fact_type = $3
-                      AND mu.id != ALL($1::uuid[])
-                ),
-                incoming AS (
-                    -- Links FROM other facts TO seeds (reverse direction)
-                    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,
-                           ml.weight
-                    FROM {fq_table("memory_links")} ml
-                    JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                    WHERE ml.to_unit_id = ANY($1::uuid[])
-                      AND ml.link_type IN ('semantic', 'temporal', 'entity')
-                      AND ml.weight >= $2
-                      AND mu.fact_type = $3
-                      AND mu.id != ALL($1::uuid[])
-                ),
-                combined AS (
-                    SELECT * FROM outgoing
-                    UNION ALL
-                    SELECT * FROM incoming
-                )
-                SELECT DISTINCT ON (id)
-                    id, text, context, event_date, occurred_start,
-                    occurred_end, mentioned_at,
-                    fact_type, document_id, chunk_id, tags,
-                    (MAX(weight) * 0.5) AS score
-                FROM combined
-                GROUP BY id, text, context, event_date, occurred_start,
-                         occurred_end, mentioned_at,
-                         fact_type, document_id, chunk_id, tags
-                ORDER BY id, score DESC
-                LIMIT $4
-                """,
-                seed_ids,
-                self.causal_weight_threshold,
-                fact_type,
-                budget,
-            )
 
             timings.edge_load_time = time.time() - query_start
-            timings.db_queries = 3
+            timings.db_queries = 1
-            timings.edge_count = len(entity_rows) + len(causal_rows) + len(fallback_rows)
+            timings.edge_count = len(entity_rows) + len(semantic_rows) + len(causal_rows)
 
-        # Merge results, taking max score per fact
+        # Merge results with additive intra-score: entity + semantic + causal ∈ [0, 3].
-        # Priority: entity links (unit_entities) > causal links > fallback links
+        #
-        score_map: dict[str, float] = {}
+        # Entity score: tanh(count × 0.5) maps shared-entity count to [0, 1]:
+        #   1 entity → 0.46,  2 → 0.76,  3 → 0.91,  4 → 0.96  (saturates naturally)
+        # Semantic score: similarity weight, already ∈ [0.7, 1.0].
+        # Causal score:   link weight, already ∈ [0, 1].
+        #
+        # Facts appearing in multiple signals accumulate higher scores, rewarding
+        # convergent evidence. The outer RRF uses rank position from this sorted list.
+        entity_scores: dict[str, float] = {}
+        semantic_scores: dict[str, float] = {}
+        causal_scores: dict[str, float] = {}
         row_map: dict[str, dict] = {}
 
         for row in entity_rows:
             fact_id = str(row["id"])
-            score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"])
+            entity_scores[fact_id] = math.tanh(row["score"] * 0.5)
             row_map[fact_id] = dict(row)
 
+        for row in semantic_rows:
+            fact_id = str(row["id"])
+            semantic_scores[fact_id] = max(semantic_scores.get(fact_id, 0.0), row["score"])
+            row_map.setdefault(fact_id, dict(row))
+
         for row in causal_rows:
             fact_id = str(row["id"])
-            score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"])
+            causal_scores[fact_id] = max(causal_scores.get(fact_id, 0.0), row["score"])
-            if fact_id not in row_map:
+            row_map.setdefault(fact_id, dict(row))
-                row_map[fact_id] = dict(row)
 
-        for row in fallback_rows:
+        all_ids = set(entity_scores) | set(semantic_scores) | set(causal_scores)
-            fact_id = str(row["id"])
+        score_map = {
-            score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"])
+            fid: entity_scores.get(fid, 0.0) + semantic_scores.get(fid, 0.0) + causal_scores.get(fid, 0.0)
-            if fact_id not in row_map:
+            for fid in all_ids
-                row_map[fact_id] = dict(row)
+        }
 
-        # Sort by score and limit
         sorted_ids = sorted(score_map.keys(), key=lambda x: score_map[x], reverse=True)[:budget]
         rows = [row_map[fact_id] for fact_id in sorted_ids]
 
-        # Convert to results
         results = []
         for row in rows:
             result = RetrievalResult.from_db_row(dict(row))
             result.activation = row["score"]
             results.append(result)
 
-        # Apply tags filtering (graph expansion may reach untagged memories)
         if tags:
             results = filter_results_by_tags(results, tags, match=tags_match)
 
@@ -389,3 +230,253 @@ class LinkExpansionRetriever(GraphRetriever):
         )
 
         return results, timings
+
+    async def _expand_combined(
+        self,
+        conn,
+        seed_ids: list,
+        fact_type: str,
+        budget: int,
+    ) -> tuple[list, list, list]:
+        """
+        Single-roundtrip CTE query combining entity, semantic, and causal expansions.
+
+        Uses a `source` discriminator column so the caller can apply per-signal
+        score transformations.  The three CTEs share one connection slot — important
+        for asyncpg which does not allow concurrent queries on the same connection.
+
+        Index coverage (requires migration d2e3f4a5b6c7):
+          entity:   idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
+                    WHERE link_type = 'entity'  → index-only scan, no heap reads
+          semantic incoming:
+                    idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
+                    → replaces costly BitmapAnd of two separate scans
+        """
+        ml = fq_table("memory_links")
+        mu = fq_table("memory_units")
+        all_rows = await conn.fetch(
+            f"""
+            WITH entity_expanded AS (
+                -- Entity co-occurrence: seeds → their precomputed entity-link neighbors.
+                -- Score = distinct shared entities (bounded at retain time to
+                -- MAX_LINKS_PER_ENTITY=50).  GROUP BY mu.id is sufficient because mu.id
+                -- is the primary key and functionally determines all other mu columns.
+                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 ml.entity_id)::float AS score,
+                    'entity'::text AS source
+                FROM {ml} ml
+                JOIN {mu} mu ON mu.id = ml.to_unit_id
+                WHERE ml.from_unit_id = ANY($1::uuid[])
+                  AND ml.link_type = 'entity'
+                  AND mu.fact_type = $2
+                  AND mu.id != ALL($1::uuid[])
+                GROUP BY mu.id
+                ORDER BY score DESC
+                LIMIT $3
+            ),
+            semantic_expanded AS (
+                -- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
+                -- incoming (facts inserted after seeds that found seeds as kNN).
+                -- Score = max similarity weight across both directions.
+                SELECT
+                    id, text, context, event_date, occurred_start,
+                    occurred_end, mentioned_at,
+                    fact_type, document_id, chunk_id, tags,
+                    MAX(weight) AS score,
+                    'semantic'::text AS source
+                FROM (
+                    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,
+                        ml.weight
+                    FROM {ml} ml
+                    JOIN {mu} mu ON mu.id = ml.to_unit_id
+                    WHERE ml.from_unit_id = ANY($1::uuid[])
+                      AND ml.link_type = 'semantic'
+                      AND mu.fact_type = $2
+                      AND mu.id != ALL($1::uuid[])
+                    UNION ALL
+                    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,
+                        ml.weight
+                    FROM {ml} ml
+                    JOIN {mu} mu ON mu.id = ml.from_unit_id
+                    WHERE ml.to_unit_id = ANY($1::uuid[])
+                      AND ml.link_type = 'semantic'
+                      AND mu.fact_type = $2
+                      AND mu.id != ALL($1::uuid[])
+                ) sem_raw
+                GROUP BY id, text, context, event_date, occurred_start,
+                         occurred_end, mentioned_at,
+                         fact_type, document_id, chunk_id, tags
+                ORDER BY score DESC
+                LIMIT $3
+            ),
+            causal_expanded AS (
+                -- Causal chains: explicit causes/enables/prevents links from seeds.
+                -- DISTINCT ON handles the case where a seed has multiple causal links
+                -- to the same target; best weight wins.
+                SELECT DISTINCT ON (mu.id)
+                    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,
+                    ml.weight AS score,
+                    'causal'::text AS source
+                FROM {ml} ml
+                JOIN {mu} mu ON ml.to_unit_id = mu.id
+                WHERE ml.from_unit_id = ANY($1::uuid[])
+                  AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
+                  AND ml.weight >= $4
+                  AND mu.fact_type = $2
+                ORDER BY mu.id, ml.weight DESC
+                LIMIT $3
+            )
+            SELECT * FROM entity_expanded
+            UNION ALL
+            SELECT * FROM semantic_expanded
+            UNION ALL
+            SELECT * FROM causal_expanded
+            """,
+            seed_ids,
+            fact_type,
+            budget,
+            self.causal_weight_threshold,
+        )
+
+        entity_rows = [r for r in all_rows if r["source"] == "entity"]
+        semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
+        causal_rows = [r for r in all_rows if r["source"] == "causal"]
+        return entity_rows, semantic_rows, causal_rows
+
+    async def _expand_observations(
+        self,
+        conn,
+        seed_ids: list,
+        budget: int,
+    ) -> tuple[list, list, list]:
+        """
+        Observation-specific expansion.
+
+        Observations don't have direct entity links in memory_links (they're created
+        by consolidation, not retain).  Instead, traverse source_memory_ids → world
+        facts → entities → other world facts → their observations.
+
+        Semantic and causal expansions run as a second combined CTE query.
+        """
+        source_ids_found: list = []
+        if logger.isEnabledFor(logging.DEBUG):
+            debug_rows = await conn.fetch(
+                f"""
+                SELECT id, source_memory_ids
+                FROM {fq_table("memory_units")}
+                WHERE id = ANY($1::uuid[])
+                """,
+                seed_ids,
+            )
+            for row in debug_rows:
+                if row["source_memory_ids"]:
+                    source_ids_found.extend(row["source_memory_ids"])
+            logger.debug(
+                f"[LinkExpansion] observation graph: {len(seed_ids)} seeds, "
+                f"{len(source_ids_found)} source_memory_ids found"
+            )
+
+        entity_rows = await conn.fetch(
+            f"""
+            WITH seed_sources AS (
+                SELECT DISTINCT unnest(source_memory_ids) AS source_id
+                FROM {fq_table("memory_units")}
+                WHERE id = ANY($1::uuid[])
+                  AND source_memory_ids IS NOT NULL
+            ),
+            source_entities AS (
+                SELECT DISTINCT ue.entity_id
+                FROM seed_sources ss
+                JOIN {fq_table("unit_entities")} ue ON ss.source_id = ue.unit_id
+            ),
+            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
+            )
+            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]
+            WHERE mu.fact_type = 'observation'
+              AND mu.id != ALL($1::uuid[])
+            GROUP BY mu.id
+            ORDER BY score DESC
+            LIMIT $2
+            """,
+            seed_ids,
+            budget,
+        )
+        logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
+
+        # Semantic + causal for observations in one query
+        ml = fq_table("memory_links")
+        mu = fq_table("memory_units")
+        sem_causal_rows = await conn.fetch(
+            f"""
+            WITH semantic_expanded AS (
+                SELECT
+                    id, text, context, event_date, occurred_start,
+                    occurred_end, mentioned_at,
+                    fact_type, document_id, chunk_id, tags,
+                    MAX(weight) AS score,
+                    'semantic'::text AS source
+                FROM (
+                    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, ml.weight
+                    FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
+                    WHERE ml.from_unit_id = ANY($1::uuid[])
+                      AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
+                      AND mu.id != ALL($1::uuid[])
+                    UNION ALL
+                    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, ml.weight
+                    FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
+                    WHERE ml.to_unit_id = ANY($1::uuid[])
+                      AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
+                      AND mu.id != ALL($1::uuid[])
+                ) sem_raw
+                GROUP BY id, text, context, event_date, occurred_start, occurred_end,
+                         mentioned_at, fact_type, document_id, chunk_id, tags
+                ORDER BY score DESC LIMIT $2
+            ),
+            causal_expanded AS (
+                SELECT DISTINCT ON (mu.id)
+                    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, ml.weight AS score, 'causal'::text AS source
+                FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
+                WHERE ml.from_unit_id = ANY($1::uuid[])
+                  AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
+                  AND ml.weight >= $3 AND mu.fact_type = 'observation'
+                ORDER BY mu.id, ml.weight DESC LIMIT $2
+            )
+            SELECT * FROM semantic_expanded
+            UNION ALL
+            SELECT * FROM causal_expanded
+            """,
+            seed_ids,
+            budget,
+            self.causal_weight_threshold,
+        )
+
+        semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
+        causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
+        return entity_rows, semantic_rows, causal_rows

hindsight-api/hindsight_api/engine/search/retrieval.py

@@ -297,13 +297,20 @@ async def retrieve_temporal_combined(
     if tags:
         params.append(tags)
 
-    # Batch query: Get entry points for ALL fact types at once with window function
+    # Two-phase entry point query:
+    # Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
+    #   the temporal window. This lets the planner use date indexes for filtering.
+    # Phase 2 (sim_ranked): join back to memory_units for only the top-50-per-type candidates
+    #   and compute embedding similarity for that small set (≤ 50 × len(fact_types) rows).
+    # This avoids computing embedding distances for potentially thousands of date-range rows.
     entry_points = await conn.fetch(
         f"""
-        WITH ranked_entries AS (
+        WITH date_ranked AS MATERIALIZED (
-            SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags,
+            SELECT id, fact_type,
-                   1 - (embedding <=> $1::vector) AS similarity,
+                   ROW_NUMBER() OVER (
-                   ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, embedding <=> $1::vector) AS rn
+                       PARTITION BY fact_type
+                       ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC NULLS LAST
+                   ) AS rn
             FROM {fq_table("memory_units")}
             WHERE bank_id = $2
               AND fact_type = ANY($3)
@@ -318,12 +325,20 @@ async def retrieve_temporal_combined(
                   OR
                   (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
               )
-              AND (1 - (embedding <=> $1::vector)) >= $6
               {tags_clause}
+        ),
+        sim_ranked AS (
+            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,
+                   1 - (mu.embedding <=> $1::vector) AS similarity,
+                   ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
+            FROM date_ranked dr
+            JOIN {fq_table("memory_units")} mu ON mu.id = dr.id
+            WHERE dr.rn <= 50
+              AND (1 - (mu.embedding <=> $1::vector)) >= $6
         )
         SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity
-        FROM ranked_entries
+        FROM sim_ranked
-        WHERE rn <= 10
+        WHERE sim_rn <= 10
         """,
         *params,
     )
@@ -387,34 +402,52 @@ async def retrieve_temporal_combined(
         frontier = list(node_scores.keys())
         budget_remaining = budget - len(ft_entry_points)
         batch_size = 20
+        # Per-source neighbor limit: lets the planner use the composite index
+        # (from_unit_id, link_type, weight DESC) with early termination, avoiding
+        # a full scan of all links from all source nodes before sorting.
+        per_source_limit = 10
+        # Safety cap on BFS iterations to prevent runaway spreading in dense graphs.
+        max_iterations = 5
+        iteration = 0
 
-        # Build tags clause for spreading (use param 6 since 1-5 are used)
+        # Build tags clause for spreading (use param 7 since 1-6 are used)
-        spreading_tags_clause = build_tags_where_clause_simple(tags, 6, table_alias="mu.", match=tags_match)
+        spreading_tags_clause = build_tags_where_clause_simple(tags, 7, table_alias="mu.", match=tags_match)
 
-        while frontier and budget_remaining > 0:
+        while frontier and budget_remaining > 0 and iteration < max_iterations:
+            iteration += 1
             batch_ids = frontier[:batch_size]
             frontier = frontier[batch_size:]
 
-            spreading_params = [query_emb_str, batch_ids, ft, semantic_threshold, batch_size * 10]
+            # $1=query_emb, $2=batch_ids, $3=fact_type, $4=threshold, $5=per_source_limit, $6=bank_id, $7=tags
+            spreading_params = [query_emb_str, batch_ids, ft, semantic_threshold, per_source_limit, bank_id]
             if tags:
                 spreading_params.append(tags)
 
+            # LATERAL join: for each source node, fetch top-K neighbors by weight using
+            # the existing idx_memory_links_from_type_weight index with early-exit semantics.
+            # This avoids scanning all temporal links from all source nodes before sorting.
+            # bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
             neighbors = await conn.fetch(
                 f"""
-                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,
+                SELECT src.from_unit_id, 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,
-                       ml.weight, ml.link_type, ml.from_unit_id,
+                       l.weight, l.link_type,
                        1 - (mu.embedding <=> $1::vector) AS similarity
+                FROM unnest($2::uuid[]) AS src(from_unit_id)
+                CROSS JOIN LATERAL (
+                    SELECT ml.to_unit_id, ml.weight, ml.link_type
                     FROM {fq_table("memory_links")} ml
-                JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
+                    WHERE ml.from_unit_id = src.from_unit_id
-                WHERE ml.from_unit_id = ANY($2::uuid[])
                       AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
                       AND ml.weight >= 0.1
+                    ORDER BY ml.weight DESC
+                    LIMIT $5
+                ) l
+                JOIN {fq_table("memory_units")} mu ON mu.id = l.to_unit_id
+                WHERE mu.bank_id = $6
                   AND mu.fact_type = $3
                   AND mu.embedding IS NOT NULL
                   AND (1 - (mu.embedding <=> $1::vector)) >= $4
                   {spreading_tags_clause}
-                ORDER BY ml.weight DESC
-                LIMIT $5
                 """,
                 *spreading_params,
             )

hindsight-api/hindsight_api/main.py

@@ -252,6 +252,7 @@ def main():
             retain_mission=config.retain_mission,
             retain_custom_instructions=config.retain_custom_instructions,
             retain_batch_tokens=config.retain_batch_tokens,
+            retain_entity_lookup=config.retain_entity_lookup,
             retain_batch_enabled=config.retain_batch_enabled,
             retain_batch_poll_interval_seconds=config.retain_batch_poll_interval_seconds,
             file_storage_type=config.file_storage_type,

hindsight-api/hindsight_api/migrations.py

@@ -18,6 +18,7 @@ No alembic.ini required - all configuration is done programmatically.
 import hashlib
 import logging
 import os
+import time
 from pathlib import Path
 
 from alembic import command
@@ -220,13 +221,40 @@ def run_migrations(
         lock_id = _get_schema_lock_id(schema) if schema else MIGRATION_LOCK_ID
         schema_name = schema or "public"
 
-        # Use PostgreSQL advisory lock to coordinate between distributed workers
+        # Use PostgreSQL advisory lock to coordinate between distributed workers.
+        #
+        # IMPORTANT: We must avoid holding an open transaction on the advisory-lock
+        # connection while CREATE INDEX CONCURRENTLY runs inside a migration.
+        # CONCURRENTLY waits for ALL active transactions to finish before the index
+        # becomes valid.  If the advisory-lock connection (or any waiting worker's
+        # connection) holds an open transaction, CONCURRENTLY deadlocks:
+        #   - migration worker waits for other workers' transactions to close
+        #   - other workers wait for the advisory lock to be released
+        #
+        # Fix:
+        #   1. Use pg_try_advisory_lock (non-blocking) in a poll loop instead of
+        #      blocking pg_advisory_lock, so we can COMMIT the transaction between
+        #      retries.  Between retries the connection holds no open transaction.
+        #   2. After acquiring the lock, COMMIT the transaction on the advisory-lock
+        #      connection itself before running migrations.  pg_advisory_lock is
+        #      session-level, so the lock survives the COMMIT.
         engine = create_engine(database_url)
         with engine.connect() as conn:
-            # pg_advisory_lock blocks until the lock is acquired
-            # The lock is automatically released when the connection closes
             logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
-            conn.execute(text(f"SELECT pg_advisory_lock({lock_id})"))
+            while True:
+                acquired = conn.execute(text(f"SELECT pg_try_advisory_lock({lock_id})")).scalar()
+                if acquired:
+                    break
+                # Commit the transaction so this connection holds no open snapshot
+                # while waiting.  This prevents blocking CREATE INDEX CONCURRENTLY
+                # that may be running in the migration worker.
+                conn.commit()
+                time.sleep(0.5)
+
+            # Commit AFTER acquiring the lock too.  pg_advisory_lock is session-level
+            # and survives the COMMIT, but the open transaction on this connection
+            # would otherwise block any CREATE INDEX CONCURRENTLY in the migration.
+            conn.commit()
             logger.debug("Migration advisory lock acquired")
 
             try:
@@ -347,6 +375,13 @@ def run_migrations(
                                     "Please install it with: CREATE EXTENSION vectorscale CASCADE;"
                                 ) from e
 
+                # Commit any pending transaction on the advisory-lock connection
+                # before running migrations.  Some code paths above (e.g., the
+                # pgvector extension check) may have started a transaction via
+                # SQLAlchemy's autobegin.  If we leave it open, CREATE INDEX
+                # CONCURRENTLY inside a migration will deadlock waiting for it.
+                conn.commit()
+
                 # Run migrations while holding the lock
                 _run_migrations_internal(database_url, script_location, schema=schema)
             finally:

hindsight-api/tests/test_worker.py

@@ -48,11 +48,11 @@ async def pool(pg0_db_url):
 @pytest_asyncio.fixture
 async def clean_operations(pool):
     """Clean up async_operations table before and after tests."""
-    # Clean before test
+    # Clean before test - covers both 'test-worker-' and 'test_worker_recovery' patterns
-    await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'")
+    await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%' OR bank_id LIKE 'test_worker_%'")
     yield
     # Clean after test
-    await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'")
+    await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%' OR bank_id LIKE 'test_worker_%'")
 
 
 class TestBrokerTaskBackend:

hindsight-dev/benchmarks/perf/recall_perf.py

@@ -0,0 +1,942 @@
+"""
+Large-bank recall load test (no LLM).
+
+Populates a synthetic bank using the real retain pipeline with a mocked LLM
+for fact extraction, then benchmarks recall latency at production scale.
+
+Usage (run from hindsight-api/):
+    cd hindsight-api
+
+    # Generate a small bank (~10K memory units):
+    uv run python ../hindsight-dev/benchmarks/perf/recall_perf.py generate \\
+        --bank-id recall-perf-small2 --scale small
+
+    # Benchmark recall:
+    uv run python ../hindsight-dev/benchmarks/perf/recall_perf.py benchmark \\
+        --bank-id recall-perf-small --query "database migration" --iterations 5
+
+    # Clean up:
+    uv run python ../hindsight-dev/benchmarks/perf/recall_perf.py clean \\
+        --bank-id recall-perf-small
+"""
+
+import argparse
+import asyncio
+import os
+import statistics
+import time
+from typing import Any
+
+from rich.console import Console
+from rich.progress import BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from rich.table import Table
+
+console = Console()
+
+# ---------------------------------------------------------------------------
+# Fact corpus
+# ---------------------------------------------------------------------------
+
+# ~200 entity names spanning people, technologies, and places.
+ENTITIES = [
+    # People
+    "Alice Chen",
+    "Bob Martinez",
+    "Carol Thompson",
+    "David Kim",
+    "Eva Rodriguez",
+    "Frank Johnson",
+    "Grace Liu",
+    "Henry Park",
+    "Irene Nguyen",
+    "James Wilson",
+    "Karen Davis",
+    "Leo Brown",
+    "Mia Patel",
+    "Nathan Clark",
+    "Olivia Walker",
+    "Paul Harris",
+    "Quinn Lewis",
+    "Rachel Young",
+    "Sam Scott",
+    "Tina Adams",
+    "Uma Turner",
+    "Victor Hall",
+    "Wendy Allen",
+    "Xavier Hill",
+    "Yara Wright",
+    "Zoe King",
+    "Aaron Green",
+    "Beth Baker",
+    "Chris Nelson",
+    "Diana Carter",
+    "Ethan Mitchell",
+    "Fiona Perez",
+    "George Roberts",
+    "Hannah Turner",
+    "Ivan Phillips",
+    "Julia Campbell",
+    "Kevin Parker",
+    "Laura Evans",
+    "Mike Edwards",
+    "Nina Collins",
+    "Oscar Stewart",
+    "Penny Sanchez",
+    "Ryan Morris",
+    "Sandra Rogers",
+    "Tom Reed",
+    # Technologies
+    "PostgreSQL",
+    "Redis",
+    "Kubernetes",
+    "Docker",
+    "Python",
+    "Rust",
+    "TypeScript",
+    "React",
+    "FastAPI",
+    "GraphQL",
+    "gRPC",
+    "Kafka",
+    "Elasticsearch",
+    "Prometheus",
+    "Grafana",
+    "Terraform",
+    "Ansible",
+    "Nginx",
+    "SQLite",
+    "MongoDB",
+    "Cassandra",
+    "RabbitMQ",
+    "Celery",
+    "Pandas",
+    "NumPy",
+    "PyTorch",
+    "TensorFlow",
+    "OpenAI API",
+    "Anthropic API",
+    "LangChain",
+    "ChromaDB",
+    "Pinecone",
+    "Weaviate",
+    "pgvector",
+    "Alembic",
+    "SQLAlchemy",
+    "asyncpg",
+    "Pydantic",
+    "pytest",
+    "Ruff",
+    "GitHub Actions",
+    "CircleCI",
+    "AWS S3",
+    "AWS Lambda",
+    "GCP BigQuery",
+    "Azure DevOps",
+    "Datadog",
+    "Sentry",
+    "OpenTelemetry",
+    "Jaeger",
+    # Places / teams / projects
+    "San Francisco",
+    "New York",
+    "Seattle",
+    "Austin",
+    "London",
+    "Berlin",
+    "Tokyo",
+    "Singapore",
+    "Platform Team",
+    "Infrastructure Team",
+    "Data Science Team",
+    "Frontend Team",
+    "Backend Team",
+    "Security Team",
+    "DevOps Team",
+    "ML Platform Team",
+    "Project Orion",
+    "Project Helios",
+    "Project Atlas",
+    "Project Nexus",
+    "Project Titan",
+    "Project Echo",
+    "Project Phoenix",
+    "Hindsight",
+    "Memory Engine",
+    "Control Plane",
+    "Data Warehouse",
+    "API Gateway",
+    "Auth Service",
+    "Billing Service",
+    "Search Service",
+    "Notification Service",
+    "Analytics Dashboard",
+    "Admin Console",
+    "CI Pipeline",
+    "Staging Environment",
+    "Production Environment",
+    "Development Environment",
+    "Load Balancer",
+    "Service Mesh",
+    "Feature Flag Service",
+    "Observability Stack",
+    "Data Lake",
+    "Event Bus",
+    "Message Queue",
+    "Cache Layer",
+    "CDN",
+    "DNS",
+    "VPN",
+    "SSO",
+]
+
+# ~300 fact templates.  Placeholders {E0}..{E4} are filled with entity names
+# drawn from ENTITIES using a Zipf-like distribution so that ~20 entities
+# recur frequently across templates.
+FACT_TEMPLATES = [
+    "{E0} deployed a new version of {E1} to {E2} on {E3}.",
+    "{E0} reported a performance regression in {E1} affecting {E2}.",
+    "{E0} and {E1} completed the migration of {E2} from {E3} to {E4}.",
+    "{E0} updated the {E1} configuration to use {E2} for caching.",
+    "{E0} reviewed {E1}'s pull request for the {E2} integration.",
+    "{E0} noticed that {E1} latency increased after the {E2} upgrade.",
+    "{E0} created a dashboard in {E1} to monitor {E2} performance.",
+    "{E0} onboarded {E1} to the {E2} platform.",
+    "{E0} resolved the {E1} incident that caused {E2} downtime.",
+    "{E0} led a design review for {E1} with {E2} and {E3}.",
+    "{E0} wrote unit tests for the {E1} module in {E2}.",
+    "{E0} configured {E1} rate limiting in {E2} for {E3} endpoints.",
+    "{E0} set up {E1} alerts in {E2} for the {E3} service.",
+    "{E0} refactored the {E1} layer to use {E2} instead of {E3}.",
+    "{E0} mentored {E1} on {E2} best practices.",
+    "{E0} presented the {E1} roadmap to {E2} leadership in {E3}.",
+    "{E0} integrated {E1} with {E2} for real-time streaming.",
+    "{E0} opened a ticket for {E1} memory leak in {E2}.",
+    "{E0} benchmarked {E1} against {E2} for the {E3} use case.",
+    "{E0} scheduled a migration window for {E1} maintenance in {E2}.",
+    "{E0} documented the {E1} API changes for {E2}.",
+    "{E0} paired with {E1} to debug the {E2} timeout issue.",
+    "{E0} upgraded {E1} from version 3.1 to 4.0 in {E2}.",
+    "{E0} provisioned new {E1} instances in {E2} to handle load.",
+    "{E0} fixed a data race in {E1} caused by {E2} concurrency.",
+    "{E0} added {E1} tracing spans to the {E2} service.",
+    "{E0} rotated the {E1} credentials used by {E2}.",
+    "{E0} enabled {E1} compression in {E2} to reduce storage costs.",
+    "{E0} flagged a security issue in {E1} shared with {E2}.",
+    "{E0} ran chaos tests against {E1} in {E2}.",
+    "{E0} opened a feature request for {E1} pagination in {E2}.",
+    "{E0} trained the {E1} model on data from {E2}.",
+    "{E0} synced {E1} state to {E2} using {E3}.",
+    "{E0} tuned {E1} connection pool size for {E2} workloads.",
+    "{E0} evaluated {E1} vs {E2} for the {E3} project.",
+    "{E0} set up {E1} CI pipeline for {E2}.",
+    "{E0} migrated {E1} secrets from {E2} to {E3}.",
+    "{E0} identified a {E1} bottleneck in the {E2} hot path.",
+    "{E0} implemented {E1} retry logic in {E2}.",
+    "{E0} reviewed the {E1} threat model with {E2}.",
+    "{E0} enabled {E1} audit logging in {E2}.",
+    "{E0} scaled {E1} to handle {E2} throughput requirements.",
+    "{E0} added {E1} health checks for {E2}.",
+    "{E0} published the {E1} release notes for {E2}.",
+    "{E0} ran a load test against {E1} using {E2}.",
+    "{E0} submitted a PR to add {E1} support to {E2}.",
+    "{E0} kicked off {E1} data backfill in {E2}.",
+    "{E0} filed a post-mortem for the {E1} outage affecting {E2}.",
+    "{E0} enabled {E1} feature flag for {E2} users.",
+    "{E0} profiled {E1} memory usage in the {E2} environment.",
+    "{E0} opened a {E1} security advisory for {E2}.",
+    "{E0} created a {E1} runbook for {E2} on-call rotation.",
+    "{E0} shipped a hotfix for {E1} parsing bug in {E2}.",
+    "{E0} demoed the {E1} prototype to {E2} stakeholders.",
+    "{E0} archived old {E1} indexes in {E2} to free storage.",
+    "{E0} configured {E1} TLS termination at {E2}.",
+    "{E0} added {E1} request validation to the {E2} API.",
+    "{E0} updated {E1} Helm charts for {E2} deployment.",
+    "{E0} set up {E1} blue-green deployment for {E2}.",
+    "{E0} onboarded {E1} as a dependency for {E2}.",
+    "{E0} improved {E1} query performance by 40% in {E2}.",
+    "{E0} compiled a report on {E1} adoption across {E2}.",
+    "{E0} fixed {E1} deadlock under high concurrency in {E2}.",
+    "{E0} restarted {E1} to clear stale state in {E2}.",
+    "{E0} paired {E1} with {E2} to unblock {E3} migration.",
+    "{E0} shipped the {E1} v2 API for {E2} consumers.",
+    "{E0} reduced {E1} cold start time by optimizing {E2} imports.",
+    "{E0} enabled {E1} dark launch for {E2} traffic.",
+    "{E0} set up {E1} canary release for {E2}.",
+    "{E0} resolved {E1} certificate expiry alert for {E2}.",
+    "{E0} added {E1} caching layer to reduce {E2} load.",
+    "{E0} restructured {E1} schema in {E2} for performance.",
+    "{E0} published {E1} metrics to {E2} for alerting.",
+    "{E0} opened a discussion on {E1} alternatives for {E2}.",
+    "{E0} implemented {E1} circuit breaker pattern in {E2}.",
+    "{E0} rolled back {E1} after breaking changes in {E2}.",
+    "{E0} enabled {E1} distributed tracing across {E2} services.",
+    "{E0} ran a {E1} audit to identify {E2} vulnerabilities.",
+    "{E0} deprecated {E1} endpoints in {E2} for {E3}.",
+    "{E0} tested {E1} failover behavior in {E2}.",
+    "{E0} enabled {E1} CORS configuration in {E2}.",
+    "{E0} created {E1} cost allocation tags in {E2}.",
+    "{E0} optimized {E1} batch processing in {E2}.",
+    "{E0} added {E1} pagination to the {E2} listing endpoint.",
+    "{E0} integrated {E1} SSO with {E2}.",
+    "{E0} shipped {E1} feature for {E2} enterprise customers.",
+    "{E0} updated {E1} dependencies to fix {E2} vulnerabilities.",
+    "{E0} presented {E1} capacity plan for {E2} growth.",
+    "{E0} configured {E1} auto-scaling for {E2}.",
+    "{E0} generated {E1} API client for {E2} consumers.",
+    "{E0} added {E1} soft-delete support to {E2}.",
+    "{E0} resolved {E1} DNS resolution failure in {E2}.",
+    "{E0} deployed {E1} across three regions starting with {E2}.",
+    "{E0} benchmarked {E1} embedding throughput for {E2}.",
+    "{E0} set up {E1} read replicas in {E2} to offload load.",
+    "{E0} cleaned up {E1} orphaned resources in {E2}.",
+    "{E0} merged the {E1} feature branch into {E2} main.",
+    "{E0} performed a {E1} code review for {E2} security standards.",
+    "{E0} generated synthetic data using {E1} for {E2} tests.",
+    "{E0} integrated {E1} observability into {E2} pipeline.",
+    "{E0} added {E1} retry budget to the {E2} client.",
+    "{E0} compressed {E1} backups stored in {E2}.",
+    "{E0} enabled {E1} slow query logging in {E2}.",
+    "{E0} resolved {E1} config drift between {E2} environments.",
+    "{E0} automated {E1} provisioning with {E2} scripts.",
+    "{E0} fixed {E1} pagination bug in the {E2} API.",
+    "{E0} increased {E1} timeout from 30s to 60s in {E2}.",
+    "{E0} set up {E1} cross-region replication for {E2}.",
+    "{E0} migrated {E1} workloads from {E2} to {E3}.",
+    "{E0} reviewed {E1} schema change proposal for {E2}.",
+    "{E0} shipped {E1} structured logging for {E2}.",
+    "{E0} ran {E1} end-to-end tests against {E2}.",
+    "{E0} added {E1} graceful shutdown to {E2} workers.",
+    "{E0} investigated {E1} anomaly detected in {E2} metrics.",
+    "{E0} enabled {E1} write-ahead logging in {E2}.",
+    "{E0} decommissioned legacy {E1} in favor of {E2}.",
+    "{E0} added {E1} vector index to {E2} for similarity search.",
+    "{E0} trained {E1} classifier on {E2} labeled dataset.",
+    "{E0} profiled {E1} CPU usage spike in {E2}.",
+    "{E0} set up {E1} webhook integration for {E2} events.",
+    "{E0} implemented {E1} rate limiting using {E2}.",
+    "{E0} exported {E1} traces to {E2} for analysis.",
+    "{E0} configured {E1} connection pooling for {E2}.",
+    "{E0} updated {E1} documentation with {E2} examples.",
+    "{E0} resolved {E1} type mismatch between {E2} versions.",
+    "{E0} shipped {E1} bulk import feature for {E2}.",
+    "{E0} enabled {E1} query caching in {E2}.",
+    "{E0} ran security scan on {E1} container images for {E2}.",
+    "{E0} added {E1} idempotency keys to {E2} endpoints.",
+    "{E0} migrated {E1} logs to {E2} for centralized search.",
+    "{E0} opened discussion on {E1} retention policy in {E2}.",
+    "{E0} tagged {E1} release candidate for {E2} deployment.",
+    "{E0} configured {E1} resource quotas in {E2} namespace.",
+    "{E0} opened a {E1} incident for {E2} degradation.",
+    "{E0} validated {E1} schema migration on {E2} staging.",
+    "{E0} submitted {E1} change request for {E2} production.",
+    "{E0} identified {E1} as critical dependency for {E2}.",
+    "{E0} refactored {E1} plugin interface for {E2}.",
+    "{E0} enabled {E1} delta compression for {E2} exports.",
+    "{E0} built {E1} smoke tests for {E2} deployment checks.",
+    "{E0} resolved {E1} clock skew issue in {E2} cluster.",
+    "{E0} set up {E1} chaos mesh for {E2} resilience testing.",
+    "{E0} ran {E1} migration dry-run against {E2} production data.",
+    "{E0} added {E1} custom metrics to {E2} dashboards.",
+    "{E0} enabled {E1} multi-region failover for {E2}.",
+    "{E0} tested {E1} rollback procedure for {E2}.",
+    "{E0} audited {E1} access logs for {E2}.",
+    "{E0} updated {E1} routing rules in {E2}.",
+    "{E0} added {E1} API versioning to {E2}.",
+    "{E0} created {E1} architecture diagram for {E2}.",
+    "{E0} configured {E1} resource limits for {E2} pods.",
+    "{E0} shipped {E1} streaming response for {E2} endpoints.",
+    "{E0} added {E1} content negotiation to {E2}.",
+    "{E0} enabled {E1} request signing for {E2}.",
+    "{E0} fixed {E1} goroutine leak in {E2}.",
+    "{E0} tuned {E1} GC parameters for {E2}.",
+    "{E0} deployed {E1} hotfix to unblock {E2}.",
+    "{E0} ran capacity review for {E1} ahead of {E2} launch.",
+    "{E0} configured {E1} alerting thresholds for {E2}.",
+    "{E0} migrated {E1} config to environment variables in {E2}.",
+    "{E0} hardened {E1} container image for {E2}.",
+    "{E0} added {E1} exponential backoff to {E2} client.",
+    "{E0} enabled {E1} connection keep-alive in {E2}.",
+    "{E0} built {E1} synthetic monitor for {E2}.",
+    "{E0} added {E1} compression to {E2} API responses.",
+    "{E0} enabled {E1} query explain plans in {E2}.",
+    "{E0} reduced {E1} package size in {E2} by tree-shaking.",
+    "{E0} added {E1} dark mode support to {E2}.",
+    "{E0} configured {E1} output caching in {E2}.",
+    "{E0} submitted {E1} benchmarks comparing {E2} and {E3}.",
+    "{E0} enabled {E1} mTLS between {E2} and {E3}.",
+    "{E0} integrated {E1} error tracking into {E2}.",
+    "{E0} enabled {E1} feature gates for {E2} beta users.",
+    "{E0} shipped {E1} analytics events for {E2} funnel.",
+    "{E0} identified {E1} as cause of {E2} tail latency.",
+    "{E0} standardized {E1} logging format across {E2}.",
+    "{E0} ran {E1} regression tests before {E2} release.",
+    "{E0} set up {E1} PR preview environments for {E2}.",
+    "{E0} fixed {E1} index missing in {E2} query.",
+    "{E0} enabled {E1} statement timeout in {E2}.",
+    "{E0} reviewed {E1} data model with {E2} data team.",
+    "{E0} refactored {E1} middleware stack in {E2}.",
+    "{E0} ran {E1} fuzzing tests against {E2}.",
+    "{E0} set {E1} memory limits for {E2} workers.",
+    "{E0} added {E1} observability hooks to {E2}.",
+    "{E0} resolved {E1} permission issue between {E2} and {E3}.",
+    "{E0} enabled {E1} auto-vacuum in {E2}.",
+    "{E0} shipped {E1} batch delete API for {E2}.",
+    "{E0} added {E1} soft-delete flag to {E2} records.",
+    "{E0} created {E1} integration test suite for {E2}.",
+    "{E0} configured {E1} load shedding in {E2}.",
+    "{E0} set up {E1} incident response playbook for {E2}.",
+    "{E0} fixed {E1} N+1 query in {E2} listing endpoint.",
+    "{E0} validated {E1} SLOs for {E2} over the past quarter.",
+    "{E0} enabled {E1} continuous profiling in {E2}.",
+    "{E0} configured {E1} service discovery for {E2}.",
+    "{E0} deployed {E1} update with zero downtime to {E2}.",
+    "{E0} ran {E1} penetration test against {E2}.",
+    "{E0} added {E1} DKIM signing for {E2} emails.",
+    "{E0} resolved {E1} OOM kill in {E2} under load.",
+    "{E0} tuned {E1} work queue parallelism for {E2}.",
+    "{E0} deployed {E1} read replica for {E2} reporting.",
+    "{E0} shipped {E1} event replay feature for {E2}.",
+    "{E0} added {E1} schema registry support to {E2}.",
+    "{E0} configured {E1} dead letter queue for {E2}.",
+    "{E0} implemented {E1} RBAC for {E2} admin endpoints.",
+    "{E0} enabled {E1} mutual authentication for {E2}.",
+    "{E0} profiled {E1} GC pressure in {E2}.",
+    "{E0} ran {E1} compliance audit for {E2} data.",
+    "{E0} added {E1} multi-tenant isolation to {E2}.",
+    "{E0} configured {E1} network policies for {E2}.",
+    "{E0} shipped {E1} webhook retry logic for {E2}.",
+    "{E0} enabled {E1} prepared statements in {E2}.",
+    "{E0} deployed {E1} for zero-trust networking in {E2}.",
+    "{E0} enabled {E1} query result caching in {E2}.",
+    "{E0} added {E1} circuit breaker to {E2} outbound calls.",
+    "{E0} ran {E1} disaster recovery drill for {E2}.",
+    "{E0} tuned {E1} thread pool for {E2} workload.",
+    "{E0} shipped {E1} export endpoint for {E2} data.",
+    "{E0} configured {E1} anomaly detection in {E2}.",
+    "{E0} set up {E1} chaos experiment for {E2} resilience.",
+    "{E0} added {E1} distributed lock to {E2} cron jobs.",
+    "{E0} migrated {E1} codebase from {E2} to {E3}.",
+    "{E0} enabled {E1} structured error responses in {E2}.",
+    "{E0} shipped {E1} async processing for {E2} heavy tasks.",
+    "{E0} verified {E1} data integrity after {E2} migration.",
+    "{E0} added {E1} request deduplication to {E2}.",
+    "{E0} configured {E1} log sampling in {E2}.",
+    "{E0} enabled {E1} query parallelism in {E2}.",
+    "{E0} shipped {E1} data masking for {E2} PII fields.",
+    "{E0} reviewed {E1} deployment procedure for {E2}.",
+    "{E0} fixed {E1} connection leak under {E2} high load.",
+    "{E0} set up {E1} fan-out pattern for {E2} events.",
+    "{E0} enabled {E1} request coalescing in {E2}.",
+    "{E0} shipped {E1} GraphQL federation for {E2}.",
+    "{E0} enabled {E1} distributed caching for {E2}.",
+    "{E0} ran {E1} smoke test suite after {E2} deploy.",
+    "{E0} resolved {E1} config injection issue in {E2}.",
+    "{E0} shipped {E1} multi-region write support for {E2}.",
+]
+
+# Scale configuration: number of content items to submit as a single async batch
+SCALES = {
+    "tiny": 1,
+    "mini": 50,
+    "small": 2_000,
+    "medium": 10_000,
+    "large": 33_000,
+    "very-large": 100_000,
+}
+
+# ---------------------------------------------------------------------------
+# Zipf-like entity selector
+# ---------------------------------------------------------------------------
+
+
+def _make_entity_selector(seed: int = 42) -> "Callable[[int], list[str]]":
+    """Return a function that draws N entity names with Zipf-like distribution."""
+    import random
+
+    rng = random.Random(seed)
+    n = len(ENTITIES)
+    # Weights: entity i gets weight 1/(i+1)
+    weights = [1.0 / (i + 1) for i in range(n)]
+
+    def pick(count: int) -> list[str]:
+        seen = set()
+        result = []
+        while len(result) < count:
+            choice = rng.choices(ENTITIES, weights=weights, k=1)[0]
+            if choice not in seen:
+                seen.add(choice)
+                result.append(choice)
+        return result
+
+    return pick
+
+
+_pick_entities = _make_entity_selector()
+
+
+def _fill_template(template: str) -> str:
+    """Replace {E0}..{E4} placeholders in a template with entity names."""
+    placeholders = [f"{{E{i}}}" for i in range(5)]
+    needed = sum(1 for p in placeholders if p in template)
+    if needed == 0:
+        return template
+    entities = _pick_entities(needed)
+    result = template
+    for i, entity in enumerate(entities):
+        result = result.replace(f"{{E{i}}}", entity)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Mock LLM callback
+# ---------------------------------------------------------------------------
+
+from collections.abc import Callable  # noqa: E402  (after stdlib)
+
+
+def _make_fact_callback() -> tuple[Callable[[list[dict], str], Any], list[int]]:
+    """
+    Return (callback, call_counter) where call_counter[0] tracks invocations.
+
+    The callback cycles through FACT_TEMPLATES and returns a valid
+    FactExtractionResponse-compatible dict for the retain pipeline.
+    """
+    call_counter = [0]
+
+    def callback(messages: list[dict], scope: str) -> Any:
+        if scope == "retain_extract_facts":
+            idx = call_counter[0] % len(FACT_TEMPLATES)
+            call_counter[0] += 1
+            template = FACT_TEMPLATES[idx]
+            fact_text = _fill_template(template)
+            # Extract a few entity names from the filled text to populate entities
+            # field (very rough — enough to drive entity link creation)
+            entity_names = [e for e in ENTITIES if e in fact_text][:3]
+            entities = [{"text": e} for e in entity_names]
+            return {
+                "facts": [
+                    {
+                        "what": fact_text,
+                        "when": "N/A",
+                        "where": "N/A",
+                        "who": "N/A",
+                        "why": "N/A",
+                        "fact_type": "world",
+                        "entities": entities,
+                    }
+                ]
+            }
+        # All other scopes (entity resolution, etc.) — return empty
+        return {"facts": []}
+
+    return callback, call_counter
+
+
+# ---------------------------------------------------------------------------
+# Engine helpers
+# ---------------------------------------------------------------------------
+
+
+def _build_engine(*, disable_observations: bool = False) -> "Any":
+    """Create a MemoryEngine using mock LLM and DB from env."""
+    from hindsight_api import MemoryEngine
+
+    db_url = os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0")
+    if disable_observations:
+        os.environ["HINDSIGHT_API_ENABLE_OBSERVATIONS"] = "false"
+    engine = MemoryEngine(
+        db_url=db_url,
+        memory_llm_provider="mock",
+        memory_llm_api_key="",
+        memory_llm_model="mock",
+        skip_llm_verification=True,
+        db_command_timeout=600,  # Long timeout needed for large-bank inserts
+    )
+    return engine
+
+
+# ---------------------------------------------------------------------------
+# Subcommand: generate
+# ---------------------------------------------------------------------------
+
+
+async def _wait_for_operation(pool: Any, operation_id: str, timeout: float = 86400.0) -> str:
+    """
+    Poll async_operations every second until the parent reaches completed or failed.
+
+    Raises immediately if:
+    - the parent itself reaches 'failed'
+    - any direct child operation reaches 'failed'
+    """
+    import uuid
+
+    from hindsight_api.engine.task_backend import fq_table
+
+    table = fq_table("async_operations")
+    deadline = asyncio.get_event_loop().time() + timeout
+    parent_uuid = uuid.UUID(operation_id)
+
+    while asyncio.get_event_loop().time() < deadline:
+        # Check parent status
+        row = await pool.fetchrow(
+            f"SELECT status, error_message FROM {table} WHERE operation_id = $1",
+            parent_uuid,
+        )
+        if row:
+            if row["status"] == "completed":
+                return "completed"
+            if row["status"] == "failed":
+                raise RuntimeError(f"Operation {operation_id} failed: {row['error_message'] or 'unknown error'}")
+
+        # Fast-fail: any direct child that has already failed
+        failed_child = await pool.fetchrow(
+            f"""
+            SELECT operation_id, error_message
+            FROM {table}
+            WHERE result_metadata::jsonb->>'parent_operation_id' = $1
+              AND status = 'failed'
+            LIMIT 1
+            """,
+            operation_id,
+        )
+        if failed_child:
+            err = failed_child["error_message"] or "unknown error"
+            raise RuntimeError(f"Child operation {failed_child['operation_id']} failed: {err}")
+
+        await asyncio.sleep(1.0)
+
+    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:
+    """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
+
+    total_items = SCALES[scale]
+    console.print(
+        f"\n[bold cyan]Generate[/bold cyan] bank=[bold]{bank_id}[/bold] scale=[bold]{scale}[/bold] workers=[bold]{workers}[/bold]"
+    )
+    console.print(f"  Total items : {total_items:,}\n")
+
+    engine = _build_engine(disable_observations=True)
+    await engine.initialize()
+
+    # Attach mock callback to retain LLM config
+    callback, call_counter = _make_fact_callback()
+    engine._retain_llm_config.set_response_callback(callback)
+    engine._llm_config.set_response_callback(callback)
+
+    # Build all content items upfront
+    all_contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)])} for i in range(total_items)]
+
+    # Submit the whole batch as a single async operation (auto-splits by token budget)
+    console.print(f"  Submitting {total_items:,} items as async batch…")
+    result = await engine.submit_async_retain(
+        bank_id=bank_id,
+        contents=all_contents,
+        request_context=RequestContext(),
+    )
+    operation_id = result["operation_id"]
+    console.print(f"  Operation  : {operation_id}")
+
+    # Start in-process worker to drain the queue
+    pool = await engine._get_pool()
+    poller = WorkerPoller(
+        pool=pool,
+        worker_id="recall-perf-worker",
+        executor=engine.execute_task,
+        poll_interval_ms=200,
+        max_slots=workers,
+        consolidation_max_slots=0,
+        max_retries=20,
+    )
+    poller_task = asyncio.create_task(poller.run())
+    console.print("  Worker     : started\n")
+
+    # Wait for the parent operation to reach a terminal state
+    t0 = time.perf_counter()
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TimeElapsedColumn(),
+        console=console,
+    ) as progress:
+        progress.add_task("Processing async retain…")
+        final_status = await _wait_for_operation(pool, operation_id)
+
+    elapsed = time.perf_counter() - t0
+
+    # Graceful shutdown: stop accepting new tasks and wait for all in-flight
+    # tasks to complete (including post-transaction flush_pending_stats).
+    await poller.shutdown_graceful(timeout=60.0)
+    poller_task.cancel()
+    try:
+        await poller_task
+    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] "
+        f"in {elapsed:.1f}s  ({total_items / elapsed:.0f} items/s)"
+    )
+    console.print(f"LLM callback invoked {call_counter[0]:,} times.")
+
+
+# ---------------------------------------------------------------------------
+# RRF-only reranker (bypasses cross-encoder for DB-focused benchmarking)
+# ---------------------------------------------------------------------------
+
+
+class _RRFReranker:
+    """
+    Drop-in replacement for CrossEncoderReranker that uses RRF scores only.
+
+    Eliminates cross-encoder (CPU-bound ML inference) so recall timings
+    reflect pure DB interaction costs.
+    """
+
+    async def ensure_initialized(self) -> None:
+        pass
+
+    async def rerank(self, query: str, candidates: list) -> list:
+        from hindsight_api.engine.search.types import ScoredResult
+
+        scored = [ScoredResult(candidate=c, weight=c.rrf_score) for c in candidates]
+        scored.sort(key=lambda x: x.weight, reverse=True)
+        return scored
+
+
+# ---------------------------------------------------------------------------
+# Subcommand: benchmark
+# ---------------------------------------------------------------------------
+
+
+async def cmd_benchmark(bank_id: str, query: str, iterations: int, concurrency: int, reranker: str) -> None:
+    """Run recall in parallel and report p50/p95/p99 timings with per-step breakdown."""
+    from hindsight_api.models import RequestContext
+
+    console.print(f"\n[bold cyan]Benchmark[/bold cyan] bank=[bold]{bank_id}[/bold]")
+    console.print(f"  Query       : {query}")
+    console.print(f"  Iterations  : {iterations}  (total recall calls)")
+    console.print(f"  Concurrency : {concurrency}")
+    console.print(f"  Reranker    : {reranker}\n")
+
+    engine = _build_engine()
+    await engine.initialize()
+
+    if reranker == "rrf":
+        engine._cross_encoder_reranker = _RRFReranker()
+
+    request_context = RequestContext()
+    durations: list[float] = []
+    all_phase_timings: dict[str, list[float]] = {}
+
+    async def recall_one() -> float:
+        t0 = time.perf_counter()
+        result = await engine.recall_async(
+            bank_id=bank_id,
+            query=query,
+            max_tokens=4096,
+            enable_trace=True,
+            request_context=request_context,
+            _quiet=True,
+        )
+        elapsed = time.perf_counter() - t0
+        if result.trace:
+            summary = result.trace.get("summary", {})
+            for pm in summary.get("phase_metrics", []):
+                name = pm["phase_name"]
+                dur = pm["duration_seconds"]
+                all_phase_timings.setdefault(name, []).append(dur)
+        return elapsed
+
+    # Run in parallel batches of `concurrency` until `iterations` total calls are done
+    remaining = iterations
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        console=console,
+    ) as progress:
+        task = progress.add_task("Running recall…", total=iterations)
+
+        while remaining > 0:
+            batch_size = min(concurrency, remaining)
+            batch = await asyncio.gather(*[recall_one() for _ in range(batch_size)])
+            durations.extend(batch)
+            remaining -= batch_size
+            progress.advance(task, batch_size)
+
+    pool = await engine._get_pool()
+    await pool.close()
+
+    # Compute percentiles
+    sorted_d = sorted(durations)
+    n = len(sorted_d)
+
+    def pct(p: float) -> float:
+        idx = min(int(p / 100 * n), n - 1)
+        return sorted_d[idx]
+
+    table = Table(title=f"Recall Latency — bank={bank_id!r}  query={query!r}")
+    table.add_column("Metric", style="cyan")
+    table.add_column("Value", style="green", justify="right")
+
+    table.add_row("Total calls", str(n))
+    table.add_row("Concurrency", str(concurrency))
+    table.add_row("Mean", f"{statistics.mean(durations):.3f}s")
+    table.add_row("p50", f"{pct(50):.3f}s")
+    table.add_row("p95", f"{pct(95):.3f}s")
+    table.add_row("p99", f"{pct(99):.3f}s")
+    table.add_row("Max", f"{max(durations):.3f}s")
+    table.add_row("Min", f"{min(durations):.3f}s")
+
+    console.print("\n")
+    console.print(table)
+
+    if all_phase_timings:
+        phase_table = Table(title="Per-Step Timing Breakdown (across all calls)")
+        phase_table.add_column("Step", style="cyan")
+        phase_table.add_column("Mean", style="green", justify="right")
+        phase_table.add_column("p50", style="green", justify="right")
+        phase_table.add_column("p95", style="yellow", justify="right")
+        phase_table.add_column("Max", style="red", justify="right")
+
+        # Sort by mean duration descending so the bottleneck is at the top
+        sorted_phases = sorted(all_phase_timings.items(), key=lambda x: statistics.mean(x[1]), reverse=True)
+        for name, times in sorted_phases:
+            st = sorted(times)
+            m = statistics.mean(times)
+            p50_v = st[min(int(0.5 * len(st)), len(st) - 1)]
+            p95_v = st[min(int(0.95 * len(st)), len(st) - 1)]
+            mx = max(times)
+            phase_table.add_row(name, f"{m:.3f}s", f"{p50_v:.3f}s", f"{p95_v:.3f}s", f"{mx:.3f}s")
+
+        console.print(phase_table)
+
+
+# ---------------------------------------------------------------------------
+# Subcommand: stats
+# ---------------------------------------------------------------------------
+
+
+async def cmd_stats(bank_ids: list[str]) -> None:
+    """Print memory / entity / link counts for one or more banks."""
+    from hindsight_api.models import RequestContext
+
+    engine = _build_engine()
+    await engine.initialize()
+    ctx = RequestContext()
+
+    table = Table(title="Bank Statistics")
+    table.add_column("Bank ID", style="cyan")
+    table.add_column("Units", style="green", justify="right")
+    table.add_column("Links (total)", style="yellow", justify="right")
+    table.add_column("Links by type", style="white")
+
+    for bank_id in bank_ids:
+        try:
+            stats = await engine.get_bank_stats(bank_id=bank_id, request_context=ctx)
+            total_units = sum(stats.get("node_counts", {}).values())
+            total_links = sum(stats.get("link_counts", {}).values())
+            links_detail = "  ".join(f"{k}={v:,}" for k, v in sorted(stats.get("link_counts", {}).items()))
+            table.add_row(bank_id, f"{total_units:,}", "-", f"{total_links:,}", links_detail)
+        except Exception as e:
+            table.add_row(bank_id, "ERROR", "", "", str(e))
+
+    pool = await engine._get_pool()
+    await pool.close()
+
+    console.print("\n")
+    console.print(table)
+
+
+# ---------------------------------------------------------------------------
+# Subcommand: clean
+# ---------------------------------------------------------------------------
+
+
+async def cmd_clean(bank_id: str) -> None:
+    """Delete all data for a bank."""
+    from hindsight_api.models import RequestContext
+
+    console.print(f"\n[bold red]Clean[/bold red] bank=[bold]{bank_id}[/bold]")
+
+    engine = _build_engine()
+    await engine.initialize()
+
+    result = await engine.delete_bank(bank_id=bank_id, request_context=RequestContext())
+
+    pool = await engine._get_pool()
+    await pool.close()
+
+    table = Table(title=f"Deleted from bank={bank_id!r}")
+    table.add_column("Table", style="cyan")
+    table.add_column("Rows deleted", style="red", justify="right")
+    for k, v in result.items():
+        table.add_row(k, str(v))
+    console.print("\n")
+    console.print(table)
+    console.print("\n[green]Done.[/green]")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Large-bank recall load test (no LLM)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    # generate
+    gen = sub.add_parser("generate", help="Populate a synthetic bank")
+    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)")
+
+    # benchmark
+    bm = sub.add_parser("benchmark", help="Run recall and report latency")
+    bm.add_argument("--bank-id", required=True)
+    bm.add_argument("--query", required=True)
+    bm.add_argument("--iterations", type=int, default=10, help="Total number of recall calls (default: 10)")
+    bm.add_argument("--concurrency", type=int, default=1, help="Parallel recalls per batch (default: 1)")
+    bm.add_argument(
+        "--reranker",
+        choices=["rrf", "cross-encoder"],
+        default="rrf",
+        help="Reranker to use: rrf=RRF scores only (no ML), cross-encoder=neural reranker (default: rrf)",
+    )
+
+    # stats
+    st = sub.add_parser("stats", help="Print memory/entity/link counts for banks")
+    st.add_argument("bank_ids", nargs="+", metavar="BANK_ID")
+
+    # clean
+    cl = sub.add_parser("clean", help="Delete all data for a bank")
+    cl.add_argument("--bank-id", required=True)
+
+    args = parser.parse_args()
+
+    if args.cmd == "generate":
+        asyncio.run(cmd_generate(args.bank_id, args.scale, workers=args.workers))
+    elif args.cmd == "benchmark":
+        asyncio.run(cmd_benchmark(args.bank_id, args.query, args.iterations, args.concurrency, args.reranker))
+    elif args.cmd == "stats":
+        asyncio.run(cmd_stats(args.bank_ids))
+    elif args.cmd == "clean":
+        asyncio.run(cmd_clean(args.bank_id))
+
+
+if __name__ == "__main__":
+    main()

package-lock.json

@@ -13,7 +13,7 @@
     },
     "hindsight-clients/typescript": {
       "name": "@vectorize-io/hindsight-client",
-      "version": "0.4.11",
+      "version": "0.4.14",
       "license": "MIT",
       "devDependencies": {
         "@hey-api/openapi-ts": "0.88.0",
@@ -131,7 +131,7 @@
     },
     "hindsight-control-plane": {
       "name": "@vectorize-io/hindsight-control-plane",
-      "version": "0.4.11",
+      "version": "0.4.14",
       "license": "ISC",
       "dependencies": {
         "@radix-ui/react-alert-dialog": "^1.1.15",