From b09eca16987284601bb5fad4cc0b29ad9d6b3154 Mon Sep 17 00:00:00 2001 From: phernandez Date: Wed, 25 Feb 2026 11:44:50 -0600 Subject: [PATCH] feat: add EmbeddingStatus schema and get_embedding_status() service method MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add EmbeddingStatus model to project_info schemas and wire it into ProjectInfoResponse. ProjectService.get_embedding_status() queries vector tables for chunk/embedding counts, detects orphaned chunks and missing embeddings, and recommends reindex when appropriate. Handles both SQLite and Postgres backends. 🔍 Includes 6 unit tests covering: disabled search, missing vector tables, entities without chunks, orphaned chunks, healthy state, and integration with get_project_info(). Co-Authored-By: Claude Opus 4.6 Signed-off-by: phernandez --- src/basic_memory/schemas/__init__.py | 2 + src/basic_memory/schemas/project_info.py | 27 ++ src/basic_memory/services/project_service.py | 162 +++++++++++ .../test_project_service_embedding_status.py | 267 ++++++++++++++++++ 4 files changed, 458 insertions(+) create mode 100644 tests/services/test_project_service_embedding_status.py diff --git a/src/basic_memory/schemas/__init__.py b/src/basic_memory/schemas/__init__.py index 6d4feecc..c821bf9f 100644 --- a/src/basic_memory/schemas/__init__.py +++ b/src/basic_memory/schemas/__init__.py @@ -41,6 +41,7 @@ from basic_memory.schemas.project_info import ( ProjectStatistics, ActivityMetrics, SystemStatus, + EmbeddingStatus, ProjectInfoResponse, ) @@ -78,6 +79,7 @@ __all__ = [ "ProjectStatistics", "ActivityMetrics", "SystemStatus", + "EmbeddingStatus", "ProjectInfoResponse", # Directory "DirectoryNode", diff --git a/src/basic_memory/schemas/project_info.py b/src/basic_memory/schemas/project_info.py index 80dd8122..c71def13 100644 --- a/src/basic_memory/schemas/project_info.py +++ b/src/basic_memory/schemas/project_info.py @@ -79,6 +79,28 @@ class SystemStatus(BaseModel): timestamp: datetime = Field(description="Timestamp when the information was collected") +class EmbeddingStatus(BaseModel): + """Embedding/vector index status for a project.""" + + # Config + semantic_search_enabled: bool + embedding_provider: Optional[str] = None + embedding_model: Optional[str] = None + embedding_dimensions: Optional[int] = None + + # Counts + total_indexed_entities: int = 0 + total_entities_with_chunks: int = 0 + total_chunks: int = 0 + total_embeddings: int = 0 + orphaned_chunks: int = 0 + vector_tables_exist: bool = False + + # Derived + reindex_recommended: bool = False + reindex_reason: Optional[str] = None + + class ProjectInfoResponse(BaseModel): """Response for the project_info tool.""" @@ -99,6 +121,11 @@ class ProjectInfoResponse(BaseModel): # System status system: SystemStatus = Field(description="System and service status information") + # Embedding status + embedding_status: Optional[EmbeddingStatus] = Field( + default=None, description="Embedding/vector index status" + ) + class ProjectInfoRequest(BaseModel): """Request model for switching projects.""" diff --git a/src/basic_memory/services/project_service.py b/src/basic_memory/services/project_service.py index 36d072cb..a5358405 100644 --- a/src/basic_memory/services/project_service.py +++ b/src/basic_memory/services/project_service.py @@ -16,6 +16,7 @@ from basic_memory.models import Project from basic_memory.repository.project_repository import ProjectRepository from basic_memory.schemas import ( ActivityMetrics, + EmbeddingStatus, ProjectInfoResponse, ProjectStatistics, SystemStatus, @@ -597,6 +598,9 @@ class ProjectService: # Get activity metrics for the specified project activity = await self.get_activity_metrics(db_project.id) + # Get embedding status for the specified project + embedding_status = await self.get_embedding_status(db_project.id) + # Get system status system = self.get_system_status() @@ -650,6 +654,7 @@ class ProjectService: statistics=statistics, activity=activity, system=system, + embedding_status=embedding_status, ) async def get_statistics(self, project_id: int) -> ProjectStatistics: @@ -918,6 +923,163 @@ class ProjectService: monthly_growth=monthly_growth, ) + async def get_embedding_status(self, project_id: int) -> EmbeddingStatus: + """Get embedding/vector index status for the specified project. + + Reports config, counts, and whether a reindex is recommended. + """ + config = self.config_manager.config + semantic_enabled = config.semantic_search_enabled + + # When semantic search is disabled, return minimal status + if not semantic_enabled: + return EmbeddingStatus(semantic_search_enabled=False) + + provider = config.semantic_embedding_provider + model = config.semantic_embedding_model + dimensions = config.semantic_embedding_dimensions + + is_postgres = config.database_backend == DatabaseBackend.POSTGRES + + # --- Check vector table existence --- + if is_postgres: + table_check_sql = text( + "SELECT COUNT(*) FROM information_schema.tables " + "WHERE table_name = 'search_vector_chunks'" + ) + else: + table_check_sql = text( + "SELECT COUNT(*) FROM sqlite_master " + "WHERE type = 'table' AND name = 'search_vector_chunks'" + ) + + table_result = await self.repository.execute_query(table_check_sql, {}) + vector_tables_exist = (table_result.scalar() or 0) > 0 + + if not vector_tables_exist: + # Count distinct entities in search index for the recommendation message + si_result = await self.repository.execute_query( + text( + "SELECT COUNT(DISTINCT entity_id) FROM search_index " + "WHERE project_id = :project_id" + ), + {"project_id": project_id}, + ) + total_indexed_entities = si_result.scalar() or 0 + + return EmbeddingStatus( + semantic_search_enabled=True, + embedding_provider=provider, + embedding_model=model, + embedding_dimensions=dimensions, + total_indexed_entities=total_indexed_entities, + vector_tables_exist=False, + reindex_recommended=True, + reindex_reason=( + "Vector tables not initialized — run: bm reindex --embeddings" + ), + ) + + # --- Count queries (tables exist) --- + si_result = await self.repository.execute_query( + text( + "SELECT COUNT(DISTINCT entity_id) FROM search_index " + "WHERE project_id = :project_id" + ), + {"project_id": project_id}, + ) + total_indexed_entities = si_result.scalar() or 0 + + chunks_result = await self.repository.execute_query( + text("SELECT COUNT(*) FROM search_vector_chunks WHERE project_id = :project_id"), + {"project_id": project_id}, + ) + total_chunks = chunks_result.scalar() or 0 + + entities_with_chunks_result = await self.repository.execute_query( + text( + "SELECT COUNT(DISTINCT entity_id) FROM search_vector_chunks " + "WHERE project_id = :project_id" + ), + {"project_id": project_id}, + ) + total_entities_with_chunks = entities_with_chunks_result.scalar() or 0 + + # Embeddings count — join pattern differs between SQLite and Postgres + if is_postgres: + embeddings_sql = text( + "SELECT COUNT(*) FROM search_vector_chunks c " + "JOIN search_vector_embeddings e ON e.chunk_id = c.id " + "WHERE c.project_id = :project_id" + ) + else: + embeddings_sql = text( + "SELECT COUNT(*) FROM search_vector_chunks c " + "JOIN search_vector_embeddings e ON e.rowid = c.id " + "WHERE c.project_id = :project_id" + ) + + embeddings_result = await self.repository.execute_query( + embeddings_sql, {"project_id": project_id} + ) + total_embeddings = embeddings_result.scalar() or 0 + + # Orphaned chunks (chunks without embeddings — indicates interrupted indexing) + if is_postgres: + orphan_sql = text( + "SELECT COUNT(*) FROM search_vector_chunks c " + "LEFT JOIN search_vector_embeddings e ON e.chunk_id = c.id " + "WHERE c.project_id = :project_id AND e.chunk_id IS NULL" + ) + else: + orphan_sql = text( + "SELECT COUNT(*) FROM search_vector_chunks c " + "LEFT JOIN search_vector_embeddings e ON e.rowid = c.id " + "WHERE c.project_id = :project_id AND e.rowid IS NULL" + ) + + orphan_result = await self.repository.execute_query( + orphan_sql, {"project_id": project_id} + ) + orphaned_chunks = orphan_result.scalar() or 0 + + # --- Reindex recommendation logic (priority order) --- + reindex_recommended = False + reindex_reason = None + + if total_indexed_entities > 0 and total_chunks == 0: + reindex_recommended = True + reindex_reason = ( + "Embeddings have never been built — run: bm reindex --embeddings" + ) + elif orphaned_chunks > 0: + reindex_recommended = True + reindex_reason = ( + f"{orphaned_chunks} orphaned chunks found (interrupted indexing) " + "— run: bm reindex --embeddings" + ) + elif total_indexed_entities > total_entities_with_chunks: + missing = total_indexed_entities - total_entities_with_chunks + reindex_recommended = True + reindex_reason = ( + f"{missing} entities missing embeddings — run: bm reindex --embeddings" + ) + + return EmbeddingStatus( + semantic_search_enabled=True, + embedding_provider=provider, + embedding_model=model, + embedding_dimensions=dimensions, + total_indexed_entities=total_indexed_entities, + total_entities_with_chunks=total_entities_with_chunks, + total_chunks=total_chunks, + total_embeddings=total_embeddings, + orphaned_chunks=orphaned_chunks, + vector_tables_exist=True, + reindex_recommended=reindex_recommended, + reindex_reason=reindex_reason, + ) + def get_system_status(self) -> SystemStatus: """Get system status information.""" import basic_memory diff --git a/tests/services/test_project_service_embedding_status.py b/tests/services/test_project_service_embedding_status.py new file mode 100644 index 00000000..fce57397 --- /dev/null +++ b/tests/services/test_project_service_embedding_status.py @@ -0,0 +1,267 @@ +"""Tests for ProjectService.get_embedding_status().""" + +from unittest.mock import patch + +import pytest + +from basic_memory.schemas.project_info import EmbeddingStatus +from basic_memory.services.project_service import ProjectService + + +@pytest.mark.asyncio +async def test_embedding_status_semantic_disabled(project_service: ProjectService, test_project): + """When semantic search is disabled, return minimal status with zero counts.""" + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=False) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert isinstance(status, EmbeddingStatus) + assert status.semantic_search_enabled is False + assert status.reindex_recommended is False + assert status.total_chunks == 0 + assert status.total_embeddings == 0 + + +@pytest.mark.asyncio +async def test_embedding_status_vector_tables_missing( + project_service: ProjectService, test_graph, test_project +): + """When vector tables don't exist, recommend reindex.""" + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + # Vector tables are not created by the standard test fixtures + # If they don't exist, status should flag it + assert status.semantic_search_enabled is True + assert status.embedding_provider == "fastembed" + assert status.embedding_model == "bge-small-en-v1.5" + + if not status.vector_tables_exist: + assert status.reindex_recommended is True + assert "Vector tables not initialized" in (status.reindex_reason or "") + + +@pytest.mark.asyncio +async def test_embedding_status_entities_without_chunks( + project_service: ProjectService, test_graph, test_project +): + """When entities have search_index rows but no chunks, recommend reindex.""" + # Create vector tables (empty) so the table-existence check passes + from sqlalchemy import text + + await project_service.repository.execute_query( + text( + "CREATE TABLE IF NOT EXISTS search_vector_chunks (" + " id INTEGER PRIMARY KEY AUTOINCREMENT," + " entity_id INTEGER NOT NULL," + " project_id INTEGER NOT NULL," + " chunk_key TEXT NOT NULL," + " chunk_text TEXT NOT NULL," + " source_hash TEXT NOT NULL," + " updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP" + ")" + ), + {}, + ) + + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert status.semantic_search_enabled is True + assert status.vector_tables_exist is True + # test_graph creates entities indexed in search_index, but no vector chunks + assert status.total_indexed_entities > 0 + assert status.total_chunks == 0 + assert status.reindex_recommended is True + assert "never been built" in (status.reindex_reason or "") + + +@pytest.mark.asyncio +async def test_embedding_status_orphaned_chunks( + project_service: ProjectService, test_graph, test_project +): + """When chunks exist without matching embeddings, recommend reindex.""" + from sqlalchemy import text + + # Create vector tables + await project_service.repository.execute_query( + text( + "CREATE TABLE IF NOT EXISTS search_vector_chunks (" + " id INTEGER PRIMARY KEY AUTOINCREMENT," + " entity_id INTEGER NOT NULL," + " project_id INTEGER NOT NULL," + " chunk_key TEXT NOT NULL," + " chunk_text TEXT NOT NULL," + " source_hash TEXT NOT NULL," + " updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP" + ")" + ), + {}, + ) + + # Insert a chunk row (no matching embedding = orphan) + # Get a real entity_id from the test graph + entity_result = await project_service.repository.execute_query( + text("SELECT id FROM entity WHERE project_id = :project_id LIMIT 1"), + {"project_id": test_project.id}, + ) + entity_id = entity_result.scalar() + + await project_service.repository.execute_query( + text( + "INSERT INTO search_vector_chunks " + "(entity_id, project_id, chunk_key, chunk_text, source_hash) " + "VALUES (:entity_id, :project_id, 'chunk-1', 'test text', 'abc123')" + ), + {"entity_id": entity_id, "project_id": test_project.id}, + ) + + # Create a minimal search_vector_embeddings table (not sqlite-vec virtual table) + # so the LEFT JOIN works and finds the orphan + await project_service.repository.execute_query( + text( + "CREATE TABLE IF NOT EXISTS search_vector_embeddings (" + " rowid INTEGER PRIMARY KEY" + ")" + ), + {}, + ) + + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert status.vector_tables_exist is True + assert status.total_chunks == 1 + assert status.orphaned_chunks == 1 + assert status.reindex_recommended is True + assert "orphaned chunks" in (status.reindex_reason or "") + + +@pytest.mark.asyncio +async def test_embedding_status_healthy( + project_service: ProjectService, test_graph, test_project +): + """When all entities have embeddings, no reindex recommended.""" + from sqlalchemy import text + + # Create vector chunks table + await project_service.repository.execute_query( + text( + "CREATE TABLE IF NOT EXISTS search_vector_chunks (" + " id INTEGER PRIMARY KEY AUTOINCREMENT," + " entity_id INTEGER NOT NULL," + " project_id INTEGER NOT NULL," + " chunk_key TEXT NOT NULL," + " chunk_text TEXT NOT NULL," + " source_hash TEXT NOT NULL," + " updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP" + ")" + ), + {}, + ) + + # Drop any existing virtual table (may have been created by search_service init) + # and recreate as a simple regular table for testing the join logic + await project_service.repository.execute_query( + text("DROP TABLE IF EXISTS search_vector_embeddings"), {} + ) + await project_service.repository.execute_query( + text( + "CREATE TABLE search_vector_embeddings (" + " rowid INTEGER PRIMARY KEY" + ")" + ), + {}, + ) + + # Insert a chunk + matching embedding for every search_index entity + entity_result = await project_service.repository.execute_query( + text( + "SELECT DISTINCT entity_id FROM search_index WHERE project_id = :project_id" + ), + {"project_id": test_project.id}, + ) + entity_ids = [row[0] for row in entity_result.fetchall()] + + chunk_id = 1 + for eid in entity_ids: + await project_service.repository.execute_query( + text( + "INSERT INTO search_vector_chunks " + "(id, entity_id, project_id, chunk_key, chunk_text, source_hash) " + "VALUES (:id, :entity_id, :project_id, :key, 'text', 'hash')" + ), + { + "id": chunk_id, + "entity_id": eid, + "project_id": test_project.id, + "key": f"chunk-{chunk_id}", + }, + ) + await project_service.repository.execute_query( + text("INSERT INTO search_vector_embeddings (rowid) VALUES (:rowid)"), + {"rowid": chunk_id}, + ) + chunk_id += 1 + + with patch.object( + type(project_service), + "config_manager", + new_callable=lambda: property( + lambda self: _config_manager_with(semantic_search_enabled=True) + ), + ): + status = await project_service.get_embedding_status(test_project.id) + + assert status.vector_tables_exist is True + assert status.total_chunks > 0 + assert status.total_embeddings == status.total_chunks + assert status.orphaned_chunks == 0 + assert status.reindex_recommended is False + assert status.reindex_reason is None + + +@pytest.mark.asyncio +async def test_get_project_info_includes_embedding_status( + project_service: ProjectService, test_graph, test_project +): + """get_project_info() response includes embedding_status field.""" + info = await project_service.get_project_info(test_project.name) + assert info.embedding_status is not None + assert isinstance(info.embedding_status, EmbeddingStatus) + + +# --- Helper --- + + +def _config_manager_with(semantic_search_enabled: bool): + """Create a ConfigManager whose config has the given semantic_search_enabled value.""" + from basic_memory.config import ConfigManager + + cm = ConfigManager() + # Patch the config object in-place + cm.config.semantic_search_enabled = semantic_search_enabled + return cm