mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
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:
@@ -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.
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user