holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

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

Browse source
This commit is contained in:
Nicolò Boschi 2026-03-03 13:35:22 +01:00 committed by GitHub
parent 5aff8e0c70
commit 7942f181c2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 2028 additions and 822 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
hindsight-api/hindsight_api/alembic/versions/b3c4d5e6f7g8_add_temporal_date_indexes.py Normal file
View file

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

46
hindsight-api/hindsight_api/alembic/versions/c1a2b3d4e5f6_enable_pg_trgm_and_entities_trgm_index.py Normal file
View file

@ -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

83
hindsight-api/hindsight_api/alembic/versions/d2e3f4a5b6c7_add_memory_links_expansion_indexes.py Normal file
View file

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

154
hindsight-api/hindsight_api/api/http.py
View file

@ -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())
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=total_nodes,
total_links=total_links,
total_documents=total_documents,
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=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,
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):

4
hindsight-api/hindsight_api/config.py
View file

@ -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(

1
hindsight-api/hindsight_api/engine/db_utils.py
View file

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

361
hindsight-api/hindsight_api/engine/entity_resolver.py
View file

@ -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))
# 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,
entities_to_create.append(
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
)
# Batch create new entities using COPY + INSERT for maximum speed
# This handles duplicates via ON CONFLICT and returns all IDs
# 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)
# 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]:
# 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]:

8
hindsight-api/hindsight_api/engine/llm_wrapper.py
View file

@ -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

184
hindsight-api/hindsight_api/engine/memory_engine.py
View file

@ -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(

26
hindsight-api/hindsight_api/engine/providers/mock_llm.py
View file

@ -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.

8
hindsight-api/hindsight_api/engine/query_analyzer.py
View file

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

3
hindsight-api/hindsight_api/engine/retain/__init__.py
View file

@ -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",

85
hindsight-api/hindsight_api/engine/retain/deduplication.py
View file

@ -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]

3
hindsight-api/hindsight_api/engine/retain/embedding_utils.py
View file

@ -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,
)

147
hindsight-api/hindsight_api/engine/retain/link_utils.py
View file

@ -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)
if not embedding_arrays:
existing_embeddings = np.array([])
elif len(embedding_arrays) == 1:
# Single embedding: reshape to (1, dim)
existing_embeddings = embedding_arrays[0].reshape(1, -1)
else:
# Multiple embeddings: vstack
existing_embeddings = np.vstack(embedding_arrays)
# For each new unit, compute similarities with ALL existing units
for unit_id, new_embedding in zip(unit_ids, embeddings):
new_emb_array = np.array(new_embedding)
emb_str = str(list(new_embedding) if not isinstance(new_embedding, list) else new_embedding)
rows = await conn.fetch(
f"""
SELECT id::text,
1 - (embedding <=> $1::vector) AS similarity
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND embedding IS NOT NULL
AND id != ALL($3::uuid[])
ORDER BY embedding <=> $1::vector
LIMIT $4
""",
emb_str,
bank_id,
exclude_uuids,
top_k,
)
for row in rows:
sim = float(min(1.0, max(0.0, row["similarity"])))
if sim >= threshold:
all_links.append((unit_id, str(row["id"]), "semantic", sim, None))
# Compute cosine similarities (dot product for normalized vectors)
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"""
total_rows = len(records)
chunks = 0
for chunk_start in range(0, total_rows, chunk_size):
chunk_end = chunk_start + chunk_size
await conn.execute(
f"""
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
FROM _temp_entity_links
WHERE _row_num > $1 AND _row_num <= $2
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""")
logger.debug(f" [9.5] INSERT from temp table: {time_mod.time() - insert_start:.3f}s")
""",
chunk_start,
chunk_end,
)
chunks += 1
logger.debug(f" [9.5] INSERT {total_rows} rows in {chunks} chunks: {time_mod.time() - insert_start:.3f}s")
logger.debug(f" [9.TOTAL] Entity links batch insert: {time_mod.time() - total_start:.3f}s")

51
hindsight-api/hindsight_api/engine/retain/orchestrator.py
View file

@ -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

521
hindsight-api/hindsight_api/engine/search/link_expansion_retrieval.py
View file

@ -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

69
hindsight-api/hindsight_api/engine/search/retrieval.py
View file

@ -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 unnest($2::uuid[]) AS src(from_unit_id)
CROSS JOIN LATERAL (
SELECT ml.to_unit_id, ml.weight, ml.link_type
FROM {fq_table("memory_links")} ml
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($2::uuid[])
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,
)

1
hindsight-api/hindsight_api/main.py
View file

@ -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,

43
hindsight-api/hindsight_api/migrations.py
View file

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

6
hindsight-api/tests/test_worker.py
View file

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

942
hindsight-dev/benchmarks/perf/recall_perf.py Normal file
View file

@ -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()

4
package-lock.json generated
View file

@ -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",