From 5488ff49d1ce98923965eaba3447d8d41e52809f Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 4 Jan 2025 20:14:36 -0600 Subject: [PATCH] add search_service.reindex_all() and api route --- src/basic_memory/api/routers/search_router.py | 10 +- src/basic_memory/deps.py | 4 +- src/basic_memory/services/search_service.py | 76 +++++++++++++-- tests/api/test_search_router.py | 88 ++++++++++++++--- tests/conftest.py | 8 +- tests/services/test_search_service.py | 95 ++++++++++++++++++- 6 files changed, 250 insertions(+), 31 deletions(-) diff --git a/src/basic_memory/api/routers/search_router.py b/src/basic_memory/api/routers/search_router.py index 28f7349f..cb046e14 100644 --- a/src/basic_memory/api/routers/search_router.py +++ b/src/basic_memory/api/routers/search_router.py @@ -3,6 +3,7 @@ from fastapi import APIRouter, Depends, BackgroundTasks from typing import List +from loguru import logger from basic_memory.services.search_service import SearchService from basic_memory.schemas.search import SearchQuery, SearchResult from basic_memory.deps import get_search_service @@ -22,6 +23,9 @@ async def reindex( background_tasks: BackgroundTasks, search_service: SearchService = Depends(get_search_service) ): - """Recreate the search index.""" - await search_service.init_search_index() - return {"status": "ok"} \ No newline at end of file + """Recreate and populate the search index.""" + await search_service.reindex_all(background_tasks=background_tasks) + return { + "status": "ok", + "message": "Reindex initiated" + } \ No newline at end of file diff --git a/src/basic_memory/deps.py b/src/basic_memory/deps.py index 84e6251f..ccd2df51 100644 --- a/src/basic_memory/deps.py +++ b/src/basic_memory/deps.py @@ -165,10 +165,10 @@ DocumentServiceDep = Annotated[DocumentService, Depends(get_document_service)] async def get_search_service( - search_repository: SearchRepositoryDep, + search_repository: SearchRepositoryDep, entity_service: EntityServiceDep, document_service: DocumentServiceDep ) -> SearchService: """Create SearchService with dependencies.""" - return SearchService(search_repository) + return SearchService(search_repository, document_service, entity_service) SearchServiceDep = Annotated[SearchService, Depends(get_search_service)] diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index 465e5e04..da9f3d07 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -1,28 +1,72 @@ """Service for search operations.""" -from typing import List, Optional +from typing import List, Optional, Any + +from fastapi import BackgroundTasks +from loguru import logger from basic_memory.repository.search_repository import SearchRepository +from basic_memory.services.document_service import DocumentService +from basic_memory.services.entity_service import EntityService from basic_memory.schemas.search import SearchQuery, SearchResult, SearchItemType class SearchService: """Service for search operations.""" - def __init__(self, search_repository: SearchRepository): + def __init__( + self, + search_repository: SearchRepository, + document_service: DocumentService, + entity_service: EntityService, + ): self.repository = search_repository + self.document_service = document_service + self.entity_service = entity_service async def init_search_index(self): """Create FTS5 virtual table if it doesn't exist.""" await self.repository.init_search_index() + + async def reindex_all( + self, + background_tasks: Optional[BackgroundTasks] = None + ) -> None: + """Reindex all content from database.""" + logger.info("Starting full reindex") + + # Clear and recreate search index + await self.init_search_index() + + # Reindex all entities + logger.debug("Indexing entities") + entities = await self.entity_service.get_all() + for entity in entities: + await self.index_entity(entity, background_tasks) + + # Reindex all documents + logger.debug("Indexing documents") + documents = await self.document_service.list_documents() + for doc in documents: + # Read content for each document + doc_with_content = await self.document_service.read_document_by_path_id(doc.path_id) + await self.index_document(doc, doc_with_content[1], background_tasks) + + logger.info("Reindex complete") async def search( - self, query: SearchQuery, context: Optional[List[str]] = None + self, + query: SearchQuery, + context: Optional[List[str]] = None ) -> List[SearchResult]: """Search across all indexed content.""" return await self.repository.search(query, context) - async def index_entity(self, entity, background_tasks=None): + async def index_entity( + self, + entity: Any, # Could be more specific if we have an Entity type + background_tasks: Optional[BackgroundTasks] = None + ) -> None: """Index an entity and its components.""" # Build searchable content content = "\n".join( @@ -64,7 +108,12 @@ class SearchService: metadata=metadata, ) - async def index_document(self, document, content: str, background_tasks=None): + async def index_document( + self, + document: Any, # Could be more specific if we have a Document type + content: str, + background_tasks: Optional[BackgroundTasks] = None + ) -> None: """Index a document and its content.""" metadata = { **document.doc_metadata, @@ -91,6 +140,19 @@ class SearchService: metadata=metadata, ) - async def _do_index(self, **kwargs): + async def _do_index( + self, + content: str, + path_id: str, + file_path: str, + type: SearchItemType, + metadata: dict + ) -> None: """Actually perform the indexing.""" - await self.repository.index_item(**kwargs) \ No newline at end of file + await self.repository.index_item( + content=content, + path_id=path_id, + file_path=file_path, + type=type, + metadata=metadata + ) \ No newline at end of file diff --git a/tests/api/test_search_router.py b/tests/api/test_search_router.py index bef8aa13..fac3079d 100644 --- a/tests/api/test_search_router.py +++ b/tests/api/test_search_router.py @@ -3,13 +3,12 @@ from datetime import datetime, timezone import pytest -import pytest_asyncio -from basic_memory.schemas.search import SearchQuery -from basic_memory.services.search_service import SearchService +from sqlalchemy import text +from basic_memory import db +from basic_memory.schemas.search import SearchQuery, SearchItemType - -@pytest_asyncio.fixture +@pytest.fixture def test_entity(): """Create a test entity.""" class Entity: @@ -26,13 +25,37 @@ def test_entity(): return Entity() -@pytest_asyncio.fixture +@pytest.fixture +def test_document(): + """Create a test document.""" + class Document: + id = 1 + path_id = "docs/test_doc.md" + file_path = "docs/test_doc.md" + doc_metadata = { + "title": "Test Document", + "type": "technical" + } + created_at = datetime.now(timezone.utc) + updated_at = datetime.now(timezone.utc) + return Document() + + +@pytest.fixture async def indexed_entity(test_entity, search_service): """Create an entity and index it.""" await search_service.index_entity(test_entity) return test_entity +@pytest.fixture +async def indexed_document(test_document, search_service): + """Create a document and index it.""" + content = "Test document content for search" + await search_service.index_document(test_document, content) + return test_document, content + + @pytest.mark.asyncio async def test_search_basic(client, indexed_entity): """Test basic text search.""" @@ -56,7 +79,7 @@ async def test_search_with_type_filter(client, indexed_entity): "/search/", json={ "text": "test", - "types": ["entity"] + "types": [SearchItemType.ENTITY.value] } ) assert response.status_code == 200 @@ -68,7 +91,7 @@ async def test_search_with_type_filter(client, indexed_entity): "/search/", json={ "text": "test", - "types": ["document"] + "types": [SearchItemType.DOCUMENT.value] } ) assert response.status_code == 200 @@ -166,11 +189,48 @@ async def test_search_empty(search_service, client): @pytest.mark.asyncio -async def test_reindex(search_service, client): +async def test_reindex( + client, + search_service, + entity_service, + document_service, + test_entity, + test_document, + session_maker +): """Test reindex endpoint.""" - response = await client.post("/search/reindex") - assert response.status_code == 200 - assert response.json()["status"] == "ok" + # Create test entity and document + await entity_service.create_entity(test_entity) + await document_service.create_document( + test_document.path_id, + "Test content", + test_document.doc_metadata + ) + + # Clear search index + async with db.scoped_session(session_maker) as session: + await session.execute(text("DELETE FROM search_index")) + await session.commit() + + # Verify nothing is searchable + response = await client.post( + "/search/", + json={"text": "test"} + ) + assert len(response.json()) == 0 + + # Trigger reindex + reindex_response = await client.post("/search/reindex") + assert reindex_response.status_code == 200 + assert reindex_response.json()["status"] == "ok" + + # Verify content is searchable again + search_response = await client.post( + "/search/", + json={"text": "test"} + ) + results = search_response.json() + assert len(results) == 2 # Both entity and document should be found @pytest.mark.asyncio @@ -180,7 +240,7 @@ async def test_multiple_filters(client, indexed_entity): "/search/", json={ "text": "test", - "types": ["entity"], + "types": [SearchItemType.ENTITY.value], "entity_types": ["component"], "after_date": datetime(2020, 1, 1, tzinfo=timezone.utc).isoformat() } @@ -190,5 +250,5 @@ async def test_multiple_filters(client, indexed_entity): assert len(results) == 1 result = results[0] assert result["path_id"] == indexed_entity.path_id - assert result["type"] == "entity" + assert result["type"] == SearchItemType.ENTITY.value assert result["metadata"]["entity_type"] == "component" \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index bcb68c5b..6566f036 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -220,9 +220,13 @@ async def search_repository(session_maker): @pytest_asyncio.fixture -async def search_service(search_repository: SearchRepository): +async def search_service( + search_repository: SearchRepository, + entity_service: EntityService, + document_service: DocumentService, +) -> SearchService: """Create and initialize search service""" - service = SearchService(search_repository) + service = SearchService(search_repository, document_service, entity_service) await service.init_search_index() return service diff --git a/tests/services/test_search_service.py b/tests/services/test_search_service.py index e12b44cb..f79b8aa6 100644 --- a/tests/services/test_search_service.py +++ b/tests/services/test_search_service.py @@ -6,8 +6,7 @@ from sqlalchemy import text from basic_memory import db from basic_memory.repository.search_repository import SearchRepository -from basic_memory.schemas.search import SearchQuery, SearchItemType -from basic_memory.services.search_service import SearchService +from basic_memory.schemas.search import SearchQuery, SearchResult, SearchItemType @pytest.fixture @@ -153,4 +152,94 @@ async def test_update_document_index(search_service, test_document): # Search for new terms results = await search_service.search(SearchQuery(text="new terms")) - assert len(results) == 1 \ No newline at end of file + assert len(results) == 1 + + +@pytest.mark.asyncio +async def test_reindex_all( + search_service, + entity_service, + document_service, + test_entity, + test_document, + session_maker +): + """Test reindexing all content.""" + # Create test entities and documents + entity = await entity_service.create_entity( + test_entity + ) + + document_content = "Test document content" + document = await document_service.create_document( + test_document.path_id, + document_content, + test_document.doc_metadata + ) + + # Clear the search index + async with db.scoped_session(session_maker) as session: + await session.execute(text("DELETE FROM search_index")) + await session.commit() + + # Verify nothing is searchable + results = await search_service.search(SearchQuery(text="test")) + assert len(results) == 0 + + # Reindex everything + await search_service.reindex_all() + + # Verify entity is searchable + entity_results = await search_service.search( + SearchQuery(text="TestComponent", types=[SearchItemType.ENTITY]) + ) + assert len(entity_results) == 1 + assert entity_results[0].path_id == test_entity.path_id + assert entity_results[0].type == SearchItemType.ENTITY + + # Verify document is searchable + doc_results = await search_service.search( + SearchQuery(text="document content", types=[SearchItemType.DOCUMENT]) + ) + assert len(doc_results) == 1 + assert doc_results[0].path_id == test_document.path_id + assert doc_results[0].type == SearchItemType.DOCUMENT + + +@pytest.mark.asyncio +async def test_reindex_with_background_tasks( + search_service, + entity_service, + document_service, + test_entity, + test_document, + session_maker +): + """Test reindexing with background tasks.""" + from fastapi import BackgroundTasks + + # Create test data + entity = await entity_service.create_entity(test_entity) + document = await document_service.create_document( + test_document.path_id, + "Test content", + test_document.doc_metadata + ) + + # Clear index + async with db.scoped_session(session_maker) as session: + await session.execute(text("DELETE FROM search_index")) + await session.commit() + + # Create background tasks + background_tasks = BackgroundTasks() + + # Reindex with background tasks + await search_service.reindex_all(background_tasks=background_tasks) + + # Execute background tasks + await background_tasks() + + # Verify everything was indexed + all_results = await search_service.search(SearchQuery(text="test")) + assert len(all_results) == 2 # Both entity and document should be found \ No newline at end of file