Commit graph

2 commits

Author SHA1 Message Date
Nicolò Boschi
914ba7962c
perf: 3-phase retain pipeline — fix deadlocks, cap temporal links, query-time entity expansion (#722)
* perf: 3-phase retain pipeline — fix deadlocks, cap temporal links, query-time entity expansion

Major retain pipeline overhaul addressing deadlocks, write amplification,
and TimeoutErrors. Restructures retain into three phases:

Phase 1: Entity resolution on separate connection (read-heavy)
Phase 2: Core write transaction (atomic) — facts, unit_entities, links
Phase 3: Best-effort display data (error-isolated) — entity viz links, stats

Key changes:
- Sorted bulk INSERT FROM unnest() prevents deadlocks
- Temporal links capped to top-20 per unit (95% reduction)
- Batched semantic ANN via temp table + LATERAL
- Query-time entity expansion via unit_entities self-join
- Entity viz links moved to Phase 3 (post-transaction)
- HINDSIGHT_API_RETAIN_MAX_CONCURRENT config (default: 32)

* fix: increase semantic link top_k from 5 to 20

The hardcoded top_k=5 was artificially limiting semantic link creation.
Link expansion retrieval can consume up to budget (50-200) semantic
neighbors per seed set, but each fact only had 5 outgoing edges — making
the bidirectional graph very sparse.

Increasing to 20 gives retrieval 4x more edges to work with. The ANN
probe cost is unchanged (same HNSW traversal per fact, just returning
more rows). INSERT cost is negligible (~14k rows via bulk INSERT).

Also: all 18 TimeoutErrors in the latest benchmark (beam-1m-u20) were
from Gemini LLM calls, zero from the database — confirming the entity
resolution split eliminated DB timeouts entirely.

* perf: move semantic ANN search to Phase 1 to avoid transaction timeouts

The batched LATERAL ANN query (700 HNSW probes) was the last remaining
source of DB TimeoutErrors — all 29 in the latest benchmark were from
create_semantic_links_batch inside the Phase 2 write transaction.

Split semantic link creation into three phases:
- Phase 1 (separate conn, autocommit): ANN search via temp table + LATERAL.
  No transaction locks, no contention with concurrent writers.
- Phase 2 (write transaction): within-batch numpy similarities (instant) +
  INSERT of both within-batch and Phase 1 ANN results. No DB reads.
- Phase 3 (flush_pending_stats): future hook point for re-checking ANN
  results after commit to catch links missed by concurrent batches.

Also adds 7 unit tests for compute_semantic_links_within_batch covering
empty input, identical/orthogonal embeddings, threshold filtering, top_k
cap, and tuple structure validation.

* fix: handle placeholder unit_ids in Phase 1 ANN search (not valid UUIDs)

* test: add Phase 1 ANN cross-batch test + configurable test PG port

- New test_semantic_links_phase1_ann_cross_batch verifies that the Phase 1
  ANN search with placeholder unit IDs correctly creates cross-batch
  semantic links after remapping to real IDs.
- Test PG port now configurable via HINDSIGHT_TEST_PG_PORT env var
  (default: 5556) to avoid conflicts with running benchmark daemons.

* perf: remove retry_with_backoff from retain, set semaphore default to 4

Remove retry_with_backoff from _run_db_work and _run_delta_db_work:
- Deadlocks are prevented by sorted bulk INSERT (no need for retry)
- Transient timeouts are handled by the worker poller's task-level retry
  (3 attempts, 60s spacing) which is better than rapid internal retries
  that amplify I/O pressure during contention storms

Set HINDSIGHT_API_RETAIN_MAX_CONCURRENT default from 32 to 4:
- The semaphore gates Phase 1 (ANN + entity resolution) + Phase 2 (writes)
- At 4 concurrent, HNSW index I/O is manageable; at 10+ concurrent the
  probes saturate disk and cause cascading timeouts
- LLM extraction still runs at full parallelism (semaphore acquired after)

* fix: add fact_type filter to Phase 1 ANN query to use per-bank HNSW indexes

The LATERAL ANN query was falling back to sequential scan + sort (90ms/probe)
because the per-bank HNSW indexes are partial indexes filtered on fact_type.
Without fact_type in the WHERE clause, PostgreSQL couldn't use them.

Fix: iterate over ('world', 'experience') and run one HNSW-indexed ANN per
type. EXPLAIN shows 8ms/probe (was 90ms) — 11x faster.

700 probes × 8ms × 2 types = ~11s total (was ~63s via seq scan).

* fix: scope temporal links by fact_type + add integration tests

Temporal links now filter by fact_type in the LATERAL query — world facts
only link to world facts, experience to experience. This matches how
retrieval filters results and avoids wasted cross-type link rows.

New integration tests:
- test_semantic_ann_uses_hnsw_index: verifies Phase 1 ANN creates
  cross-batch semantic links (tests fact_type filter + placeholder remap)
- test_temporal_links_scoped_by_fact_type: verifies world facts get
  temporal links to other world facts but NOT to experience facts

* fix: tolerate individual chunk LLM failures instead of failing entire batch

Changed asyncio.gather(*tasks) to asyncio.gather(*tasks, return_exceptions=True)
in both chunk-level and content-level fact extraction. A single chunk timeout
(e.g., Gemini >90s) no longer discards all other successfully extracted facts.

For a 50MB document with 17k chunks, even a 2% chunk failure rate previously
caused 0 completions (entire batch discarded). Now 16,700 facts are extracted
and only the 300 failed chunks are skipped with a warning log.

* fix: batch temporal LATERAL query for large documents (16k+ chunks)

The LATERAL query for temporal links passed all unit_ids at once into
unnest(), causing PostgreSQL timeouts on documents with 16k+ chunks.
Split into batches of 500 units per query to keep each under the
command_timeout.

Also identified: HNSW index creation on shared pg0 instances with
50k+ existing units exceeds the 60s command_timeout. This is a
test infrastructure issue (shared pg0 accumulates data) but also
affects production when creating new banks on large instances.

* feat: streaming chunk batching for large documents (RETAIN_CHUNK_BATCH_SIZE)

Process chunks in mini-batches of N (default 500), committing each batch
to the DB before starting the next. This prevents OOM kills on large
documents (50MB / 17k+ chunks) by keeping only ~500 facts + embeddings
in memory at a time instead of 50k+.

Each mini-batch goes through the full Phase 1 → 2 → 3 pipeline
independently, sharing the same document_id. On recovery (process dies
mid-way), delta retain detects already-committed chunks via content_hash
and skips them — only remaining chunks get re-extracted.

Config: HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE (default: 500, 0 to disable)
Per-bank configurable via the hierarchical config system.

Tests:
- test_streaming_chunk_batching_produces_same_facts
- test_streaming_chunk_batching_recovery (delta retain skips committed chunks)
- test_streaming_disabled_for_small_docs

* perf(retain): producer-consumer pipeline + deferred semantic ANN

Replace the sequential streaming loop with a producer-consumer pipeline:
- LLM producer fires concurrent chunk extractions (semaphore-bounded)
- DB consumer drains queue in batches, runs Phase 1+2+3 per batch
- LLM and DB work overlap instead of running sequentially

Defer semantic links to a single final ANN pass after all batches commit:
- Remove within-batch semantic links from Phase 2 (was 2.6s/batch)
- Run parallel ANN (4 connections) after all facts committed
- top_k reduced from 50 to 20 (recall uses at most 20 neighbors)
- Recovery via operation result_metadata checkpoint

Additional optimizations:
- skip_exists_check on temporal/causal link INSERT (saves ~0.5s/batch)
- WHERE EXISTS guard on semantic link INSERT (handles document upsert)
- timeout=300s on ANN queries and bulk INSERT for large banks
- Demote [ANN] debug logs to logger.debug()
- Fix docstring typos (agent_id → bank_id)
- Fix content_index remapping in producer-consumer batches
- Fix delta retain passing contents vs delta_contents

50MB benchmark (mock LLM): 9.2 min (was 23 min) — 2.5x faster.
BEAM 10m benchmark: zero deadlocks, zero DB errors.

* refactor(retain): remove legacy fallback code paths

- Remove process_entities_batch (legacy single-connection entity processing)
- Remove extract_entities_batch_optimized (only caller was the above)
- Remove fallback entity processing inside Phase 2 transaction
- Remove legacy ANN inline fallback in create_semantic_links_batch
- Remove fallback entity_links direct-insert path in Phase 3
- Make resolved_entity_ids/entity_to_unit/unit_to_entity_ids required params

* refactor(retain): replace tuple returns with dataclasses, remove dead code

- Add EntityResolutionResult and Phase1Result dataclasses in types.py
- Replace 4-tuple return from _pre_resolve_phase1 with Phase1Result
- Remove dead `entity_links = []` variables in retain_batch and _try_delta_retain
- Remove unused `confidence_score` parameter from orchestrator.retain_batch
  and _retain_batch_async_internal (was accepted but never used)

* fix(entity-resolver): remove LIKE full-scan fallbacks, use index-only trigram matching

The entity resolution query had LIKE '%...' substring conditions that bypassed
the GIN trigram index, causing full sequential scans of the entities table.
On banks with 10k+ entities, this caused TimeoutErrors (observed in BEAM 10m).

Changes:
- Remove LIKE fallbacks, use trigram % operator only (GIN index-based)
- Lower similarity threshold from 0.3 to 0.15 to catch substring relationships
- Use LOWER() on both sides for case-insensitive matching
- Migration: recreate GIN trigram index on LOWER(canonical_name)

* fix: remove schema prefix from index names in trigram migration

* fix(delta-retain): use same chunk_size as streaming path (3000 vs 120000)

_chunk_contents_for_delta defaulted to chunk_size=120000 while the streaming
path used 3000. On retry, delta re-chunked the document with different
boundaries, found 0 matching chunks, and fell through to full re-extraction.
This wasted all LLM calls on already-committed chunks.

Fix: use the same default (3000) so chunk hashes match on recovery.

* fix(retain): persist generated document_id in operation metadata for retry recovery

When no document_id is provided, retain generates a UUID. On retry, a new UUID
was generated, making delta retain and streaming chunk-hash recovery unable to
find previously committed chunks. All LLM extraction was wasted on retry.

Fix: resolve document_id early in retain_batch (before delta), persist it to
operation result_metadata, and recover it on retry. Both delta and streaming
paths now see the same document_id across attempts.

* refactor(retain): unify into single streaming pipeline, remove non-streaming path

All retains now go through the producer-consumer streaming pipeline,
regardless of document size. Small documents are processed as a single batch.
This eliminates the maintenance burden of two separate code paths.

Also fix document upsert: compare content hash to distinguish recovery
(same content, partially committed) from update (different content, needs
cascade-delete). Previously, existing chunks always triggered recovery mode.

* refactor(retain): remove dead code, replace raw dicts with Phase3Context dataclass

- Remove dead _handle_zero_facts_documents (no callers after path unification)
- Remove unused imports: defaultdict, EntityLink
- Replace raw dict phase3_context with typed Phase3Context dataclass
- Update _build_and_insert_entity_links_phase3 to use typed parameter
2026-04-01 12:52:49 +02:00
Nicolò Boschi
aefb3fcf4d
fix: improve async batch retain with large payloads (#366)
* fix: improve async batch retain with large payloads

* fix: improve async batch retain with large payloads

* api

* api

* api

* api

* api

* Clean up perf benchmark: keep only Python files

- Remove README.md and PERFORMANCE_FINDINGS.md
- Remove results/ JSON files (gitignored)
- Remove test_data/ directory
- Keep only __init__.py and retain_perf.py

* docs: explain automatic batch optimization for async retain

- Add section explaining Hindsight automatically handles batch sizing
- Users don't need to manually tune batch sizes with async mode
- Hindsight splits large batches (>10k tokens) into optimized sub-batches
- Include example showing best practices

* docs: remove emojis and code example from performance page

* fix: correct OperationDetails type to match API response

- Change optional fields to use | null instead of ?
- Fixes TypeScript compilation error in control plane build

* fix: use discriminated union for OperationDetails type

- Support both success and error states properly
- Fixes TypeScript error when setting error state

* fix: use unique document_ids in batch retain examples

- Each item in a batch must have unique document_id
- Update both Python and JavaScript examples
- Fixes test-doc-examples CI failure

* chore: trigger CI

* fix: test mocking and duplicate document_ids in examples

- Mock _get_pool() in test_async_retain_tags.py to avoid _initialized error
- Set _initialized = True on mocked MemoryEngine instances
- Fix duplicate document_ids in retain.py and retain.mjs examples

* fix: properly mock async pool/connection and fix more duplicate document_ids

- Use AsyncMock for pool.acquire() to fix 'can't be used in await' error
- Fix duplicate document_ids in retain-async examples (retain.py and retain.mjs)
- Remove batch-level document_id parameter that caused duplicates

* ci: collect all doc example failures and show summary

- Run all Python/Node.js/CLI examples regardless of individual failures
- Collect failure list and display summary at the end
- Show pass/fail count and list of failed files
- Exit with failure only after running all examples

* refactor: extract doc example testing to standalone script

- Create scripts/test-doc-examples.sh to run all examples
- Collects logs of failed examples separately
- Shows full error logs only for failures at the end
- Clean summary with pass/fail counts
- Proper exit codes
- Replaces inline bash in CI workflow

* fix: doc examples - duplicate document_ids and error handling

- retain.py: move document_id to item level to avoid duplicates
- documents.mjs: add error handling for getDocument to show clear error message

* fix: update tests for duplicate document_id validation

- test_async_retain_tags: verify operation structure instead of exact UUID
- test_delete_bank: use unique document_ids (team-doc-1, team-doc-2)
2026-02-16 12:51:42 +01:00