From aac9946676bb8d18da629be32f9872e7e23040ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 26 May 2026 17:35:46 +0200 Subject: [PATCH 1/4] chore: regenerate docs skill (sync Tigris S3 config notes) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drift picked up by the generate-docs-skill pre-commit hook — keeps skills/hindsight-docs/ in sync with the upstream hindsight-docs/ sources. --- .../references/developer/configuration.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/skills/hindsight-docs/references/developer/configuration.md b/skills/hindsight-docs/references/developer/configuration.md index d3ee797dd3..089f5f14ca 100644 --- a/skills/hindsight-docs/references/developer/configuration.md +++ b/skills/hindsight-docs/references/developer/configuration.md @@ -1071,10 +1071,12 @@ export HINDSIGHT_API_FILE_STORAGE_TYPE=native |----------|-------------|---------| | `HINDSIGHT_API_FILE_STORAGE_S3_BUCKET` | S3 bucket name | - | | `HINDSIGHT_API_FILE_STORAGE_S3_REGION` | AWS region | - | -| `HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT` | Custom endpoint URL (for S3-compatible stores like MinIO, Cloudflare R2) | AWS default | +| `HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT` | Custom endpoint URL (for S3-compatible stores like MinIO, Cloudflare R2, Tigris) | AWS default | | `HINDSIGHT_API_FILE_STORAGE_S3_ACCESS_KEY_ID` | AWS access key ID | - | | `HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY` | AWS secret access key | - | +For S3-compatible providers that don't expose AWS-style regions (MinIO, Cloudflare R2, Tigris), set `HINDSIGHT_API_FILE_STORAGE_S3_REGION=auto`. The value is required for SigV4 request signing but is ignored by the service. + ```bash # AWS S3 export HINDSIGHT_API_FILE_STORAGE_TYPE=s3 @@ -1086,9 +1088,18 @@ export HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPx # S3-compatible (MinIO, Cloudflare R2, etc.) export HINDSIGHT_API_FILE_STORAGE_TYPE=s3 export HINDSIGHT_API_FILE_STORAGE_S3_BUCKET=my-bucket +export HINDSIGHT_API_FILE_STORAGE_S3_REGION=auto export HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT=https://your-minio.example.com export HINDSIGHT_API_FILE_STORAGE_S3_ACCESS_KEY_ID=minioadmin export HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY=minioadmin + +# Tigris (S3-compatible, single global endpoint) +export HINDSIGHT_API_FILE_STORAGE_TYPE=s3 +export HINDSIGHT_API_FILE_STORAGE_S3_BUCKET=my-hindsight-bucket +export HINDSIGHT_API_FILE_STORAGE_S3_REGION=auto +export HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT=https://t3.storage.dev +export HINDSIGHT_API_FILE_STORAGE_S3_ACCESS_KEY_ID=tid_your_access_key +export HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY=tsec_your_secret_key ``` #### Google Cloud Storage From 4432d9858b8f02018b78c04e037dda4002e8514d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 26 May 2026 17:36:34 +0200 Subject: [PATCH 2/4] perf(api): derive entity edges from unit_entities instead of materializing them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stop writing link_type='entity' rows to memory_links and derive entity edges on demand in the /graph endpoint (from the unit_entities self-join recall already uses) and in /stats (by replicating the historical writer cap). Why: on the recall-perf-medium bench bank (10k units), entity rows were 53% of all memory_links — 345k rows, ~190 MB of table+index — and recall never read them (entity expansion in link_expansion_retrieval.py uses unit_entities, not memory_links). Retain was running a synchronous pairwise loop per shared entity to write rows nothing read; per-unit entity degree was uncapped (max 326 outgoing on a single unit), and overall per-unit total degree averaged 130 with a p99 of 462. Changes: - Drop Phase 3 entity-link build/insert from retain orchestrator. Keep entity_resolver.flush_pending_stats() so entity_cooccurrences (which feeds /entities/graph) still updates. - Delete build_entity_links_from_resolved, insert_entity_links_batch, MAX_LINKS_PER_ENTITY, EntityLink, Phase3Context, and the now-dead fetch_entity_unit_fanout op (PG + Oracle). - /graph: filter memory_links query to link_type <> 'entity'; broaden the existing observation-inferred entity-pair loop to cover all visible units; cap at 10 units per entity to bound hot entities. - /stats: split link_breakdown into a memory_links query (non-entity) and a unit_entities-based derivation for entity, sized to the historical writer cap so link_counts.entity stays in the same magnitude. - Migration e9b2c7d1f3a4: drop idx_memory_links_entity_covering and chunk-delete existing entity rows (PG + Oracle paths). - Tests: rewrite test_entity_links_creation and test_all_link_types_together to assert via /graph + /stats; assert no entity rows in memory_links. API response shapes (graph edges, stats link_counts/links_breakdown) are unchanged at the boundary, so SDKs and the control plane do not need to be regenerated. --- .../e9b2c7d1f3a4_drop_entity_memory_links.py | 90 ++++++++++ .../hindsight_api/engine/db/ops.py | 15 -- .../hindsight_api/engine/db/ops_oracle.py | 24 --- .../hindsight_api/engine/db/ops_postgresql.py | 23 --- .../hindsight_api/engine/memory_engine.py | 88 +++++++--- .../engine/retain/entity_processing.py | 72 +------- .../hindsight_api/engine/retain/link_utils.py | 155 ------------------ .../engine/retain/orchestrator.py | 91 ++-------- .../hindsight_api/engine/retain/types.py | 30 ---- hindsight-api-slim/tests/test_retain.py | 113 +++++-------- 10 files changed, 215 insertions(+), 486 deletions(-) create mode 100644 hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py diff --git a/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py b/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py new file mode 100644 index 0000000000..b6e6b01df8 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py @@ -0,0 +1,90 @@ +"""Drop materialized entity rows from memory_links. + +Entity edges are no longer stored in ``memory_links``. The /graph endpoint +derives them on demand from ``unit_entities``, and recall already used the +``unit_entities`` self-join. Storing entity rows duplicated state we never +read from the link table — on a 10k-unit benchmark bank, entity rows were +53% of all link rows (~190 MB after indexes) and recall never touched them. + +This migration: + +1. Drops ``idx_memory_links_entity_covering`` (the partial index targeting + ``WHERE link_type = 'entity'`` rows that no longer exist). +2. Deletes ``memory_links`` rows with ``link_type = 'entity'``. + +Revision ID: e9b2c7d1f3a4 +Revises: 86f7a033d372 +Create Date: 2026-05-26 +""" + +from collections.abc import Sequence + +from alembic import context, op + +from hindsight_api.alembic._dialect import run_for_dialect + +revision: str = "e9b2c7d1f3a4" +down_revision: str | Sequence[str] | None = "86f7a033d372" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _pg_schema_prefix() -> str: + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + +def _pg_upgrade() -> None: + schema = _pg_schema_prefix() + + # Drop the partial covering index first so the bulk DELETE doesn't churn it. + # CREATE/DROP INDEX CONCURRENTLY must run outside a transaction block. + op.execute("COMMIT") + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering") + + # Delete entity rows. Chunked to keep individual transactions small on + # large banks (the perf-medium bench had ~345k entity rows; production + # banks can be much larger). + op.execute( + f""" + DO $$ + DECLARE + deleted INTEGER; + BEGIN + LOOP + DELETE FROM {schema}memory_links + WHERE ctid IN ( + SELECT ctid FROM {schema}memory_links + WHERE link_type = 'entity' + LIMIT 50000 + ); + GET DIAGNOSTICS deleted = ROW_COUNT; + EXIT WHEN deleted = 0; + COMMIT; + END LOOP; + END$$; + """ + ) + + +def _pg_downgrade() -> None: + # Cannot reconstruct deleted entity links — the writer was path-dependent + # on retain order. New retains will not produce entity rows either, so the + # partial index would stay empty. Leave both no-op. + pass + + +def _oracle_upgrade() -> None: + op.execute("DELETE FROM memory_links WHERE link_type = 'entity'") + + +def _oracle_downgrade() -> None: + pass + + +def upgrade() -> None: + run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade) + + +def downgrade() -> None: + run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade) diff --git a/hindsight-api-slim/hindsight_api/engine/db/ops.py b/hindsight-api-slim/hindsight_api/engine/db/ops.py index 2d4d880870..99ee197b39 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/ops.py +++ b/hindsight-api-slim/hindsight_api/engine/db/ops.py @@ -166,21 +166,6 @@ async def bulk_insert_unit_entities( # -- LATERAL / fan-out queries --------------------------------------- - @abstractmethod - async def fetch_entity_unit_fanout( - self, - conn: DatabaseConnection, - ue_table: str, - entity_id_list: list[UUID], - limit_per_entity: int, - ) -> list[ResultRow]: - """Fetch unit_ids for a list of entities with per-entity row cap. - - PG uses unnest + CROSS JOIN LATERAL with LIMIT. - Non-PG queries each entity individually. - """ - ... - @abstractmethod async def fetch_unit_dates( self, diff --git a/hindsight-api-slim/hindsight_api/engine/db/ops_oracle.py b/hindsight-api-slim/hindsight_api/engine/db/ops_oracle.py index 2c2cc7a8a1..bc877b8ad5 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/ops_oracle.py +++ b/hindsight-api-slim/hindsight_api/engine/db/ops_oracle.py @@ -215,30 +215,6 @@ async def bulk_insert_unit_entities( list(zip(unit_ids, entity_ids)), ) - async def fetch_entity_unit_fanout( - self, - conn: DatabaseConnection, - ue_table: str, - entity_id_list: list[UUID], - limit_per_entity: int, - ) -> list[ResultRow]: - # Query each entity individually - rows: list[ResultRow] = [] - for eid in entity_id_list: - entity_rows = await conn.fetch( - f""" - SELECT $1 AS entity_id, ue.unit_id - FROM {ue_table} ue - WHERE ue.entity_id = $1 - ORDER BY ue.unit_id DESC - LIMIT $2 - """, - eid, - limit_per_entity, - ) - rows.extend(entity_rows) - return rows - async def fetch_unit_dates( self, conn: DatabaseConnection, diff --git a/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py b/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py index cd5d9b178e..eeccd947b0 100644 --- a/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py +++ b/hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py @@ -290,29 +290,6 @@ async def bulk_insert_unit_entities( entity_ids, ) - async def fetch_entity_unit_fanout( - self, - conn: DatabaseConnection, - ue_table: str, - entity_id_list: list[UUID], - limit_per_entity: int, - ) -> list[ResultRow]: - return await conn.fetch( - f""" - SELECT e.entity_id, n.unit_id - FROM unnest($1::uuid[]) AS e(entity_id) - CROSS JOIN LATERAL ( - SELECT ue.unit_id - FROM {ue_table} ue - WHERE ue.entity_id = e.entity_id - ORDER BY ue.unit_id DESC - LIMIT $2 - ) n - """, - entity_id_list, - limit_per_entity, - ) - async def fetch_unit_dates( self, conn: DatabaseConnection, diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index e215dc9c31..bd6f310e23 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -4881,7 +4881,9 @@ async def get_graph_data( source_memory_ids.extend(unit["source_memory_ids"]) source_memory_ids = list(set(source_memory_ids)) # Deduplicate - # Fetch links where BOTH endpoints are in the visible set (or source memories) + # Fetch non-entity links where BOTH endpoints are in the visible set (or + # source memories). Entity edges are derived below from unit_entities so + # we don't materialize them in memory_links anymore. # Cap at 10k edges — the UI can't usefully render more, and uncapped queries # on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow. max_edges = 10000 @@ -4893,10 +4895,11 @@ async def get_graph_data( ml.to_unit_id, ml.link_type, ml.weight, - e.canonical_name as entity_name + NULL::text AS entity_name FROM {fq_table("memory_links")} ml - LEFT JOIN {fq_table("entities")} e ON ml.entity_id = e.id - WHERE ml.from_unit_id = ANY($1::uuid[]) AND ml.to_unit_id = ANY($1::uuid[]) + WHERE ml.link_type <> 'entity' + AND ml.from_unit_id = ANY($1::uuid[]) + AND ml.to_unit_id = ANY($1::uuid[]) ORDER BY ml.weight DESC NULLS LAST LIMIT $2 """, @@ -5034,16 +5037,19 @@ async def get_graph_data( } ) - # Build observation-inferred links from inherited entities and shared source memories. - # Observations never have direct memory_links rows, so all their links must be derived. + # Build derived links: entity edges for all visible units (from unit_entities) + # and observation semantic edges via shared source memories. + # Observations never have direct memory_links rows, so all their links are derived. observation_units = [unit for unit in units if unit["fact_type"] == "observation"] - observation_ids = {unit["id"] for unit in observation_units} - # Entity links: pair observations that share at least one inherited entity - entity_to_observations: dict[str, list] = {} - for obs_id in observation_ids: - for entity_name in entity_map.get(obs_id, []): - entity_to_observations.setdefault(entity_name, []).append(obs_id) + # Entity links: pair any visible units that share at least one entity. + # Cap per entity to bound hot entities (e.g. one entity in 500 visible units + # would otherwise produce ~125k edges by itself). + max_units_per_entity = 10 + entity_to_units_visible: dict[str, list] = {} + for unit_id in unit_ids: + for entity_name in entity_map.get(unit_id, []): + entity_to_units_visible.setdefault(entity_name, []).append(unit_id) # Semantic links: pair observations that share at least one source memory source_to_obs_for_semantic: dict = {} @@ -5055,16 +5061,17 @@ async def get_graph_data( observation_inferred_links = [] seen_inferred: set[tuple] = set() - for entity_name, obs_ids in entity_to_observations.items(): - for i, obs_a in enumerate(obs_ids): - for obs_b in obs_ids[i + 1 :]: - pair = (min(str(obs_a), str(obs_b)), max(str(obs_a), str(obs_b)), "entity", entity_name) + for entity_name, ent_unit_ids in entity_to_units_visible.items(): + capped = ent_unit_ids[:max_units_per_entity] + for i, unit_a in enumerate(capped): + for unit_b in capped[i + 1 :]: + pair = (min(str(unit_a), str(unit_b)), max(str(unit_a), str(unit_b)), "entity", entity_name) if pair not in seen_inferred: seen_inferred.add(pair) observation_inferred_links.append( { - "from_unit_id": obs_a, - "to_unit_id": obs_b, + "from_unit_id": unit_a, + "to_unit_id": unit_b, "link_type": "entity", "weight": 1.0, "entity_name": entity_name, @@ -6996,19 +7003,56 @@ async def get_bank_stats( bank_id, ) - # Link stats — filter on ml.bank_id directly instead of joining through mu.bank_id. - # link_counts and link_counts_by_fact_type are derived in Python from the breakdown. - link_breakdown_stats = await conn.fetch( + # Link stats — filter on ml.bank_id (indexed) instead of joining through mu.bank_id. + # With the idx_memory_links_bank_link_type index this turns a full-table hash join + # into an indexed scan + PK lookups. link_counts and link_counts_by_fact_type are + # derived in Python from the breakdown. + non_entity_breakdown = 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 ml.bank_id = $1 + WHERE ml.bank_id = $1 AND ml.link_type <> 'entity' GROUP BY mu.fact_type, ml.link_type """, bank_id, ) + # Entity links are derived from unit_entities (no longer stored in memory_links). + # Replicate the historical writer cap: each unit linked bidirectionally to up to + # MAX_LINKS_PER_ENTITY others sharing each entity. Count per (unit, entity) the + # outgoing rows the writer would have produced — by fact_type of the source unit. + max_links_per_entity = 10 + entity_breakdown = await conn.fetch( + f""" + WITH per_entity AS ( + SELECT ue.entity_id, COUNT(*) AS n + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id + WHERE mu.bank_id = $1 + GROUP BY ue.entity_id + ) + SELECT mu.fact_type, SUM(LEAST(pe.n - 1, $2))::bigint AS count + FROM {fq_table("unit_entities")} ue + JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id + JOIN per_entity pe ON pe.entity_id = ue.entity_id + WHERE mu.bank_id = $1 + GROUP BY mu.fact_type + """, + bank_id, + max_links_per_entity, + ) + + link_breakdown_stats = [ + {"fact_type": row["fact_type"], "link_type": row["link_type"], "count": row["count"]} + for row in non_entity_breakdown + ] + link_breakdown_stats.extend( + {"fact_type": row["fact_type"], "link_type": "entity", "count": int(row["count"] or 0)} + for row in entity_breakdown + if (row["count"] or 0) > 0 + ) + link_counts: dict[str, int] = {} link_counts_by_fact_type: dict[str, int] = {} for row in link_breakdown_stats: diff --git a/hindsight-api-slim/hindsight_api/engine/retain/entity_processing.py b/hindsight-api-slim/hindsight_api/engine/retain/entity_processing.py index bf77999b27..c89aca0a60 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/entity_processing.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/entity_processing.py @@ -1,13 +1,13 @@ """ Entity processing for retain pipeline. -Handles entity extraction, resolution, and link creation for stored facts. +Handles entity extraction and resolution for stored facts. """ import logging from . import link_utils -from .types import EntityLink, ProcessedFact +from .types import ProcessedFact logger = logging.getLogger(__name__) @@ -76,8 +76,7 @@ async def resolve_entities( entity_labels: Optional entity label taxonomy Returns: - Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids) - to pass to build_entity_links(). + Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids). """ if not unit_ids or not facts: return [], [], {} @@ -99,68 +98,3 @@ async def resolve_entities( log_buffer, entity_labels=entity_labels, ) - - -async def build_entity_links( - entity_resolver, - conn, - bank_id: str, - unit_ids: list[str], - resolved_entity_ids: list[str], - entity_to_unit: list[tuple], - unit_to_entity_ids: dict[str, list[str]], - log_buffer: list[str] = None, - skip_unit_entities_insert: bool = False, - ops=None, -) -> list[EntityLink]: - """ - Build entity links for UI graph visualization. - - Queries unit_entities to find shared entities between new and existing units, - then generates EntityLink objects. When called from Phase 3 (post-transaction), - set skip_unit_entities_insert=True since unit_entities were already inserted - in Phase 2. - - Args: - entity_resolver: EntityResolver instance - conn: Database connection - bank_id: Bank identifier - unit_ids: Actual unit IDs (must already be inserted in the DB) - resolved_entity_ids: From resolve_entities() - entity_to_unit: From resolve_entities() - unit_to_entity_ids: From resolve_entities() - log_buffer: Optional buffer for detailed logging - skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2) - ops: DataAccessOps instance (from backend.ops) - - Returns: - List of EntityLink objects for batch insertion - """ - return await link_utils.build_entity_links_from_resolved( - entity_resolver, - conn, - bank_id, - unit_ids, - resolved_entity_ids, - entity_to_unit, - unit_to_entity_ids, - log_buffer, - skip_unit_entities_insert=skip_unit_entities_insert, - ops=ops, - ) - - -async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None: - """ - Insert entity links in batch. - - Args: - conn: Database connection - entity_links: List of EntityLink objects - bank_id: Bank identifier (stored directly on memory_links for fast filtering) - ops: DataAccessOps instance (from backend.ops) - """ - if not entity_links: - return - - await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py index 3c0b7e8c1c..cc00b7455f 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py @@ -5,10 +5,8 @@ import logging import time from datetime import UTC, datetime, timedelta -from uuid import UUID from ..memory_engine import fq_table -from .types import EntityLink logger = logging.getLogger(__name__) @@ -366,136 +364,6 @@ async def resolve_entities_only( return resolved_entity_ids, entity_to_unit, unit_to_entity_ids -async def build_entity_links_from_resolved( - entity_resolver, - conn, - bank_id: str, - unit_ids: list[str], - resolved_entity_ids: list[str], - entity_to_unit: list[tuple], - unit_to_entity_ids: dict[str, list[str]], - log_buffer: list[str] = None, - skip_unit_entities_insert: bool = False, - ops=None, -) -> list["EntityLink"]: - """ - Build entity links between units that share entities. - - Queries unit_entities to find which existing units share entities with the - new units, then generates EntityLink objects for UI graph visualization. - - Args: - entity_resolver: EntityResolver instance - conn: Database connection - bank_id: Bank identifier - unit_ids: Actual unit IDs (must already be inserted in the DB) - resolved_entity_ids: Entity IDs from resolve_entities_only - entity_to_unit: Mapping from resolve_entities_only - unit_to_entity_ids: Mapping from resolve_entities_only - log_buffer: Optional logging buffer - skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2) - - Returns: - List of EntityLink objects for batch insertion - """ - if not resolved_entity_ids: - return [] - - if not skip_unit_entities_insert: - # Insert unit-entity links (used in fallback path where Phase 2 didn't do this) - substep_start = time.time() - unit_entity_pairs = [] - for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit): - # Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred - # reflects the event timeline, not the ingest moment. - unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date)) - - await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn) - _log( - log_buffer, - f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s", - level="debug", - ) - - # Build entity links between units that share entities - substep_start = time.time() - all_entity_ids = set() - for entity_ids_list in unit_to_entity_ids.values(): - all_entity_ids.update(entity_ids_list) - - _log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug") - - MAX_LINKS_PER_ENTITY = 10 - - entity_to_units = {} - if all_entity_ids: - query_start = time.time() - import uuid - - entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids] - limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap - - rows = await ops.fetch_entity_unit_fanout( - conn, - fq_table("unit_entities"), - entity_id_list, - limit_per_entity, - ) - _log( - log_buffer, - f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s", - level="debug", - ) - - group_start = time.time() - for row in rows: - entity_id = row["entity_id"] - if entity_id not in entity_to_units: - entity_to_units[entity_id] = [] - entity_to_units[entity_id].append(row["unit_id"]) - _log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug") - link_gen_start = time.time() - links: list[EntityLink] = [] - new_unit_set = set(unit_ids) - - def to_uuid(val) -> UUID: - return UUID(val) if isinstance(val, str) else val - - for entity_id, units_with_entity in entity_to_units.items(): - entity_uuid = to_uuid(entity_id) - new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set] - existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set] - - new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units - for i, unit_id_1 in enumerate(new_units_to_link): - for unit_id_2 in new_units_to_link[i + 1 :]: - links.append( - EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid) - ) - links.append( - EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid) - ) - - existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:] - for new_unit in new_units: - for existing_unit in existing_to_link: - links.append( - EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid) - ) - links.append( - EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid) - ) - - _log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug") - _log( - log_buffer, - f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s", - level="debug", - ) - - return links - - async def create_temporal_links_batch_per_fact( conn, bank_id: str, @@ -889,29 +757,6 @@ async def create_semantic_links_batch( raise -async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None): - """ - Bulk-insert entity links via sorted INSERT FROM unnest(). - - Args: - conn: Database connection - links: List of EntityLink objects - bank_id: Bank identifier (stored directly on memory_links for fast filtering) - chunk_size: Number of rows per INSERT chunk (default 5000) - """ - if not links: - return - - import time as time_mod - - total_start = time_mod.time() - tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links] - await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops) - logger.debug( - f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s" - ) - - async def create_causal_links_batch( conn, bank_id: str, diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index c62878f661..d7ffd238d4 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -100,7 +100,6 @@ def parse_datetime_flexible(value: Any) -> datetime: ChunkMetadata, EntityResolutionResult, Phase1Result, - Phase3Context, ProcessedFact, RetainContent, RetainContentDict, @@ -259,30 +258,27 @@ async def _insert_facts_and_links( skip_semantic_links: bool = False, outbox_callback=None, ops=None, -) -> tuple[list[list[str]], Phase3Context]: +) -> list[list[str]]: """ Phase 2 of the retain pipeline: insert facts and retrieval-critical links. Runs inside a single database transaction to ensure atomicity of the data that retrieval depends on (facts, unit_entities, temporal/semantic/causal links). - Entity link generation and insertion for UI visualization are NOT done here — - only the unit_entities INSERT (FK to memory_units) stays in the transaction. - Entity link building is deferred to Phase 3 (post-transaction, best-effort). + Entity edges for UI graph visualization are derived on demand from + unit_entities by the /graph endpoint, so no entity rows are written to + memory_links here. """ set_stage("retain.phase2.insert_facts") unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops) step_start = time.time() log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s") - # Context for Phase 3 entity link building (after transaction commits) - phase3_context = Phase3Context() - if unit_ids: # Entity resolution was done in Phase 1 (separate connection). # Remap placeholder IDs to actual unit IDs. step_start = time.time() - remapped_entity_to_unit, remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results( + remapped_entity_to_unit, _remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results( resolved_entity_ids, entity_to_unit, unit_to_entity_ids, semantic_ann_links or [], unit_ids ) # Update semantic_ann_links with remapped IDs for Phase 2 @@ -296,13 +292,6 @@ async def _insert_facts_and_links( ] await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn) log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s") - # Save context for Phase 3 entity link building (after commit) - phase3_context = Phase3Context( - unit_ids=unit_ids, - resolved_entity_ids=resolved_entity_ids, - entity_to_unit=remapped_entity_to_unit, - unit_to_entity_ids=remapped_unit_to_entity_ids, - ) # Create temporal links step_start = time.time() @@ -346,49 +335,7 @@ async def _insert_facts_and_links( if outbox_callback is not None: await outbox_callback(conn) - return result_unit_ids, phase3_context - - -async def _build_and_insert_entity_links_phase3( - pool: Any, - entity_resolver, - bank_id: str, - phase3_ctx: Phase3Context, - log_buffer: list[str], -) -> None: - """ - Phase 3 helper: build entity links from resolved data and insert them. - - Runs on a fresh connection after the main transaction has committed. - Entity links are for UI graph visualization only — retrieval uses - the unit_entities self-join instead. - """ - set_stage("retain.phase3.entity_links") - p3_unit_ids = phase3_ctx.unit_ids - p3_resolved = phase3_ctx.resolved_entity_ids - p3_entity_to_unit = phase3_ctx.entity_to_unit - p3_unit_to_entity_ids = phase3_ctx.unit_to_entity_ids - - if not p3_unit_ids or not p3_resolved: - return - - async with acquire_with_retry(pool) as conn: - step_start = time.time() - entity_links = await entity_processing.build_entity_links( - entity_resolver, - conn, - bank_id, - p3_unit_ids, - p3_resolved, - p3_entity_to_unit, - p3_unit_to_entity_ids, - log_buffer, - skip_unit_entities_insert=True, # Already inserted in Phase 2 - ops=pool.ops, - ) - if entity_links: - await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id, ops=pool.ops) - log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s") + return result_unit_ids async def _extract_and_embed( @@ -1200,7 +1147,6 @@ async def _run_mini_batch_db_work() -> None: p2_start = time.time() batch_result_ids = None - phase3_ctx = None async with acquire_with_retry(pool) as conn: async with conn.transaction(): # --- Document ownership gate --- @@ -1289,7 +1235,7 @@ async def _run_mini_batch_db_work() -> None: # Insert facts and links — skip semantic links entirely in streaming # mode; they are created in a single final ANN pass after all batches. - batch_result_ids, phase3_ctx = await _insert_facts_and_links( + batch_result_ids = await _insert_facts_and_links( conn, entity_resolver, bank_id, @@ -1309,15 +1255,13 @@ async def _run_mini_batch_db_work() -> None: logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s") - # Best-effort: entity viz + stats (fast, not semantic ANN) - if phase3_ctx is not None: - try: - await entity_resolver.flush_pending_stats() - await _build_and_insert_entity_links_phase3( - pool, entity_resolver, bank_id, phase3_ctx, log_buffer - ) - except Exception: - logger.warning(f"Phase 3 stats (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True) + # Best-effort: flush entity_cooccurrences and other deferred stats. + try: + await entity_resolver.flush_pending_stats() + except Exception: + logger.warning( + f"Entity stats flush (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True + ) logger.info( f"[streaming] Consumer batch {consumer_batch_idx + 1} total " @@ -1775,7 +1719,7 @@ async def _run_delta_db_work() -> None: # Insert facts and retrieval-critical links. # Use delta_contents (the changed/new chunks) as the content list, # since extracted_facts have content_index relative to delta_contents. - result_unit_ids, phase3_ctx = await _insert_facts_and_links( + result_unit_ids = await _insert_facts_and_links( conn, entity_resolver, bank_id, @@ -1792,12 +1736,11 @@ async def _run_delta_db_work() -> None: ops=pool.ops, ) - # PHASE 3 — Best-Effort Display Data (post-transaction) + # Flush deferred entity_cooccurrences stats (post-transaction, best-effort). try: await entity_resolver.flush_pending_stats() - await _build_and_insert_entity_links_phase3(pool, entity_resolver, bank_id, phase3_ctx, log_buffer) except Exception: - logger.warning("Phase 3 (best-effort display data) failed — retrieval unaffected", exc_info=True) + logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True) total_time = time.time() - start_time log_buffer.append(f"{'=' * 60}") diff --git a/hindsight-api-slim/hindsight_api/engine/retain/types.py b/hindsight-api-slim/hindsight_api/engine/retain/types.py index ca9736417b..b4c1c6a426 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/types.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/types.py @@ -224,21 +224,6 @@ def from_extracted_fact( ) -@dataclass -class Phase3Context: - """ - Data passed from Phase 2 to Phase 3 for entity link building. - - Contains the unit IDs and entity resolution data needed to build - entity links for UI graph visualization after the write transaction commits. - """ - - unit_ids: list[str] = field(default_factory=list) - resolved_entity_ids: list[str] = field(default_factory=list) - entity_to_unit: list[tuple] = field(default_factory=list) - unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict) - - @dataclass class EntityResolutionResult: """ @@ -263,21 +248,6 @@ class Phase1Result: semantic_ann_links: list[tuple] -@dataclass -class EntityLink: - """ - Link between two memory units through a shared entity. - - Used for entity-based graph connections in the memory graph. - """ - - from_unit_id: UUID - to_unit_id: UUID - entity_id: UUID - link_type: str = "entity" - weight: float = 1.0 - - @dataclass class RetainBatch: """ diff --git a/hindsight-api-slim/tests/test_retain.py b/hindsight-api-slim/tests/test_retain.py index bc2e31dcc0..61bbfeb82a 100644 --- a/hindsight-api-slim/tests/test_retain.py +++ b/hindsight-api-slim/tests/test_retain.py @@ -1585,39 +1585,32 @@ async def test_semantic_links_creation(memory, request_context): @pytest.mark.asyncio async def test_entity_links_creation(memory, request_context): """ - Test that entity links are created between facts that mention the same entities. - - Entity links connect facts that reference the same person, place, or concept. - This is core functionality and should work consistently. + Test that entity edges surface in the /graph response between facts that + mention the same entities, and that the /stats endpoint reports a non-zero + entity link count. Entity edges are derived on demand from unit_entities; + no rows of link_type='entity' are written to memory_links. """ bank_id = f"test_entity_links_{datetime.now(timezone.utc).timestamp()}" try: - # Store facts that mention the same entities unit_ids_1 = await memory.retain_async( bank_id=bank_id, content="Alice joined Google as a software engineer in 2020.", context="career history", request_context=request_context, ) - - # Mentions same entity (Alice) - should create entity link unit_ids_2 = await memory.retain_async( bank_id=bank_id, content="Alice led the development of the new authentication system.", context="project updates", request_context=request_context, ) - - # Mentions same entity (Google) - should create entity link unit_ids_3 = await memory.retain_async( bank_id=bank_id, content="Google announced new cloud services at their annual conference.", context="tech news", request_context=request_context, ) - - # Different entities - no entity link expected unit_ids_4 = await memory.retain_async( bank_id=bank_id, content="Bob works at Meta on machine learning infrastructure.", @@ -1626,57 +1619,36 @@ async def test_entity_links_creation(memory, request_context): ) assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0 and len(unit_ids_4) > 0 + all_unit_ids = {str(uid) for uid in unit_ids_1 + unit_ids_2 + unit_ids_3 + unit_ids_4} - logger.info(f"Created {len(unit_ids_1) + len(unit_ids_2) + len(unit_ids_3) + len(unit_ids_4)} facts") - - # Query the memory_links table to verify entity links exist + # memory_links should NOT contain entity rows — they are derived on read. async with memory._pool.acquire() as conn: - all_unit_ids = unit_ids_1 + unit_ids_2 + unit_ids_3 + unit_ids_4 - - entity_links = await conn.fetch( + stored_entity_links = await conn.fetchval( """ - SELECT from_unit_id, to_unit_id, link_type, weight, entity_id - FROM memory_links - WHERE from_unit_id::text = ANY($1) - AND link_type = 'entity' - ORDER BY from_unit_id, to_unit_id + SELECT COUNT(*) FROM memory_links + WHERE bank_id = $1 AND link_type = 'entity' """, - all_unit_ids + bank_id, ) + assert stored_entity_links == 0, "Entity edges must not be materialized in memory_links" + + # /graph derives entity edges from unit_entities — pairs that share an entity + # should appear as link_type='entity' edges. + graph = await memory.get_graph_data(bank_id=bank_id, request_context=request_context) + entity_edges = [edge["data"] for edge in graph["edges"] if edge["data"].get("linkType") == "entity"] + assert entity_edges, "Graph should surface entity edges for facts sharing an entity" + for edge in entity_edges: + assert edge["source"] != edge["target"] + # At least one edge must connect two facts we retained directly (proves the + # derivation works for non-observation units, not just for inherited entities). + retained_pair_edges = [ + edge for edge in entity_edges if edge["source"] in all_unit_ids and edge["target"] in all_unit_ids + ] + assert retained_pair_edges, "Expected at least one entity edge between two directly retained units" - logger.info(f"Found {len(entity_links)} entity links") - - # Entity extraction is core functionality and should work - assert len(entity_links) > 0, "Should have created entity links between facts with shared entities (Alice, Google)" - - # Verify link properties - entities_seen = set() - for link in entity_links: - entity_id = link['entity_id'] - entities_seen.add(str(entity_id)) - from_id = str(link['from_unit_id']) - to_id = str(link['to_unit_id']) - logger.info(f" Link: {from_id[:8]}... -> {to_id[:8]}... via entity {str(entity_id)[:8]}...") - assert link['link_type'] == 'entity', "Link type should be 'entity'" - assert link['weight'] == 1.0, "Entity links should have weight 1.0" - assert entity_id is not None, "Entity links must reference an entity_id" - - logger.info(f"Entity links created successfully for {len(entities_seen)} unique entities") - - # Verify bidirectional links (entity links should be bidirectional) - link_pairs = set() - for link in entity_links: - from_id = str(link['from_unit_id']) - to_id = str(link['to_unit_id']) - entity_id = str(link['entity_id']) - link_pairs.add((from_id, to_id, entity_id)) - - # Check that for each (A -> B) link, there's a (B -> A) link with same entity - for from_id, to_id, entity_id in link_pairs: - reverse_exists = (to_id, from_id, entity_id) in link_pairs - assert reverse_exists, f"Entity links should be bidirectional: missing reverse link for {from_id[:8]} -> {to_id[:8]}" - - logger.info("Entity links are properly bidirectional") + # /stats should report a non-zero entity count under the same key as before. + stats = await memory.get_bank_stats(bank_id=bank_id, request_context=request_context) + assert stats["link_counts"].get("entity", 0) > 0, "Stats must report a non-zero entity link count" finally: await memory.delete_bank(bank_id, request_context=request_context) @@ -2002,11 +1974,12 @@ async def test_all_link_types_together(memory, request_context): logger.info(f"Created {len(unit_ids_1) + len(unit_ids_2) + len(unit_ids_3)} facts") - # Query for all link types + # Temporal/semantic/causal links live in memory_links; entity links are + # derived from unit_entities at read time, so check both surfaces. async with memory._pool.acquire() as conn: all_unit_ids = unit_ids_1 + unit_ids_2 + unit_ids_3 - all_links = await conn.fetch( + stored_links = await conn.fetch( """ SELECT link_type, COUNT(*) as count FROM memory_links @@ -2014,24 +1987,16 @@ async def test_all_link_types_together(memory, request_context): GROUP BY link_type ORDER BY link_type """, - all_unit_ids + all_unit_ids, ) + stored_by_type = {row["link_type"]: row["count"] for row in stored_links} + + assert "temporal" in stored_by_type, "Should have temporal links (facts with nearby dates)" + assert "semantic" in stored_by_type, "Should have semantic links (similar content about Python/auth)" + assert stored_by_type.get("entity", 0) == 0, "Entity links must not be materialized in memory_links" - logger.info("Link types created:") - link_types_found = {} - for row in all_links: - link_type = row['link_type'] - count = row['count'] - link_types_found[link_type] = count - logger.info(f" - {link_type}: {count} links") - - # Should have temporal, semantic, and entity links - assert 'temporal' in link_types_found, "Should have temporal links (facts with nearby dates)" - assert 'semantic' in link_types_found, "Should have semantic links (similar content about Python/auth)" - assert 'entity' in link_types_found, "Should have entity links (all mention Alice)" - - logger.info(f"Successfully created {len(link_types_found)} different link types") - logger.info("All major link types (temporal, semantic, entity) are working correctly") + stats = await memory.get_bank_stats(bank_id=bank_id, request_context=request_context) + assert stats["link_counts"].get("entity", 0) > 0, "Stats must report entity links (all facts mention Alice)" finally: await memory.delete_bank(bank_id, request_context=request_context) From 3017c4423532f3b7c19bc3df249e65f58ee76e0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 26 May 2026 18:03:56 +0200 Subject: [PATCH 3/4] fix(graph): cap entity edges per unit, not per entity list MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous derivation kept only the first 10 units per entity before pairing, so any unit beyond #10 for a hot entity had zero entity edges in /graph — even though it shared the entity with many visible units. Switch to a sliding window: each unit links to its next N neighbors in the per-entity list. Every unit that shares an entity with another visible unit gets edges (its successors directly, predecessors via their pairs), and total edges stay bounded at ~N * cap per entity instead of N². Adds a regression test that retains 15 facts mentioning the same person and asserts every retained unit appears in at least one entity edge in /graph. --- .../hindsight_api/engine/memory_engine.py | 19 ++++++--- hindsight-api-slim/tests/test_retain.py | 40 +++++++++++++++++++ 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index bd6f310e23..86bd5033dd 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -5043,9 +5043,11 @@ async def get_graph_data( observation_units = [unit for unit in units if unit["fact_type"] == "observation"] # Entity links: pair any visible units that share at least one entity. - # Cap per entity to bound hot entities (e.g. one entity in 500 visible units - # would otherwise produce ~125k edges by itself). - max_units_per_entity = 10 + # Each unit links to up to max_neighbors_per_unit subsequent units in the + # per-entity list, so every unit that shares an entity with another visible + # unit gets edges (matches the historical writer cap, which was per-unit). + # Bounds total edges to ~N * cap per entity instead of N² for hot entities. + max_neighbors_per_unit = 10 entity_to_units_visible: dict[str, list] = {} for unit_id in unit_ids: for entity_name in entity_map.get(unit_id, []): @@ -5062,9 +5064,14 @@ async def get_graph_data( seen_inferred: set[tuple] = set() for entity_name, ent_unit_ids in entity_to_units_visible.items(): - capped = ent_unit_ids[:max_units_per_entity] - for i, unit_a in enumerate(capped): - for unit_b in capped[i + 1 :]: + n = len(ent_unit_ids) + for i, unit_a in enumerate(ent_unit_ids): + # Sliding window: link unit_a to its next max_neighbors_per_unit + # in the list. Each pair is also "incoming" for the later unit, + # so every unit ends up with up to ~2*max_neighbors_per_unit edges + # for this entity (its successors + its predecessors via their pairs). + for j in range(i + 1, min(i + 1 + max_neighbors_per_unit, n)): + unit_b = ent_unit_ids[j] pair = (min(str(unit_a), str(unit_b)), max(str(unit_a), str(unit_b)), "entity", entity_name) if pair not in seen_inferred: seen_inferred.add(pair) diff --git a/hindsight-api-slim/tests/test_retain.py b/hindsight-api-slim/tests/test_retain.py index 61bbfeb82a..e93abcad25 100644 --- a/hindsight-api-slim/tests/test_retain.py +++ b/hindsight-api-slim/tests/test_retain.py @@ -1654,6 +1654,46 @@ async def test_entity_links_creation(memory, request_context): await memory.delete_bank(bank_id, request_context=request_context) +@pytest.mark.asyncio +async def test_graph_entity_edges_cover_all_visible_units(memory, request_context): + """ + Regression test: when a hot entity is shared by more than the per-entity + cap (default 10) visible units, every visible unit mentioning that entity + must still appear in at least one entity edge. The previous implementation + capped the per-entity list to the first 10 units before pairing, leaving + units #11+ without any entity edges in /graph. + """ + bank_id = f"test_graph_entity_coverage_{datetime.now(timezone.utc).timestamp()}" + + try: + # 15 facts all mentioning the same person — more than max_neighbors_per_unit=10. + contents = [ + {"content": f"Alice completed task #{i} in the authentication module.", "context": "sprint log"} + for i in range(15) + ] + retained_lists = await memory.retain_batch_async( + bank_id=bank_id, + contents=contents, + request_context=request_context, + ) + retained_unit_ids = {str(uid) for sublist in retained_lists for uid in sublist} + assert len(retained_unit_ids) >= 15, f"Expected >=15 units, got {len(retained_unit_ids)}" + + graph = await memory.get_graph_data(bank_id=bank_id, request_context=request_context) + entity_edges = [edge["data"] for edge in graph["edges"] if edge["data"].get("linkType") == "entity"] + assert entity_edges, "Graph should have entity edges for facts sharing an entity" + + units_in_entity_edges = {edge["source"] for edge in entity_edges} | {edge["target"] for edge in entity_edges} + missing = retained_unit_ids - units_in_entity_edges + assert not missing, ( + f"{len(missing)}/{len(retained_unit_ids)} retained units have no entity edges in /graph. " + f"The per-entity edge cap must not exclude units beyond the first N from pairing." + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + @pytest.mark.asyncio async def test_people_name_extraction(memory, request_context): """ From 205901ac6cf10ae600207476a8578ab7befdaa67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Tue, 26 May 2026 18:21:55 +0200 Subject: [PATCH 4/4] fix(migration): re-parent entity-link drop after e1b2c3d4f5a6 landed on main #1762 landed e1b2c3d4f5a6_drop_unused_indexes between this PR opening and CI run, which also drops idx_memory_links_entity_covering. Our migration's down_revision still pointed at the prior head, leaving Alembic with two heads and tripping test_alembic_dag.test_single_head. Re-parent to e1b2c3d4f5a6 to unify the head. The DROP INDEX IF EXISTS line becomes a defensive no-op (since #1762 already dropped it), but is retained in case this migration runs against a snapshot taken before #1762. --- .../e9b2c7d1f3a4_drop_entity_memory_links.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py b/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py index b6e6b01df8..344d5aa65b 100644 --- a/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py +++ b/hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py @@ -6,14 +6,13 @@ read from the link table — on a 10k-unit benchmark bank, entity rows were 53% of all link rows (~190 MB after indexes) and recall never touched them. -This migration: - -1. Drops ``idx_memory_links_entity_covering`` (the partial index targeting - ``WHERE link_type = 'entity'`` rows that no longer exist). -2. Deletes ``memory_links`` rows with ``link_type = 'entity'``. +This migration deletes ``memory_links`` rows with ``link_type = 'entity'``. +``idx_memory_links_entity_covering`` was already dropped by migration +``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case +this migration runs against an older snapshot that predates that one. Revision ID: e9b2c7d1f3a4 -Revises: 86f7a033d372 +Revises: e1b2c3d4f5a6 Create Date: 2026-05-26 """ @@ -24,7 +23,7 @@ from hindsight_api.alembic._dialect import run_for_dialect revision: str = "e9b2c7d1f3a4" -down_revision: str | Sequence[str] | None = "86f7a033d372" +down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None