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