From dbe9cea261809cff4b8e38efe5b78c99a4c3aede Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 4 Jan 2025 20:54:52 -0600 Subject: [PATCH] add search indexing to api routes --- .../api/routers/documents_router.py | 44 ++++-- .../api/routers/knowledge_router.py | 75 ++++++++-- .../repository/search_repository.py | 5 +- src/basic_memory/services/search_service.py | 13 +- tests/api/test_documents_router.py | 91 ++++++++++++ tests/api/test_knowledge_router.py | 130 +++++++++++++++++- tests/api/test_search_router.py | 5 +- tests/conftest.py | 5 + tests/mcp/conftest.py | 6 + 9 files changed, 336 insertions(+), 38 deletions(-) diff --git a/src/basic_memory/api/routers/documents_router.py b/src/basic_memory/api/routers/documents_router.py index 733810fa..d97a8de2 100644 --- a/src/basic_memory/api/routers/documents_router.py +++ b/src/basic_memory/api/routers/documents_router.py @@ -2,9 +2,9 @@ from typing import List -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, BackgroundTasks, Depends -from basic_memory.deps import DocumentServiceDep +from basic_memory.deps import DocumentServiceDep, get_search_service from basic_memory.schemas.request import DocumentRequest, DocumentPathId from basic_memory.schemas.response import DocumentResponse, DocumentCreateResponse from basic_memory.services.document_service import ( @@ -20,22 +20,23 @@ router = APIRouter(prefix="/documents", tags=["documents"]) @router.post("/create", response_model=DocumentCreateResponse, status_code=201) async def create_document( doc: DocumentRequest, + background_tasks: BackgroundTasks, service: DocumentServiceDep, + search_service = Depends(get_search_service) ) -> DocumentCreateResponse: - """Create a new document. - - The document will be created with appropriate frontmatter including: - - Generated ID - - Creation timestamp - - Last modified timestamp - - Any provided doc_metadata - """ + """Create a new document with search indexing.""" try: document = await service.create_document( path_id=doc.path_id, content=doc.content, metadata=doc.doc_metadata, ) + # Index the new document + await search_service.index_document( + document, + doc.content, + background_tasks=background_tasks + ) return DocumentCreateResponse.model_validate(document.__dict__) except DocumentError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -71,9 +72,11 @@ async def get_document( async def update_document( path_id: DocumentPathId, doc: DocumentRequest, + background_tasks: BackgroundTasks, service: DocumentServiceDep, + search_service = Depends(get_search_service) ) -> DocumentResponse: - """Update a document by ID.""" + """Update a document by ID with search indexing.""" # Verify FilePaths match if doc.path_id != path_id: raise HTTPException( @@ -86,10 +89,16 @@ async def update_document( content=doc.content, metadata=doc.doc_metadata, ) + # Update search index + await search_service.index_document( + document, + doc.content, + background_tasks=background_tasks + ) doc_dict = document.__dict__ | {"content": doc.content} return DocumentResponse.model_validate(doc_dict) except DocumentNotFoundError: - raise HTTPException(status_code=404, detail=f"Document not found: {id}") + raise HTTPException(status_code=404, detail=f"Document not found: {path_id}") except DocumentWriteError as e: raise HTTPException(status_code=400, detail=str(e)) @@ -97,12 +106,17 @@ async def update_document( @router.delete("/{path_id:path}", status_code=204) async def delete_document( path_id: DocumentPathId, + background_tasks: BackgroundTasks, service: DocumentServiceDep, + search_service = Depends(get_search_service) ) -> None: - """Delete a document by ID.""" + """Delete a document by ID and remove from search index.""" try: + # Delete from storage await service.delete_document_by_path_id(path_id) + # Remove from search index (in background) + background_tasks.add_task(search_service.delete_by_path_id, path_id) except DocumentNotFoundError: - raise HTTPException(status_code=404, detail=f"Document not found: {id}") + raise HTTPException(status_code=404, detail=f"Document not found: {path_id}") except DocumentWriteError as e: - raise HTTPException(status_code=400, detail=str(e)) + raise HTTPException(status_code=400, detail=str(e)) \ No newline at end of file diff --git a/src/basic_memory/api/routers/knowledge_router.py b/src/basic_memory/api/routers/knowledge_router.py index e58b006c..e19d5818 100644 --- a/src/basic_memory/api/routers/knowledge_router.py +++ b/src/basic_memory/api/routers/knowledge_router.py @@ -1,11 +1,12 @@ """Router for knowledge graph operations.""" -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, BackgroundTasks, Depends from loguru import logger from basic_memory.deps import ( EntityServiceDep, KnowledgeServiceDep, + get_search_service, ) from basic_memory.schemas import ( CreateEntityRequest, @@ -31,10 +32,18 @@ router = APIRouter(prefix="/knowledge", tags=["knowledge"]) @router.post("/entities", response_model=EntityListResponse) async def create_entities( - data: CreateEntityRequest, knowledge_service: KnowledgeServiceDep + data: CreateEntityRequest, + background_tasks: BackgroundTasks, + knowledge_service: KnowledgeServiceDep, + search_service = Depends(get_search_service) ) -> EntityListResponse: - """Create new entities in the knowledge graph.""" + """Create new entities in the knowledge graph and index them.""" entities = await knowledge_service.create_entities(data.entities) + + # Index each entity + for entity in entities: + await search_service.index_entity(entity, background_tasks=background_tasks) + return EntityListResponse( entities=[EntityResponse.model_validate(entity) for entity in entities] ) @@ -42,10 +51,18 @@ async def create_entities( @router.post("/relations", response_model=EntityListResponse) async def create_relations( - data: CreateRelationsRequest, knowledge_service: KnowledgeServiceDep + data: CreateRelationsRequest, + background_tasks: BackgroundTasks, + knowledge_service: KnowledgeServiceDep, + search_service = Depends(get_search_service), ) -> EntityListResponse: - """Create relations between entities.""" + """Create relations between entities and update search index.""" updated_entities = await knowledge_service.create_relations(data.relations) + + # Reindex updated entities since relations have changed + for entity in updated_entities: + await search_service.index_entity(entity, background_tasks=background_tasks) + return EntityListResponse( entities=[EntityResponse.model_validate(entity) for entity in updated_entities] ) @@ -53,13 +70,20 @@ async def create_relations( @router.post("/observations", response_model=EntityResponse) async def add_observations( - data: AddObservationsRequest, knowledge_service: KnowledgeServiceDep + data: AddObservationsRequest, + background_tasks: BackgroundTasks, + knowledge_service: KnowledgeServiceDep, + search_service = Depends(get_search_service) ) -> EntityResponse: - """Add observations to an entity.""" + """Add observations to an entity and update search index.""" logger.debug(f"Adding observations to entity: {data.path_id}") updated_entity = await knowledge_service.add_observations( data.path_id, data.observations, data.context ) + + # Reindex the entity with new observations + await search_service.index_entity(updated_entity, background_tasks=background_tasks) + return EntityResponse.model_validate(updated_entity) @@ -104,29 +128,52 @@ async def open_nodes(data: OpenNodesRequest, entity_service: EntityServiceDep) - @router.post("/entities/delete", response_model=DeleteEntitiesResponse) async def delete_entities( - data: DeleteEntitiesRequest, knowledge_service: KnowledgeServiceDep + data: DeleteEntitiesRequest, + background_tasks: BackgroundTasks, + knowledge_service: KnowledgeServiceDep, + search_service = Depends(get_search_service) ) -> DeleteEntitiesResponse: - """Delete a specific entity by PathId.""" + """Delete entities and remove from search index.""" deleted = await knowledge_service.delete_entities(data.path_ids) + + # Remove each deleted entity from search index + for path_id in data.path_ids: + background_tasks.add_task(search_service.delete_by_path_id, path_id) + return DeleteEntitiesResponse(deleted=deleted) @router.post("/observations/delete", response_model=EntityResponse) async def delete_observations( - data: DeleteObservationsRequest, knowledge_service: KnowledgeServiceDep + data: DeleteObservationsRequest, + background_tasks: BackgroundTasks, + knowledge_service: KnowledgeServiceDep, + search_service = Depends(get_search_service) ) -> EntityResponse: - """Delete observations from an entity.""" + """Delete observations and update search index.""" path_id = data.path_id updated_entity = await knowledge_service.delete_observations(path_id, data.observations) + + # Reindex the entity since observations changed + await search_service.index_entity(updated_entity, background_tasks=background_tasks) + return EntityResponse.model_validate(updated_entity) @router.post("/relations/delete", response_model=EntityListResponse) async def delete_relations( - data: DeleteRelationsRequest, knowledge_service: KnowledgeServiceDep + data: DeleteRelationsRequest, + background_tasks: BackgroundTasks, + knowledge_service: KnowledgeServiceDep, + search_service = Depends(get_search_service) ) -> EntityListResponse: - """Delete relations between entities.""" + """Delete relations and update search index.""" updated_entities = await knowledge_service.delete_relations(data.relations) + + # Reindex entities since relations changed + for entity in updated_entities: + await search_service.index_entity(entity, background_tasks=background_tasks) + return EntityListResponse( entities=[EntityResponse.model_validate(entity) for entity in updated_entities] - ) + ) \ No newline at end of file diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 88f117b0..91b1ebb2 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -3,6 +3,8 @@ import json from typing import List, Optional from datetime import datetime + +from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -120,9 +122,10 @@ class SearchRepository(): "metadata": json.dumps(metadata) } ) + logger.debug(f"indexed {path_id}") await session.commit() - async def delete_by_path(self, path_id: str): + async def delete_by_path_id(self, path_id: str): """Delete an item from the search index.""" async with db.scoped_session(self.session_maker) as session: await session.execute( diff --git a/src/basic_memory/services/search_service.py b/src/basic_memory/services/search_service.py index da9f3d07..b8e3f7f8 100644 --- a/src/basic_memory/services/search_service.py +++ b/src/basic_memory/services/search_service.py @@ -5,6 +5,7 @@ from typing import List, Optional, Any from fastapi import BackgroundTasks from loguru import logger +from basic_memory.models import Document, Entity from basic_memory.repository.search_repository import SearchRepository from basic_memory.services.document_service import DocumentService from basic_memory.services.entity_service import EntityService @@ -64,7 +65,7 @@ class SearchService: async def index_entity( self, - entity: Any, # Could be more specific if we have an Entity type + entity: Entity, background_tasks: Optional[BackgroundTasks] = None ) -> None: """Index an entity and its components.""" @@ -110,13 +111,13 @@ class SearchService: async def index_document( self, - document: Any, # Could be more specific if we have a Document type + document: Document, content: str, background_tasks: Optional[BackgroundTasks] = None ) -> None: """Index a document and its content.""" metadata = { - **document.doc_metadata, + **(document.doc_metadata or {}), "created_at": document.created_at.isoformat(), "updated_at": document.updated_at.isoformat(), } @@ -155,4 +156,8 @@ class SearchService: file_path=file_path, type=type, metadata=metadata - ) \ No newline at end of file + ) + + async def delete_by_path_id(self, path_id: str): + """Delete an item from the search index.""" + await self.repository.delete_by_path_id(path_id) \ No newline at end of file diff --git a/tests/api/test_documents_router.py b/tests/api/test_documents_router.py index 1ee8ab4a..74aea128 100644 --- a/tests/api/test_documents_router.py +++ b/tests/api/test_documents_router.py @@ -6,6 +6,97 @@ import pytest from httpx import AsyncClient from basic_memory.config import ProjectConfig +from basic_memory.schemas.search import SearchItemType + + +@pytest.mark.asyncio +async def test_document_indexing(client: AsyncClient, test_config): + """Test document creation includes search indexing.""" + test_doc = { + "path_id": "test.md", + "content": "# Test\nThis is a test document with unique searchable content.", + "doc_metadata": {"type": "test", "tags": ["documentation", "test"]}, + } + + # Create document + response = await client.post("/documents/create", json=test_doc) + assert response.status_code == 201 + + # Verify it's searchable + search_response = await client.post( + "/search/", + json={"text": "unique searchable content", "types": [SearchItemType.DOCUMENT.value]}, + ) + assert search_response.status_code == 200 + results = search_response.json() + assert len(results) == 1 + assert results[0]["path_id"] == "test.md" + assert results[0]["type"] == SearchItemType.DOCUMENT.value + + +@pytest.mark.asyncio +async def test_document_update_indexing(client: AsyncClient): + """Test document updates are reflected in search index.""" + # Create initial document + test_doc = { + "path_id": "test.md", + "content": "Original content without special terms.", + "doc_metadata": {"type": "test", "status": "draft"}, + } + create_response = await client.post("/documents/create", json=test_doc) + assert create_response.status_code == 201 + + # Update document with new content + update_doc = { + "path_id": "test.md", + "content": "Updated content with special sphinx terms.", + "doc_metadata": {"type": "test", "status": "final"}, + } + update_response = await client.put(f"/documents/{test_doc["path_id"]}", json=update_doc) + assert update_response.status_code == 200 + + # Search for new terms + search_response = await client.post( + "/search/", json={"text": "sphinx", "types": [SearchItemType.DOCUMENT.value]} + ) + results = search_response.json() + assert len(results) == 1 + assert results[0]["path_id"] == "test.md" + + # Original terms shouldn't be found + search_response = await client.post( + "/search/", json={"text": "without special", "types": [SearchItemType.DOCUMENT.value]} + ) + assert len(search_response.json()) == 0 + + +@pytest.mark.asyncio +async def test_document_delete_indexing(client: AsyncClient): + """Test deleted documents are removed from search index.""" + # Create document + test_doc = { + "path_id": "test.md", + "content": "Searchable content that should disappear.", + "doc_metadata": {"type": "test"}, + } + create_response = await client.post("/documents/create", json=test_doc) + assert create_response.status_code == 201 + + # Verify it's initially searchable + search_response = await client.post( + "/search/", json={"text": "should disappear", "types": [SearchItemType.DOCUMENT.value]} + ) + assert len(search_response.json()) == 1 + + # Delete document + delete_response = await client.delete(f"/documents/{test_doc["path_id"]}") + assert delete_response.status_code == 204 + + # Verify it's no longer searchable + search_response = await client.post( + "/search/", json={"text": "should disappear", "types": [SearchItemType.DOCUMENT.value]} + ) + assert len(search_response.json()) == 0 @pytest.mark.asyncio diff --git a/tests/api/test_knowledge_router.py b/tests/api/test_knowledge_router.py index 5f9b7304..291d444c 100644 --- a/tests/api/test_knowledge_router.py +++ b/tests/api/test_knowledge_router.py @@ -12,6 +12,7 @@ from basic_memory.schemas import ( ObservationResponse, RelationResponse, ) +from basic_memory.schemas.search import SearchItemType async def create_entity(client) -> EntityResponse: @@ -47,7 +48,7 @@ async def add_observations(client, path_id: str) -> List[ObservationResponse]: {"content": "First observation", "category": "tech"}, {"content": "Second observation", "category": "note"}, ], - "context": "something special" + "context": "something special", }, ) # Verify observations were added @@ -417,7 +418,7 @@ async def test_full_knowledge_flow(client: AsyncClient): {"content": "Connected to first related entity", "category": "tech"}, {"content": "Connected to second related entity", "category": "note"}, ], - "context": "testing the flow" + "context": "testing the flow", }, ) @@ -449,3 +450,128 @@ async def test_full_knowledge_flow(client: AsyncClient): path_id = quote("test/MainEntity") response = await client.get(f"/knowledge/entities/{path_id}") assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_entity_indexing(client: AsyncClient): + """Test entity creation includes search indexing.""" + data = { + "name": "SearchTest", + "entity_type": "test", + "observations": ["Unique searchable observation"], + } + + # Create entity + response = await client.post("/knowledge/entities", json={"entities": [data]}) + assert response.status_code == 200 + + # Verify it's searchable + search_response = await client.post( + "/search/", json={"text": "unique searchable", "types": [SearchItemType.ENTITY.value]} + ) + assert search_response.status_code == 200 + results = search_response.json() + assert len(results) == 1 + assert results[0]["path_id"] == "test/search_test" + assert results[0]["type"] == SearchItemType.ENTITY.value + + +@pytest.mark.asyncio +async def test_observation_update_indexing(client: AsyncClient): + """Test observation changes are reflected in search.""" + # Create entity + data = { + "name": "TestEntity", + "entity_type": "test", + "observations": ["Initial observation"], + } + response = await client.post("/knowledge/entities", json={"entities": [data]}) + assert response.status_code == 200 + entity = response.json()["entities"][0] + + # Add new observation + await client.post( + "/knowledge/observations", + json={ + "path_id": entity["path_id"], + "observations": [{"content": "Unique sphinx observation", "category": "tech"}], + }, + ) + + # Search for new observation + search_response = await client.post( + "/search/", json={"text": "sphinx", "types": [SearchItemType.ENTITY.value]} + ) + results = search_response.json() + assert len(results) == 1 + assert results[0]["path_id"] == entity["path_id"] + + +@pytest.mark.asyncio +async def test_entity_delete_indexing(client: AsyncClient): + """Test deleted entities are removed from search index.""" + data = { + "name": "DeleteTest", + "entity_type": "test", + "observations": ["Searchable observation that should be removed"], + } + + # Create entity + response = await client.post("/knowledge/entities", json={"entities": [data]}) + assert response.status_code == 200 + entity = response.json()["entities"][0] + + # Verify it's initially searchable + search_response = await client.post( + "/search/", json={"text": "should be removed", "types": [SearchItemType.ENTITY.value]} + ) + assert len(search_response.json()) == 1 + + # Delete entity + delete_response = await client.post( + "/knowledge/entities/delete", json={"path_ids": [entity["path_id"]]} + ) + assert delete_response.status_code == 200 + + # Verify it's no longer searchable + search_response = await client.post( + "/search/", json={"text": "should be removed", "types": [SearchItemType.ENTITY.value]} + ) + assert len(search_response.json()) == 0 + + +@pytest.mark.asyncio +async def test_relation_indexing(client: AsyncClient): + """Test relations are included in search index.""" + # Create entities + entities = [ + {"name": "SourceTest", "entity_type": "test"}, + {"name": "TargetTest", "entity_type": "test"}, + ] + create_response = await client.post("/knowledge/entities", json={"entities": entities}) + assert create_response.status_code == 200 + + # Create relation with unique description + response = await client.post( + "/knowledge/relations", + json={ + "relations": [ + { + "from_id": "test/source_test", + "to_id": "test/target_test", + "relation_type": "sphinx_relation", + "context": "Unique sphinx relation context", + } + ] + }, + ) + assert response.status_code == 200 + + # Search should find both entities through relation + search_response = await client.post( + "/search/", json={"text": "sphinx relation", "types": [SearchItemType.ENTITY.value]} + ) + results = search_response.json() + assert len(results) == 2 # Both source and target entities + path_ids = {r["path_id"] for r in results} + assert path_ids == {"test/source_test", "test/target_test"} diff --git a/tests/api/test_search_router.py b/tests/api/test_search_router.py index fac3079d..693b7147 100644 --- a/tests/api/test_search_router.py +++ b/tests/api/test_search_router.py @@ -3,6 +3,7 @@ from datetime import datetime, timezone import pytest +import pytest_asyncio from sqlalchemy import text from basic_memory import db from basic_memory.schemas.search import SearchQuery, SearchItemType @@ -41,8 +42,8 @@ def test_document(): return Document() -@pytest.fixture -async def indexed_entity(test_entity, search_service): +@pytest_asyncio.fixture +async def indexed_entity(init_search_index, test_entity, search_service): """Create an entity and index it.""" await search_service.index_entity(test_entity) return test_entity diff --git a/tests/conftest.py b/tests/conftest.py index 6566f036..905ee094 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -218,6 +218,11 @@ async def search_repository(session_maker): """Create SearchRepository instance""" return SearchRepository(session_maker) +@pytest_asyncio.fixture(autouse=True) +async def init_search_index(search_service): + await search_service.init_search_index() + + @pytest_asyncio.fixture async def search_service( diff --git a/tests/mcp/conftest.py b/tests/mcp/conftest.py index adb3d59d..377e6b6e 100644 --- a/tests/mcp/conftest.py +++ b/tests/mcp/conftest.py @@ -6,6 +6,7 @@ from httpx import AsyncClient, ASGITransport from basic_memory.api.app import app as fastapi_app from basic_memory.deps import get_project_config, get_engine_factory +from basic_memory.services.search_service import SearchService @pytest_asyncio.fixture @@ -56,3 +57,8 @@ def test_directory_entity_data(): } ] } + +@pytest_asyncio.fixture(autouse=True) +async def init_search_index(search_service: SearchService): + await search_service.init_search_index() +