From 914ba7962c46b4b826a12ce973aa558bc2e698d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Wed, 1 Apr 2026 12:52:49 +0200 Subject: [PATCH] =?UTF-8?q?perf:=203-phase=20retain=20pipeline=20=E2=80=94?= =?UTF-8?q?=20fix=20deadlocks,=20cap=20temporal=20links,=20query-time=20en?= =?UTF-8?q?tity=20expansion=20(#722)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf: 3-phase retain pipeline — fix deadlocks, cap temporal links, query-time entity expansion Major retain pipeline overhaul addressing deadlocks, write amplification, and TimeoutErrors. Restructures retain into three phases: Phase 1: Entity resolution on separate connection (read-heavy) Phase 2: Core write transaction (atomic) — facts, unit_entities, links Phase 3: Best-effort display data (error-isolated) — entity viz links, stats Key changes: - Sorted bulk INSERT FROM unnest() prevents deadlocks - Temporal links capped to top-20 per unit (95% reduction) - Batched semantic ANN via temp table + LATERAL - Query-time entity expansion via unit_entities self-join - Entity viz links moved to Phase 3 (post-transaction) - HINDSIGHT_API_RETAIN_MAX_CONCURRENT config (default: 32) * fix: increase semantic link top_k from 5 to 20 The hardcoded top_k=5 was artificially limiting semantic link creation. Link expansion retrieval can consume up to budget (50-200) semantic neighbors per seed set, but each fact only had 5 outgoing edges — making the bidirectional graph very sparse. Increasing to 20 gives retrieval 4x more edges to work with. The ANN probe cost is unchanged (same HNSW traversal per fact, just returning more rows). INSERT cost is negligible (~14k rows via bulk INSERT). Also: all 18 TimeoutErrors in the latest benchmark (beam-1m-u20) were from Gemini LLM calls, zero from the database — confirming the entity resolution split eliminated DB timeouts entirely. * perf: move semantic ANN search to Phase 1 to avoid transaction timeouts The batched LATERAL ANN query (700 HNSW probes) was the last remaining source of DB TimeoutErrors — all 29 in the latest benchmark were from create_semantic_links_batch inside the Phase 2 write transaction. Split semantic link creation into three phases: - Phase 1 (separate conn, autocommit): ANN search via temp table + LATERAL. No transaction locks, no contention with concurrent writers. - Phase 2 (write transaction): within-batch numpy similarities (instant) + INSERT of both within-batch and Phase 1 ANN results. No DB reads. - Phase 3 (flush_pending_stats): future hook point for re-checking ANN results after commit to catch links missed by concurrent batches. Also adds 7 unit tests for compute_semantic_links_within_batch covering empty input, identical/orthogonal embeddings, threshold filtering, top_k cap, and tuple structure validation. * fix: handle placeholder unit_ids in Phase 1 ANN search (not valid UUIDs) * test: add Phase 1 ANN cross-batch test + configurable test PG port - New test_semantic_links_phase1_ann_cross_batch verifies that the Phase 1 ANN search with placeholder unit IDs correctly creates cross-batch semantic links after remapping to real IDs. - Test PG port now configurable via HINDSIGHT_TEST_PG_PORT env var (default: 5556) to avoid conflicts with running benchmark daemons. * perf: remove retry_with_backoff from retain, set semaphore default to 4 Remove retry_with_backoff from _run_db_work and _run_delta_db_work: - Deadlocks are prevented by sorted bulk INSERT (no need for retry) - Transient timeouts are handled by the worker poller's task-level retry (3 attempts, 60s spacing) which is better than rapid internal retries that amplify I/O pressure during contention storms Set HINDSIGHT_API_RETAIN_MAX_CONCURRENT default from 32 to 4: - The semaphore gates Phase 1 (ANN + entity resolution) + Phase 2 (writes) - At 4 concurrent, HNSW index I/O is manageable; at 10+ concurrent the probes saturate disk and cause cascading timeouts - LLM extraction still runs at full parallelism (semaphore acquired after) * fix: add fact_type filter to Phase 1 ANN query to use per-bank HNSW indexes The LATERAL ANN query was falling back to sequential scan + sort (90ms/probe) because the per-bank HNSW indexes are partial indexes filtered on fact_type. Without fact_type in the WHERE clause, PostgreSQL couldn't use them. Fix: iterate over ('world', 'experience') and run one HNSW-indexed ANN per type. EXPLAIN shows 8ms/probe (was 90ms) — 11x faster. 700 probes × 8ms × 2 types = ~11s total (was ~63s via seq scan). * fix: scope temporal links by fact_type + add integration tests Temporal links now filter by fact_type in the LATERAL query — world facts only link to world facts, experience to experience. This matches how retrieval filters results and avoids wasted cross-type link rows. New integration tests: - test_semantic_ann_uses_hnsw_index: verifies Phase 1 ANN creates cross-batch semantic links (tests fact_type filter + placeholder remap) - test_temporal_links_scoped_by_fact_type: verifies world facts get temporal links to other world facts but NOT to experience facts * fix: tolerate individual chunk LLM failures instead of failing entire batch Changed asyncio.gather(*tasks) to asyncio.gather(*tasks, return_exceptions=True) in both chunk-level and content-level fact extraction. A single chunk timeout (e.g., Gemini >90s) no longer discards all other successfully extracted facts. For a 50MB document with 17k chunks, even a 2% chunk failure rate previously caused 0 completions (entire batch discarded). Now 16,700 facts are extracted and only the 300 failed chunks are skipped with a warning log. * fix: batch temporal LATERAL query for large documents (16k+ chunks) The LATERAL query for temporal links passed all unit_ids at once into unnest(), causing PostgreSQL timeouts on documents with 16k+ chunks. Split into batches of 500 units per query to keep each under the command_timeout. Also identified: HNSW index creation on shared pg0 instances with 50k+ existing units exceeds the 60s command_timeout. This is a test infrastructure issue (shared pg0 accumulates data) but also affects production when creating new banks on large instances. * feat: streaming chunk batching for large documents (RETAIN_CHUNK_BATCH_SIZE) Process chunks in mini-batches of N (default 500), committing each batch to the DB before starting the next. This prevents OOM kills on large documents (50MB / 17k+ chunks) by keeping only ~500 facts + embeddings in memory at a time instead of 50k+. Each mini-batch goes through the full Phase 1 → 2 → 3 pipeline independently, sharing the same document_id. On recovery (process dies mid-way), delta retain detects already-committed chunks via content_hash and skips them — only remaining chunks get re-extracted. Config: HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE (default: 500, 0 to disable) Per-bank configurable via the hierarchical config system. Tests: - test_streaming_chunk_batching_produces_same_facts - test_streaming_chunk_batching_recovery (delta retain skips committed chunks) - test_streaming_disabled_for_small_docs * perf(retain): producer-consumer pipeline + deferred semantic ANN Replace the sequential streaming loop with a producer-consumer pipeline: - LLM producer fires concurrent chunk extractions (semaphore-bounded) - DB consumer drains queue in batches, runs Phase 1+2+3 per batch - LLM and DB work overlap instead of running sequentially Defer semantic links to a single final ANN pass after all batches commit: - Remove within-batch semantic links from Phase 2 (was 2.6s/batch) - Run parallel ANN (4 connections) after all facts committed - top_k reduced from 50 to 20 (recall uses at most 20 neighbors) - Recovery via operation result_metadata checkpoint Additional optimizations: - skip_exists_check on temporal/causal link INSERT (saves ~0.5s/batch) - WHERE EXISTS guard on semantic link INSERT (handles document upsert) - timeout=300s on ANN queries and bulk INSERT for large banks - Demote [ANN] debug logs to logger.debug() - Fix docstring typos (agent_id → bank_id) - Fix content_index remapping in producer-consumer batches - Fix delta retain passing contents vs delta_contents 50MB benchmark (mock LLM): 9.2 min (was 23 min) — 2.5x faster. BEAM 10m benchmark: zero deadlocks, zero DB errors. * refactor(retain): remove legacy fallback code paths - Remove process_entities_batch (legacy single-connection entity processing) - Remove extract_entities_batch_optimized (only caller was the above) - Remove fallback entity processing inside Phase 2 transaction - Remove legacy ANN inline fallback in create_semantic_links_batch - Remove fallback entity_links direct-insert path in Phase 3 - Make resolved_entity_ids/entity_to_unit/unit_to_entity_ids required params * refactor(retain): replace tuple returns with dataclasses, remove dead code - Add EntityResolutionResult and Phase1Result dataclasses in types.py - Replace 4-tuple return from _pre_resolve_phase1 with Phase1Result - Remove dead `entity_links = []` variables in retain_batch and _try_delta_retain - Remove unused `confidence_score` parameter from orchestrator.retain_batch and _retain_batch_async_internal (was accepted but never used) * fix(entity-resolver): remove LIKE full-scan fallbacks, use index-only trigram matching The entity resolution query had LIKE '%...' substring conditions that bypassed the GIN trigram index, causing full sequential scans of the entities table. On banks with 10k+ entities, this caused TimeoutErrors (observed in BEAM 10m). Changes: - Remove LIKE fallbacks, use trigram % operator only (GIN index-based) - Lower similarity threshold from 0.3 to 0.15 to catch substring relationships - Use LOWER() on both sides for case-insensitive matching - Migration: recreate GIN trigram index on LOWER(canonical_name) * fix: remove schema prefix from index names in trigram migration * fix(delta-retain): use same chunk_size as streaming path (3000 vs 120000) _chunk_contents_for_delta defaulted to chunk_size=120000 while the streaming path used 3000. On retry, delta re-chunked the document with different boundaries, found 0 matching chunks, and fell through to full re-extraction. This wasted all LLM calls on already-committed chunks. Fix: use the same default (3000) so chunk hashes match on recovery. * fix(retain): persist generated document_id in operation metadata for retry recovery When no document_id is provided, retain generates a UUID. On retry, a new UUID was generated, making delta retain and streaming chunk-hash recovery unable to find previously committed chunks. All LLM extraction was wasted on retry. Fix: resolve document_id early in retain_batch (before delta), persist it to operation result_metadata, and recover it on retry. Both delta and streaming paths now see the same document_id across attempts. * refactor(retain): unify into single streaming pipeline, remove non-streaming path All retains now go through the producer-consumer streaming pipeline, regardless of document size. Small documents are processed as a single batch. This eliminates the maintenance burden of two separate code paths. Also fix document upsert: compare content hash to distinguish recovery (same content, partially committed) from update (different content, needs cascade-delete). Previously, existing chunks always triggered recovery mode. * refactor(retain): remove dead code, replace raw dicts with Phase3Context dataclass - Remove dead _handle_zero_facts_documents (no callers after path unification) - Remove unused imports: defaultdict, EntityLink - Replace raw dict phase3_context with typed Phase3Context dataclass - Update _build_and_insert_entity_links_phase3 to use typed parameter --- ...c1_case_insensitive_entities_trgm_index.py | 45 + hindsight-api-slim/hindsight_api/config.py | 11 + .../hindsight_api/engine/entity_resolver.py | 29 +- .../hindsight_api/engine/memory_engine.py | 83 +- .../engine/retain/entity_processing.py | 136 ++- .../engine/retain/fact_extraction.py | 98 +- .../engine/retain/link_creation.py | 15 +- .../hindsight_api/engine/retain/link_utils.py | 1080 +++++++++------- .../engine/retain/orchestrator.py | 1081 +++++++++++++---- .../hindsight_api/engine/retain/types.py | 39 + .../engine/search/link_expansion_retrieval.py | 71 +- hindsight-api-slim/tests/conftest.py | 2 +- hindsight-api-slim/tests/test_link_utils.py | 136 ++- hindsight-api-slim/tests/test_retain.py | 540 ++++++++ hindsight-dev/benchmarks/perf/retain_perf.py | 380 +++++- .../docs/developer/configuration.md | 1 + 16 files changed, 2908 insertions(+), 839 deletions(-) create mode 100644 hindsight-api-slim/hindsight_api/alembic/versions/d6e7f8a9b0c1_case_insensitive_entities_trgm_index.py diff --git a/hindsight-api-slim/hindsight_api/alembic/versions/d6e7f8a9b0c1_case_insensitive_entities_trgm_index.py b/hindsight-api-slim/hindsight_api/alembic/versions/d6e7f8a9b0c1_case_insensitive_entities_trgm_index.py new file mode 100644 index 00000000..682a8520 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/alembic/versions/d6e7f8a9b0c1_case_insensitive_entities_trgm_index.py @@ -0,0 +1,45 @@ +"""Recreate entities trigram index on LOWER(canonical_name) for case-insensitive matching + +The previous GIN trigram index on canonical_name was case-sensitive, causing +"Alice" and "alice" to have different trigram sets. This recreates it on +LOWER(canonical_name) so the % operator matches case-insensitively. + +Revision ID: d6e7f8a9b0c1 +Revises: c5d6e7f8a9b0 +Create Date: 2026-03-31 +""" + +from collections.abc import Sequence + +from alembic import context, op + +revision: str = "d6e7f8a9b0c1" +down_revision: str | Sequence[str] | None = "c5d6e7f8a9b0" +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() + # Drop the old case-sensitive trigram index + op.execute(f"DROP INDEX IF EXISTS entities_canonical_name_trgm_idx") + # Create case-insensitive trigram index on LOWER(canonical_name) + op.execute( + f"CREATE INDEX IF NOT EXISTS entities_canonical_name_lower_trgm_idx " + f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops)" + ) + + +def downgrade() -> None: + op.execute(f"DROP INDEX IF EXISTS entities_canonical_name_lower_trgm_idx") + schema = _get_schema_prefix() + # Restore original case-sensitive index + op.execute( + f"CREATE INDEX IF NOT EXISTS entities_canonical_name_trgm_idx " + f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)" + ) diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index 53dc6fa5..d1878130 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -277,6 +277,7 @@ 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" +ENV_RETAIN_CHUNK_BATCH_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE" # File storage configuration ENV_FILE_STORAGE_TYPE = "HINDSIGHT_API_FILE_STORAGE_TYPE" @@ -340,6 +341,7 @@ ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES" ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT" ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS" ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS" +ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT" # Reflect agent settings ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS" @@ -466,6 +468,9 @@ DEFAULT_RETAIN_MISSION = None # Declarative spec of what to retain (injected in DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom") DEFAULT_RETAIN_DEFAULT_STRATEGY = None # Default strategy name (None = no strategy override) DEFAULT_RETAIN_STRATEGIES: dict | None = None # Named retain strategies (dict of name → config overrides) +DEFAULT_RETAIN_CHUNK_BATCH_SIZE = ( + 100 # Max chunks per streaming batch. Each chunk produces ~17 facts, so 100 chunks = ~1700 facts/batch. +) 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) @@ -513,6 +518,7 @@ DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker +DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention. # Reflect agent settings DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response @@ -759,6 +765,7 @@ class HindsightConfig: retain_batch_enabled: bool retain_batch_poll_interval_seconds: int retain_entity_lookup: str # "full" or "trigram" + retain_chunk_batch_size: int # Max chunks per streaming batch (0 = disabled) # File storage (static - server-level only) file_storage_type: str # "native" (PostgreSQL) or "s3" (S3-compatible) @@ -830,6 +837,7 @@ class HindsightConfig: worker_http_port: int worker_max_slots: int worker_consolidation_max_slots: int + retain_max_concurrent: int # Reflect agent settings reflect_max_iterations: int @@ -896,6 +904,7 @@ class HindsightConfig: "retain_custom_instructions", "retain_default_strategy", "retain_strategies", + "retain_chunk_batch_size", # Entity labels (controlled vocabulary for entity classification) "entity_labels", "entities_allow_free_form", @@ -1249,6 +1258,7 @@ class HindsightConfig: retain_batch_poll_interval_seconds=int( os.getenv(ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS, str(DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS)) ), + retain_chunk_batch_size=int(os.getenv(ENV_RETAIN_CHUNK_BATCH_SIZE, str(DEFAULT_RETAIN_CHUNK_BATCH_SIZE))), # File storage file_storage_type=os.getenv(ENV_FILE_STORAGE_TYPE, DEFAULT_FILE_STORAGE_TYPE), file_storage_s3_bucket=os.getenv(ENV_FILE_STORAGE_S3_BUCKET) or None, @@ -1330,6 +1340,7 @@ class HindsightConfig: worker_consolidation_max_slots=int( os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS)) ), + retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))), # Reflect agent settings reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))), reflect_max_context_tokens=int( diff --git a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py index 47ee5be3..3c1e64fb 100644 --- a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py +++ b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py @@ -317,8 +317,13 @@ class EntityResolver: 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. + # Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive + # similarity lookup. Previous version also had LIKE '%...' substring fallbacks, + # but those forced full sequential scans of the entities table and caused + # TimeoutErrors on banks with 10k+ entities. Lowering the similarity threshold + # to 0.15 (from default 0.3) catches most substring relationships while + # staying fully index-based. + await conn.execute("SET pg_trgm.similarity_threshold = 0.15") rows = await conn.fetch( f""" SELECT DISTINCT ON (e.id) @@ -327,16 +332,13 @@ class EntityResolver: 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) || '%' - ) + AND LOWER(e.canonical_name) % LOWER(q.query_text) ) """, bank_id, entity_texts, ) + await conn.execute("RESET pg_trgm.similarity_threshold") # Group candidates by query_text all_candidates: dict[str, list] = {t: [] for t in entity_texts} @@ -808,14 +810,19 @@ class EntityResolver: return await self._link_units_to_entities_batch_impl(conn, unit_entity_pairs) async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str]]): - # Batch insert all unit-entity links - await conn.executemany( + # Sorted bulk insert to prevent deadlocks from inconsistent lock ordering + # across concurrent transactions on the unit_entities unique index. + sorted_pairs = sorted(unit_entity_pairs) + unit_ids = [p[0] for p in sorted_pairs] + entity_ids = [p[1] for p in sorted_pairs] + await conn.execute( f""" INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id) - VALUES ($1, $2) + SELECT u, e FROM unnest($1::uuid[], $2::uuid[]) AS t(u, e) ON CONFLICT DO NOTHING """, - unit_entity_pairs, + unit_ids, + entity_ids, ) # Build map of unit -> entities for co-occurrence calculation diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 2d87176a..f5e73cc2 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -532,10 +532,11 @@ class MemoryEngine(MemoryEngineInterface): # Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50) self._search_semaphore = asyncio.Semaphore(get_config().recall_max_concurrent) - # Backpressure for put operations: limit concurrent puts to prevent database contention - # Each put_batch holds a connection for the entire transaction, so we limit to 5 - # concurrent puts to avoid connection pool exhaustion and reduce write contention - self._put_semaphore = asyncio.Semaphore(5) + # Backpressure for retain DB writes: limit concurrent transactions to prevent contention + # on entity/link tables. Acquired in the orchestrator *after* LLM extraction completes, + # so LLM calls run in full parallelism while only the DB-heavy phase is throttled. + # Configurable via HINDSIGHT_API_RETAIN_MAX_CONCURRENT (default: 4). + self._put_semaphore = asyncio.Semaphore(get_config().retain_max_concurrent) # initialize encoding eagerly to avoid delaying the first time _get_tiktoken_encoding() @@ -2213,7 +2214,6 @@ class MemoryEngine(MemoryEngineInterface): document_id=document_id, is_first_batch=i == 1, # Only upsert on first batch fact_type_override=fact_type_override, - confidence_score=confidence_score, document_tags=document_tags, operation_id=operation_id, strategy=strategy, @@ -2238,7 +2238,6 @@ class MemoryEngine(MemoryEngineInterface): document_id=document_id, is_first_batch=True, fact_type_override=fact_type_override, - confidence_score=confidence_score, document_tags=document_tags, operation_id=operation_id, strategy=strategy, @@ -2290,7 +2289,6 @@ class MemoryEngine(MemoryEngineInterface): document_id: str | None = None, is_first_batch: bool = True, fact_type_override: str | None = None, - confidence_score: float | None = None, document_tags: list[str] | None = None, operation_id: str | None = None, outbox_callback: "Callable[[asyncpg.Connection], Awaitable[None]] | None" = None, @@ -2311,54 +2309,51 @@ class MemoryEngine(MemoryEngineInterface): document_id: Optional document ID (always upserts if exists) is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch) fact_type_override: Override fact type for all facts - confidence_score: Confidence score for opinions document_tags: Tags applied to all items in this batch Returns: Tuple of (unit ID lists, token usage for fact extraction) """ - # Backpressure: limit concurrent retains to prevent database contention - async with self._put_semaphore: - # Use the new modular orchestrator - from .retain import orchestrator + # Use the new modular orchestrator + from .retain import orchestrator - pool = await self._get_pool() + pool = await self._get_pool() - # Resolve bank-specific config for this operation - resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context) + # Resolve bank-specific config for this operation + resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context) - # Force chunks mode when LLM provider is "none" (no LLM available for fact extraction) - if self._llm_config.provider == "none": - resolved_config.retain_extraction_mode = "chunks" - resolved_config.enable_observations = False + # Force chunks mode when LLM provider is "none" (no LLM available for fact extraction) + if self._llm_config.provider == "none": + resolved_config.retain_extraction_mode = "chunks" + resolved_config.enable_observations = False - # Apply strategy overrides: explicit strategy > bank default strategy - from hindsight_api.config_resolver import apply_strategy + # Apply strategy overrides: explicit strategy > bank default strategy + from hindsight_api.config_resolver import apply_strategy - effective_strategy = strategy or resolved_config.retain_default_strategy - if effective_strategy: - resolved_config = apply_strategy(resolved_config, effective_strategy) + effective_strategy = strategy or resolved_config.retain_default_strategy + if effective_strategy: + resolved_config = apply_strategy(resolved_config, effective_strategy) - # Create parent span for retain operation - with create_operation_span("retain", bank_id): - return await orchestrator.retain_batch( - pool=pool, - embeddings_model=self.embeddings, - llm_config=self._retain_llm_config.with_config(resolved_config), - entity_resolver=self.entity_resolver, - format_date_fn=self._format_readable_date, - bank_id=bank_id, - contents_dicts=contents, - document_id=document_id, - is_first_batch=is_first_batch, - fact_type_override=fact_type_override, - confidence_score=confidence_score, - document_tags=document_tags, - config=resolved_config, - operation_id=operation_id, - schema=_current_schema.get(), - outbox_callback=outbox_callback, - ) + # Create parent span for retain operation + with create_operation_span("retain", bank_id): + return await orchestrator.retain_batch( + pool=pool, + embeddings_model=self.embeddings, + llm_config=self._retain_llm_config.with_config(resolved_config), + entity_resolver=self.entity_resolver, + format_date_fn=self._format_readable_date, + bank_id=bank_id, + contents_dicts=contents, + document_id=document_id, + is_first_batch=is_first_batch, + fact_type_override=fact_type_override, + document_tags=document_tags, + config=resolved_config, + operation_id=operation_id, + schema=_current_schema.get(), + outbox_callback=outbox_callback, + db_semaphore=self._put_semaphore, + ) def recall( self, 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 828c4bb8..286eb79c 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/entity_processing.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/entity_processing.py @@ -12,61 +12,27 @@ from .types import EntityLink, ProcessedFact logger = logging.getLogger(__name__) -async def process_entities_batch( - entity_resolver, - conn, - bank_id: str, - unit_ids: list[str], +def _prepare_facts_for_entity_processing( facts: list[ProcessedFact], - log_buffer: list[str] = None, - user_entities_per_content: dict[int, list[dict]] = None, - entity_labels: list | None = None, -) -> list[EntityLink]: + user_entities_per_content: dict[int, list[dict]] | None = None, +) -> tuple[list[str], list, list[list[dict]]]: """ - Process entities for all facts and create entity links. - - This function: - 1. Extracts entity mentions from fact texts - 2. Merges user-provided entities with LLM-extracted entities - 3. Resolves entity names to canonical entities - 4. Creates entity records in the database - 5. Returns entity links ready for insertion - - Args: - entity_resolver: EntityResolver instance for entity resolution - conn: Database connection - bank_id: Bank identifier - unit_ids: List of unit IDs (same length as facts) - facts: List of ProcessedFact objects - log_buffer: Optional buffer for detailed logging - user_entities_per_content: Dict mapping content_index to list of user-provided entities + Extract fact texts, dates, and merged entity lists from ProcessedFact objects. Returns: - List of EntityLink objects for batch insertion + Tuple of (fact_texts, fact_dates, entities_per_fact) """ - if not unit_ids or not facts: - return [] - - if len(unit_ids) != len(facts): - raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})") - user_entities_per_content = user_entities_per_content or {} - # Extract data for link_utils function fact_texts = [fact.fact_text for fact in facts] - # Use occurred_start if available, otherwise use mentioned_at for entity timestamps fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts] - # Convert EntityRef objects to dict format and merge with user-provided entities entities_per_fact = [] for fact in facts: - # Start with LLM-extracted entities llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])] - # Get user entities for this content (use content_index from fact) user_entities = user_entities_per_content.get(fact.content_index, []) - # Merge with case-insensitive deduplication seen_texts = {e["text"].lower() for e in llm_entities} for user_entity in user_entities: if user_entity["text"].lower() not in seen_texts: @@ -80,8 +46,48 @@ async def process_entities_batch( entities_per_fact.append(llm_entities) - # Use existing link_utils function for entity processing - entity_links = await link_utils.extract_entities_batch_optimized( + return fact_texts, fact_dates, entities_per_fact + + +async def resolve_entities( + entity_resolver, + conn, + bank_id: str, + unit_ids: list[str], + facts: list[ProcessedFact], + log_buffer: list[str] = None, + user_entities_per_content: dict[int, list[dict]] = None, + entity_labels: list | None = None, +) -> tuple[list[str], list[tuple], dict[str, list[str]]]: + """ + Phase 1: Resolve entity names to canonical IDs (read-heavy). + + Should be called on a SEPARATE connection OUTSIDE the main write transaction + to avoid holding the transaction open during expensive trigram scans. + + Args: + entity_resolver: EntityResolver instance + conn: Database connection (separate from the main write transaction) + bank_id: Bank identifier + unit_ids: Placeholder unit IDs (used only for grouping) + facts: List of ProcessedFact objects + log_buffer: Optional buffer for detailed logging + user_entities_per_content: Dict mapping content_index to user-provided 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(). + """ + if not unit_ids or not facts: + return [], [], {} + + if len(unit_ids) != len(facts): + raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})") + + fact_texts, fact_dates, entities_per_fact = _prepare_facts_for_entity_processing(facts, user_entities_per_content) + + return await link_utils.resolve_entities_only( entity_resolver, conn, bank_id, @@ -90,11 +96,55 @@ async def process_entities_batch( "", # context (not used in current implementation) fact_dates, entities_per_fact, - log_buffer, # Pass log_buffer for detailed logging + log_buffer, entity_labels=entity_labels, ) - return entity_links + +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, +) -> 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) + + 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, + ) async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str) -> None: diff --git a/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py b/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py index ac46c04b..f4f7363e 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py @@ -1464,28 +1464,76 @@ async def extract_facts_from_text( f"chunk_size={config.retain_chunk_size:,}) - starting parallel LLM extraction" ) - tasks = [ - _extract_facts_with_auto_split( - chunk=chunk, - chunk_index=i, - total_chunks=len(chunks), - event_date=event_date, - context=context, - llm_config=llm_config, - config=config, - agent_name=agent_name, - metadata=metadata, - ) - for i, chunk in enumerate(chunks) - ] - chunk_results = await asyncio.gather(*tasks) + # Per-chunk retry wrapper: each chunk gets up to MAX_CHUNK_RETRIES attempts. + # This handles transient LLM failures (timeouts, rate limits, malformed responses) + # without discarding the entire batch. If a chunk still fails after all retries, + # the ENTIRE retain fails — we do not accept partial extraction. + MAX_CHUNK_RETRIES = 3 + CHUNK_RETRY_BASE_DELAY = 2.0 # seconds, doubles each retry + + async def _extract_chunk_with_retry(chunk: str, chunk_index: int) -> tuple: + """Extract facts from a single chunk with retries on failure.""" + last_exception = None + for attempt in range(MAX_CHUNK_RETRIES): + try: + return await _extract_facts_with_auto_split( + chunk=chunk, + chunk_index=chunk_index, + total_chunks=len(chunks), + event_date=event_date, + context=context, + llm_config=llm_config, + config=config, + agent_name=agent_name, + metadata=metadata, + ) + except Exception as e: + last_exception = e + if attempt < MAX_CHUNK_RETRIES - 1: + delay = CHUNK_RETRY_BASE_DELAY * (2**attempt) + logger.warning( + f"Chunk {chunk_index}/{len(chunks)} extraction failed " + f"(attempt {attempt + 1}/{MAX_CHUNK_RETRIES}): " + f"{type(e).__name__}. Retrying in {delay:.0f}s..." + ) + await asyncio.sleep(delay) + else: + logger.error( + f"Chunk {chunk_index}/{len(chunks)} extraction failed after " + f"{MAX_CHUNK_RETRIES} attempts: {type(e).__name__}: {e}" + ) + raise last_exception + + tasks = [_extract_chunk_with_retry(chunk, i) for i, chunk in enumerate(chunks)] + + # return_exceptions=True so we can collect all results even if some chunks + # exhausted their retries. We check for failures below and fail the retain + # if ANY chunk could not be extracted — partial extraction is not acceptable. + chunk_results = await asyncio.gather(*tasks, return_exceptions=True) + all_facts = [] chunk_metadata = [] # [(chunk_text, fact_count), ...] total_usage = TokenUsage() - for chunk, (chunk_facts, chunk_usage) in zip(chunks, chunk_results): + failed_chunks = [] + for i, (chunk, result) in enumerate(zip(chunks, chunk_results)): + if isinstance(result, Exception): + failed_chunks.append((i, result)) + continue + chunk_facts, chunk_usage = result all_facts.extend(chunk_facts) chunk_metadata.append((chunk, len(chunk_facts))) total_usage = total_usage + chunk_usage + + if failed_chunks: + # Fail the entire retain — partial extraction is not acceptable. + # All successfully extracted facts are discarded because the transaction + # hasn't committed yet. The worker poller will retry the entire task. + failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5]) + raise RuntimeError( + f"Fact extraction failed: {len(failed_chunks)}/{len(chunks)} chunks failed " + f"after {MAX_CHUNK_RETRIES} retries each. First failures: {failed_summary}" + ) + return all_facts, chunk_metadata, total_usage @@ -2055,8 +2103,9 @@ async def extract_facts_from_contents( ) fact_extraction_tasks.append(task) - # Step 2: Wait for all fact extractions to complete - all_fact_results = await asyncio.gather(*fact_extraction_tasks) + # Step 2: Wait for all fact extractions to complete. + # Use return_exceptions=True so one content item failure doesn't discard the rest. + all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True) # Step 3: Flatten and convert to typed objects extracted_facts: list[ExtractedFactType] = [] @@ -2066,9 +2115,16 @@ async def extract_facts_from_contents( global_chunk_idx = 0 global_fact_idx = 0 - for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate( - zip(contents, all_fact_results) - ): + # Filter out failed content items + valid_results = [] + for content, result in zip(contents, all_fact_results): + if isinstance(result, Exception): + logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}") + valid_results.append((content, ([], [], TokenUsage()))) + else: + valid_results.append((content, result)) + + for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results): total_usage = total_usage + content_usage chunk_start_idx = global_chunk_idx diff --git a/hindsight-api-slim/hindsight_api/engine/retain/link_creation.py b/hindsight-api-slim/hindsight_api/engine/retain/link_creation.py index 416f057b..82c6bbdb 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/link_creation.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/link_creation.py @@ -32,17 +32,26 @@ async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str]) - return await link_utils.create_temporal_links_batch_per_fact(conn, bank_id, unit_ids, log_buffer=[]) -async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], embeddings: list[list[float]]) -> int: +async def create_semantic_links_batch( + conn, + bank_id: str, + unit_ids: list[str], + embeddings: list[list[float]], + pre_computed_ann_links: list[tuple] | None = None, +) -> int: """ Create semantic links between facts. Links facts that are semantically similar based on embeddings. + When pre_computed_ann_links are provided (from Phase 1), they are used + instead of running ANN queries inside the transaction. Args: conn: Database connection bank_id: Bank identifier unit_ids: List of unit IDs to create links for embeddings: List of embedding vectors (same length as unit_ids) + pre_computed_ann_links: Pre-computed ANN results from Phase 1 Returns: Number of semantic links created @@ -53,7 +62,9 @@ async def create_semantic_links_batch(conn, bank_id: str, unit_ids: list[str], e if len(unit_ids) != len(embeddings): raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})") - return await link_utils.create_semantic_links_batch(conn, bank_id, unit_ids, embeddings, log_buffer=[]) + return await link_utils.create_semantic_links_batch( + conn, bank_id, unit_ids, embeddings, log_buffer=[], pre_computed_ann_links=pre_computed_ann_links + ) async def create_causal_links_batch(conn, bank_id: str, unit_ids: list[str], facts: list[ProcessedFact]) -> int: 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 ed41e52d..1d80bcaa 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py @@ -12,6 +12,114 @@ from .types import EntityLink logger = logging.getLogger(__name__) +# Sentinel UUID used in the unique index to represent NULL entity_id +_NIL_ENTITY_UUID = "00000000-0000-0000-0000-000000000000" + +# Maximum number of temporal links to keep per unit (from_unit_id). +# Retrieval only reads top 10-20 per unit via LATERAL join, so keeping +# more is wasted storage and write amplification. +MAX_TEMPORAL_LINKS_PER_UNIT = 20 + + +def _cap_links_per_unit(links: list[tuple], max_per_unit: int = MAX_TEMPORAL_LINKS_PER_UNIT) -> list[tuple]: + """Keep only the top-N links per from_unit_id, ranked by weight descending. + + Args: + links: List of (from_unit_id, to_unit_id, link_type, weight, entity_id) tuples. + max_per_unit: Maximum number of links to retain per from_unit_id. + + Returns: + Filtered list of link tuples. + """ + if not links: + return links + + # Group by from_unit_id (index 0) + groups: dict[str, list[tuple]] = {} + for link in links: + key = str(link[0]) + if key not in groups: + groups[key] = [] + groups[key].append(link) + + # For each group, sort by weight (index 3) descending and keep top N + result: list[tuple] = [] + for group_links in groups.values(): + group_links.sort(key=lambda lnk: lnk[3], reverse=True) + result.extend(group_links[:max_per_unit]) + + return result + + +async def _bulk_insert_links( + conn, + links: list[tuple], + bank_id: str = "", + chunk_size: int = 5000, + skip_exists_check: bool = False, +) -> None: + """Insert links into memory_links using sorted bulk INSERT FROM unnest(). + + Sorting by (from_unit_id, to_unit_id) ensures all concurrent transactions + acquire index locks in the same order, eliminating circular-wait deadlocks. + + A single INSERT ... SELECT FROM unnest() is also faster than executemany + (one round-trip vs N), and acquires all locks within one statement execution + rather than interleaving with other transactions between rows. + + Args: + conn: Database connection (must be inside a transaction). + links: List of (from_unit_id, to_unit_id, link_type, weight, entity_id) tuples. + bank_id: Bank identifier stored on memory_links for fast filtering. + chunk_size: Max rows per INSERT statement to avoid query timeouts on + very large tables (100M+ rows). + skip_exists_check: Skip WHERE EXISTS checks on memory_units. Use when + all referenced unit IDs are guaranteed to exist (e.g., within + the same transaction that inserted them). + """ + if not links: + return + + # Sort by (from_unit_id, to_unit_id) to guarantee consistent lock ordering + # across concurrent transactions — prevents deadlocks. + sorted_links = sorted(links, key=lambda lnk: (str(lnk[0]), str(lnk[1]))) + + from_ids = [lnk[0] for lnk in sorted_links] + to_ids = [lnk[1] for lnk in sorted_links] + types = [lnk[2] for lnk in sorted_links] + weights = [lnk[3] for lnk in sorted_links] + entity_ids = [lnk[4] for lnk in sorted_links] + + exists_clause = "" + if not skip_exists_check: + exists_clause = ( + f"WHERE EXISTS (SELECT 1 FROM {fq_table('memory_units')} mu WHERE mu.id = f)" + f" AND EXISTS (SELECT 1 FROM {fq_table('memory_units')} mu WHERE mu.id = t)" + ) + + for chunk_start in range(0, len(sorted_links), chunk_size): + chunk_end = min(chunk_start + chunk_size, len(sorted_links)) + await conn.execute( + f""" + INSERT INTO {fq_table("memory_links")} + (from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id) + SELECT f, t, tp, w, e, $6 + FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[]) + AS t(f, t, tp, w, e) + {exists_clause} + ON CONFLICT (from_unit_id, to_unit_id, link_type, + COALESCE(entity_id, '{_NIL_ENTITY_UUID}'::uuid)) + DO NOTHING + """, + from_ids[chunk_start:chunk_end], + to_ids[chunk_start:chunk_end], + types[chunk_start:chunk_end], + weights[chunk_start:chunk_end], + entity_ids[chunk_start:chunk_end], + bank_id, + timeout=300, + ) + def _normalize_datetime(dt): """Normalize datetime to be timezone-aware (UTC) for consistent comparison.""" @@ -78,7 +186,7 @@ def compute_temporal_links( weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours)) links.append((unit_id, str(recent_id), "temporal", weight, None)) - return links + return _cap_links_per_unit(links) def compute_temporal_query_bounds( @@ -140,7 +248,70 @@ def _log(log_buffer, message, level="info"): logger.log(logging.WARNING if level == "warning" else logging.ERROR, message) -async def extract_entities_batch_optimized( +def _prepare_entities_for_resolution( + unit_ids: list[str], + sentences: list[str], + fact_dates: list, + llm_entities: list[list[dict]], + log_buffer: list[str] = None, +) -> tuple[list[dict], list[list[dict]], list[tuple]]: + """ + Convert LLM entities into the flat format expected by entity resolver. + + Returns: + Tuple of (all_entities_flat, all_entities, entity_to_unit) where: + - all_entities_flat: flat list of entity dicts ready for resolve_entities_batch + - all_entities: per-unit formatted entity lists + - entity_to_unit: maps flat index to (unit_id, local_index, fact_date) + """ + substep_start = time.time() + all_entities = [] + for entity_list in llm_entities: + formatted_entities = [] + for ent in entity_list: + if hasattr(ent, "text"): + formatted_entities.append({"text": ent.text, "type": "CONCEPT"}) + elif isinstance(ent, dict): + formatted_entities.append({"text": ent.get("text", ""), "type": ent.get("type", "CONCEPT")}) + all_entities.append(formatted_entities) + + total_entities = sum(len(ents) for ents in all_entities) + _log( + log_buffer, + f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s", + level="debug", + ) + + substep_start = time.time() + all_entities_flat = [] + entity_to_unit: list[tuple] = [] + + for unit_id, entities, fact_date in zip(unit_ids, all_entities, fact_dates): + if not entities: + continue + for local_idx, entity in enumerate(entities): + all_entities_flat.append( + { + "text": entity["text"], + "type": entity["type"], + "nearby_entities": entities, + } + ) + entity_to_unit.append((unit_id, local_idx, fact_date)) + _log( + log_buffer, + f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_start:.3f}s", + level="debug", + ) + + # Attach per-entity dates + for idx, (_unit_id, _local_idx, fact_date) in enumerate(entity_to_unit): + all_entities_flat[idx]["event_date"] = fact_date + + return all_entities_flat, all_entities, entity_to_unit + + +async def resolve_entities_only( entity_resolver, conn, bank_id: str, @@ -151,244 +322,205 @@ async def extract_entities_batch_optimized( llm_entities: list[list[dict]], log_buffer: list[str] = None, entity_labels: list | None = None, -) -> list[tuple]: +) -> tuple[list[str], list[tuple], dict[str, list[str]]]: """ - Process LLM-extracted entities for ALL facts in batch. + Phase 1 of entity processing: resolve entity names to canonical IDs. - Uses entities provided by the LLM (no spaCy needed), then resolves - and links them in bulk. + Runs the expensive read-heavy trigram search, co-occurrence fetch, and scoring + OUTSIDE the main write transaction. Also INSERTs new entities (idempotent + DO NOTHING) so that IDs are available for the subsequent write phase. Args: - entity_resolver: EntityResolver instance for entity resolution - conn: Database connection - agent_id: bank IDentifier - unit_ids: List of unit IDs - sentences: List of fact sentences + entity_resolver: EntityResolver instance + conn: Database connection (separate from the main write transaction) + bank_id: Bank identifier + unit_ids: Placeholder unit IDs (used only for grouping, not yet inserted) + sentences: Fact texts context: Context string - fact_dates: List of fact dates - llm_entities: List of entity lists from LLM extraction - log_buffer: Optional buffer for logging + fact_dates: Per-fact dates + llm_entities: Per-fact entity lists from LLM extraction + log_buffer: Optional logging buffer + entity_labels: Optional entity label taxonomy Returns: - List of tuples for batch insertion: (from_unit_id, to_unit_id, link_type, weight, entity_id) + Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids) where: + - resolved_entity_ids: list of entity IDs in same order as flattened entities + - entity_to_unit: maps flat index to (unit_id, local_index, fact_date) + - unit_to_entity_ids: maps unit_id to list of resolved entity IDs """ - try: - # Step 1: Convert LLM entities to the format expected by entity resolver - substep_start = time.time() - all_entities = [] - for entity_list in llm_entities: - # Convert List[Entity] or List[dict] to List[Dict] format - formatted_entities = [] - for ent in entity_list: - # Handle both Entity objects and dicts - if hasattr(ent, "text"): - # Entity objects only have 'text', default type to 'CONCEPT' - formatted_entities.append({"text": ent.text, "type": "CONCEPT"}) - elif isinstance(ent, dict): - formatted_entities.append({"text": ent.get("text", ""), "type": ent.get("type", "CONCEPT")}) - all_entities.append(formatted_entities) + all_entities_flat, _all_entities, entity_to_unit = _prepare_entities_for_resolution( + unit_ids, sentences, fact_dates, llm_entities, log_buffer + ) - total_entities = sum(len(ents) for ents in all_entities) + if not all_entities_flat: + _log(log_buffer, " [6.2] Entity resolution (batched): 0 entities", level="debug") + return [], [], {} + + step_start = time.time() + resolved_entity_ids = await entity_resolver.resolve_entities_batch( + bank_id=bank_id, + entities_data=all_entities_flat, + context=context, + unit_event_date=None, + conn=conn, + entity_labels=entity_labels, + ) + _log( + log_buffer, + f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in single batch in {time.time() - step_start:.3f}s", + level="debug", + ) + + # Build unit_to_entity_ids mapping + unit_to_entity_ids: dict[str, list[str]] = {} + for idx, (unit_id, _local_idx, _fact_date) in enumerate(entity_to_unit): + if unit_id not in unit_to_entity_ids: + unit_to_entity_ids[unit_id] = [] + unit_to_entity_ids[unit_id].append(resolved_entity_ids[idx]) + + _log( + log_buffer, + f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_start:.3f}s", + level="debug", + ) + + 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, +) -> 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): + unit_entity_pairs.append((unit_id, resolved_entity_ids[idx])) + + await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn) _log( log_buffer, - f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s", + f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s", level="debug", ) - # Step 2: Resolve entities in BATCH (much faster!) - substep_start = time.time() - step_6_2_start = time.time() + # 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) - # [6.2.1] Prepare all entities for batch resolution - substep_6_2_1_start = time.time() - all_entities_flat = [] - entity_to_unit = [] # Maps flat index to (unit_id, local_index) + _log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug") - for unit_id, entities, fact_date in zip(unit_ids, all_entities, fact_dates): - if not entities: - continue + MAX_LINKS_PER_ENTITY = 10 - for local_idx, entity in enumerate(entities): - all_entities_flat.append( - { - "text": entity["text"], - "type": entity["type"], - "nearby_entities": entities, - } + 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] + # Use LATERAL with LIMIT to cap rows fetched per entity at the SQL level, + # avoiding transfer of thousands of rows for high-cardinality entities. + rows = 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 {fq_table("unit_entities")} ue + WHERE ue.entity_id = e.entity_id + ORDER BY ue.unit_id DESC + LIMIT $2 + ) n + """, + entity_id_list, + MAX_LINKS_PER_ENTITY + len(unit_ids), # room for new units + existing cap + ) + _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) ) - entity_to_unit.append((unit_id, local_idx, fact_date)) - _log( - log_buffer, - f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s", - level="debug", - ) - # Resolve ALL entities in one batch call - if all_entities_flat: - # [6.2.2] Batch resolve entities - single call with per-entity dates - substep_6_2_2_start = time.time() + 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) + ) - # Add per-entity dates to entity data for batch resolution - for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit): - all_entities_flat[idx]["event_date"] = fact_date + _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", + ) - # Resolve ALL entities in ONE batch call (much faster than sequential buckets) - # INSERT ... ON CONFLICT handles any race conditions at the DB level - resolved_entity_ids = await entity_resolver.resolve_entities_batch( - bank_id=bank_id, - entities_data=all_entities_flat, - context=context, - unit_event_date=None, # Not used when per-entity dates provided - conn=conn, # Use main transaction connection - entity_labels=entity_labels, - ) - - _log( - log_buffer, - f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in single batch in {time.time() - substep_6_2_2_start:.3f}s", - level="debug", - ) - - # [6.2.3] Create unit-entity links in BATCH - substep_6_2_3_start = time.time() - # Map resolved entities back to units and collect all (unit, entity) pairs - unit_to_entity_ids = {} - unit_entity_pairs = [] - for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit): - if unit_id not in unit_to_entity_ids: - unit_to_entity_ids[unit_id] = [] - - entity_id = resolved_entity_ids[idx] - unit_to_entity_ids[unit_id].append(entity_id) - unit_entity_pairs.append((unit_id, entity_id)) - - # Batch insert all unit-entity links (MUCH faster!) - 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_6_2_3_start:.3f}s", - level="debug", - ) - - _log( - log_buffer, - f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s", - level="debug", - ) - else: - unit_to_entity_ids = {} - _log( - log_buffer, - f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s", - level="debug", - ) - - # Step 3: Create entity links between units that share entities - substep_start = time.time() - # Collect all unique entity IDs - all_entity_ids = set() - for entity_ids in unit_to_entity_ids.values(): - all_entity_ids.update(entity_ids) - - _log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug") - - # Find all units that reference these entities (ONE batched query) - 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] - rows = await conn.fetch( - f""" - SELECT entity_id, unit_id - FROM {fq_table("unit_entities")} - WHERE entity_id = ANY($1::uuid[]) - """, - entity_id_list, - ) - _log( - log_buffer, - f" [6.3.1] Query unit_entities: {len(rows)} rows in {time.time() - query_start:.3f}s", - level="debug", - ) - - # Group by entity_id - 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") - - # Create bidirectional links between units that share entities - # OPTIMIZATION: Limit links per entity to avoid N² explosion - # Only link each new unit to the most recent MAX_LINKS_PER_ENTITY units - MAX_LINKS_PER_ENTITY = 50 # Limit to prevent explosion when entity appears in many facts - link_gen_start = time.time() - links: list[EntityLink] = [] - new_unit_set = set(unit_ids) # Units from this batch - - 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) - # Separate new units (from this batch) and existing units - 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] - - # Link new units to each other (within batch) - also limited - # For very common entities, limit within-batch links too - 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 - ) - ) - - # Link new units to LIMITED existing units (most recent) - existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:] # Take most recent - 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 - - except Exception as e: - logger.error(f"Failed to extract entities in batch: {str(e)}") - import traceback - - traceback.print_exc() - raise + return links async def create_temporal_links_batch_per_fact( @@ -406,7 +538,7 @@ async def create_temporal_links_batch_per_fact( Args: conn: Database connection - agent_id: bank IDentifier + bank_id: Bank identifier unit_ids: List of unit IDs time_window_hours: Time window in hours for temporal links log_buffer: Optional buffer for logging @@ -424,63 +556,118 @@ async def create_temporal_links_batch_per_fact( fetch_dates_start = time_mod.time() rows = await conn.fetch( f""" - SELECT id, event_date + SELECT id, event_date, fact_type FROM {fq_table("memory_units")} WHERE id::text = ANY($1) """, unit_ids, ) - new_units = {str(row["id"]): row["event_date"] for row in rows} + new_units = {str(row["id"]): (row["event_date"], row["fact_type"]) for row in rows} _log( log_buffer, f" [7.1] Fetch event_dates for {len(unit_ids)} units: {time_mod.time() - fetch_dates_start:.3f}s", ) - # Fetch ALL potential temporal neighbors in ONE query (much faster!) - # Get time range across all units with overflow protection - min_date, max_date = compute_temporal_query_bounds(new_units, time_window_hours) - + # Use LATERAL push-down to fetch only top-N temporal neighbors per new unit, + # avoiding transfer of the entire time-window result set (could be 50k+ rows). fetch_neighbors_start = time_mod.time() - if min_date is not None and max_date is not None: - all_candidates = await conn.fetch( - f""" - SELECT id, event_date - FROM {fq_table("memory_units")} - WHERE bank_id = $1 - AND event_date BETWEEN $2 AND $3 - AND id::text != ALL($4) - ORDER BY event_date DESC - """, - bank_id, - min_date, - max_date, - unit_ids, - ) + + # Build arrays of new unit IDs, event dates, and fact types for the LATERAL query + new_unit_entries = [(uid, edate, ftype) for uid, (edate, ftype) in new_units.items() if edate is not None] + if new_unit_entries: + import uuid as uuid_mod + + lateral_unit_ids = [ + uuid_mod.UUID(uid) if isinstance(uid, str) else uid for uid in [e[0] for e in new_unit_entries] + ] + lateral_event_dates = [_normalize_datetime(e[1]) for e in new_unit_entries] + lateral_fact_types = [e[2] for e in new_unit_entries] + # Bidirectional index scan: instead of scanning all units in the 24h + # window (O(N) — 164k rows at scale) and sorting by proximity, we scan + # the nearest K units in each direction using the B-tree index on + # (bank_id, fact_type, event_date). This reads only 2×K rows per probe + # regardless of bank size — 120x faster at 164k units (0.6ms vs 74ms). + TEMPORAL_LATERAL_BATCH = 500 + half_limit = MAX_TEMPORAL_LINKS_PER_UNIT # fetch K in each direction, take top K combined + mu = fq_table("memory_units") + rows = [] + for batch_start in range(0, len(new_unit_entries), TEMPORAL_LATERAL_BATCH): + batch_end = batch_start + TEMPORAL_LATERAL_BATCH + batch_rows = await conn.fetch( + f""" + SELECT from_id, id, event_date, time_diff_hours FROM ( + SELECT src.unit_id::text AS from_id, combined.*, + ROW_NUMBER() OVER ( + PARTITION BY src.unit_id + ORDER BY combined.time_diff_hours + ) AS rn + FROM unnest($1::uuid[], $2::timestamptz[], $3::text[]) + AS src(unit_id, event_date, fact_type) + CROSS JOIN LATERAL ( + -- Scan backward (older events) using index order + (SELECT mu.id, mu.event_date, + ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours + FROM {mu} mu + WHERE mu.bank_id = $4 + AND mu.fact_type = src.fact_type + AND mu.event_date <= src.event_date + AND mu.id != src.unit_id + ORDER BY mu.event_date DESC + LIMIT $5) + UNION ALL + -- Scan forward (newer events) using index order + (SELECT mu.id, mu.event_date, + ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours + FROM {mu} mu + WHERE mu.bank_id = $4 + AND mu.fact_type = src.fact_type + AND mu.event_date > src.event_date + AND mu.id != src.unit_id + ORDER BY mu.event_date ASC + LIMIT $5) + ) combined + ) ranked + WHERE rn <= $5 + """, + lateral_unit_ids[batch_start:batch_end], + lateral_event_dates[batch_start:batch_end], + lateral_fact_types[batch_start:batch_end], + bank_id, + half_limit, + ) + rows.extend(batch_rows) else: - all_candidates = [] + rows = [] + _log( log_buffer, - f" [7.2] Fetch {len(all_candidates)} candidate neighbors (1 query): {time_mod.time() - fetch_neighbors_start:.3f}s", + f" [7.2] Fetch {len(rows)} candidate neighbors (LATERAL): {time_mod.time() - fetch_neighbors_start:.3f}s", ) - # Filter and create links in memory (much faster than N queries) + # Build links directly from the LATERAL results (already per-unit limited) link_gen_start = time_mod.time() - links = compute_temporal_links(new_units, all_candidates, time_window_hours) + links = [] + for row in rows: + time_diff_h = float(row["time_diff_hours"]) + weight = max(0.3, 1.0 - (time_diff_h / time_window_hours)) + links.append((row["from_id"], str(row["id"]), "temporal", weight, None)) # Also compute temporal links WITHIN the new batch (new units to each other) if len(new_units) > 1: # Convert new_units dict to candidate format for within-batch linking new_unit_items = list(new_units.items()) - for i, (unit_id, event_date) in enumerate(new_unit_items): + for i, (unit_id, (event_date, fact_type)) in enumerate(new_unit_items): if event_date is None: continue # Skip units without event_date for temporal linking unit_event_date_norm = _normalize_datetime(event_date) # Compare with other new units (only those after this one to avoid duplicates) for j in range(i + 1, len(new_unit_items)): - other_id, other_event_date = new_unit_items[j] + other_id, (other_event_date, other_fact_type) = new_unit_items[j] if other_event_date is None: continue # Skip units without event_date + if fact_type != other_fact_type: + continue # Only link facts of the same type other_event_date_norm = _normalize_datetime(other_event_date) # Check if within time window @@ -491,23 +678,15 @@ async def create_temporal_links_batch_per_fact( links.append((unit_id, other_id, "temporal", weight, None)) links.append((other_id, unit_id, "temporal", weight, None)) + # Cap temporal links per unit to avoid write amplification; + # retrieval only reads top 10-20 per unit anyway. + links = _cap_links_per_unit(links) + _log(log_buffer, f" [7.3] Generate {len(links)} temporal links: {time_mod.time() - link_gen_start:.3f}s") if links: insert_start = time_mod.time() - # Add bank_id to each tuple for direct filtering (avoids expensive JOIN in stats) - links_with_bank = [(*link, bank_id) for link in links] - # Batch inserts to avoid timeout on large batches - BATCH_SIZE = 1000 - for batch_start in range(0, len(links_with_bank), BATCH_SIZE): - await conn.executemany( - f""" - INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id) - VALUES ($1, $2, $3, $4, $5, $6) - ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING - """, - links_with_bank[batch_start : batch_start + BATCH_SIZE], - ) + await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True) _log(log_buffer, f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s") return len(links) @@ -520,28 +699,203 @@ async def create_temporal_links_batch_per_fact( raise +async def compute_semantic_links_ann( + conn, + bank_id: str, + unit_ids: list[str], + embeddings: list[list[float]], + fact_types: list[str] | None = None, + top_k: int = 50, + threshold: float = 0.7, + log_buffer: list[str] = None, +) -> list[tuple]: + """ + Phase 1: ANN search for semantic neighbors among existing units. + + Runs on a separate connection OUTSIDE the write transaction to avoid + holding locks during expensive HNSW index probes. Uses a temp table + + LATERAL join to batch all probes in a single query. + + Queries are split by fact_type so PostgreSQL uses the per-bank partial + HNSW indexes (idx_mu_emb_worl_*, idx_mu_emb_expr_*). Without the + fact_type filter, the planner falls back to sequential scan (~50x slower). + + Args: + conn: Database connection (separate from write transaction, autocommit) + bank_id: Bank identifier + unit_ids: Placeholder unit IDs (real IDs not yet created) + embeddings: Embedding vectors for each unit + fact_types: Per-unit fact types (same length as unit_ids). Used to + query only the matching HNSW index per seed. + top_k: Max neighbors per unit + threshold: Minimum cosine similarity + log_buffer: Optional logging buffer + + Returns: + List of (from_id, to_id, "semantic", similarity, None) tuples + where from_id uses placeholder IDs. + """ + if not unit_ids or not embeddings: + return [] + + import time as time_mod + import uuid as uuid_mod + + ann_start = time_mod.time() + links = [] + + # Lower ef_search for retain ANN — default 400 is tuned for recall precision + # but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms/probe + # (35x faster) with sufficient accuracy for top-50 semantic link creation. + # Reset after to avoid polluting the connection pool for recall queries. + await conn.execute("SET hnsw.ef_search = 60") + + logger.debug(f"[ANN] Starting: {len(unit_ids)} seeds, top_k={top_k}") + + # Build per-unit fact_types (default to 'world' if not provided) + if fact_types is None: + fact_types = ["world"] * len(unit_ids) + + # No exclude_uuids — large exclusion lists (8k+ UUIDs) force PostgreSQL to + # sequential-scan every HNSW probe result against the array, destroying + # performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO + # NOTHING handles duplicates in memory_links). + t_setup = time_mod.time() + await conn.execute("CREATE TEMP TABLE IF NOT EXISTS _ann_seeds (unit_id text, emb_text text, fact_type text)") + await conn.execute("TRUNCATE _ann_seeds") + + records = [ + (uid, emb if isinstance(emb, str) else str(emb), ft) for uid, emb, ft in zip(unit_ids, embeddings, fact_types) + ] + await conn.copy_records_to_table("_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"]) + logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)") + + # Run one ANN query per fact_type so each uses the right HNSW index. + rows = [] + active_types = set(fact_types) + for fact_type in active_types: + t_query = time_mod.time() + seed_count = sum(1 for ft in fact_types if ft == fact_type) + logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds") + ft_rows = await conn.fetch( + f""" + SELECT s.unit_id AS from_id, + n.id::text AS to_id, + n.similarity + FROM _ann_seeds s + CROSS JOIN LATERAL ( + SELECT mu.id, + 1 - (mu.embedding <=> s.emb_text::vector) AS similarity + FROM {fq_table("memory_units")} mu + WHERE mu.bank_id = $1 + AND mu.fact_type = $2 + AND mu.embedding IS NOT NULL + ORDER BY mu.embedding <=> s.emb_text::vector + LIMIT $3 + ) n + WHERE s.fact_type = $2 + """, + bank_id, + fact_type, + top_k, + timeout=300, # ANN on large banks can take minutes + ) + logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s") + rows.extend(ft_rows) + + # Clean up temp table (no ON COMMIT DROP since we're not in a transaction) + await conn.execute("DROP TABLE IF EXISTS _ann_seeds") + + # Reset ef_search to default so the pooled connection doesn't affect recall queries + await conn.execute("RESET hnsw.ef_search") + + for row in rows: + sim = float(min(1.0, max(0.0, row["similarity"]))) + if sim >= threshold: + links.append((row["from_id"], row["to_id"], "semantic", sim, None)) + + _log( + log_buffer, + f" [8.1] ANN search (Phase 1): {len(unit_ids)} units → {len(links)} links in {time_mod.time() - ann_start:.3f}s", + ) + + return links + + +def compute_semantic_links_within_batch( + unit_ids: list[str], + embeddings: list[list[float]], + top_k: int = 50, + threshold: float = 0.7, +) -> list[tuple]: + """ + Compute semantic links between units within the same batch (no DB needed). + + Uses numpy dot product on embeddings already in memory — instant. + + Args: + unit_ids: Unit IDs (real IDs from insert_facts_batch) + embeddings: Embedding vectors + top_k: Max neighbors per unit + threshold: Minimum cosine similarity + + Returns: + List of (from_id, to_id, "semantic", similarity, None) tuples + """ + if len(unit_ids) < 2: + return [] + + import numpy as np + + links = [] + new_embeddings_matrix = np.array(embeddings) + + for i, unit_id in enumerate(unit_ids): + other_indices = [j for j in range(len(unit_ids)) if j != i] + if not other_indices: + continue + + other_embeddings = new_embeddings_matrix[other_indices] + similarities = np.dot(other_embeddings, new_embeddings_matrix[i]) + + above_threshold = np.where(similarities >= threshold)[0] + if len(above_threshold) > 0: + sorted_local_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k] + for local_idx in sorted_local_indices: + other_idx = other_indices[local_idx] + other_id = unit_ids[other_idx] + similarity = float(min(1.0, max(0.0, similarities[local_idx]))) + links.append((unit_id, other_id, "semantic", similarity, None)) + + return links + + async def create_semantic_links_batch( conn, bank_id: str, unit_ids: list[str], embeddings: list[list[float]], - top_k: int = 5, + top_k: int = 50, threshold: float = 0.7, log_buffer: list[str] = None, + pre_computed_ann_links: list[tuple] | None = None, ) -> int: """ - Create semantic links for multiple units efficiently. + Phase 2: Create semantic links (within-batch + pre-computed ANN results). - For each unit, finds similar units and creates links. + Within-batch similarities are computed in Python (numpy, instant). + ANN results from Phase 1 are passed in via pre_computed_ann_links and + inserted alongside the within-batch links. Args: - conn: Database connection - agent_id: bank IDentifier - unit_ids: List of unit IDs - embeddings: List of embedding vectors - top_k: Number of top similar units to link - threshold: Minimum similarity threshold - log_buffer: Optional buffer for logging + conn: Database connection (inside write transaction) + bank_id: Bank identifier + unit_ids: Real unit IDs (from insert_facts_batch) + embeddings: Embedding vectors + top_k: Max neighbors per unit + threshold: Minimum cosine similarity + log_buffer: Optional logging buffer + pre_computed_ann_links: ANN results from Phase 1 (already remapped to real IDs) Returns: Number of semantic links created @@ -552,96 +906,28 @@ async def create_semantic_links_batch( try: import time as time_mod - import numpy as np - - # 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 = [] - # Build UUID exclude list once for all ANN queries - import uuid as uuid_mod + # Within-batch similarities (numpy, no DB) + batch_start = time_mod.time() + within_batch_links = compute_semantic_links_within_batch(unit_ids, embeddings, top_k, threshold) + all_links.extend(within_batch_links) + _log( + log_buffer, + f" [8.1] Within-batch semantic: {len(within_batch_links)} links in {time_mod.time() - batch_start:.3f}s", + ) - exclude_uuids = [uuid_mod.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids] - - for unit_id, new_embedding in zip(unit_ids, embeddings): - emb_str = str(list(new_embedding) if not isinstance(new_embedding, list) else new_embedding) - rows = await conn.fetch( - f""" - SELECT id::text, - 1 - (embedding <=> $1::vector) AS similarity - FROM {fq_table("memory_units")} - WHERE bank_id = $2 - AND embedding IS NOT NULL - AND id != ALL($3::uuid[]) - ORDER BY embedding <=> $1::vector - LIMIT $4 - """, - emb_str, - bank_id, - exclude_uuids, - top_k, + # Add pre-computed ANN links from Phase 1 + if pre_computed_ann_links: + all_links.extend(pre_computed_ann_links) + _log( + log_buffer, + f" [8.2] Pre-computed ANN: {len(pre_computed_ann_links)} links", ) - 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)) - - _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 - if len(unit_ids) > 1: - new_embeddings_matrix = np.array(embeddings) - - for i, unit_id in enumerate(unit_ids): - # Compute similarities with all OTHER new units - other_indices = [j for j in range(len(unit_ids)) if j != i] - if not other_indices: - continue - - other_embeddings = new_embeddings_matrix[other_indices] - similarities = np.dot(other_embeddings, new_embeddings_matrix[i]) - - # Find top-k above threshold (same logic as existing units) - above_threshold = np.where(similarities >= threshold)[0] - - if len(above_threshold) > 0: - # Sort by similarity (descending) and take top-k - sorted_local_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k] - - for local_idx in sorted_local_indices: - other_idx = other_indices[local_idx] - other_id = unit_ids[other_idx] - # Clamp to [0, 1] to handle floating point precision issues - similarity = float(min(1.0, max(0.0, similarities[local_idx]))) - all_links.append((unit_id, other_id, "semantic", similarity, None)) - - _log( - log_buffer, - f" [8.2] Within-batch similarities added {len(all_links)} total semantic links", - ) if all_links: insert_start = time_mod.time() - # Add bank_id to each tuple for direct filtering (avoids expensive JOIN in stats) - all_links_with_bank = [(*link, bank_id) for link in all_links] - # Batch inserts to avoid timeout on large batches - BATCH_SIZE = 1000 - for batch_start in range(0, len(all_links_with_bank), BATCH_SIZE): - await conn.executemany( - f""" - INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id) - VALUES ($1, $2, $3, $4, $5, $6) - ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING - """, - all_links_with_bank[batch_start : batch_start + BATCH_SIZE], - ) + await _bulk_insert_links(conn, all_links, bank_id=bank_id) _log( log_buffer, f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s" ) @@ -658,11 +944,7 @@ async def create_semantic_links_batch( async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000): """ - Insert all entity links using COPY to temp table + chunked INSERT for reliability. - - 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). + Insert entity links into memory_links via sorted bulk INSERT FROM unnest(). Args: conn: Database connection @@ -676,63 +958,11 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, import time as time_mod total_start = time_mod.time() - - # 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, - weight float, - entity_id uuid, - bank_id text - ) ON COMMIT DROP - """) - logger.debug(f" [9.1] Create temp table: {time_mod.time() - create_start:.3f}s") - - # Clear any existing data in temp table - truncate_start = time_mod.time() - await conn.execute("TRUNCATE _temp_entity_links") - logger.debug(f" [9.2] Truncate temp table: {time_mod.time() - truncate_start:.3f}s") - - # Convert EntityLink objects to tuples for COPY - convert_start = time_mod.time() - records = [ - (link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id, bank_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) - copy_start = time_mod.time() - await conn.copy_records_to_table( - "_temp_entity_links", - records=records, - columns=["from_unit_id", "to_unit_id", "link_type", "weight", "entity_id", "bank_id"], + 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) + logger.debug( + f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s" ) - logger.debug(f" [9.4] COPY {len(records)} records to temp table: {time_mod.time() - copy_start:.3f}s") - - # Insert from temp table in chunks to avoid single-query timeouts on large tables - insert_start = time_mod.time() - 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, bank_id) - SELECT from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id - FROM _temp_entity_links - WHERE _row_num > $1 AND _row_num <= $2 - ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING - """, - chunk_start, - chunk_end, - ) - chunks += 1 - logger.debug(f" [9.5] INSERT {total_rows} rows in {chunks} chunks: {time_mod.time() - insert_start:.3f}s") - logger.debug(f" [9.TOTAL] Entity links batch insert: {time_mod.time() - total_start:.3f}s") async def create_causal_links_batch( @@ -804,28 +1034,12 @@ async def create_causal_links_batch( # Add the causal link # link_type is the relation_type (e.g., "causes", "caused_by") # weight is the strength of the relationship - links.append((from_unit_id, to_unit_id, relation_type, strength, None, bank_id)) + links.append((from_unit_id, to_unit_id, relation_type, strength, None)) if links: insert_start = time_mod.time() - try: - await conn.executemany( - f""" - INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id) - VALUES ($1, $2, $3, $4, $5, $6) - ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING - """, - links, - ) - except Exception as db_error: - # Log the actual data being inserted for debugging - logger.error(f"Database insert failed for causal links. Error: {db_error}") - logger.error(f"Attempted to insert {len(links)} links. First few:") - for i, link in enumerate(links[:3]): - logger.error( - f" Link {i}: from={link[0]}, to={link[1]}, type='{link[2]}' (repr={repr(link[2])}), weight={link[3]}, entity={link[4]}" - ) - raise + await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True) + logger.debug(f" [10.1] Insert {len(links)} causal links: {time_mod.time() - insert_start:.3f}s") return len(links) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index acc53d6b..992ab2c5 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -4,15 +4,18 @@ Main orchestrator for the retain pipeline. Coordinates all retain pipeline modules to store memories efficiently. """ +import asyncio +import hashlib +import json import logging import time import uuid -from collections import defaultdict from collections.abc import Awaitable, Callable from datetime import UTC, datetime from typing import Any -from ..db_utils import acquire_with_retry, retry_with_backoff +from ..db_utils import acquire_with_retry +from ..memory_engine import fq_table from . import bank_utils @@ -65,7 +68,16 @@ from . import ( fact_storage, link_creation, ) -from .types import ChunkMetadata, EntityLink, ExtractedFact, ProcessedFact, RetainContent, RetainContentDict +from .types import ( + ChunkMetadata, + EntityResolutionResult, + ExtractedFact, + Phase1Result, + Phase3Context, + ProcessedFact, + RetainContent, + RetainContentDict, +) logger = logging.getLogger(__name__) @@ -101,6 +113,105 @@ def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None): return retain_params, merged_tags +async def _pre_resolve_phase1( + pool, + entity_resolver, + bank_id: str, + contents: list[RetainContent], + processed_facts: list[ProcessedFact], + config, + log_buffer: list[str], + skip_semantic_ann: bool = False, +) -> Phase1Result: + """ + Phase 1: Run expensive read-heavy operations on a separate connection + OUTSIDE the write transaction. + + - Entity resolution: trigram GIN scan + co-occurrence fetch + scoring + - Semantic ANN: HNSW index probes to find similar existing units + + Running these outside the transaction avoids holding row locks during + slow reads, eliminating TimeoutErrors under concurrent load. + """ + from .link_utils import compute_semantic_links_ann + + user_entities_per_content = {idx: content.entities for idx, content in enumerate(contents) if content.entities} + + # Use placeholder unit_ids for grouping during resolution. The actual + # unit_ids are created later by insert_facts_batch inside the transaction, + # but entity resolution and ANN search only need them as grouping keys. + placeholder_unit_ids = [str(i) for i in range(len(processed_facts))] + embeddings = [fact.embedding for fact in processed_facts] + + async with acquire_with_retry(pool) as resolve_conn: + resolved_entity_ids, entity_to_unit, unit_to_entity_ids = await entity_processing.resolve_entities( + entity_resolver, + resolve_conn, + bank_id, + placeholder_unit_ids, + processed_facts, + log_buffer, + user_entities_per_content=user_entities_per_content, + entity_labels=getattr(config, "entity_labels", None), + ) + + # Semantic ANN search on the same connection (autocommit, no transaction). + # Skipped in streaming mode — deferred to Phase 3 to avoid O(bank_size) + # scaling bottleneck that makes later streaming batches progressively slower. + semantic_ann_links = [] + if not skip_semantic_ann: + fact_types = [fact.fact_type for fact in processed_facts] + semantic_ann_links = await compute_semantic_links_ann( + resolve_conn, bank_id, placeholder_unit_ids, embeddings, fact_types=fact_types, log_buffer=log_buffer + ) + + return Phase1Result( + entities=EntityResolutionResult( + resolved_entity_ids=resolved_entity_ids, + entity_to_unit=entity_to_unit, + unit_to_entity_ids=unit_to_entity_ids, + ), + semantic_ann_links=semantic_ann_links, + ) + + +def _remap_phase1_results( + resolved_entity_ids: list[str], + entity_to_unit: list[tuple], + unit_to_entity_ids: dict[str, list[str]], + semantic_ann_links: list[tuple], + actual_unit_ids: list[str], +) -> tuple[list[tuple], dict[str, list[str]], list[tuple]]: + """ + Remap Phase 1 results from placeholder unit IDs to actual unit IDs. + + During Phase 1 we use str(fact_index) as placeholder unit IDs. + After insert_facts_batch creates real UUIDs, this function replaces the + placeholders so that all rows reference the correct memory_units. + """ + # Build placeholder -> actual mapping + placeholder_to_actual = {str(i): actual_id for i, actual_id in enumerate(actual_unit_ids)} + + # Remap entity_to_unit tuples + remapped_entity_to_unit = [ + (placeholder_to_actual.get(unit_id, unit_id), local_idx, fact_date) + for unit_id, local_idx, fact_date in entity_to_unit + ] + + # Remap unit_to_entity_ids keys + remapped_unit_to_entity_ids: dict[str, list[str]] = {} + for placeholder_id, entity_ids in unit_to_entity_ids.items(): + actual_id = placeholder_to_actual.get(placeholder_id, placeholder_id) + remapped_unit_to_entity_ids[actual_id] = entity_ids + + # Remap semantic ANN links (from_id uses placeholder) + remapped_semantic = [ + (placeholder_to_actual.get(lnk[0], lnk[0]), lnk[1], lnk[2], lnk[3], lnk[4]) for lnk in semantic_ann_links + ] + + return remapped_entity_to_unit, remapped_unit_to_entity_ids, remapped_semantic + + async def _insert_facts_and_links( conn, entity_resolver, @@ -110,56 +221,78 @@ async def _insert_facts_and_links( processed_facts: list[ProcessedFact], config, log_buffer: list[str], + resolved_entity_ids: list[str], + entity_to_unit: list[tuple], + unit_to_entity_ids: dict[str, list[str]], + semantic_ann_links: list[tuple], + skip_semantic_links: bool = False, outbox_callback=None, -) -> list[list[str]]: +) -> tuple[list[list[str]], Phase3Context]: """ - Shared pipeline: insert facts, process entities, create all link types. + Phase 2 of the retain pipeline: insert facts and retrieval-critical links. - Used by both the full retain and delta retain paths. + Runs inside a single database transaction to ensure atomicity of the data + that retrieval depends on (facts, unit_entities, temporal/semantic/causal links). - Returns: - List of unit ID lists mapped back to original content items. + 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). """ unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts) 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: - # Process entities + # Entity resolution was done in Phase 1 (separate connection). + # Remap placeholder IDs to actual unit IDs. step_start = time.time() - user_entities_per_content = {idx: content.entities for idx, content in enumerate(contents) if content.entities} - entity_links = await entity_processing.process_entities_batch( - entity_resolver, - conn, - bank_id, - unit_ids, - processed_facts, - log_buffer, - user_entities_per_content=user_entities_per_content, - entity_labels=getattr(config, "entity_labels", None), + 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 + semantic_ann_links = remapped_semantic + # INSERT unit_entities (FK to memory_units, must be in transaction) + unit_entity_pairs = [ + (unit_id, resolved_entity_ids[idx]) + for idx, (unit_id, _local_idx, _fact_date) in enumerate(remapped_entity_to_unit) + ] + 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, ) - log_buffer.append(f" Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s") # Create temporal links step_start = time.time() temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids) log_buffer.append(f" Temporal links: {temporal_link_count} links in {time.time() - step_start:.3f}s") - # Create semantic links - step_start = time.time() - embeddings_for_links = [fact.embedding for fact in processed_facts] - semantic_link_count = await link_creation.create_semantic_links_batch( - conn, bank_id, unit_ids, embeddings_for_links - ) - log_buffer.append(f" Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s") + # Create semantic links (within-batch + pre-computed ANN from Phase 1) + if skip_semantic_links: + log_buffer.append(" Semantic links: skipped (deferred to final ANN pass)") + semantic_link_count = 0 + else: + step_start = time.time() + embeddings_for_links = [fact.embedding for fact in processed_facts] + semantic_link_count = await link_creation.create_semantic_links_batch( + conn, + bank_id, + unit_ids, + embeddings_for_links, + pre_computed_ann_links=semantic_ann_links, + ) + log_buffer.append(f" Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s") - # Insert entity links - step_start = time.time() - if entity_links: - await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id) - log_buffer.append( - f" Entity links: {len(entity_links) if entity_links else 0} links in {time.time() - step_start:.3f}s" - ) + # NOTE: Entity links are NOT inserted here. They are deferred to + # Phase 3 (post-transaction, best-effort) since retrieval uses the + # unit_entities self-join instead. Entity links only serve UI visualization. # Create causal links step_start = time.time() @@ -172,7 +305,47 @@ async def _insert_facts_and_links( if outbox_callback: await outbox_callback(conn) - return result_unit_ids + return result_unit_ids, phase3_context + + +async def _build_and_insert_entity_links_phase3( + pool, + 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. + """ + 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 + ) + if entity_links: + await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id) + log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s") async def _extract_and_embed( @@ -232,11 +405,11 @@ async def retain_batch( document_id: str | None = None, is_first_batch: bool = True, fact_type_override: str | None = None, - confidence_score: float | None = None, document_tags: list[str] | None = None, operation_id: str | None = None, schema: str | None = None, outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None, + db_semaphore: "asyncio.Semaphore | None" = None, ) -> tuple[list[list[str]], TokenUsage]: """ Process a batch of content through the retain pipeline. @@ -261,6 +434,48 @@ async def retain_batch( # Convert dicts to RetainContent objects contents = _build_contents(contents_dicts, document_tags) + # Resolve effective document_id early so both delta and streaming paths + # can find existing chunks from a prior attempt. On retry, the generated + # document_id is recovered from operation result_metadata. + effective_doc_id = document_id + if not effective_doc_id: + doc_ids = {item.get("document_id") for item in contents_dicts if item.get("document_id")} + if len(doc_ids) == 1: + effective_doc_id = doc_ids.pop() + if not effective_doc_id and operation_id: + try: + async with acquire_with_retry(pool) as conn: + row = await conn.fetchrow( + f"SELECT result_metadata FROM {fq_table('async_operations')} WHERE operation_id = $1", + uuid.UUID(operation_id), + ) + if row and row["result_metadata"]: + meta = ( + row["result_metadata"] + if isinstance(row["result_metadata"], dict) + else json.loads(row["result_metadata"]) + ) + effective_doc_id = meta.get("generated_document_id") + except Exception: + pass + if not effective_doc_id: + effective_doc_id = str(uuid.uuid4()) + # Persist so retries reuse the same document_id + if operation_id: + try: + async with acquire_with_retry(pool) as conn: + await conn.execute( + f""" + UPDATE {fq_table("async_operations")} + SET result_metadata = result_metadata || $1::jsonb, updated_at = now() + WHERE operation_id = $2 + """, + json.dumps({"generated_document_id": effective_doc_id}), + uuid.UUID(operation_id), + ) + except Exception: + logger.warning("Failed to persist generated document_id", exc_info=True) + # --- Delta retain: check if we can skip unchanged chunks --- if is_first_batch: delta_result = await _try_delta_retain( @@ -273,7 +488,7 @@ async def retain_batch( contents_dicts, contents, config, - document_id, + effective_doc_id, fact_type_override, document_tags, agent_name, @@ -282,163 +497,583 @@ async def retain_batch( operation_id, schema, outbox_callback, + db_semaphore, ) if delta_result is not None: return delta_result - # --- Full retain path --- - extracted_facts, processed_facts, chunks, usage = await _extract_and_embed( - contents, - llm_config, - agent_name, - config, - embeddings_model, - format_date_fn, - fact_type_override, - log_buffer, - pool, - operation_id, - schema, + # --- Always use the streaming pipeline (producer-consumer batching) --- + # Even small documents go through the same path — they just end up as a + # single batch. This eliminates the maintenance burden of two separate + # retain code paths. + chunk_batch_size = getattr(config, "retain_chunk_batch_size", 100) + chunk_size = getattr(config, "retain_chunk_size", 3000) + all_pre_chunks = [] + for content in contents: + content_chunks = fact_extraction.chunk_text(content.content, chunk_size) + all_pre_chunks.extend(content_chunks) + + total_pre_chunks = len(all_pre_chunks) + num_batches = (total_pre_chunks + chunk_batch_size - 1) // chunk_batch_size if total_pre_chunks > 0 else 1 + log_buffer.append( + f"[streaming] {total_pre_chunks} chunks, batch_size {chunk_batch_size} — " + f"{num_batches} batch{'es' if num_batches != 1 else ''}" ) - if not extracted_facts: - await _handle_zero_facts_documents( - pool, + return await _streaming_retain_batch( + pool=pool, + embeddings_model=embeddings_model, + llm_config=llm_config, + entity_resolver=entity_resolver, + format_date_fn=format_date_fn, + bank_id=bank_id, + contents_dicts=contents_dicts, + contents=contents, + config=config, + document_id=effective_doc_id, + is_first_batch=is_first_batch, + fact_type_override=fact_type_override, + document_tags=document_tags, + agent_name=agent_name, + log_buffer=log_buffer, + start_time=start_time, + all_pre_chunks=all_pre_chunks, + chunk_batch_size=chunk_batch_size, + operation_id=operation_id, + schema=schema, + outbox_callback=outbox_callback, + db_semaphore=db_semaphore, + ) + + +# --------------------------------------------------------------------------- +# Final semantic ANN pass (post-commit) +# --------------------------------------------------------------------------- + +_ANN_CHUNK_SIZE = 1000 # Max seeds per ANN query — smaller chunks avoid timeouts +_ANN_PARALLELISM = 4 # Max concurrent ANN chunks to avoid pool saturation + + +async def _run_final_semantic_ann( + pool, + bank_id: str, + unit_ids: list[str], + log_buffer: list[str], +) -> None: + """ + Create semantic links for all committed units in a single pass. + + Called after all streaming batches have committed. Loads embeddings and + fact_types from the database, then runs ANN in chunks of _ANN_CHUNK_SIZE + seeds. This replaces per-batch within-batch + fire-and-forget ANN with + one efficient pass that sees the full bank. + """ + from .link_utils import _bulk_insert_links, compute_semantic_links_ann + + if not unit_ids: + return + + # Load embeddings and fact_types for all committed units + load_start = time.time() + async with acquire_with_retry(pool) as conn: + rows = await conn.fetch( + f""" + SELECT id::text, embedding::text, fact_type + FROM {fq_table("memory_units")} + WHERE bank_id = $1 AND id = ANY($2::uuid[]) + ORDER BY id + """, bank_id, - contents_dicts, - contents, - config, - document_id, - is_first_batch, - document_tags, - chunks, - log_buffer, - start_time, + unit_ids, ) - return [[] for _ in contents], usage - # Group contents by document_id - contents_by_doc = defaultdict(list) - for idx, content_dict in enumerate(contents_dicts): - doc_id = content_dict.get("document_id") - contents_by_doc[doc_id].append((idx, content_dict)) + if not rows: + log_buffer.append("[streaming] Final ANN: no units found in DB (unexpected)") + return - # Database transaction (retried on deadlock) - result_unit_ids: list[list[str]] = [] - log_buffer_pre_db = len(log_buffer) + # Build lookup: unit_id -> (embedding_text, fact_type) + unit_map: dict[str, tuple[str, str]] = {} + for row in rows: + unit_map[row["id"]] = (row["embedding"], row["fact_type"]) - async def _run_db_work() -> None: - nonlocal result_unit_ids - del log_buffer[log_buffer_pre_db:] - document_ids_added: list[str] = [] - for pf in processed_facts: - pf.document_id = None - pf.chunk_id = None - entity_resolver.discard_pending_stats() + # Filter to units that have embeddings + ann_unit_ids = [] + ann_embeddings = [] + ann_fact_types = [] + for uid in unit_ids: + if uid in unit_map and unit_map[uid][0] is not None: + ann_unit_ids.append(uid) + ann_embeddings.append(unit_map[uid][0]) # embedding as text (for temp table) + ann_fact_types.append(unit_map[uid][1]) - async with acquire_with_retry(pool) as conn: - async with conn.transaction(): - # Handle document tracking - step_start = time.time() - doc_id_mapping = {} + log_buffer.append( + f"[streaming] Final ANN: loaded {len(ann_unit_ids)} units with embeddings in {time.time() - load_start:.3f}s" + ) - if document_id: - combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) - retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) - await fact_storage.handle_document_tracking( - conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags - ) - document_ids_added.append(document_id) - doc_id_mapping[None] = document_id - else: - has_any_doc_ids = any(item.get("document_id") for item in contents_dicts) - if has_any_doc_ids or chunks: - for original_doc_id, doc_contents in contents_by_doc.items(): - actual_doc_id = original_doc_id - should_create_doc = (original_doc_id is not None) or chunks - if should_create_doc: - if actual_doc_id is None: - actual_doc_id = str(uuid.uuid4()) - doc_id_mapping[original_doc_id] = actual_doc_id - combined_content = "\n".join([c.get("content", "") for _, c in doc_contents]) - retain_params, merged_tags = _build_retain_params( - contents_dicts, document_tags, doc_contents=doc_contents - ) - await fact_storage.handle_document_tracking( - conn, - bank_id, - actual_doc_id, - combined_content, - is_first_batch, - retain_params, - merged_tags, - ) - document_ids_added.append(actual_doc_id) + if not ann_unit_ids: + return - if document_ids_added: - log_buffer.append( - f" Document tracking: {len(document_ids_added)} documents in {time.time() - step_start:.3f}s" - ) + # Process in parallel chunks — each chunk runs ANN query + INSERT on its own connection. + # Parallelism bounded by _ANN_PARALLELISM to avoid saturating the connection pool. + num_chunks = (len(ann_unit_ids) + _ANN_CHUNK_SIZE - 1) // _ANN_CHUNK_SIZE + ann_semaphore = asyncio.Semaphore(_ANN_PARALLELISM) + chunk_link_counts: list[int] = [0] * num_chunks - # Store chunks and map to facts - step_start = time.time() - chunk_id_map_by_doc = {} - if chunks: - chunks_by_doc = defaultdict(list) - for chunk in chunks: - original_doc_id = contents_dicts[chunk.content_index].get("document_id") - actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) - if actual_doc_id is None and document_id: - actual_doc_id = document_id - chunks_by_doc[actual_doc_id].append(chunk) + async def _process_ann_chunk(chunk_idx: int) -> None: + chunk_start = chunk_idx * _ANN_CHUNK_SIZE + chunk_end = min(chunk_start + _ANN_CHUNK_SIZE, len(ann_unit_ids)) + chunk_ids = ann_unit_ids[chunk_start:chunk_end] + chunk_embs = ann_embeddings[chunk_start:chunk_end] + chunk_ftypes = ann_fact_types[chunk_start:chunk_end] - for doc_id, doc_chunks in chunks_by_doc.items(): - chunk_id_map = await chunk_storage.store_chunks_batch(conn, bank_id, doc_id, doc_chunks) - for chunk_idx, chunk_id in chunk_id_map.items(): - chunk_id_map_by_doc[(doc_id, chunk_idx)] = chunk_id - - log_buffer.append( - f" Store chunks: {len(chunks)} chunks for {len(chunks_by_doc)} documents " - f"in {time.time() - step_start:.3f}s" - ) - - # Map chunk_ids and document_ids to facts - for fact, processed_fact in zip(extracted_facts, processed_facts): - original_doc_id = contents_dicts[fact.content_index].get("document_id") - actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id) - if actual_doc_id is None and document_id: - actual_doc_id = document_id - processed_fact.document_id = actual_doc_id - if chunks and fact.chunk_index is not None: - chunk_id = chunk_id_map_by_doc.get((actual_doc_id, fact.chunk_index)) - if chunk_id: - processed_fact.chunk_id = chunk_id - - # Insert facts and create all links (shared pipeline) - result_unit_ids = await _insert_facts_and_links( + async with ann_semaphore: + t0 = time.time() + async with acquire_with_retry(pool) as conn: + await conn.execute("SET statement_timeout = '300s'") + ann_links = await compute_semantic_links_ann( conn, - entity_resolver, bank_id, - contents, - extracted_facts, - processed_facts, - config, - log_buffer, - outbox_callback, + chunk_ids, + chunk_embs, + fact_types=chunk_ftypes, + top_k=20, # Recall uses at most 20 neighbors + log_buffer=log_buffer, ) + if ann_links: + await _bulk_insert_links(conn, ann_links, bank_id=bank_id) + chunk_link_counts[chunk_idx] = len(ann_links) + await conn.execute("RESET statement_timeout") + logger.info( + f"[streaming] Final ANN chunk {chunk_idx + 1}/{num_chunks}: " + f"{len(ann_links)} links in {time.time() - t0:.3f}s" + ) - await entity_resolver.flush_pending_stats() + await asyncio.gather(*[_process_ann_chunk(i) for i in range(num_chunks)]) + total_links = sum(chunk_link_counts) + log_buffer.append(f"[streaming] Final ANN: {total_links} total semantic links") - total_time = time.time() - start_time - log_buffer.append(f"{'=' * 60}") - log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(processed_facts)} units in {total_time:.3f}s") - if document_ids_added: - log_buffer.append(f"Documents: {', '.join(document_ids_added)}") - log_buffer.append(f"{'=' * 60}") - logger.info("\n" + "\n".join(log_buffer) + "\n") - await retry_with_backoff(_run_db_work) - return result_unit_ids, usage +# --------------------------------------------------------------------------- +# Streaming chunk batching +# --------------------------------------------------------------------------- + + +async def _streaming_retain_batch( + pool, + embeddings_model, + llm_config, + entity_resolver, + format_date_fn, + bank_id: str, + contents_dicts: list[RetainContentDict], + contents: list[RetainContent], + config, + document_id: str | None, + is_first_batch: bool, + fact_type_override: str | None, + document_tags: list[str] | None, + agent_name: str, + log_buffer: list[str], + start_time: float, + all_pre_chunks: list[str], + chunk_batch_size: int, + operation_id: str | None = None, + schema: str | None = None, + outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None, + db_semaphore: "asyncio.Semaphore | None" = None, +) -> tuple[list[list[str]], TokenUsage]: + """ + Process a large document in streaming mini-batches to bound memory usage. + + Instead of extracting facts from ALL chunks at once (which can OOM for 17k+ + chunk documents), this splits the pre-chunked content into batches of + ``chunk_batch_size`` chunks. Each mini-batch goes through the full + extract -> embed -> Phase 1/2/3 pipeline and commits to the DB before the + next batch starts, so memory is released between batches. + + All mini-batches share the same ``document_id`` so that: + - Delta retain can detect already-committed chunks on retry + - The document row tracks the full content + - Chunks are associated with the correct document + """ + total_chunks = len(all_pre_chunks) + total_usage = TokenUsage() + all_unit_ids: list[str] = [] + + # document_id is already resolved by retain_batch (includes recovery from + # operation result_metadata on retry). + effective_doc_id = document_id + + # Use the first content item as the template for metadata (context, event_date, etc.) + template_content = contents[0] if contents else RetainContent(content="") + + # Load existing chunk hashes BEFORE document tracking to detect recovery. + # If chunks exist AND the document content hash matches, this is a retry of + # the same content — preserve existing data. If content differs, this is an + # update — cascade-delete old data and start fresh. + existing_chunk_hashes: set[str] = set() + combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) + new_content_hash = hashlib.sha256(combined_content.encode()).hexdigest() + is_recovery = False + + try: + async with acquire_with_retry(pool) as conn: + # Check if document exists with matching content hash + doc_row = await conn.fetchrow( + f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2", + effective_doc_id, + bank_id, + ) + if doc_row and doc_row["content_hash"] == new_content_hash: + # Same content — load chunk hashes for recovery skip + existing_rows = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id) + existing_chunk_hashes = {c.content_hash for c in existing_rows if c.content_hash} + if existing_chunk_hashes: + is_recovery = True + log_buffer.append( + f"[streaming] RECOVERY: found {len(existing_chunk_hashes)} already-committed chunks — " + f"will skip matching and preserve existing data" + ) + except Exception: + pass # If we can't load, just process all chunks + + # Create/update the document row. + retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) + async with acquire_with_retry(pool) as conn: + async with conn.transaction(): + if is_recovery: + # Recovery: same content, partially committed — preserve existing data + await fact_storage.upsert_document_metadata( + conn, bank_id, effective_doc_id, combined_content, retain_params, merged_tags + ) + log_buffer.append( + f"[streaming] Document {effective_doc_id} updated (recovery, preserving existing chunks)" + ) + else: + # Fresh or update: cascade-delete old data if document exists + await fact_storage.handle_document_tracking( + conn, bank_id, effective_doc_id, combined_content, is_first_batch, retain_params, merged_tags + ) + log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)") + + # --------------------------------------------------------------------------- + # Producer-consumer pipeline: LLM extraction runs concurrently with DB writes + # --------------------------------------------------------------------------- + num_batches = (total_chunks + chunk_batch_size - 1) // chunk_batch_size + + # Queue for enriched chunks (extracted facts + embeddings). + # Buffer up to 2x batch_size items so the producer can stay ahead of the consumer. + chunk_queue: asyncio.Queue = asyncio.Queue(maxsize=chunk_batch_size * 2) + + # Shared mutable state for the producer to report skipped chunks and usage + producer_error: list[BaseException] = [] + + # ---- LLM Producer ---- + # Fires all chunk extractions as concurrent tasks (bounded by the LLM + # semaphore inside fact_extraction to 32 concurrent). As each completes + # it pushes the enriched result into the queue for the DB consumer. + async def _llm_producer() -> None: + async def _extract_one(global_idx: int, chunk_text: str) -> None: + content = RetainContent( + content=chunk_text, + context=template_content.context, + event_date=template_content.event_date, + metadata=template_content.metadata, + entities=template_content.entities, + tags=template_content.tags, + observation_scopes=template_content.observation_scopes, + ) + extracted, processed, chunk_meta, usage = await _extract_and_embed( + [content], + llm_config, + agent_name, + config, + embeddings_model, + format_date_fn, + fact_type_override, + log_buffer, + pool, + operation_id, + schema, + ) + await chunk_queue.put((global_idx, content, extracted, processed, chunk_meta, usage)) + + tasks: list[asyncio.Task] = [] + skipped_total = 0 + for i, chunk_text in enumerate(all_pre_chunks): + chunk_hash = chunk_storage.compute_chunk_hash(chunk_text) + if chunk_hash in existing_chunk_hashes: + skipped_total += 1 + continue + tasks.append(asyncio.create_task(_extract_one(i, chunk_text))) + + if skipped_total > 0: + log_buffer.append(f"[streaming] Producer: skipped {skipped_total}/{total_chunks} already-committed chunks") + + # Wait for all extractions; collect exceptions + results = await asyncio.gather(*tasks, return_exceptions=True) + for r in results: + if isinstance(r, BaseException): + producer_error.append(r) + + # Signal the consumer that production is done + await chunk_queue.put(None) + + # ---- DB Consumer ---- + # Drains enriched chunks from the queue in batches and runs + # Phase 1 (entity resolution) -> Phase 2 (write txn) -> Phase 3 (ANN fire-and-forget). + async def _db_consumer() -> None: + batch: list[tuple] = [] + global_chunk_offset = 0 + consumer_batch_idx = 0 + + while True: + item = await chunk_queue.get() + if item is None: + # Process any remaining items + if batch: + await _process_db_batch( + batch, + global_chunk_offset, + consumer_batch_idx, + is_last=True, + ) + break + + batch.append(item) + + if len(batch) >= chunk_batch_size: + await _process_db_batch( + batch, + global_chunk_offset, + consumer_batch_idx, + is_last=False, + ) + global_chunk_offset += len(batch) + consumer_batch_idx += 1 + batch = [] + + async def _process_db_batch( + batch: list[tuple], + global_chunk_offset: int, + consumer_batch_idx: int, + is_last: bool, + ) -> None: + """Run Phase 1 + Phase 2 + Phase 3 for a batch of pre-extracted chunks.""" + # Combine results from individual chunk extractions + batch_contents: list[RetainContent] = [] + batch_extracted: list = [] + batch_processed: list[ProcessedFact] = [] + batch_chunk_meta: list[ChunkMetadata] = [] + batch_usage = TokenUsage() + + for global_idx, content, extracted, processed, chunk_meta, usage in batch: + content_idx_in_batch = len(batch_contents) + # Adjust chunk indices to global offsets and remap content_index + for fact in extracted: + fact.content_index = content_idx_in_batch + if fact.chunk_index is not None: + fact.chunk_index = global_chunk_offset + content_idx_in_batch + for pf in processed: + pf.content_index = content_idx_in_batch + for cm in chunk_meta: + cm.chunk_index = global_chunk_offset + content_idx_in_batch + + batch_contents.append(content) + batch_extracted.extend(extracted) + batch_processed.extend(processed) + batch_chunk_meta.extend(chunk_meta) + batch_usage = batch_usage + usage + + nonlocal total_usage + total_usage = total_usage + batch_usage + + if not batch_extracted: + log_buffer.append( + f"[streaming] Consumer batch {consumer_batch_idx + 1}: " + f"0 facts extracted from {len(batch)} chunks, skipping" + ) + return + + log_buffer.append( + f"[streaming] Consumer batch {consumer_batch_idx + 1}: " + f"processing {len(batch_extracted)} facts from {len(batch)} chunks" + ) + + async def _run_mini_batch_db_work() -> None: + entity_resolver.discard_pending_stats() + mb_start = time.time() + + # Phase 1 — Entity Resolution only (no ANN — deferred to Phase 3) + p1_start = time.time() + phase1 = await _pre_resolve_phase1( + pool, + entity_resolver, + bank_id, + batch_contents, + batch_processed, + config, + log_buffer, + skip_semantic_ann=True, + ) + + logger.info(f"[streaming] Phase 1 (entity resolution): {time.time() - p1_start:.3f}s") + + # Phase 2 — Write transaction (within-batch semantic links only) + p2_start = time.time() + async with acquire_with_retry(pool) as conn: + async with conn.transaction(): + # Store chunks with correct global indices + step_start = time.time() + chunk_id_map = {} + if batch_chunk_meta: + chunk_id_map = await chunk_storage.store_chunks_batch( + conn, bank_id, effective_doc_id, batch_chunk_meta + ) + log_buffer.append( + f" Store chunks: {len(batch_chunk_meta)} chunks in {time.time() - step_start:.3f}s" + ) + + # Map document_id and chunk_id to processed facts + for fact, processed_fact in zip(batch_extracted, batch_processed): + processed_fact.document_id = effective_doc_id + if batch_chunk_meta and fact.chunk_index is not None: + chunk_id = chunk_id_map.get(fact.chunk_index) + if chunk_id: + processed_fact.chunk_id = chunk_id + + # 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( + conn, + entity_resolver, + bank_id, + batch_contents, + batch_extracted, + batch_processed, + config, + log_buffer, + resolved_entity_ids=phase1.entities.resolved_entity_ids, + entity_to_unit=phase1.entities.entity_to_unit, + unit_to_entity_ids=phase1.entities.unit_to_entity_ids, + semantic_ann_links=[], + skip_semantic_links=True, + outbox_callback=outbox_callback if is_last else None, + ) + + logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s") + + # Best-effort: entity viz + stats (fast, not semantic ANN) + 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) + + logger.info( + f"[streaming] Consumer batch {consumer_batch_idx + 1} total " + f"(excluding fire-and-forget): {time.time() - mb_start:.3f}s" + ) + + # Collect unit_ids from this batch + for content_ids in batch_result_ids: + all_unit_ids.extend(content_ids) + + if db_semaphore is not None: + async with db_semaphore: + await _run_mini_batch_db_work() + else: + await _run_mini_batch_db_work() + + # --------------------------------------------------------------------------- + # Check if facts are already committed (recovery from previous crash). + # If so, skip extraction+writes and jump straight to final ANN pass. + # --------------------------------------------------------------------------- + facts_already_committed = False + if operation_id: + try: + async with acquire_with_retry(pool) as conn: + row = await conn.fetchrow( + f"SELECT result_metadata FROM {fq_table('async_operations')} WHERE operation_id = $1", + uuid.UUID(operation_id), + ) + if row and row["result_metadata"]: + meta = ( + row["result_metadata"] + if isinstance(row["result_metadata"], dict) + else json.loads(row["result_metadata"]) + ) + if meta.get("facts_committed"): + facts_already_committed = True + log_buffer.append( + f"[streaming] Recovery: facts already committed ({meta.get('unit_ids_count', '?')} units), " + f"skipping to final ANN pass" + ) + except Exception: + logger.warning("Failed to check operation recovery state", exc_info=True) + + if not facts_already_committed: + # Run producer and consumer concurrently + await asyncio.gather(_llm_producer(), _db_consumer()) + + # Propagate producer errors (e.g. LLM failures) + if producer_error: + raise producer_error[0] + + # Mark facts as committed in operation metadata (crash recovery checkpoint) + if operation_id and all_unit_ids: + try: + async with acquire_with_retry(pool) as conn: + await conn.execute( + f""" + UPDATE {fq_table("async_operations")} + SET result_metadata = result_metadata || $1::jsonb, updated_at = now() + WHERE operation_id = $2 + """, + json.dumps({"facts_committed": True, "unit_ids_count": len(all_unit_ids)}), + uuid.UUID(operation_id), + ) + log_buffer.append(f"[streaming] Checkpoint: {len(all_unit_ids)} facts committed, ANN pass next") + except Exception: + logger.warning("Failed to save facts_committed checkpoint", exc_info=True) + else: + # Recovery path: load committed unit IDs from DB + async with acquire_with_retry(pool) as conn: + rows = await conn.fetch( + f""" + SELECT id::text FROM {fq_table("memory_units")} + WHERE bank_id = $1 AND document_id = $2 + ORDER BY created_at + """, + bank_id, + effective_doc_id, + ) + all_unit_ids = [row["id"] for row in rows] + log_buffer.append(f"[streaming] Recovery: loaded {len(all_unit_ids)} unit IDs from DB") + + # --------------------------------------------------------------------------- + # Final ANN pass: create semantic links for ALL committed units at once. + # This replaces per-batch within-batch + fire-and-forget ANN with a single + # efficient pass after all facts are in the database. + # --------------------------------------------------------------------------- + if all_unit_ids: + ann_start = time.time() + await _run_final_semantic_ann(pool, bank_id, all_unit_ids, log_buffer) + log_buffer.append(f"[streaming] Final ANN pass: {time.time() - ann_start:.3f}s for {len(all_unit_ids)} units") + + total_time = time.time() - start_time + log_buffer.append(f"{'=' * 60}") + log_buffer.append( + f"STREAMING RETAIN COMPLETE: {len(all_unit_ids)} units across {num_batches} batches in {total_time:.3f}s" + ) + log_buffer.append(f"Document: {effective_doc_id}") + log_buffer.append(f"{'=' * 60}") + logger.info("\n" + "\n".join(log_buffer) + "\n") + + # Map all unit_ids back to the original content items. + # For streaming mode with a single document, all units belong to content 0. + result_unit_ids = [all_unit_ids] + [[] for _ in contents[1:]] + return result_unit_ids, total_usage # --------------------------------------------------------------------------- @@ -465,6 +1100,7 @@ async def _try_delta_retain( operation_id, schema, outbox_callback, + db_semaphore: "asyncio.Semaphore | None" = None, ): """ Attempt delta retain for a document upsert. Returns result tuple if delta @@ -582,6 +1218,12 @@ async def _try_delta_retain( pf.chunk_id = None entity_resolver.discard_pending_stats() + # PHASE 1 — Entity Resolution + Semantic ANN (separate connection, read-heavy) + phase1 = await _pre_resolve_phase1( + pool, entity_resolver, bank_id, delta_contents, processed_facts, config, log_buffer + ) + + # PHASE 2 — Core Write Transaction (atomic) async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Update document metadata (no delete) @@ -652,20 +1294,31 @@ async def _try_delta_retain( if chunk_id: pf.chunk_id = chunk_id - # Insert facts and create all links (shared pipeline) - result_unit_ids = await _insert_facts_and_links( + # 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( conn, entity_resolver, bank_id, - contents, + delta_contents, extracted_facts, processed_facts, config, log_buffer, - outbox_callback, + resolved_entity_ids=phase1.entities.resolved_entity_ids, + entity_to_unit=phase1.entities.entity_to_unit, + unit_to_entity_ids=phase1.entities.unit_to_entity_ids, + semantic_ann_links=phase1.semantic_ann_links, + outbox_callback=outbox_callback, ) - await entity_resolver.flush_pending_stats() + # PHASE 3 — Best-Effort Display Data (post-transaction) + 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) total_time = time.time() - start_time log_buffer.append(f"{'=' * 60}") @@ -677,7 +1330,11 @@ async def _try_delta_retain( log_buffer.append(f"{'=' * 60}") logger.info("\n" + "\n".join(log_buffer) + "\n") - await retry_with_backoff(_run_delta_db_work) + if db_semaphore is not None: + async with db_semaphore: + await _run_delta_db_work() + else: + await _run_delta_db_work() return result_unit_ids, usage @@ -747,75 +1404,19 @@ def _build_contents(contents_dicts: list[RetainContentDict], document_tags: list return contents -async def _handle_zero_facts_documents( - pool, - bank_id, - contents_dicts, - contents, - config, - document_id, - is_first_batch, - document_tags, - chunks, - log_buffer, - start_time, -): - """Handle document tracking when zero facts were extracted.""" - docs_tracked = 0 - async with acquire_with_retry(pool) as conn: - async with conn.transaction(): - contents_by_doc = defaultdict(list) - for idx, content_dict in enumerate(contents_dicts): - doc_id = content_dict.get("document_id") - contents_by_doc[doc_id].append((idx, content_dict)) - - if document_id: - combined_content = "\n".join([c.get("content", "") for c in contents_dicts]) - retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags) - await fact_storage.handle_document_tracking( - conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags - ) - docs_tracked += 1 - else: - has_any_doc_ids = any(item.get("document_id") for item in contents_dicts) - if has_any_doc_ids or chunks: - for original_doc_id, doc_contents in contents_by_doc.items(): - should_create_doc = (original_doc_id is not None) or chunks - if not should_create_doc: - continue - actual_doc_id = original_doc_id or str(uuid.uuid4()) - combined_content = "\n".join([c.get("content", "") for _, c in doc_contents]) - retain_params, merged_tags = _build_retain_params( - contents_dicts, document_tags, doc_contents=doc_contents - ) - await fact_storage.handle_document_tracking( - conn, - bank_id, - actual_doc_id, - combined_content, - is_first_batch, - retain_params, - merged_tags, - ) - docs_tracked += 1 - - total_time = time.time() - start_time - doc_status = f"{docs_tracked} document(s) tracked" if docs_tracked > 0 else "no document tracked" - logger.info( - f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents " - f"in {total_time:.3f}s ({doc_status}, no facts)" - ) - - def _chunk_contents_for_delta(contents: list[RetainContent], config) -> dict[int, str]: """ - Chunk contents the same way fact_extraction does, returning a map of + Chunk contents the same way the streaming path does, returning a map of global_chunk_index -> chunk_text. + + Must use the same chunk_size as the streaming path (default 3000) so that + chunk boundaries match and delta can detect unchanged chunks. + Previously defaulted to 120000, causing all chunks to appear changed on retry. """ result = {} global_chunk_idx = 0 for content in contents: - chunk_size = getattr(config, "retain_chunk_size", 120000) + chunk_size = getattr(config, "retain_chunk_size", 3000) chunks = fact_extraction.chunk_text(content.content, chunk_size) for chunk_text in chunks: result[global_chunk_idx] = chunk_text diff --git a/hindsight-api-slim/hindsight_api/engine/retain/types.py b/hindsight-api-slim/hindsight_api/engine/retain/types.py index e328cece..6cf70608 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/types.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/types.py @@ -221,6 +221,45 @@ class ProcessedFact: ) +@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: + """ + Result of Phase 1 entity resolution. + + Contains resolved entity IDs and the mapping data needed to remap + placeholder unit IDs to real IDs after fact insertion in Phase 2. + """ + + resolved_entity_ids: list[str] + entity_to_unit: list[tuple] + unit_to_entity_ids: dict[str, list[str]] + + +@dataclass +class Phase1Result: + """ + Full result of Phase 1 (entity resolution + optional semantic ANN). + """ + + entities: EntityResolutionResult + semantic_ann_links: list[tuple] + + @dataclass class EntityLink: """ diff --git a/hindsight-api-slim/hindsight_api/engine/search/link_expansion_retrieval.py b/hindsight-api-slim/hindsight_api/engine/search/link_expansion_retrieval.py index caa0aa0a..3b3f434a 100644 --- a/hindsight-api-slim/hindsight_api/engine/search/link_expansion_retrieval.py +++ b/hindsight-api-slim/hindsight_api/engine/search/link_expansion_retrieval.py @@ -4,9 +4,9 @@ Link Expansion graph retrieval. Expands from semantic/temporal seeds through three parallel, first-class signals stored in memory_links: -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. +1. Entity links — query-time self-join through unit_entities. Score = number of distinct + shared entities between the seed set and each candidate, computed via + COUNT(DISTINCT entity_id). More accurate than precomputed entity links. 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. @@ -264,29 +264,33 @@ class LinkExpansionRetriever(GraphRetriever): """ 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' + ue = fq_table("unit_entities") + + entity_cte = f""" + entity_expanded AS ( + -- Entity co-occurrence via unit_entities self-join. + -- Finds units sharing entities with seeds at query time — more accurate + -- than precomputed entity links (no stale 50-neighbor cap). + -- Score = COUNT(DISTINCT shared entities), mapped to [0,1] via tanh. + 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 ue_seed.entity_id)::float AS score, + 'entity'::text AS source + FROM {ue} ue_seed + JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id + JOIN {mu} mu ON mu.id = ue_target.unit_id + WHERE ue_seed.unit_id = ANY($1::uuid[]) + AND ue_target.unit_id != ALL($1::uuid[]) AND mu.fact_type = $2 - AND mu.id != ALL($1::uuid[]) GROUP BY mu.id ORDER BY score DESC LIMIT $3 - ), + )""" + + all_rows = await conn.fetch( + f""" + WITH {entity_cte}, semantic_expanded AS ( -- Semantic kNN: both outgoing (seeds → their kNN at insert time) and -- incoming (facts inserted after seeds that found seeds as kNN). @@ -397,6 +401,19 @@ class LinkExpansionRetriever(GraphRetriever): f"{len(source_ids_found)} source_memory_ids found" ) + ue = fq_table("unit_entities") + + connected_sources_cte = f""" + connected_sources AS ( + -- Find sources sharing entities with seed observation sources + -- via unit_entities self-join (query-time, no precomputed links needed). + SELECT DISTINCT ue_target.unit_id AS source_id + FROM seed_sources ss + JOIN {ue} ue_seed ON ue_seed.unit_id = ss.source_id + JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id + WHERE ue_target.unit_id != ss.source_id + )""" + entity_rows = await conn.fetch( f""" WITH seed_sources AS ( @@ -405,15 +422,7 @@ class LinkExpansionRetriever(GraphRetriever): WHERE id = ANY($1::uuid[]) AND source_memory_ids IS NOT NULL ), - connected_sources AS ( - -- Mirror the non-observation entity expansion: follow pre-bounded entity - -- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time). - -- Score = number of distinct shared entities, same as the non-obs path. - SELECT DISTINCT ml.to_unit_id AS source_id - FROM seed_sources ss - JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id - WHERE ml.link_type = 'entity' - ), + {connected_sources_cte}, connected_array AS ( SELECT array_agg(source_id) AS source_ids FROM connected_sources ) diff --git a/hindsight-api-slim/tests/conftest.py b/hindsight-api-slim/tests/conftest.py index 083de070..07278037 100644 --- a/hindsight-api-slim/tests/conftest.py +++ b/hindsight-api-slim/tests/conftest.py @@ -17,7 +17,7 @@ from hindsight_api.pg0 import EmbeddedPostgres # Default pg0 instance configuration for tests DEFAULT_PG0_INSTANCE_NAME = "hindsight-test" -DEFAULT_PG0_PORT = None +DEFAULT_PG0_PORT = int(os.environ.get("HINDSIGHT_TEST_PG_PORT", "5556")) # Load environment variables from .env at the start of test session diff --git a/hindsight-api-slim/tests/test_link_utils.py b/hindsight-api-slim/tests/test_link_utils.py index 969ca70f..c5d6262d 100644 --- a/hindsight-api-slim/tests/test_link_utils.py +++ b/hindsight-api-slim/tests/test_link_utils.py @@ -1,11 +1,15 @@ -"""Tests for link_utils datetime handling and temporal link computation.""" +"""Tests for link_utils datetime handling, temporal link computation, and semantic link splitting.""" +import numpy as np import pytest from datetime import datetime, timezone, timedelta from hindsight_api.engine.retain.link_utils import ( _normalize_datetime, + _cap_links_per_unit, compute_temporal_links, compute_temporal_query_bounds, + compute_semantic_links_within_batch, + MAX_TEMPORAL_LINKS_PER_UNIT, ) @@ -254,3 +258,133 @@ class TestComputeTemporalLinks: assert len(links) == 1 assert links[0][3] >= 0.3 + + +class TestCapLinksPerUnit: + """Tests for the _cap_links_per_unit helper function.""" + + def test_empty_links(self): + assert _cap_links_per_unit([]) == [] + + def test_under_cap_unchanged(self): + links = [ + ("unit_a", "unit_x", "temporal", 0.9, None), + ("unit_a", "unit_y", "temporal", 0.8, None), + ] + result = _cap_links_per_unit(links, max_per_unit=5) + assert len(result) == 2 + + def test_caps_to_max_per_unit(self): + # Create 30 links from the same unit with descending weights + links = [("unit_a", f"unit_{i}", "temporal", 1.0 - i * 0.01, None) for i in range(30)] + result = _cap_links_per_unit(links, max_per_unit=10) + assert len(result) == 10 + # Should keep the highest-weight links + weights = [lnk[3] for lnk in result] + assert weights == sorted(weights, reverse=True) + assert weights[0] == 1.0 # Highest weight kept + + def test_caps_independently_per_unit(self): + links_a = [("unit_a", f"target_{i}", "temporal", 0.9 - i * 0.01, None) for i in range(10)] + links_b = [("unit_b", f"target_{i}", "temporal", 0.8 - i * 0.01, None) for i in range(10)] + result = _cap_links_per_unit(links_a + links_b, max_per_unit=5) + # 5 from unit_a + 5 from unit_b + assert len(result) == 10 + from_a = [lnk for lnk in result if lnk[0] == "unit_a"] + from_b = [lnk for lnk in result if lnk[0] == "unit_b"] + assert len(from_a) == 5 + assert len(from_b) == 5 + + def test_default_max_is_temporal_constant(self): + links = [("unit_a", f"target_{i}", "temporal", 1.0 - i * 0.01, None) for i in range(50)] + result = _cap_links_per_unit(links) + assert len(result) == MAX_TEMPORAL_LINKS_PER_UNIT + + def test_preserves_tuple_structure(self): + links = [("from_id", "to_id", "temporal", 0.95, "entity_id")] + result = _cap_links_per_unit(links, max_per_unit=5) + assert result[0] == ("from_id", "to_id", "temporal", 0.95, "entity_id") + + +class TestComputeSemanticLinksWithinBatch: + """Tests for compute_semantic_links_within_batch. + + This function computes semantic links between units in the same batch + using numpy dot product (no DB access). It runs in Phase 2 (write + transaction) while the expensive ANN search against existing units runs + in Phase 1 on a separate connection to avoid TimeoutErrors from HNSW + index contention under concurrent load. + """ + + def test_empty_returns_empty(self): + assert compute_semantic_links_within_batch([], []) == [] + + def test_single_unit_returns_empty(self): + emb = [np.random.randn(384).tolist()] + assert compute_semantic_links_within_batch(["u1"], emb) == [] + + def test_identical_embeddings_produce_links(self): + """Two identical embeddings should have similarity=1.0 (above 0.7 threshold).""" + emb = [0.1] * 384 + links = compute_semantic_links_within_batch(["u1", "u2"], [emb, emb]) + assert len(links) == 2 # bidirectional: u1→u2, u2→u1 + from_ids = {lnk[0] for lnk in links} + to_ids = {lnk[1] for lnk in links} + assert from_ids == {"u1", "u2"} + assert to_ids == {"u1", "u2"} + for lnk in links: + assert lnk[2] == "semantic" + assert lnk[3] >= 0.99 # near-1.0 similarity + assert lnk[4] is None # no entity_id + + def test_orthogonal_embeddings_no_links(self): + """Orthogonal embeddings should have similarity=0 (below 0.7 threshold).""" + emb1 = [1.0] + [0.0] * 383 + emb2 = [0.0] + [1.0] + [0.0] * 382 + links = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2]) + assert len(links) == 0 + + def test_respects_threshold(self): + """Links below threshold should be excluded.""" + emb1 = np.random.randn(384).tolist() + # Create a slightly similar embedding (add noise) + emb2 = [x + np.random.randn() * 0.5 for x in emb1] + # Normalize both + norm1 = np.linalg.norm(emb1) + norm2 = np.linalg.norm(emb2) + emb1 = [x / norm1 for x in emb1] + emb2 = [x / norm2 for x in emb2] + + links_low = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2], threshold=0.0) + links_high = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2], threshold=0.99) + # Low threshold should have more links than high threshold + assert len(links_low) >= len(links_high) + + def test_top_k_limits_per_unit(self): + """Each unit should link to at most top_k other units.""" + n = 10 + # Create similar embeddings (all close to the same vector) + base = np.random.randn(384) + base = base / np.linalg.norm(base) + embs = [(base + np.random.randn(384) * 0.01).tolist() for _ in range(n)] + unit_ids = [f"u{i}" for i in range(n)] + + links = compute_semantic_links_within_batch(unit_ids, embs, top_k=3, threshold=0.5) + # Each unit should have at most 3 outgoing links + from collections import Counter + from_counts = Counter(lnk[0] for lnk in links) + for count in from_counts.values(): + assert count <= 3 + + def test_link_tuple_structure(self): + """Verify the tuple format matches what _bulk_insert_links expects.""" + emb = [0.1] * 384 + links = compute_semantic_links_within_batch(["u1", "u2"], [emb, emb]) + for lnk in links: + assert len(lnk) == 5 + from_id, to_id, link_type, weight, entity_id = lnk + assert isinstance(from_id, str) + assert isinstance(to_id, str) + assert link_type == "semantic" + assert 0.0 <= weight <= 1.0 + assert entity_id is None diff --git a/hindsight-api-slim/tests/test_retain.py b/hindsight-api-slim/tests/test_retain.py index 35b04dec..67c41055 100644 --- a/hindsight-api-slim/tests/test_retain.py +++ b/hindsight-api-slim/tests/test_retain.py @@ -2044,6 +2044,72 @@ async def test_semantic_links_within_same_batch(memory, request_context): await memory.delete_bank(bank_id, request_context=request_context) +@pytest.mark.asyncio +async def test_semantic_links_phase1_ann_cross_batch(memory, request_context): + """ + Test that Phase 1 ANN search creates semantic links between facts from + DIFFERENT retain batches. + + The semantic ANN search runs in Phase 1 on a separate connection (outside + the write transaction) using placeholder unit IDs to avoid TimeoutErrors + from HNSW index contention under concurrent load. This test verifies that: + 1. Phase 1 ANN with placeholder IDs works correctly + 2. Placeholder IDs are remapped to real unit IDs before insertion + 3. Cross-batch semantic links are created between similar facts + """ + bank_id = f"test_semantic_phase1_{datetime.now(timezone.utc).timestamp()}" + + try: + # First batch: store some facts about Python + await memory.retain_async( + bank_id=bank_id, + content="Alice is an expert Python developer who builds web applications using FastAPI.", + context="team skills", + request_context=request_context, + ) + + # Second batch: store similar facts — Phase 1 ANN should find the first batch's + # facts via HNSW index and create cross-batch semantic links + unit_ids_2 = await memory.retain_async( + bank_id=bank_id, + content="Bob specializes in Python programming and creates REST APIs with FastAPI.", + context="team skills", + request_context=request_context, + ) + + assert len(unit_ids_2) > 0 + + # Verify cross-batch semantic links exist + async with memory._pool.acquire() as conn: + cross_batch_links = await conn.fetch( + """ + SELECT from_unit_id, to_unit_id, weight + FROM memory_links + WHERE from_unit_id::text = ANY($1) + AND link_type = 'semantic' + AND to_unit_id::text != ALL($1) + """, + unit_ids_2, + ) + + logger.info(f"Cross-batch semantic links from batch 2: {len(cross_batch_links)}") + for link in cross_batch_links: + logger.info( + f" {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... " + f"(weight: {link['weight']:.3f})" + ) + + # Phase 1 ANN should have found similar facts from batch 1 + assert len(cross_batch_links) > 0, ( + "Phase 1 ANN search should create semantic links between similar facts " + "from different retain batches. This tests that placeholder unit IDs are " + "correctly remapped to real IDs after insert_facts_batch." + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + @pytest.mark.asyncio async def test_temporal_links_within_same_batch(memory, request_context): """ @@ -2793,3 +2859,477 @@ async def test_named_strategy_applied_end_to_end(memory, request_context): finally: await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_semantic_ann_uses_hnsw_index(memory, request_context): + """ + Test that Phase 1 ANN semantic search creates links between similar world + facts across batches. This exercises the per-fact_type partial HNSW index + and the placeholder-ID remap logic. + """ + bank_id = f"test_sem_ann_hnsw_{datetime.now(timezone.utc).timestamp()}" + + try: + # Batch 1: world facts about machine learning + unit_ids_1 = await memory.retain_async( + bank_id=bank_id, + content=( + "Deep learning models require large amounts of training data. " + "Gradient descent is the primary optimization algorithm used in neural networks." + ), + context="ML knowledge base", + event_date=datetime(2024, 3, 1, tzinfo=timezone.utc), + request_context=request_context, + ) + assert len(unit_ids_1) > 0, "Batch 1 should produce facts" + + # Batch 2: similar ML world facts — Phase 1 ANN should link to batch 1 + unit_ids_2 = await memory.retain_async( + bank_id=bank_id, + content=( + "Neural networks learn by adjusting weights through backpropagation. " + "Training deep learning models requires GPUs for fast gradient computation." + ), + context="ML knowledge base", + event_date=datetime(2024, 3, 2, tzinfo=timezone.utc), + request_context=request_context, + ) + assert len(unit_ids_2) > 0, "Batch 2 should produce facts" + + logger.info(f"Batch 1: {len(unit_ids_1)} facts, Batch 2: {len(unit_ids_2)} facts") + + # Verify cross-batch semantic links exist + async with memory._pool.acquire() as conn: + cross_links = await conn.fetch( + """ + SELECT from_unit_id, to_unit_id, weight + FROM memory_links + WHERE from_unit_id::text = ANY($1) + AND link_type = 'semantic' + AND to_unit_id::text = ANY($2) + """, + unit_ids_2, + unit_ids_1, + ) + + logger.info(f"Cross-batch semantic links (batch2 -> batch1): {len(cross_links)}") + for link in cross_links: + logger.info( + f" {str(link['from_unit_id'])[:8]}... -> " + f"{str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})" + ) + + assert len(cross_links) > 0, ( + "Phase 1 ANN should create semantic links between similar world facts " + "from different batches via the HNSW index with placeholder-ID remap." + ) + + # All weights must meet the similarity threshold + for link in cross_links: + assert link["weight"] >= 0.7, ( + f"Semantic link weight {link['weight']:.3f} below threshold 0.7" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_temporal_links_scoped_by_fact_type(memory, request_context): + """ + Test that temporal links only connect facts of the SAME fact_type. + + World facts should not get temporal links to experience facts even when + their event dates fall within the time window. + """ + bank_id = f"test_temporal_scope_{datetime.now(timezone.utc).timestamp()}" + + try: + base_date = datetime(2024, 5, 10, 12, 0, 0, tzinfo=timezone.utc) + + # Store a world fact + world_ids = await memory.retain_async( + bank_id=bank_id, + content="Python 3.12 was released with significant performance improvements for the interpreter.", + context="tech news", + event_date=base_date, + fact_type_override="world", + request_context=request_context, + ) + assert len(world_ids) > 0, "Should create world fact(s)" + + # Store an experience fact at a nearby timestamp (same hour) + experience_ids = await memory.retain_async( + bank_id=bank_id, + content="I upgraded all my projects to Python 3.12 and benchmarked the speed improvements.", + context="personal log", + event_date=base_date + timedelta(hours=1), + fact_type_override="experience", + request_context=request_context, + ) + assert len(experience_ids) > 0, "Should create experience fact(s)" + + # Store another world fact at a nearby timestamp so we can confirm + # same-type temporal links ARE created + world_ids_2 = await memory.retain_async( + bank_id=bank_id, + content="The Python Software Foundation announced long-term support plans for Python 3.12.", + context="tech news", + event_date=base_date + timedelta(hours=2), + fact_type_override="world", + request_context=request_context, + ) + assert len(world_ids_2) > 0, "Should create second world fact(s)" + + logger.info( + f"World1: {world_ids}, Experience: {experience_ids}, World2: {world_ids_2}" + ) + + async with memory._pool.acquire() as conn: + # Check that world facts DO have temporal links to each other + world_all = world_ids + world_ids_2 + world_temporal = await conn.fetch( + """ + SELECT from_unit_id, to_unit_id, weight + FROM memory_links + WHERE from_unit_id::text = ANY($1) + AND to_unit_id::text = ANY($1) + AND link_type = 'temporal' + """, + world_all, + ) + logger.info(f"World-to-world temporal links: {len(world_temporal)}") + assert len(world_temporal) > 0, ( + "World facts with nearby dates should have temporal links to each other" + ) + + # Check that world facts do NOT have temporal links to experience facts + cross_type_links = await conn.fetch( + """ + SELECT from_unit_id, to_unit_id, weight + FROM memory_links + WHERE ( + (from_unit_id::text = ANY($1) AND to_unit_id::text = ANY($2)) + OR + (from_unit_id::text = ANY($2) AND to_unit_id::text = ANY($1)) + ) + AND link_type = 'temporal' + """, + world_all, + experience_ids, + ) + logger.info(f"Cross-type temporal links (world<->experience): {len(cross_type_links)}") + assert len(cross_type_links) == 0, ( + f"Temporal links should NOT cross fact types, but found {len(cross_type_links)} " + f"world<->experience links" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +# --------------------------------------------------------------------------- +# Streaming chunk batching tests +# --------------------------------------------------------------------------- + +import json +import uuid +from unittest.mock import patch + +import pytest_asyncio + +from hindsight_api.engine.llm_wrapper import TokenUsage +from hindsight_api.engine.memory_engine import MemoryEngine +from hindsight_api.engine.task_backend import SyncTaskBackend + + +def _make_mock_llm_call(): + """Create a mock LLM call function that returns deterministic facts.""" + + async def mock_llm_call(*args, **kwargs): + from hindsight_api.engine.consolidation.consolidator import _ConsolidationBatchResponse + + if kwargs.get("scope") == "consolidation": + return_usage = kwargs.get("return_usage", False) + if return_usage: + return _ConsolidationBatchResponse(), TokenUsage(input_tokens=0, output_tokens=0) + return _ConsolidationBatchResponse() + + messages = kwargs.get("messages", args[0] if args else []) + user_msg = messages[-1]["content"] if messages else "" + + # Extract sentences from the content to generate one fact per sentence + sentences = [s.strip() for s in user_msg.split(".") if s.strip() and len(s.strip()) > 10] + num_facts = max(1, min(len(sentences), 10)) + + facts = [] + for i in range(num_facts): + sentence = sentences[i] if i < len(sentences) else f"Fact {i}" + facts.append({ + "what": sentence[:200], + "when": "2024-06-15", + "where": "N/A", + "who": "N/A", + "why": "N/A", + "fact_type": "world", + "entities": [{"text": f"Entity{i}"}], + "causal_relations": [], + }) + + response_dict = {"facts": facts} + return_usage = kwargs.get("return_usage", False) + if return_usage: + usage = TokenUsage( + input_tokens=len(user_msg) // 4, + output_tokens=len(json.dumps(response_dict)) // 4, + ) + return response_dict, usage + return response_dict + + return mock_llm_call + + +@pytest_asyncio.fixture(scope="function") +async def memory_mock_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer): + """MemoryEngine with mock LLM for streaming tests.""" + mem = MemoryEngine( + db_url=pg0_db_url, + memory_llm_provider="openai", + memory_llm_api_key="mock-key", + memory_llm_model="gpt-4", + embeddings=embeddings, + cross_encoder=cross_encoder, + query_analyzer=query_analyzer, + pool_min_size=1, + pool_max_size=5, + run_migrations=False, + skip_llm_verification=True, + task_backend=SyncTaskBackend(), + ) + await mem.initialize() + yield mem + try: + if mem._pool and not mem._pool._closing: + await mem.close() + except Exception: + pass + + +def _generate_chunky_content(num_chunks: int, chunk_size: int = 3000) -> str: + """Generate content that will produce approximately num_chunks chunks. + + Each chunk is chunk_size characters, separated by double newlines. + """ + base_sentences = [ + "Alice works as a senior engineer at TechCorp in San Francisco.", + "Bob joined the marketing team last month from Chicago.", + "The project deadline was extended to December 15th.", + "Sarah mentioned she is planning a trip to Tokyo next month.", + "The quarterly budget review showed a 15% increase in revenue.", + "Mike suggested exploring alternative cloud providers.", + "The client feedback from beta testing was positive overall.", + "Emily started learning Rust programming language last week.", + "The new office will be located in the financial district.", + "David presented the annual technology roadmap to stakeholders.", + ] + + chunks = [] + for chunk_idx in range(num_chunks): + # Generate enough text for one chunk + lines = [] + chars = 0 + line_idx = 0 + while chars < chunk_size - 100: + sentence = f"[Chunk {chunk_idx}, Line {line_idx}] {base_sentences[line_idx % len(base_sentences)]}" + lines.append(sentence) + chars += len(sentence) + 1 + line_idx += 1 + chunks.append("\n".join(lines)) + + return "\n\n".join(chunks) + + +def _set_chunk_batch_size(memory: MemoryEngine, batch_size: int) -> None: + """Set retain_chunk_batch_size on the config resolver's global config.""" + memory._config_resolver._global_config.retain_chunk_batch_size = batch_size + + +@pytest.mark.asyncio +async def test_streaming_chunk_batching_produces_same_facts(memory_mock_llm, request_context): + """ + Retain a medium document (~10 chunks) with batch_size=3. + Verify all facts are extracted (streaming should not lose facts). + """ + memory = memory_mock_llm + _set_chunk_batch_size(memory, 3) + bank_id = f"test_streaming_{uuid.uuid4().hex[:8]}" + document_id = f"streaming_doc_{uuid.uuid4().hex[:8]}" + + # Generate content that produces ~10 chunks at default chunk_size (3000 chars) + content = _generate_chunky_content(num_chunks=10, chunk_size=3000) + + mock_llm_call = _make_mock_llm_call() + + try: + with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call): + # Retain with streaming enabled (batch_size=3, so 10 chunks -> 4 mini-batches) + result = await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "context": "streaming test", + "event_date": datetime(2024, 6, 15, tzinfo=timezone.utc), + }], + document_id=document_id, + request_context=request_context, + ) + + streaming_unit_ids = result[0] if result else [] + logger.info(f"Streaming produced {len(streaming_unit_ids)} facts") + assert len(streaming_unit_ids) > 0, "Streaming should produce facts" + + # Verify facts are in the DB + async with memory._pool.acquire() as conn: + fact_count = await conn.fetchval( + "SELECT COUNT(*) FROM memory_units WHERE bank_id = $1", + bank_id, + ) + assert fact_count == len(streaming_unit_ids), ( + f"DB has {fact_count} facts, but streaming returned {len(streaming_unit_ids)} unit_ids" + ) + + # Verify the document was tracked + doc = await conn.fetchrow( + "SELECT id FROM documents WHERE bank_id = $1 AND id = $2", + bank_id, document_id, + ) + assert doc is not None, "Document should be tracked in DB" + + # Verify chunks were stored with correct indices + chunk_count = await conn.fetchval( + "SELECT COUNT(*) FROM chunks WHERE bank_id = $1 AND document_id = $2", + bank_id, document_id, + ) + assert chunk_count > 0, "Chunks should be stored in DB" + logger.info(f"Stored {chunk_count} chunks for document {document_id}") + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_streaming_chunk_batching_recovery(memory_mock_llm, request_context): + """ + Test recovery: retain a document with streaming, then retain the same + document again. Delta retain should detect existing chunks and skip + re-extraction. Fact count should be unchanged (no duplicates). + """ + memory = memory_mock_llm + _set_chunk_batch_size(memory, 3) + bank_id = f"test_streaming_recovery_{uuid.uuid4().hex[:8]}" + document_id = f"recovery_doc_{uuid.uuid4().hex[:8]}" + + content = _generate_chunky_content(num_chunks=9, chunk_size=3000) + + mock_llm_call = _make_mock_llm_call() + + try: + # First retain — streaming mode + with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call): + result1 = await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "context": "recovery test", + "event_date": datetime(2024, 6, 15, tzinfo=timezone.utc), + }], + document_id=document_id, + request_context=request_context, + ) + + first_unit_ids = result1[0] if result1 else [] + assert len(first_unit_ids) > 0, "First retain should produce facts" + + async with memory._pool.acquire() as conn: + first_fact_count = await conn.fetchval( + "SELECT COUNT(*) FROM memory_units WHERE bank_id = $1", + bank_id, + ) + + logger.info(f"First retain: {first_fact_count} facts") + + # Second retain — same document, same content (should be a no-op via delta retain) + with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call): + result2 = await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "context": "recovery test", + "event_date": datetime(2024, 6, 15, tzinfo=timezone.utc), + }], + document_id=document_id, + request_context=request_context, + ) + + async with memory._pool.acquire() as conn: + second_fact_count = await conn.fetchval( + "SELECT COUNT(*) FROM memory_units WHERE bank_id = $1", + bank_id, + ) + + logger.info(f"Second retain: {second_fact_count} facts") + + # Fact count should be the same (delta retain skipped unchanged chunks) + assert second_fact_count == first_fact_count, ( + f"Second retain should not create duplicates: first={first_fact_count}, second={second_fact_count}" + ) + + finally: + await memory.delete_bank(bank_id, request_context=request_context) + + +@pytest.mark.asyncio +async def test_streaming_disabled_for_small_docs(memory_mock_llm, request_context): + """ + Retain a small document (2 chunks) with batch_size=500. + Verify it uses the non-streaming path (no batching overhead). + """ + memory = memory_mock_llm + _set_chunk_batch_size(memory, 500) + bank_id = f"test_streaming_small_{uuid.uuid4().hex[:8]}" + document_id = f"small_doc_{uuid.uuid4().hex[:8]}" + + # Generate content that produces ~2 chunks + content = _generate_chunky_content(num_chunks=2, chunk_size=3000) + + mock_llm_call = _make_mock_llm_call() + + try: + with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_call): + # batch_size=500 >> 2 chunks, so non-streaming path should be used + result = await memory.retain_batch_async( + bank_id=bank_id, + contents=[{ + "content": content, + "context": "small doc test", + "event_date": datetime(2024, 6, 15, tzinfo=timezone.utc), + }], + document_id=document_id, + request_context=request_context, + ) + + unit_ids = result[0] if result else [] + logger.info(f"Small doc produced {len(unit_ids)} facts") + assert len(unit_ids) > 0, "Should produce facts even through non-streaming path" + + # Verify the document was tracked + async with memory._pool.acquire() as conn: + doc = await conn.fetchrow( + "SELECT id FROM documents WHERE bank_id = $1 AND id = $2", + bank_id, document_id, + ) + assert doc is not None, "Document should be tracked" + + finally: + await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-dev/benchmarks/perf/retain_perf.py b/hindsight-dev/benchmarks/perf/retain_perf.py index ed8214cf..31142cd2 100644 --- a/hindsight-dev/benchmarks/perf/retain_perf.py +++ b/hindsight-dev/benchmarks/perf/retain_perf.py @@ -31,6 +31,19 @@ from rich.table import Table console = Console() +def _create_memory_engine(): + """Create a MemoryEngine from environment variables.""" + from hindsight_api import MemoryEngine + + return MemoryEngine( + db_url=os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0"), + memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"), + memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"), + memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-20b"), + memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None, + ) + + async def retain_via_memory_engine( bank_id: str, items: list[dict[str, Any]], @@ -45,17 +58,9 @@ async def retain_via_memory_engine( Returns: Tuple of (duration_seconds, response_data) """ - from hindsight_api import MemoryEngine from hindsight_api.models import RequestContext - # Initialize memory engine - memory = MemoryEngine( - db_url=os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0"), - memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"), - memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"), - memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-20b"), - memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None, - ) + memory = _create_memory_engine() await memory.initialize() # Measure time @@ -88,6 +93,303 @@ async def retain_via_memory_engine( await pool.close() +def _mock_fact_response(messages, scope): + """Generate fact extraction responses with heavily overlapping entities. + + Uses a tiny entity pool (5 names) so every concurrent retain touches the + same rows in the entities / unit_entities / memory_links tables, maximising + the chance of deadlocks from row-lock ordering conflicts. + """ + import hashlib + import random as _rng + + # Deterministic seed from message content so results are repeatable + content = str(messages) + seed = int(hashlib.md5(content.encode()).hexdigest()[:8], 16) + _rng.seed(seed) + + # Deliberately tiny pools → very high overlap across concurrent retains + names = ["Alice", "Bob", "Carol", "Dave", "Eve"] + places = ["New York", "London"] + + num_facts = _rng.randint(5, 12) + facts = [] + for _ in range(num_facts): + who1, who2 = _rng.sample(names, 2) + place = _rng.choice(places) + facts.append( + { + "what": f"{who1} met {who2} in {place}", + "when": "2024-01-15", + "where": place, + "who": f"{who1}, {who2}", + "why": "N/A", + "fact_kind": "conversation", + "fact_type": "world", + "entities": [{"text": who1}, {"text": who2}, {"text": place}], + } + ) + return {"facts": facts} + + +async def retain_via_memory_engine_async( + bank_id: str, + items: list[dict[str, Any]], +) -> tuple[float, dict[str, Any]]: + """ + Submit retain via submit_async_retain and let the WorkerPoller process it. + + This reproduces the real async flow: documents are split into sub-batches, + each becomes a separate worker task, and the worker processes them concurrently. + """ + from hindsight_api.models import RequestContext + from hindsight_api.worker.poller import WorkerPoller + + memory = _create_memory_engine() + await memory.initialize() + + # Configure mock LLM to return realistic facts with entities (after init) + for llm_config in [memory._llm_config, memory._retain_llm_config]: + if hasattr(llm_config, "set_response_callback"): + llm_config.set_response_callback(_mock_fact_response) + console.print(" [cyan]Mock LLM configured with entity-rich fact responses[/cyan]") + + pool = await memory._get_pool() + + # Start a WorkerPoller so tasks get picked up + poller = WorkerPoller( + pool=pool, + worker_id="bench-worker", + executor=memory.execute_task, + poll_interval_ms=100, + max_slots=50, + ) + poller_task = asyncio.create_task(poller.run()) + + start_time = time.time() + + try: + # Submit async retain (splits into sub-batches as worker tasks) + result = await memory.submit_async_retain( + bank_id=bank_id, + contents=items, + request_context=RequestContext(), + ) + operation_id = result["operation_id"] + console.print(f" Submitted operation {operation_id} ({result['items_count']} items)") + + # Poll for completion + while True: + status = await memory.get_operation_status( + bank_id=bank_id, + operation_id=operation_id, + request_context=RequestContext(), + ) + op_status = status.get("status") + if op_status in ("completed", "failed"): + if op_status == "failed": + console.print(f" [red]Operation FAILED: {status.get('error_message')}[/red]") + # Print child statuses if available + for child in status.get("child_operations", []): + if child.get("status") == "failed": + console.print(f" Child {child['operation_id']}: {child.get('error_message', '')}") + break + await asyncio.sleep(0.5) + + duration = time.time() - start_time + + response_data = { + "success": op_status == "completed", + "bank_id": bank_id, + "items_count": result["items_count"], + "async": True, + "usage": None, + } + + return duration, response_data + finally: + await poller.shutdown_graceful(timeout=5) + poller_task.cancel() + await memory.close() + + +async def stress_test_deadlocks( + concurrency: int = 20, + num_documents: int = 50, + max_retain_concurrent: int | None = None, +) -> dict[str, Any]: + """ + Fire many concurrent retains into the same bank with overlapping entities + to reproduce deadlocks on entity/link tables. + + Each document gets a unique short text, but the mock LLM always returns + facts referencing the same small set of entities — maximising row-lock + contention on the entities and unit_entities tables. + """ + import traceback + from collections import Counter + + from hindsight_api.models import RequestContext + + # Override semaphore limit if requested (before engine init reads config) + if max_retain_concurrent is not None: + os.environ["HINDSIGHT_API_RETAIN_MAX_CONCURRENT"] = str(max_retain_concurrent) + + # Enable logging so deadlock retry warnings are visible + import logging + + logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(name)s: %(message)s") + logging.getLogger("hindsight_api.engine.db_utils").setLevel(logging.DEBUG) + + # Force mock provider so LLM calls are instant — we're testing DB contention + os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock" + memory = _create_memory_engine() + await memory.initialize() + + # Wire up the mock callback so LLM calls return entity-rich facts instantly + for llm_config in [memory._llm_config, memory._retain_llm_config]: + if hasattr(llm_config, "set_response_callback"): + llm_config.set_response_callback(_mock_fact_response) + + bank_id = f"stress-deadlock-{int(time.time())}" + pool = await memory._get_pool() + + # Ensure the bank exists + # Ensure bank exists before firing concurrent retains + from hindsight_api.engine.retain.fact_storage import ensure_bank_exists + + async with pool.acquire() as conn: + await ensure_bank_exists(conn, bank_id) + + console.print("\n[bold]Stress test config:[/bold]") + console.print(f" Bank: {bank_id}") + console.print(f" Documents: {num_documents}") + console.print(f" Concurrency: {concurrency}") + console.print(f" DB semaphore: {max_retain_concurrent or 'default'}") + console.print() + + # Generate synthetic documents — large enough to produce many chunks. + # Default chunk_size is 3000 chars, so 100k content ≈ 33 chunks per doc. + # Each chunk triggers the mock LLM which returns 5-12 facts with overlapping + # entities, maximising row-lock contention across concurrent transactions. + content_size = int(os.getenv("STRESS_CONTENT_SIZE", "100000")) + console.print(f" Content/doc: ~{content_size:,} chars (~{content_size // 3000} chunks)") + + # Pre-populate the bank with seed documents so that the bank already has + # units with embeddings. Subsequent concurrent retains will create semantic + # and temporal links to these existing units — and to each other's new units + # — triggering INSERT ON CONFLICT share-lock deadlocks on memory_links. + seed_count = int(os.getenv("STRESS_SEED_DOCS", "5")) + if seed_count > 0: + console.print(f"\n[cyan]Seeding bank with {seed_count} documents (serial)...[/cyan]") + for i in range(seed_count): + await memory.retain_batch_async( + bank_id=bank_id, + contents=[ + { + "content": f"Seed document {i}: Alice discussed machine learning with Bob in New York. " + f"Carol and Dave reviewed the quarterly results in London. " + f"Eve presented blockchain research findings to Frank in Berlin." + } + ], + request_context=RequestContext(), + ) + console.print(f" [green]Seeded {seed_count} documents[/green]") + + documents = [] + for i in range(num_documents): + # Build a large document from repeated paragraphs with slight variation + paragraphs = [] + while len("\n\n".join(paragraphs)) < content_size: + j = len(paragraphs) + paragraphs.append( + f"Section {j} of document {i}: Alice and Bob met Carol in New York to discuss " + f"the progress on project Alpha. Dave joined from London via video call. " + f"Eve presented the quarterly results while Frank took notes. " + f"The team agreed to reconvene next week in San Francisco. " + f"Key topics included machine learning infrastructure, deployment pipelines, " + f"and the upcoming product launch scheduled for Q2." + ) + content = "\n\n".join(paragraphs)[:content_size] + documents.append([{"content": content}]) + + # Track outcomes per task + results: list[dict] = [] + semaphore = asyncio.Semaphore(concurrency) + + async def _retain_one(doc_idx: int, items: list[dict]) -> dict: + async with semaphore: + t0 = time.time() + try: + await memory.retain_batch_async( + bank_id=bank_id, + contents=items, + request_context=RequestContext(), + ) + return {"idx": doc_idx, "status": "ok", "duration": time.time() - t0} + except Exception as e: + return { + "idx": doc_idx, + "status": "error", + "error": type(e).__name__, + "message": str(e)[:200], + "traceback": traceback.format_exc(), + "duration": time.time() - t0, + } + + console.print("[cyan]Firing concurrent retains...[/cyan]") + start = time.time() + + tasks = [asyncio.create_task(_retain_one(i, docs)) for i, docs in enumerate(documents)] + results = await asyncio.gather(*tasks) + + total_time = time.time() - start + + # Summarise + status_counts = Counter(r["status"] for r in results) + error_types = Counter(r.get("error", "") for r in results if r["status"] == "error") + durations = [r["duration"] for r in results] + durations.sort() + + console.print(f"\n[bold]Results ({total_time:.2f}s total):[/bold]") + table = Table(title="Stress Test Results") + table.add_column("Metric", style="cyan") + table.add_column("Value", style="green") + table.add_row("Total documents", str(num_documents)) + table.add_row("Concurrency", str(concurrency)) + table.add_row("OK", str(status_counts.get("ok", 0))) + table.add_row("Errors", str(status_counts.get("error", 0))) + table.add_row("Wall time", f"{total_time:.2f}s") + table.add_row("p50 latency", f"{durations[len(durations) // 2]:.3f}s") + table.add_row("p95 latency", f"{durations[int(len(durations) * 0.95)]:.3f}s") + table.add_row("p99 latency", f"{durations[int(len(durations) * 0.99)]:.3f}s") + table.add_row("Max latency", f"{durations[-1]:.3f}s") + console.print(table) + + if error_types: + console.print("\n[bold red]Error breakdown:[/bold red]") + for err_type, count in error_types.most_common(): + console.print(f" {err_type}: {count}") + # Show first traceback for each error type + for r in results: + if r.get("error") == err_type: + console.print(f" [dim]{r.get('traceback', '(no traceback)')}[/dim]") + break + + # Check for deadlock retries in logs + deadlock_count = sum(1 for r in results if r.get("error") == "DeadlockDetectedError") + if deadlock_count: + console.print(f"\n[bold red]Deadlocks that exhausted retries: {deadlock_count}[/bold red]") + elif status_counts.get("error", 0) == 0: + console.print( + "\n[bold green]No errors — deadlocks may still have occurred but were retried successfully.[/bold green]" + ) + console.print("[dim]Check logs above for 'Deadlock detected' warnings from retry_with_backoff.[/dim]") + + await pool.close() + return {"ok": status_counts.get("ok", 0), "errors": status_counts.get("error", 0), "wall_time": total_time} + + async def retain_via_http( base_url: str, bank_id: str, @@ -309,7 +611,7 @@ Examples: parser.add_argument( "--document", - required=True, + required=False, help="Path to document file or directory (for directories, batches all .json/.txt/.md files)", ) parser.add_argument( @@ -342,9 +644,52 @@ Examples: action="store_true", help="Use in-memory MemoryEngine instead of HTTP (bypasses API server, useful for isolating performance)", ) + parser.add_argument( + "--async", + dest="use_async", + action="store_true", + help="Use async retain (submit_async_retain + worker poller). Only works with --in-memory.", + ) + parser.add_argument( + "--max-retain-concurrent", + type=int, + default=None, + help="Override HINDSIGHT_API_RETAIN_MAX_CONCURRENT for this run (default: from config)", + ) + parser.add_argument( + "--stress", + action="store_true", + help="Run deadlock stress test: fire many concurrent retains with overlapping entities into the same bank", + ) + parser.add_argument( + "--stress-concurrency", + type=int, + default=20, + help="Max concurrent retains for stress test (default: 20)", + ) + parser.add_argument( + "--stress-documents", + type=int, + default=50, + help="Number of documents to retain in stress test (default: 50)", + ) args = parser.parse_args() + # Stress test mode — standalone, doesn't need --document + if args.stress: + console.print("\n[bold cyan]Retain Deadlock Stress Test[/bold cyan]") + console.print("=" * 80) + await stress_test_deadlocks( + concurrency=args.stress_concurrency, + num_documents=args.stress_documents, + max_retain_concurrent=args.max_retain_concurrent, + ) + return + + if not args.document: + parser.error("--document is required (unless using --stress)") + console.print("\n[bold cyan]Retain Performance Benchmark[/bold cyan]") console.print("=" * 80) @@ -368,6 +713,11 @@ Examples: console.print(" ./scripts/dev/start-api.sh") sys.exit(1) + # Override retain_max_concurrent if specified + if args.max_retain_concurrent is not None: + os.environ["HINDSIGHT_API_RETAIN_MAX_CONCURRENT"] = str(args.max_retain_concurrent) + console.print(f" [cyan]Retain max concurrent: {args.max_retain_concurrent}[/cyan]") + # Load document(s) doc_path = Path(args.document) if doc_path.is_dir(): @@ -398,8 +748,14 @@ Examples: console.print(f"\n[3] {'Processing' if args.in_memory else 'Sending retain request to'} bank '{args.bank_id}'...") console.print(f" [cyan]Retaining {num_docs:,} document{'s' if num_docs > 1 else ''} in batch...[/cyan]") try: - if args.in_memory: - # In-memory mode: call MemoryEngine directly + if args.in_memory and args.use_async: + # In-memory async mode: submit_async_retain + worker poller + duration, result = await retain_via_memory_engine_async( + bank_id=args.bank_id, + items=items, + ) + elif args.in_memory: + # In-memory sync mode: call MemoryEngine directly duration, result = await retain_via_memory_engine( bank_id=args.bank_id, items=items, diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 2cd2e466..1179310b 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -612,6 +612,7 @@ Controls the retain (memory ingestion) pipeline. | `HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS` | Full prompt override for fact extraction (only used when mode is `custom`). Replaces built-in extraction rules entirely. | - | | `HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS` | Extract causal relationships between facts | `true` | | `HINDSIGHT_API_RETAIN_BATCH_ENABLED` | Use LLM Batch API for fact extraction (50% cost savings, only with async operations) | `false` | +| `HINDSIGHT_API_RETAIN_MAX_CONCURRENT` | Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention during high-concurrency ingestion. | `4` | | `HINDSIGHT_API_RETAIN_BATCH_TOKENS` | Max characters per sub-batch for async retain auto-splitting | `10000` | | `HINDSIGHT_API_RETAIN_ENTITY_LOOKUP` | Entity lookup method during retain: `full` (exact match) or `trigram` (fuzzy trigram matching) | `trigram` | | `HINDSIGHT_API_RETAIN_DEFAULT_STRATEGY` | Default retain strategy name. When set, all retain calls without an explicit `strategy` parameter use this strategy. | - |