fix: per-bank index creation respects HINDSIGHT_API_VECTOR_EXTENSION config (#755)
create_bank_hnsw_indexes() hardcoded USING hnsw regardless of the configured vector extension, causing "column cannot have more than 2000 dimensions for hnsw index" when using pgvectorscale or vchord with high-dimensional embeddings. Now reads get_config().vector_extension and uses the appropriate index type: - pgvector → USING hnsw - pgvectorscale → USING diskann - vchord → USING vchordrq Closes #738
This commit is contained in:
Nicolò Boschi
GitHub
parent
d2965e64e6
commit
6488c9bc77
4 changed files with 61 additions and 45 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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="")
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue