mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
fix(core): load sqlite-vec before vector table cleanup (#852)
Signed-off-by: Drew Cain <groksrc@gmail.com>
This commit is contained in:
@@ -620,6 +620,25 @@ class SQLiteSearchRepository(SearchRepositoryBase):
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
async def drop_vector_tables(self) -> None:
|
||||
"""Drop SQLite vector tables on a sqlite-vec-enabled connection."""
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
vector_sql_result = await session.execute(
|
||||
text(
|
||||
"SELECT sql FROM sqlite_master "
|
||||
"WHERE type = 'table' AND name = 'search_vector_embeddings'"
|
||||
)
|
||||
)
|
||||
vector_sql = vector_sql_result.scalar()
|
||||
if vector_sql and "using vec0" in vector_sql.lower():
|
||||
await self._ensure_sqlite_vec_loaded(session)
|
||||
|
||||
await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings"))
|
||||
await session.execute(text("DROP TABLE IF EXISTS search_vector_chunks"))
|
||||
await session.execute(text("DROP TABLE IF EXISTS search_vector_index"))
|
||||
await session.commit()
|
||||
self._vector_tables_initialized = False
|
||||
|
||||
async def delete_stale_vector_rows(self) -> None:
|
||||
"""Delete vector rows whose source entities no longer exist."""
|
||||
await self._ensure_vector_tables()
|
||||
|
||||
@@ -126,19 +126,23 @@ class SearchService:
|
||||
|
||||
async def reindex_all(self, background_tasks: Optional[BackgroundTasks] = None) -> None:
|
||||
"""Reindex all content from database."""
|
||||
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
|
||||
|
||||
logger.info("Starting full reindex")
|
||||
# Clear and recreate search index
|
||||
await self.repository.execute_query(text("DROP TABLE IF EXISTS search_index"), params={})
|
||||
await self.repository.execute_query(
|
||||
text("DROP TABLE IF EXISTS search_vector_embeddings"), params={}
|
||||
)
|
||||
await self.repository.execute_query(
|
||||
text("DROP TABLE IF EXISTS search_vector_chunks"), params={}
|
||||
)
|
||||
await self.repository.execute_query(
|
||||
text("DROP TABLE IF EXISTS search_vector_index"), params={}
|
||||
)
|
||||
if isinstance(self.repository, SQLiteSearchRepository):
|
||||
await self.repository.drop_vector_tables()
|
||||
else:
|
||||
await self.repository.execute_query(
|
||||
text("DROP TABLE IF EXISTS search_vector_embeddings"), params={}
|
||||
)
|
||||
await self.repository.execute_query(
|
||||
text("DROP TABLE IF EXISTS search_vector_chunks"), params={}
|
||||
)
|
||||
await self.repository.execute_query(
|
||||
text("DROP TABLE IF EXISTS search_vector_index"), params={}
|
||||
)
|
||||
await self.init_search_index()
|
||||
|
||||
# Reindex all entities
|
||||
|
||||
@@ -5,7 +5,9 @@ from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from basic_memory import db
|
||||
from basic_memory.repository import EntityRepository
|
||||
from basic_memory.repository.search_repository_base import VectorSyncBatchResult
|
||||
from basic_memory.repository.semantic_errors import (
|
||||
@@ -256,6 +258,98 @@ async def test_reindex_vectors_respects_embed_opt_out(search_service, monkeypatc
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reindex_vectors_purges_sqlite_vectors_before_sync(search_service, monkeypatch):
|
||||
"""Regression for #829: stale vec0 rows must be purged before batch sync."""
|
||||
repository = _sqlite_repo(search_service)
|
||||
repository._semantic_enabled = True
|
||||
calls: list[str] = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
search_service.entity_repository,
|
||||
"find_all",
|
||||
AsyncMock(return_value=[SimpleNamespace(id=42, entity_metadata={})]),
|
||||
)
|
||||
|
||||
async def delete_stale_vector_rows():
|
||||
calls.append("purge")
|
||||
|
||||
async def sync_entity_vectors_batch(entity_ids, progress_callback=None):
|
||||
assert entity_ids == [42]
|
||||
assert progress_callback is None
|
||||
calls.append("sync")
|
||||
return VectorSyncBatchResult(entities_total=1, entities_synced=1, entities_failed=0)
|
||||
|
||||
monkeypatch.setattr(repository, "delete_stale_vector_rows", delete_stale_vector_rows)
|
||||
monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_entity_vectors_batch)
|
||||
|
||||
stats = await search_service.reindex_vectors()
|
||||
|
||||
assert calls == ["purge", "sync"]
|
||||
assert stats == {
|
||||
"total_entities": 1,
|
||||
"embedded": 1,
|
||||
"skipped": 0,
|
||||
"errors": 0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reindex_all_uses_sqlite_vec_aware_drop(search_service, monkeypatch):
|
||||
"""Full service reindex should not drop vec0 tables through a raw connection."""
|
||||
repository = _sqlite_repo(search_service)
|
||||
executed_sql: list[str] = []
|
||||
calls: list[str] = []
|
||||
|
||||
async def execute_query(query, params=None):
|
||||
executed_sql.append(str(query))
|
||||
|
||||
async def drop_vector_tables():
|
||||
calls.append("drop_vector_tables")
|
||||
|
||||
monkeypatch.setattr(repository, "execute_query", execute_query)
|
||||
monkeypatch.setattr(repository, "drop_vector_tables", drop_vector_tables)
|
||||
monkeypatch.setattr(search_service, "init_search_index", AsyncMock())
|
||||
monkeypatch.setattr(search_service.entity_repository, "find_all", AsyncMock(return_value=[]))
|
||||
|
||||
await search_service.reindex_all()
|
||||
|
||||
assert calls == ["drop_vector_tables"]
|
||||
assert all("search_vector_embeddings" not in sql for sql in executed_sql)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drop_vector_tables_skips_sqlite_vec_load_for_plain_table(
|
||||
search_service,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Plain compatibility tables should not require sqlite-vec to be loaded."""
|
||||
repository = _sqlite_repo(search_service)
|
||||
await repository.drop_vector_tables()
|
||||
|
||||
async with db.scoped_session(repository.session_maker) as session:
|
||||
await session.execute(
|
||||
text("CREATE TABLE search_vector_embeddings (rowid INTEGER PRIMARY KEY)")
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
async def fail_if_loaded(_session):
|
||||
raise AssertionError("plain tables should not load sqlite-vec")
|
||||
|
||||
monkeypatch.setattr(repository, "_ensure_sqlite_vec_loaded", fail_if_loaded)
|
||||
|
||||
await repository.drop_vector_tables()
|
||||
|
||||
async with db.scoped_session(repository.session_maker) as session:
|
||||
result = await session.execute(
|
||||
text(
|
||||
"SELECT name FROM sqlite_master "
|
||||
"WHERE type = 'table' AND name = 'search_vector_embeddings'"
|
||||
)
|
||||
)
|
||||
assert result.scalar() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reindex_vectors_force_full_clears_project_vectors_before_resync(
|
||||
search_service, monkeypatch
|
||||
|
||||
Reference in New Issue
Block a user