diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index 1a2a5420..db1aacbe 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -3896,11 +3896,11 @@ class MemoryEngine(MemoryEngineInterface): except Exception as e: raise Exception(f"Failed to delete agent data: {str(e)}") - # Drop per-bank HNSW indexes AFTER the transaction commits to avoid + # Drop per-bank vector indexes AFTER the transaction commits to avoid # AccessExclusiveLock deadlocks with concurrent bank deletions. # (DROP INDEX on memory_units conflicts with RowExclusiveLock from DELETE inside tx) if bank_internal_id: - await bank_utils.drop_bank_hnsw_indexes(conn, bank_internal_id) + await bank_utils.drop_bank_vector_indexes(conn, bank_internal_id) if invalidated_obs > 0: await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py index b1c98774..7d203f00 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/bank_utils.py @@ -10,32 +10,47 @@ from typing import TypedDict from pydantic import BaseModel, Field +from ...config import get_config from ..db_utils import acquire_with_retry from ..memory_engine import fq_table, get_current_schema from ..response_models import DispositionTraits logger = logging.getLogger(__name__) -# Fact types that get per-bank partial HNSW indexes, mapped to their 4-char index suffix. -_HNSW_FACT_TYPES: dict[str, str] = { +# Fact types that get per-bank partial vector indexes, mapped to their 4-char index suffix. +_BANK_INDEX_FACT_TYPES: dict[str, str] = { "world": "worl", "experience": "expr", "observation": "obsv", } -def _hnsw_index_name(ft: str, internal_id: str) -> str: - """Deterministic, schema-safe HNSW index name for a (bank, fact_type) pair. +def _bank_index_name(ft: str, internal_id: str) -> str: + """Deterministic, schema-safe vector index name for a (bank, fact_type) pair. Uses the first 16 hex chars of internal_id (8 bytes of entropy) — unique enough in practice, fits comfortably within PostgreSQL's 63-char identifier limit. """ uid = str(internal_id).replace("-", "")[:16] - return f"idx_mu_emb_{_HNSW_FACT_TYPES[ft]}_{uid}" + return f"idx_mu_emb_{_BANK_INDEX_FACT_TYPES[ft]}_{uid}" -async def create_bank_hnsw_indexes(conn, bank_id: str, internal_id: str) -> None: - """Create per-(bank, fact_type) partial HNSW indexes for a newly created bank. +def _vector_index_clause() -> str: + """Return the USING clause for vector index creation based on the configured extension.""" + ext = get_config().vector_extension + if ext == "pgvectorscale": + return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)" + elif ext == "vchord": + return "USING vchordrq (embedding vector_l2_ops)" + else: # pgvector (default) + return "USING hnsw (embedding vector_cosine_ops)" + + +async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str) -> None: + """Create per-(bank, fact_type) partial vector indexes for a newly created bank. + + Respects the HINDSIGHT_API_VECTOR_EXTENSION config to use the appropriate + index type (HNSW for pgvector, DiskANN for pgvectorscale, vchordrq for vchord). Called immediately after the bank row is first inserted. Safe on empty banks (index build is instant). Idempotent via CREATE INDEX IF NOT EXISTS. @@ -43,24 +58,25 @@ async def create_bank_hnsw_indexes(conn, bank_id: str, internal_id: str) -> None """ table = fq_table("memory_units") escaped = bank_id.replace("'", "''") - for ft in _HNSW_FACT_TYPES: - idx = _hnsw_index_name(ft, internal_id) + using_clause = _vector_index_clause() + for ft in _BANK_INDEX_FACT_TYPES: + idx = _bank_index_name(ft, internal_id) await conn.execute( f"CREATE INDEX IF NOT EXISTS {idx} " - f"ON {table} USING hnsw (embedding vector_cosine_ops) " + f"ON {table} {using_clause} " f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'" ) -async def drop_bank_hnsw_indexes(conn, internal_id: str) -> None: - """Drop per-(bank, fact_type) partial HNSW indexes for a bank being deleted. +async def drop_bank_vector_indexes(conn, internal_id: str) -> None: + """Drop per-(bank, fact_type) partial vector indexes for a bank being deleted. Called before the bank row is deleted so internal_id is still known. Idempotent via DROP INDEX IF EXISTS. """ schema = get_current_schema() - for ft in _HNSW_FACT_TYPES: - idx = _hnsw_index_name(ft, internal_id) + for ft in _BANK_INDEX_FACT_TYPES: + idx = _bank_index_name(ft, internal_id) await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}") @@ -138,8 +154,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile: ) if inserted: - # Fresh insert — create per-bank HNSW indexes (instant on empty bank) - await create_bank_hnsw_indexes(conn, bank_id, str(internal_id)) + # Fresh insert — create per-bank vector indexes (instant on empty bank) + await create_bank_vector_indexes(conn, bank_id, str(internal_id)) return BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission="") diff --git a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py index 43b9e93c..a98181f7 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py @@ -10,7 +10,7 @@ import uuid from ...config import get_config from ..memory_engine import fq_table -from .bank_utils import DEFAULT_DISPOSITION, create_bank_hnsw_indexes +from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes from .fact_extraction import _sanitize_text from .types import ProcessedFact @@ -207,8 +207,8 @@ async def ensure_bank_exists(conn, bank_id: str) -> None: internal_id, ) if inserted: - # Fresh insert — create per-bank HNSW indexes - await create_bank_hnsw_indexes(conn, bank_id, str(internal_id)) + # Fresh insert — create per-bank vector indexes + await create_bank_vector_indexes(conn, bank_id, str(internal_id)) async def handle_document_tracking( diff --git a/hindsight-api-slim/tests/test_hnsw_indexes.py b/hindsight-api-slim/tests/test_hnsw_indexes.py index a84b6be0..8a94b03e 100644 --- a/hindsight-api-slim/tests/test_hnsw_indexes.py +++ b/hindsight-api-slim/tests/test_hnsw_indexes.py @@ -1,10 +1,10 @@ """ -Tests for per-bank HNSW index lifecycle and UNION ALL retrieval. +Tests for per-bank vector index lifecycle and UNION ALL retrieval. Covers: -- _hnsw_index_name deterministic naming -- Per-bank HNSW indexes created on bank creation (retain_async / ensure_bank_exists) -- Per-bank HNSW indexes dropped on bank deletion +- _bank_index_name deterministic naming +- Per-bank vector indexes created on bank creation (retain_async / ensure_bank_exists) +- Per-bank vector indexes dropped on bank deletion - retrieve_semantic_bm25_combined groups results correctly by fact_type and source """ import uuid @@ -12,7 +12,7 @@ from datetime import datetime, timezone import pytest -from hindsight_api.engine.retain.bank_utils import _HNSW_FACT_TYPES, _hnsw_index_name +from hindsight_api.engine.retain.bank_utils import _BANK_INDEX_FACT_TYPES, _bank_index_name # --------------------------------------------------------------------------- @@ -20,36 +20,36 @@ from hindsight_api.engine.retain.bank_utils import _HNSW_FACT_TYPES, _hnsw_index # --------------------------------------------------------------------------- -class TestHnswIndexName: +class TestBankIndexName: def test_deterministic(self): uid = "550e8400-e29b-41d4-a716-446655440000" - assert _hnsw_index_name("world", uid) == _hnsw_index_name("world", uid) + assert _bank_index_name("world", uid) == _bank_index_name("world", uid) def test_strips_dashes(self): uid = "550e8400-e29b-41d4-a716-446655440000" - name = _hnsw_index_name("world", uid) + name = _bank_index_name("world", uid) # uid16 should be hex chars only assert "-" not in name def test_uses_first_16_hex_chars(self): uid = "550e8400-e29b-41d4-a716-446655440000" uid16 = uid.replace("-", "")[:16] # "550e8400e29b41d4" - assert name_ends_with(name=_hnsw_index_name("world", uid), suffix=uid16) + assert name_ends_with(name=_bank_index_name("world", uid), suffix=uid16) def test_suffix_per_fact_type(self): uid = "550e8400-e29b-41d4-a716-446655440000" - names = {ft: _hnsw_index_name(ft, uid) for ft in _HNSW_FACT_TYPES} + names = {ft: _bank_index_name(ft, uid) for ft in _BANK_INDEX_FACT_TYPES} # All three names must be distinct assert len(set(names.values())) == 3 def test_all_fact_types_covered(self): - assert set(_HNSW_FACT_TYPES) == {"world", "experience", "observation"} + assert set(_BANK_INDEX_FACT_TYPES) == {"world", "experience", "observation"} def test_fits_pg_identifier_limit(self): # PostgreSQL max identifier length is 63 chars uid = "f" * 32 # simulated UUID without dashes - for ft in _HNSW_FACT_TYPES: - assert len(_hnsw_index_name(ft, uid)) <= 63 + for ft in _BANK_INDEX_FACT_TYPES: + assert len(_bank_index_name(ft, uid)) <= 63 def name_ends_with(name: str, suffix: str) -> bool: @@ -61,7 +61,7 @@ def name_ends_with(name: str, suffix: str) -> bool: # --------------------------------------------------------------------------- -async def _get_bank_hnsw_indexes(pool, bank_id: str) -> list[str]: +async def _get_bank_vector_indexes(pool, bank_id: str) -> list[str]: """Return index names for memory_units that match the per-bank pattern.""" async with pool.acquire() as conn: rows = await conn.fetch( @@ -79,8 +79,8 @@ async def _get_bank_hnsw_indexes(pool, bank_id: str) -> list[str]: @pytest.mark.asyncio -async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context): - """retain_async on a new bank must create 3 per-(bank, fact_type) HNSW indexes.""" +async def test_retain_creates_per_bank_vector_indexes(memory, request_context): + """retain_async on a new bank must create 3 per-(bank, fact_type) vector indexes.""" bank_id = f"test_hnsw_create_{uuid.uuid4().hex[:8]}" try: await memory.retain_async( @@ -88,9 +88,9 @@ async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context): content="Alice is a software engineer.", request_context=request_context, ) - indexes = await _get_bank_hnsw_indexes(memory._pool, bank_id) - assert len(indexes) == 3, f"Expected 3 per-bank HNSW indexes, got: {indexes}" - for ft_short in _HNSW_FACT_TYPES.values(): + indexes = await _get_bank_vector_indexes(memory._pool, bank_id) + assert len(indexes) == 3, f"Expected 3 per-bank vector indexes, got: {indexes}" + for ft_short in _BANK_INDEX_FACT_TYPES.values(): assert any(ft_short in idx for idx in indexes), ( f"Missing index for fact_type short '{ft_short}' in {indexes}" ) @@ -99,8 +99,8 @@ async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context): @pytest.mark.asyncio -async def test_delete_bank_drops_hnsw_indexes(memory, request_context): - """delete_bank must drop all per-bank HNSW indexes.""" +async def test_delete_bank_drops_vector_indexes(memory, request_context): + """delete_bank must drop all per-bank vector indexes.""" bank_id = f"test_hnsw_drop_{uuid.uuid4().hex[:8]}" await memory.retain_async( @@ -109,12 +109,12 @@ async def test_delete_bank_drops_hnsw_indexes(memory, request_context): request_context=request_context, ) # Verify indexes exist before deletion - indexes_before = await _get_bank_hnsw_indexes(memory._pool, bank_id) + indexes_before = await _get_bank_vector_indexes(memory._pool, bank_id) assert len(indexes_before) == 3 await memory.delete_bank(bank_id, request_context=request_context) - indexes_after = await _get_bank_hnsw_indexes(memory._pool, bank_id) + indexes_after = await _get_bank_vector_indexes(memory._pool, bank_id) assert indexes_after == [], f"Indexes should be dropped after bank deletion, got: {indexes_after}" @@ -133,7 +133,7 @@ async def test_retain_idempotent_bank_creation(memory, request_context): content="Carol joined the company in 2022.", request_context=request_context, ) - indexes = await _get_bank_hnsw_indexes(memory._pool, bank_id) + indexes = await _get_bank_vector_indexes(memory._pool, bank_id) assert len(indexes) == 3 finally: await memory.delete_bank(bank_id, request_context=request_context)