"""Tests for shared semantic search logic in SearchRepositoryBase. Covers: _compose_row_source_text, _split_text_into_chunks, _build_chunk_records, _search_hybrid entity_id fusion key, and SemanticSearchDisabledError in SQLite. """ from types import SimpleNamespace import pytest from basic_memory.repository.search_repository_base import ( MAX_VECTOR_CHUNK_CHARS, SearchRepositoryBase, _PreparedEntityVectorSync, ) from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.repository.semantic_errors import SemanticSearchDisabledError from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode # --- Helpers --- def _make_row( *, row_type: str, title: str = "Test Title", permalink: str = "test/permalink", content_stems: str = "", content_snippet: str = "", category: str = "", relation_type: str = "", row_id: int = 1, ): """Create a SimpleNamespace mimicking a search_index DB row.""" return SimpleNamespace( id=row_id, type=row_type, title=title, permalink=permalink, content_stems=content_stems, content_snippet=content_snippet, category=category, relation_type=relation_type, ) class _ConcreteRepo(SearchRepositoryBase): """Minimal concrete subclass for testing base class methods.""" _semantic_enabled = False _semantic_vector_k = 100 _embedding_provider = None _semantic_embedding_sync_batch_size = 64 _vector_dimensions = 4 _vector_tables_initialized = False def __init__(self): # Bypass parent __init__ since we don't need a real session_maker for unit tests self.session_maker = None self.project_id = 1 async def init_search_index(self): pass def _prepare_search_term(self, term, is_prefix=True): return term async def search(self, **kwargs): return [] async def _ensure_vector_tables(self): pass async def _run_vector_query(self, session, query_embedding, candidate_limit): return [] async def _write_embeddings(self, session, jobs, embeddings): pass async def _delete_entity_chunks(self, session, entity_id): pass async def _delete_stale_chunks(self, session, stale_ids, entity_id): pass async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" def _distance_to_similarity(self, distance: float) -> float: return 1.0 / (1.0 + max(distance, 0.0)) # --- _compose_row_source_text --- class TestComposeRowSourceText: """Verify _compose_row_source_text produces correct text for all row types.""" def setup_method(self): self.repo = _ConcreteRepo() def test_entity_row_uses_content_snippet_not_content_stems(self): """Entity rows should use content_snippet (human-readable) instead of content_stems.""" row = _make_row( row_type=SearchItemType.ENTITY.value, title="Auth Design", permalink="specs/auth-design", content_stems="auth login token session stems expanded variants", content_snippet="JWT authentication with session management", ) result = self.repo._compose_row_source_text(row) assert "Auth Design" in result assert "specs/auth-design" in result assert "JWT authentication with session management" in result # content_stems should NOT appear assert "stems expanded variants" not in result def test_observation_row_includes_category(self): row = _make_row( row_type=SearchItemType.OBSERVATION.value, title="Coffee Notes", permalink="notes/coffee", category="technique", content_snippet="Pour over produces cleaner cups", ) result = self.repo._compose_row_source_text(row) assert "Coffee Notes" in result assert "technique" in result assert "Pour over produces cleaner cups" in result def test_relation_row_includes_relation_type(self): row = _make_row( row_type=SearchItemType.RELATION.value, title="Brewing", permalink="notes/brewing", relation_type="relates_to", content_snippet="Coffee brewing method", ) result = self.repo._compose_row_source_text(row) assert "Brewing" in result assert "relates_to" in result assert "Coffee brewing method" in result def test_entity_row_with_none_fields(self): """Null fields should be skipped, not included as empty strings.""" row = _make_row( row_type=SearchItemType.ENTITY.value, title="Minimal", permalink="", content_snippet="", ) row.permalink = None row.content_snippet = None result = self.repo._compose_row_source_text(row) assert result == "Minimal" # --- _split_text_into_chunks --- class TestSplitTextIntoChunks: """Verify markdown-aware text splitting.""" def setup_method(self): self.repo = _ConcreteRepo() def test_short_text_returns_single_chunk(self): result = self.repo._split_text_into_chunks("Short text") assert result == ["Short text"] def test_empty_text_returns_empty(self): assert self.repo._split_text_into_chunks("") == [] assert self.repo._split_text_into_chunks(" ") == [] def test_splits_on_headers(self): # Make it long enough to actually split long_section_a = "## Section A\n" + ("A content. " * 100) long_section_b = "## Section B\n" + ("B content. " * 100) long_text = f"Intro paragraph\n\n{long_section_a}\n\n{long_section_b}" result = self.repo._split_text_into_chunks(long_text) assert len(result) >= 2 def test_paragraph_merging_within_limit(self): """Paragraphs shorter than limit should be merged.""" para1 = "First paragraph." para2 = "Second paragraph." text = f"# Header\n\n{para1}\n\n{para2}" result = self.repo._split_text_into_chunks(text) # Both paragraphs fit in one chunk assert len(result) == 1 assert para1 in result[0] assert para2 in result[0] def test_long_paragraph_uses_char_window(self): """A single paragraph longer than MAX_VECTOR_CHUNK_CHARS uses sliding window.""" long_para = "x" * (MAX_VECTOR_CHUNK_CHARS * 3) result = self.repo._split_text_into_chunks(long_para) assert len(result) >= 3 for chunk in result: assert len(chunk) <= MAX_VECTOR_CHUNK_CHARS # --- _build_chunk_records --- class TestBuildChunkRecords: def setup_method(self): self.repo = _ConcreteRepo() def test_produces_records_with_correct_keys(self): rows = [ _make_row( row_type=SearchItemType.ENTITY.value, title="Test", permalink="test", content_snippet="content", row_id=42, ) ] records = self.repo._build_chunk_records(rows) assert len(records) >= 1 for record in records: assert "chunk_key" in record assert "chunk_text" in record assert "source_hash" in record assert record["chunk_key"].startswith("entity:") def test_chunk_key_includes_row_id(self): rows = [ _make_row( row_type=SearchItemType.OBSERVATION.value, content_snippet="obs content", row_id=99, ) ] records = self.repo._build_chunk_records(rows) assert any("99" in r["chunk_key"] for r in records) def test_duplicate_rows_collapse_to_unique_chunk_keys(self): rows = [ _make_row( row_type=SearchItemType.ENTITY.value, title="Spec", permalink="spec", content_snippet="shared content", row_id=77, ), _make_row( row_type=SearchItemType.ENTITY.value, title="Spec", permalink="spec", content_snippet="shared content", row_id=77, ), ] records = self.repo._build_chunk_records(rows) assert len(records) == 1 assert records[0]["chunk_key"] == "entity:77:0" # --- SQLite SemanticSearchDisabledError --- @pytest.mark.asyncio async def test_sqlite_vector_search_raises_disabled_error(search_repository): """Vector mode on SQLite should raise SemanticSearchDisabledError when disabled.""" if not isinstance(search_repository, SQLiteSearchRepository): pytest.skip("SQLite-specific test.") search_repository._semantic_enabled = False with pytest.raises(SemanticSearchDisabledError): await search_repository.search( search_text="test query", retrieval_mode=SearchRetrievalMode.VECTOR, limit=5, offset=0, ) @pytest.mark.asyncio async def test_sqlite_hybrid_search_raises_disabled_error(search_repository): """Hybrid mode on SQLite should raise SemanticSearchDisabledError when disabled.""" if not isinstance(search_repository, SQLiteSearchRepository): pytest.skip("SQLite-specific test.") search_repository._semantic_enabled = False with pytest.raises(SemanticSearchDisabledError): await search_repository.search( search_text="test query", retrieval_mode=SearchRetrievalMode.HYBRID, limit=5, offset=0, ) @pytest.mark.asyncio async def test_sync_entity_vectors_batch_flushes_at_configured_threshold(monkeypatch): """Batch sync should flush queued jobs at semantic_embedding_sync_batch_size boundaries.""" repo = _ConcreteRepo() repo._semantic_enabled = True repo._embedding_provider = object() repo._semantic_embedding_sync_batch_size = 2 prepared_by_entity = { 1: _PreparedEntityVectorSync(1, 1.0, 1, [(101, "chunk-1")]), 2: _PreparedEntityVectorSync(2, 2.0, 1, [(102, "chunk-2")]), 3: _PreparedEntityVectorSync(3, 3.0, 1, [(103, "chunk-3")]), } flush_sizes: list[int] = [] async def _stub_prepare(entity_id: int) -> _PreparedEntityVectorSync: return prepared_by_entity[entity_id] async def _stub_flush(flush_jobs, entity_runtime, synced_entity_ids): flush_sizes.append(len(flush_jobs)) for job in flush_jobs: runtime = entity_runtime[job.entity_id] runtime.remaining_jobs -= 1 if runtime.remaining_jobs <= 0: synced_entity_ids.add(job.entity_id) entity_runtime.pop(job.entity_id, None) return (0.1, 0.2) monkeypatch.setattr(repo, "_prepare_entity_vector_jobs", _stub_prepare) monkeypatch.setattr(repo, "_flush_embedding_jobs", _stub_flush) result = await repo.sync_entity_vectors_batch([1, 2, 3]) assert flush_sizes == [2, 1] assert result.entities_total == 3 assert result.entities_synced == 3 assert result.entities_failed == 0 assert result.failed_entity_ids == [] assert result.embedding_jobs_total == 3 assert result.embed_seconds_total == pytest.approx(0.2) assert result.write_seconds_total == pytest.approx(0.4) @pytest.mark.asyncio async def test_sync_entity_vectors_batch_continue_on_error(monkeypatch): """Batch sync should continue after per-entity and per-flush failures.""" repo = _ConcreteRepo() repo._semantic_enabled = True repo._embedding_provider = object() repo._semantic_embedding_sync_batch_size = 1 async def _stub_prepare(entity_id: int) -> _PreparedEntityVectorSync: if entity_id == 2: raise RuntimeError("prepare failed") return _PreparedEntityVectorSync( entity_id, float(entity_id), 1, [(100 + entity_id, "chunk")] ) async def _stub_flush(flush_jobs, entity_runtime, synced_entity_ids): entity_id = flush_jobs[0].entity_id if entity_id == 3: raise RuntimeError("flush failed") runtime = entity_runtime[entity_id] runtime.remaining_jobs = 0 synced_entity_ids.add(entity_id) entity_runtime.pop(entity_id, None) return (0.05, 0.05) monkeypatch.setattr(repo, "_prepare_entity_vector_jobs", _stub_prepare) monkeypatch.setattr(repo, "_flush_embedding_jobs", _stub_flush) result = await repo.sync_entity_vectors_batch([1, 2, 3]) assert result.entities_total == 3 assert result.entities_synced == 1 assert result.entities_failed == 2 assert result.failed_entity_ids == [2, 3]