fix(core): load sqlite-vec for embedding-status query (#901)

Signed-off-by: phernandez <paul@basicmachines.co>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Paul Hernandez
2026-06-07 17:37:27 -05:00
committed by GitHub
parent 816ee85fb9
commit 271c883ea8
4 changed files with 220 additions and 21 deletions
@@ -5,7 +5,7 @@ from typing import Optional, Sequence, Union
from loguru import logger
from sqlalchemy import inspect as sa_inspect, select, text
from sqlalchemy import Executable, inspect as sa_inspect, select, text
from sqlalchemy.exc import NoResultFound, OperationalError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
@@ -258,6 +258,28 @@ class ProjectRepository(Repository[Project]):
logger.debug(f"Deleted Project and search rows for project_id: {entity_id}")
return True
async def scalar_vec_query(
self, query: Executable, params: Optional[dict] = None
) -> Optional[int]:
"""Run a scalar COUNT query that reads the sqlite-vec vec0 table.
Extension loading is per-connection, so the bare pooled session used by
`execute_query` cannot read vec0 virtual tables — SQLite raises
"no such module: vec0". This helper loads sqlite-vec on the session it
opens before running the query, reusing the same loader as project delete.
Returns None when sqlite-vec cannot be loaded on this Python build, so
callers can fall back to the genuinely-missing-dependency path.
"""
async with db.scoped_session(self.session_maker) as session:
# Trigger: query reads the vec0-backed search_vector_embeddings table.
# Why: vec0 modules are only visible on a connection that loaded sqlite-vec.
# Outcome: load the extension here, or signal absence so the caller degrades.
if not await _load_sqlite_vec_on_session(session):
return None
result = await session.execute(query, params)
return result.scalar()
async def update_path(self, project_id: int, new_path: str) -> Optional[Project]:
"""Update project path.
+22 -9
View File
@@ -1047,10 +1047,27 @@ class ProjectService:
f"WHERE c.project_id = :project_id {chunk_entity_exists}"
)
embeddings_result = await self.repository.execute_query(
embeddings_sql, {"project_id": project_id}
)
total_embeddings = embeddings_result.scalar() or 0
# The embeddings/orphan JOINs read search_vector_embeddings, a vec0
# virtual table. On SQLite that table is only visible on a connection
# that loaded sqlite-vec, so route these through scalar_vec_query which
# loads the extension first. Postgres has no per-connection extension
# and uses the bare pooled session.
async def _vec_scalar(vec_sql) -> int:
if is_postgres:
result = await self.repository.execute_query(
vec_sql, {"project_id": project_id}
)
return result.scalar() or 0
count = await self.repository.scalar_vec_query(vec_sql, {"project_id": project_id})
# Trigger: sqlite-vec genuinely can't load on this Python build.
# Why: without the extension the vec0 JOIN can't run at all.
# Outcome: raise the canonical error so the except block emits the
# true "sqlite-vec unavailable" fallback instead of reporting 0.
if count is None:
raise SAOperationalError(str(vec_sql), {}, Exception("no such module: vec0"))
return count
total_embeddings = await _vec_scalar(embeddings_sql)
# Orphaned chunks (chunks without embeddings — indicates interrupted indexing)
if is_postgres:
@@ -1065,11 +1082,7 @@ class ProjectService:
"LEFT JOIN search_vector_embeddings e ON e.rowid = c.id "
f"WHERE c.project_id = :project_id AND e.rowid IS NULL {chunk_entity_exists}"
)
orphan_result = await self.repository.execute_query(
orphan_sql, {"project_id": project_id}
)
orphaned_chunks = orphan_result.scalar() or 0
orphaned_chunks = await _vec_scalar(orphan_sql)
except SAOperationalError as exc:
# Trigger: sqlite_master can list vec0 virtual tables even when sqlite-vec
# is not loaded in the current Python runtime.
+168
View File
@@ -0,0 +1,168 @@
"""Integration regression test for get_embedding_status against a real vec0 table.
Regression for #658: after a successful `bm reindex --embeddings`, `bm project info`
still reported "sqlite-vec is unavailable", "Indexed 0/N", and "Chunks 0", and
recommended an unnecessary reindex.
Root cause: get_embedding_status() ran the vec0 JOIN count queries on a bare pooled
ProjectRepository session that never loaded the sqlite-vec extension, so SQLite raised
"no such module: vec0", which the except block mis-reported as "unavailable".
This test exercises the real failure path: it builds a REAL vec0 virtual table, writes a
real embedding into it via the search repository, then queries get_embedding_status through
a ProjectRepository session that did NOT pre-load the extension (mirroring the bug). The
healthy unit test substitutes a plain regular table for vec0 and therefore does not cover
this path.
"""
import os
import sqlite3
import pytest
from sqlalchemy import text
from basic_memory import db
from basic_memory.config import BasicMemoryConfig, DatabaseBackend
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.project_repository import ProjectRepository
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
from basic_memory.services.project_service import ProjectService
def _is_postgres() -> bool:
return os.environ.get("BASIC_MEMORY_TEST_POSTGRES", "").lower() in ("1", "true", "yes")
def _unit_vector(dimensions: int) -> list[float]:
"""Return a deterministic unit-norm vector for the vec0 embedding column."""
# vec0 stores float[dimensions]; the actual values don't matter for the count
# queries, but using a normalized vector keeps the row well-formed.
vec = [0.0] * dimensions
vec[0] = 1.0
return vec
@pytest.mark.asyncio
async def test_embedding_status_reads_real_vec0_table(engine_factory, test_project, config_manager):
"""get_embedding_status must report a populated real vec0 table as healthy.
Before the fix, the vec0 JOIN ran on a session without sqlite-vec loaded and
raised "no such module: vec0", which the except block mapped to
vector_tables_exist=False + reindex_recommended=True.
"""
# Trigger: Postgres test matrix executes the same suite.
# Why: vec0 + per-connection sqlite-vec loading is SQLite-specific.
# Outcome: keep the regression on the backend that can actually hit this path.
if _is_postgres():
pytest.skip("Real vec0 table handling is SQLite-specific.")
# Trigger: Python build without SQLite extension loading (#711 — python.org
# macOS / some Windows interpreters lack enable_load_extension).
# Why: this test creates a REAL vec0 virtual table during setup, which is
# impossible without loading the sqlite-vec extension.
# Outcome: skip the regression as an environment-capability gap; the codebase
# already degrades gracefully in that scenario (covered by the unit test).
_probe = sqlite3.connect(":memory:")
if not hasattr(_probe, "enable_load_extension"):
_probe.close()
pytest.skip(
"Python build does not support SQLite extension loading — "
"cannot create real vec0 tables"
)
_probe.close()
_engine, session_maker = engine_factory
project_id = test_project.id
# --- Build a REAL vec0 table via the search repository ---
# Semantic enabled with a fastembed provider so _ensure_vector_tables creates
# the vec0-backed search_vector_embeddings table (float[384]).
app_config = BasicMemoryConfig(
env="test",
database_backend=DatabaseBackend.SQLITE,
semantic_search_enabled=True,
)
search_repo = SQLiteSearchRepository(
session_maker,
project_id=project_id,
app_config=app_config,
)
await search_repo._ensure_vector_tables()
dimensions = search_repo._vector_dimensions
# --- Seed a real entity + search_index row so counts are non-zero ---
# Use the repository so model-level defaults (external_id) are applied.
entity_repo = EntityRepository(session_maker, project_id=project_id)
entity = await entity_repo.create(
{
"title": "Vec Note",
"note_type": "note",
"content_type": "text/markdown",
"project_id": project_id,
"permalink": "vec-note",
"file_path": "vec-note.md",
}
)
entity_id = entity.id
async with db.scoped_session(session_maker) as session:
await session.execute(
text(
"INSERT INTO search_index "
"(id, entity_id, project_id, type, title, permalink, content_stems, "
"content_snippet, file_path, metadata) "
"VALUES (:id, :eid, :pid, 'entity', 'Vec Note', 'vec-note', "
"'vec content', 'vec snippet', 'vec-note.md', '{}')"
),
{"id": entity_id, "eid": entity_id, "pid": project_id},
)
await session.commit()
# --- Insert a chunk + a real embedding into the vec0 table ---
# _write_embeddings writes the embedding into the vec0 virtual table keyed by
# rowid == chunk id, exactly like the reindex path.
async with db.scoped_session(session_maker) as session:
await search_repo._ensure_sqlite_vec_loaded(session)
chunk_result = await session.execute(
text(
"INSERT INTO search_vector_chunks "
"(entity_id, project_id, chunk_key, chunk_text, source_hash, "
"entity_fingerprint, embedding_model) "
"VALUES (:eid, :pid, 'chunk-1', 'vec content', 'hash', "
"'fp-hash', 'bge-small-en-v1.5') "
"RETURNING id"
),
{"eid": entity_id, "pid": project_id},
)
chunk_id = chunk_result.scalar_one()
await search_repo._write_embeddings(
session,
[(chunk_id, "vec content")],
[_unit_vector(dimensions)],
)
await session.commit()
# Evict the vec-loaded connection from the pool. sqlite-vec is loaded
# per-connection, so disposing forces get_embedding_status onto a brand-new
# connection that never loaded the extension — exactly the #658 bug condition
# (e.g. a fresh `bm project info` process after `bm reindex --embeddings`).
await _engine.dispose()
# --- Query status through a fresh ProjectRepository (no extension preloaded) ---
project_repository = ProjectRepository(session_maker)
project_service = ProjectService(project_repository)
status = await project_service.get_embedding_status(project_id)
assert status.semantic_search_enabled is True
# The vec0 JOIN must succeed, so the table is reported as present and healthy.
assert status.vector_tables_exist is True
assert status.reindex_recommended is False
assert status.reindex_reason is None
# Counts must reflect the real data, not the false "0" from the unavailable path.
assert status.total_indexed_entities == 1
assert status.total_chunks == 1
assert status.total_entities_with_chunks == 1
assert status.total_embeddings == 1
assert status.orphaned_chunks == 0
@@ -5,7 +5,6 @@ from unittest.mock import patch
import pytest
from sqlalchemy import text
from sqlalchemy.exc import OperationalError as SAOperationalError
from basic_memory.schemas.project_info import EmbeddingStatus
from basic_memory.services.project_service import ProjectService
@@ -153,20 +152,17 @@ async def test_embedding_status_orphaned_chunks(
async def test_embedding_status_handles_sqlite_vec_unavailable(
project_service: ProjectService, test_graph, test_project
):
"""Unreadable vec0 tables should degrade to unavailable status instead of crashing."""
"""When sqlite-vec can't load at all, degrade to unavailable status instead of crashing."""
# Trigger: Postgres test matrix executes the same unit suite.
# Why: sqlite-vec loading failures are specific to SQLite virtual tables, not Postgres joins.
# Outcome: keep the regression focused on the backend that can actually hit this path.
if _is_postgres():
pytest.skip("sqlite-vec unavailable handling is SQLite-specific.")
original_execute_query = project_service.repository.execute_query
async def _execute_query_with_vec0_failure(query, params):
query_text = str(query)
if "JOIN search_vector_embeddings" in query_text:
raise SAOperationalError(query_text, params, Exception("no such module: vec0"))
return await original_execute_query(query, params)
# scalar_vec_query returns None when the extension can't be loaded on this
# Python build (e.g. the python.org macOS interpreter). Simulate that here.
async def _vec_query_unavailable(query, params=None):
return None
with patch.object(
type(project_service),
@@ -177,8 +173,8 @@ async def test_embedding_status_handles_sqlite_vec_unavailable(
):
with patch.object(
project_service.repository,
"execute_query",
side_effect=_execute_query_with_vec0_failure,
"scalar_vec_query",
side_effect=_vec_query_unavailable,
):
status = await project_service.get_embedding_status(test_project.id)