diff --git a/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py index f83d0be0..69a8b6b5 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/link_utils.py @@ -739,17 +739,10 @@ async def compute_semantic_links_ann( return [] import time as time_mod - import uuid as uuid_mod ann_start = time_mod.time() links = [] - # Lower ef_search for retain ANN — default 400 is tuned for recall precision - # but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms/probe - # (35x faster) with sufficient accuracy for top-50 semantic link creation. - # Reset after to avoid polluting the connection pool for recall queries. - await conn.execute("SET hnsw.ef_search = 60") - logger.debug(f"[ANN] Starting: {len(unit_ids)} seeds, top_k={top_k}") # Build per-unit fact_types (default to 'world' if not provided) @@ -760,54 +753,75 @@ async def compute_semantic_links_ann( # sequential-scan every HNSW probe result against the array, destroying # performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO # NOTHING handles duplicates in memory_links). - t_setup = time_mod.time() - await conn.execute("CREATE TEMP TABLE IF NOT EXISTS _ann_seeds (unit_id text, emb_text text, fact_type text)") - await conn.execute("TRUNCATE _ann_seeds") + # + # The entire CREATE TEMP TABLE → COPY → SELECT sequence MUST run inside a + # single transaction. Callers may connect through pgBouncer in `transaction` + # pool mode, in which case the backend is only pinned to the client for the + # duration of a transaction. Outside a transaction, pgBouncer can rebind + # the client to a different backend between statements, and the temp table + # (which is session-scoped to its creating backend) becomes invisible. + # The observed failure mode was an intermittent + # `relation "_ann_seeds" does not exist` on the second statement. + # + # Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to + # manually drop the temp table or reset hnsw.ef_search — the transaction + # end handles both. + rows: list = [] + async with conn.transaction(): + # Transaction-local ef_search. Default 400 is tuned for recall precision + # but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms + # per probe (35x faster) with sufficient accuracy for top-50 semantic + # link creation. SET LOCAL auto-reverts at commit, so we don't pollute + # the pool for subsequent recall queries. + await conn.execute("SET LOCAL hnsw.ef_search = 60") - records = [ - (uid, emb if isinstance(emb, str) else str(emb), ft) for uid, emb, ft in zip(unit_ids, embeddings, fact_types) - ] - await conn.copy_records_to_table("_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"]) - logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)") - - # Run one ANN query per fact_type so each uses the right HNSW index. - rows = [] - active_types = set(fact_types) - for fact_type in active_types: - t_query = time_mod.time() - seed_count = sum(1 for ft in fact_types if ft == fact_type) - logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds") - ft_rows = await conn.fetch( - f""" - SELECT s.unit_id AS from_id, - n.id::text AS to_id, - n.similarity - FROM _ann_seeds s - CROSS JOIN LATERAL ( - SELECT mu.id, - 1 - (mu.embedding <=> s.emb_text::vector) AS similarity - FROM {fq_table("memory_units")} mu - WHERE mu.bank_id = $1 - AND mu.fact_type = $2 - AND mu.embedding IS NOT NULL - ORDER BY mu.embedding <=> s.emb_text::vector - LIMIT $3 - ) n - WHERE s.fact_type = $2 - """, - bank_id, - fact_type, - top_k, - timeout=300, # ANN on large banks can take minutes + t_setup = time_mod.time() + await conn.execute( + "CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP" ) - logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s") - rows.extend(ft_rows) - # Clean up temp table (no ON COMMIT DROP since we're not in a transaction) - await conn.execute("DROP TABLE IF EXISTS _ann_seeds") + records = [ + (uid, emb if isinstance(emb, str) else str(emb), ft) + for uid, emb, ft in zip(unit_ids, embeddings, fact_types) + ] + await conn.copy_records_to_table( + "_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"] + ) + logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)") - # Reset ef_search to default so the pooled connection doesn't affect recall queries - await conn.execute("RESET hnsw.ef_search") + # Run one ANN query per fact_type so each uses the right HNSW index. + active_types = set(fact_types) + for fact_type in active_types: + t_query = time_mod.time() + seed_count = sum(1 for ft in fact_types if ft == fact_type) + logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds") + ft_rows = await conn.fetch( + f""" + SELECT s.unit_id AS from_id, + n.id::text AS to_id, + n.similarity + FROM _ann_seeds s + CROSS JOIN LATERAL ( + SELECT mu.id, + 1 - (mu.embedding <=> s.emb_text::vector) AS similarity + FROM {fq_table("memory_units")} mu + WHERE mu.bank_id = $1 + AND mu.fact_type = $2 + AND mu.embedding IS NOT NULL + ORDER BY mu.embedding <=> s.emb_text::vector + LIMIT $3 + ) n + WHERE s.fact_type = $2 + """, + bank_id, + fact_type, + top_k, + timeout=300, # ANN on large banks can take minutes + ) + logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s") + rows.extend(ft_rows) + # Transaction commits here. _ann_seeds is dropped (ON COMMIT DROP). + # hnsw.ef_search reverts (SET LOCAL). for row in rows: sim = float(min(1.0, max(0.0, row["similarity"]))) diff --git a/hindsight-api-slim/tests/test_link_utils.py b/hindsight-api-slim/tests/test_link_utils.py index c5d6262d..c152e4e5 100644 --- a/hindsight-api-slim/tests/test_link_utils.py +++ b/hindsight-api-slim/tests/test_link_utils.py @@ -2,12 +2,14 @@ import numpy as np import pytest from datetime import datetime, timezone, timedelta +from unittest.mock import AsyncMock, MagicMock from hindsight_api.engine.retain.link_utils import ( _normalize_datetime, _cap_links_per_unit, compute_temporal_links, compute_temporal_query_bounds, + compute_semantic_links_ann, compute_semantic_links_within_batch, MAX_TEMPORAL_LINKS_PER_UNIT, ) @@ -388,3 +390,145 @@ class TestComputeSemanticLinksWithinBatch: assert link_type == "semantic" assert 0.0 <= weight <= 1.0 assert entity_id is None + + +class TestComputeSemanticLinksAnnPgBouncerSafety: + """Regression tests ensuring compute_semantic_links_ann stays in a single + transaction so that the `_ann_seeds` temp table remains visible when the + caller's connection goes through pgBouncer in `transaction` pool mode. + + In pgBouncer transaction mode, the backend is only pinned to the client + for the duration of an actual PostgreSQL transaction. Outside a + transaction, consecutive statements can land on different backends, and + session-scoped temp tables (which are bound to the backend that created + them) become invisible. The observed failure mode was an intermittent + `relation "_ann_seeds" does not exist` on the statement immediately + following the CREATE TEMP TABLE. + """ + + @pytest.fixture + def mock_conn(self): + """An asyncpg-like connection mock with an async `transaction()` + context manager and awaitable execute/fetch/copy helpers.""" + conn = MagicMock() + + txn_cm = MagicMock() + txn_cm.__aenter__ = AsyncMock(return_value=None) + txn_cm.__aexit__ = AsyncMock(return_value=None) + conn.transaction = MagicMock(return_value=txn_cm) + + conn.execute = AsyncMock() + conn.copy_records_to_table = AsyncMock() + conn.fetch = AsyncMock(return_value=[]) + return conn + + @pytest.mark.asyncio + async def test_empty_inputs_skip_transaction(self, mock_conn): + """No seeds -> no work, no transaction, no temp-table churn.""" + result = await compute_semantic_links_ann( + conn=mock_conn, + bank_id="bank-1", + unit_ids=[], + embeddings=[], + ) + assert result == [] + mock_conn.transaction.assert_not_called() + mock_conn.execute.assert_not_called() + + @pytest.mark.asyncio + async def test_runs_inside_a_transaction(self, mock_conn): + """The full CREATE TEMP TABLE -> COPY -> SELECT sequence must happen + inside a single `async with conn.transaction():` block.""" + emb = [0.1] * 384 + await compute_semantic_links_ann( + conn=mock_conn, + bank_id="bank-1", + unit_ids=["u1", "u2"], + embeddings=[emb, emb], + fact_types=["world", "world"], + ) + + # Transaction context manager was entered. + mock_conn.transaction.assert_called_once() + txn_cm = mock_conn.transaction.return_value + txn_cm.__aenter__.assert_awaited_once() + txn_cm.__aexit__.assert_awaited_once() + + @pytest.mark.asyncio + async def test_temp_table_uses_on_commit_drop(self, mock_conn): + """The CREATE TEMP TABLE statement must use ON COMMIT DROP so the + table is transaction-scoped. Without ON COMMIT DROP the table would + be session-scoped and would not survive pgBouncer backend rebinding + between transactions.""" + emb = [0.1] * 384 + await compute_semantic_links_ann( + conn=mock_conn, + bank_id="bank-1", + unit_ids=["u1"], + embeddings=[emb], + fact_types=["world"], + ) + + executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list] + create_statements = [s for s in executed_sql if "CREATE TEMP TABLE" in s] + assert len(create_statements) == 1, "Should create _ann_seeds exactly once" + assert "_ann_seeds" in create_statements[0] + assert "ON COMMIT DROP" in create_statements[0], ( + "CREATE TEMP TABLE must use ON COMMIT DROP so the table is cleaned " + "up at transaction end and is transaction-scoped" + ) + + # Must not use IF NOT EXISTS — the table is fresh each transaction. + assert "IF NOT EXISTS" not in create_statements[0], ( + "With ON COMMIT DROP the table is always fresh at transaction start, " + "so IF NOT EXISTS is both unnecessary and misleading (suggests the " + "table might persist across transactions)" + ) + + @pytest.mark.asyncio + async def test_no_manual_drop_or_truncate(self, mock_conn): + """With ON COMMIT DROP we must not re-add manual TRUNCATE or DROP + statements — they were the source of the original pgBouncer bug.""" + emb = [0.1] * 384 + await compute_semantic_links_ann( + conn=mock_conn, + bank_id="bank-1", + unit_ids=["u1"], + embeddings=[emb], + fact_types=["world"], + ) + + executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list] + assert not any("TRUNCATE _ann_seeds" in s for s in executed_sql), ( + "TRUNCATE is unnecessary with ON COMMIT DROP and was previously " + "the statement that failed with 'relation does not exist' when " + "pgBouncer rebound the backend" + ) + assert not any("DROP TABLE" in s and "_ann_seeds" in s for s in executed_sql), ( + "Explicit DROP is unnecessary with ON COMMIT DROP" + ) + + @pytest.mark.asyncio + async def test_uses_set_local_for_ef_search(self, mock_conn): + """hnsw.ef_search must be set with SET LOCAL so the change is scoped + to the transaction. Without SET LOCAL, the setting would leak onto + the pooled backend and affect subsequent recall queries that land + on the same backend.""" + emb = [0.1] * 384 + await compute_semantic_links_ann( + conn=mock_conn, + bank_id="bank-1", + unit_ids=["u1"], + embeddings=[emb], + fact_types=["world"], + ) + + executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list] + ef_statements = [s for s in executed_sql if "hnsw.ef_search" in s] + assert ef_statements, "ef_search must be tuned down for retain ANN" + for stmt in ef_statements: + assert stmt.strip().startswith("SET LOCAL"), ( + f"hnsw.ef_search must use SET LOCAL, got: {stmt}" + ) + # And there must not be a RESET — SET LOCAL handles it at commit. + assert not any("RESET hnsw.ef_search" in s for s in executed_sql)