diff --git a/hindsight-api/hindsight_api/alembic/versions/b3c4d5e6f7g8_add_temporal_date_indexes.py b/hindsight-api/hindsight_api/alembic/versions/b3c4d5e6f7g8_add_temporal_date_indexes.py new file mode 100644 index 00000000..4daf0a72 --- /dev/null +++ b/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") diff --git a/hindsight-api/hindsight_api/alembic/versions/c1a2b3d4e5f6_enable_pg_trgm_and_entities_trgm_index.py b/hindsight-api/hindsight_api/alembic/versions/c1a2b3d4e5f6_enable_pg_trgm_and_entities_trgm_index.py new file mode 100644 index 00000000..db7a707d --- /dev/null +++ b/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 diff --git a/hindsight-api/hindsight_api/alembic/versions/d2e3f4a5b6c7_add_memory_links_expansion_indexes.py b/hindsight-api/hindsight_api/alembic/versions/d2e3f4a5b6c7_add_memory_links_expansion_indexes.py new file mode 100644 index 00000000..8a129ff0 --- /dev/null +++ b/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") diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 49d7442f..a3597e09 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/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 - await app.state.memory._authenticate_tenant(request_context) - if app.state.memory._operation_validator: - from hindsight_api.extensions import BankReadContext - - ctx = BankReadContext(bank_id=bank_id, operation="get_bank_stats", request_context=request_context) - await app.state.memory._validate_operation( - app.state.memory._operation_validator.validate_bank_read(ctx) - ) - pool = await app.state.memory._get_pool() - async with acquire_with_retry(pool) as conn: - # 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_links=total_links, - total_documents=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, - failed_operations=failed_operations, - last_consolidated_at=(last_consolidated_at.isoformat() if last_consolidated_at else None), - pending_consolidation=pending_consolidation, - total_observations=total_observations, - ) - + stats = await app.state.memory.get_bank_stats(bank_id, request_context=request_context) + nodes_by_type = stats["node_counts"] + links_by_type = stats["link_counts"] + links_by_fact_type = stats["link_counts_by_fact_type"] + links_breakdown: dict[str, dict[str, int]] = {} + for row in stats["link_breakdown"]: + ft = row["fact_type"] + if ft not in links_breakdown: + links_breakdown[ft] = {} + links_breakdown[ft][row["link_type"]] = row["count"] + ops = stats["operations"] + return BankStatsResponse( + bank_id=bank_id, + total_nodes=sum(nodes_by_type.values()), + total_links=sum(links_by_type.values()), + 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=ops.get("pending", 0), + failed_operations=ops.get("failed", 0), + last_consolidated_at=stats["last_consolidated_at"], + pending_consolidation=stats["pending_consolidation"], + total_observations=stats["total_observations"], + ) except OperationValidationError as e: raise HTTPException(status_code=e.status_code, detail=e.reason) except (AuthenticationError, HTTPException): diff --git a/hindsight-api/hindsight_api/config.py b/hindsight-api/hindsight_api/config.py index 409485b0..0a26c16f 100644 --- a/hindsight-api/hindsight_api/config.py +++ b/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( diff --git a/hindsight-api/hindsight_api/engine/db_utils.py b/hindsight-api/hindsight_api/engine/db_utils.py index acdc5cbc..475262dc 100644 --- a/hindsight-api/hindsight_api/engine/db_utils.py +++ b/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, diff --git a/hindsight-api/hindsight_api/engine/entity_resolver.py b/hindsight-api/hindsight_api/engine/entity_resolver.py index 770b72d1..f166e8a8 100644 --- a/hindsight-api/hindsight_api/engine/entity_resolver.py +++ b/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_create = [] # (idx, entity_data, event_date) - - taxonomy_lookup = taxonomy_lookup or set() + entities_to_update: list[_EntityStat] = [] + entities_to_create: list[_EntityToCreate] = [] 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, - ) + # Existing entities: IDs already known from the candidate SELECT above. + # 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) - # Batch create new entities using COPY + INSERT for maximum speed - # This handles duplicates via ON CONFLICT and returns all IDs + # 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 - # For duplicates, we only insert once and reuse the ID, but track the count - unique_entities = {} # lowercase_name -> (entity_data, event_date, [indices]) - for idx, entity_data, event_date in entities_to_create: - name_lower = entity_data["text"].lower() - if name_lower not in unique_entities: - 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) + # Group by lowercase name — deduplicate within the batch. + @dataclass + class _NameGroup: + name: str + event_date: datetime | None + indices: list[int] = field(default_factory=list) - # Batch insert unique entities and get their IDs - # Use a single query with unnest for speed - entity_names = [] - entity_dates = [] - entity_counts = [] # Track how many times each entity appears in this batch - indices_map = [] # Maps result index -> list of original indices + groups: dict[str, _NameGroup] = {} + for e in entities_to_create: + name_lower = e.name.lower() + if name_lower not in groups: + groups[name_lower] = _NameGroup(name=e.name, event_date=e.event_date) + groups[name_lower].indices.append(e.idx) - for name_lower, (entity_data, event_date, indices) in unique_entities.items(): - entity_names.append(entity_data["text"]) - entity_dates.append(event_date) - entity_counts.append(len(indices)) # Count of occurrences in this batch - indices_map.append(indices) + # Sort by lowercase name for deterministic ordering. + sorted_groups = sorted(groups.items()) + entity_names = [g.name for _, g in sorted_groups] + entity_dates = [g.event_date for _, g in sorted_groups] - # Batch INSERT ... ON CONFLICT with RETURNING - # Uses the batch count for mention_count instead of always 1 - rows = await conn.fetch( + # INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities. + inserted_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 - FROM unnest($2::text[], $3::timestamptz[], $4::int[]) AS t(name, event_date, cnt) + SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 1 + FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date) ON CONFLICT (bank_id, LOWER(canonical_name)) - DO UPDATE SET - mention_count = {fq_table("entities")}.mention_count + EXCLUDED.mention_count, - last_seen = EXCLUDED.last_seen - RETURNING id + DO NOTHING + RETURNING id, LOWER(canonical_name) AS name_lower """, 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 - for result_idx, row in enumerate(rows): - entity_id = row["id"] - for original_idx in indices_map[result_idx]: - entity_ids[original_idx] = entity_id + # Fallback SELECT for names that conflicted (another worker won the race). + missing = [n for n, _ in sorted_groups if n not in id_by_name] + if missing: + 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) - 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 + 1, - last_cooccurred = EXCLUDED.last_cooccurred - """, - [(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs], + key = self._task_key() + self._pending_cooccurrences.setdefault(key, []).extend( + _CooccurrencePair(entity_id_1=e1, entity_id_2=e2) for e1, e2 in cooccurrence_pairs ) async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]: diff --git a/hindsight-api/hindsight_api/engine/llm_wrapper.py b/hindsight-api/hindsight_api/engine/llm_wrapper.py index a6f1f2a6..86805070 100644 --- a/hindsight-api/hindsight_api/engine/llm_wrapper.py +++ b/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 diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 84b20f2a..0251f306 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/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 - self.entity_resolver = EntityResolver(self._pool) + # Initialize entity resolver with pool and configured lookup strategy + 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 - 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 + # Single query for all link stats — avoids triple join on memory_links (can be 21M+ rows). + # link_counts and link_counts_by_fact_type are derived in Python from the breakdown. 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}, - "link_counts": {row["link_type"]: row["count"] for row in link_stats}, - "link_counts_by_fact_type": {row["fact_type"]: row["count"] for row in link_fact_type_stats}, + "node_counts": node_counts, + "link_counts": link_counts, + "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( diff --git a/hindsight-api/hindsight_api/engine/providers/mock_llm.py b/hindsight-api/hindsight_api/engine/providers/mock_llm.py index 680b5602..a8ca2b68 100644 --- a/hindsight-api/hindsight_api/engine/providers/mock_llm.py +++ b/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. diff --git a/hindsight-api/hindsight_api/engine/query_analyzer.py b/hindsight-api/hindsight_api/engine/query_analyzer.py index 5c9b568d..a78581c7 100644 --- a/hindsight-api/hindsight_api/engine/query_analyzer.py +++ b/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: """ diff --git a/hindsight-api/hindsight_api/engine/retain/__init__.py b/hindsight-api/hindsight_api/engine/retain/__init__.py index 6fd95827..71cd84ef 100644 --- a/hindsight-api/hindsight_api/engine/retain/__init__.py +++ b/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", diff --git a/hindsight-api/hindsight_api/engine/retain/deduplication.py b/hindsight-api/hindsight_api/engine/retain/deduplication.py deleted file mode 100644 index ead97001..00000000 --- a/hindsight-api/hindsight_api/engine/retain/deduplication.py +++ /dev/null @@ -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] diff --git a/hindsight-api/hindsight_api/engine/retain/embedding_utils.py b/hindsight-api/hindsight_api/engine/retain/embedding_utils.py index 53ebc762..269acb07 100644 --- a/hindsight-api/hindsight_api/engine/retain/embedding_utils.py +++ b/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, ) diff --git a/hindsight-api/hindsight_api/engine/retain/link_utils.py b/hindsight-api/hindsight_api/engine/retain/link_utils.py index 103dfc31..0f6e2a00 100644 --- a/hindsight-api/hindsight_api/engine/retain/link_utils.py +++ b/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 - fetch_start = time_mod.time() - all_existing = await conn.fetch( - f""" - SELECT id, embedding - 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() + # Use pgvector ANN search (HNSW index) for each new unit instead of fetching + # all existing embeddings into Python. At large scale (100K+ units) the old + # approach would transfer 100K × 384 floats (~150 MB) per retain call; the + # ANN query completes in <5 ms and transfers only top_k rows. + ann_start = time_mod.time() all_links = [] - if all_existing: - # Convert existing embeddings to numpy array - 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 + # Build UUID exclude list once for all ANN queries + import uuid as uuid_mod - 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 (works for numpy arrays, pgvector objects, etc.) - emb = np.array(raw_emb, dtype=np.float32) + exclude_uuids = [uuid_mod.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids] - # Ensure it's 1D - if emb.ndim != 1: - raise ValueError(f"Expected 1D embedding, got shape {emb.shape}") - embedding_arrays.append(emb) + for unit_id, new_embedding in zip(unit_ids, embeddings): + 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)) - 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) - - # Compute cosine similarities (dot product for normalized vectors) - similarities = np.dot(existing_embeddings, new_emb_array) - - # 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)) + _log( + 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", + ) # 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, - then INSERT ... ON CONFLICT from temp table. This is the fastest - method for bulk inserts with conflict handling. + Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading into a + temp table, then INSERT ... ON CONFLICT in chunks of chunk_size. Chunking + 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 = [] - for link in links: - records.append((link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id)) + records = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links] 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""" - 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 - 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") + 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 + """, + 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") diff --git a/hindsight-api/hindsight_api/engine/retain/orchestrator.py b/hindsight-api/hindsight_api/engine/retain/orchestrator.py index 3d1829a7..deca2ee8 100644 --- a/hindsight-api/hindsight_api/engine/retain/orchestrator.py +++ b/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 - 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 + non_duplicate_facts = processed_facts # 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. - - 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))} + """Map created unit IDs back to original content items.""" + facts_by_content: dict[int, list[int]] = {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]: - if not is_duplicate_flags[fact_idx]: - content_unit_ids.append(unit_ids[filtered_idx]) - filtered_idx += 1 + for _ in facts_by_content[content_index]: + content_unit_ids.append(unit_ids[unit_idx]) + unit_idx += 1 result_unit_ids.append(content_unit_ids) return result_unit_ids diff --git a/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py b/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py index eb97c914..5720565f 100644 --- a/hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py +++ b/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: -1. Entity links: Find facts sharing entities with seeds (filtered by entity frequency) -2. Causal links: Find facts causally linked to seeds (top-k by weight) +Expands from semantic/temporal seeds through three parallel, first-class signals +stored in memory_links: -Characteristics: -- 2-3 DB queries (seed finding + parallel entity/causal expansion) -- Sublinear: only touches connected facts via indexes -- No iteration, no propagation, no normalization -- Target: <100ms +1. Entity links — precomputed co-occurrence graph (created at retain time, bounded to + MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared + entities between the seed set and each candidate. +2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most + 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. - Fast and simple alternative to MPFP. + Runs three expansions through precomputed memory_links: entity co-occurrence, + 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 - causal_limit_per_seed: Max causal links to follow per seed + causal_weight_threshold: Minimum weight for causal links to follow. """ - 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 - 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") + entity_rows, semantic_rows, causal_rows = await self._expand_observations(conn, seed_ids, budget) else: - # For world/experience facts, use direct entity lookup - 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, - ) + entity_rows, semantic_rows, causal_rows = await self._expand_combined(conn, seed_ids, fact_type, budget) timings.edge_load_time = time.time() - query_start - timings.db_queries = 3 - timings.edge_count = len(entity_rows) + len(causal_rows) + len(fallback_rows) + timings.db_queries = 1 + timings.edge_count = len(entity_rows) + len(semantic_rows) + len(causal_rows) - # Merge results, taking max score per fact - # Priority: entity links (unit_entities) > causal links > fallback links - score_map: dict[str, float] = {} + # Merge results with additive intra-score: entity + semantic + causal ∈ [0, 3]. + # + # 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"]) - if fact_id not in row_map: - row_map[fact_id] = dict(row) + causal_scores[fact_id] = max(causal_scores.get(fact_id, 0.0), row["score"]) + row_map.setdefault(fact_id, dict(row)) - for row in fallback_rows: - fact_id = str(row["id"]) - score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"]) - if fact_id not in row_map: - row_map[fact_id] = dict(row) + all_ids = set(entity_scores) | set(semantic_scores) | set(causal_scores) + score_map = { + fid: entity_scores.get(fid, 0.0) + semantic_scores.get(fid, 0.0) + causal_scores.get(fid, 0.0) + for fid in all_ids + } - # 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 diff --git a/hindsight-api/hindsight_api/engine/search/retrieval.py b/hindsight-api/hindsight_api/engine/search/retrieval.py index 9682b55d..0189c731 100644 --- a/hindsight-api/hindsight_api/engine/search/retrieval.py +++ b/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 ( - SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, - 1 - (embedding <=> $1::vector) AS similarity, - ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, embedding <=> $1::vector) AS rn + WITH date_ranked AS MATERIALIZED ( + SELECT id, fact_type, + ROW_NUMBER() OVER ( + 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 - WHERE rn <= 10 + FROM sim_ranked + 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) - spreading_tags_clause = build_tags_where_clause_simple(tags, 6, table_alias="mu.", match=tags_match) + # Build tags clause for spreading (use param 7 since 1-6 are used) + 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, - ml.weight, ml.link_type, ml.from_unit_id, + 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, + l.weight, l.link_type, 1 - (mu.embedding <=> $1::vector) AS similarity - 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($2::uuid[]) - AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents') - AND ml.weight >= 0.1 + 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 + WHERE ml.from_unit_id = src.from_unit_id + 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, ) diff --git a/hindsight-api/hindsight_api/main.py b/hindsight-api/hindsight_api/main.py index d76766c0..635a1caf 100644 --- a/hindsight-api/hindsight_api/main.py +++ b/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, diff --git a/hindsight-api/hindsight_api/migrations.py b/hindsight-api/hindsight_api/migrations.py index fd7d43df..45245087 100644 --- a/hindsight-api/hindsight_api/migrations.py +++ b/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: diff --git a/hindsight-api/tests/test_worker.py b/hindsight-api/tests/test_worker.py index b4341a5b..e5e750b9 100644 --- a/hindsight-api/tests/test_worker.py +++ b/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 - await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'") + # 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-%' 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: diff --git a/hindsight-dev/benchmarks/perf/recall_perf.py b/hindsight-dev/benchmarks/perf/recall_perf.py new file mode 100644 index 00000000..711580c8 --- /dev/null +++ b/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() diff --git a/package-lock.json b/package-lock.json index b3418a0b..ebfefd44 100644 --- a/package-lock.json +++ b/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",