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
This commit is contained in:
Nicolò Boschi
GitHub
parent
6f173b10a7
commit
914ba7962c
16 changed files with 2908 additions and 839 deletions
|
|
@ -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)"
|
||||
)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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. | - |
|
||||
|
|
|
|||
Loading…
Reference in a new issue