From 4d3f2dc03daa91d67306a4c85ac5930f43a7dbf6 Mon Sep 17 00:00:00 2001 From: phernandez Date: Thu, 2 Jan 2025 18:21:37 -0600 Subject: [PATCH] entity_sync_service --- src/basic_memory/services/__init__.py | 2 - .../services/sync/entity_sync_service.py | 118 ++++++++++ .../services/sync/sync_service.py | 3 +- .../services/sync/test_entity_sync_service.py | 211 ++++++++++++++++++ 4 files changed, 331 insertions(+), 3 deletions(-) create mode 100644 src/basic_memory/services/sync/entity_sync_service.py create mode 100644 tests/services/sync/test_entity_sync_service.py diff --git a/src/basic_memory/services/__init__.py b/src/basic_memory/services/__init__.py index c77cdfb6..3578eb7b 100644 --- a/src/basic_memory/services/__init__.py +++ b/src/basic_memory/services/__init__.py @@ -2,7 +2,6 @@ from .document_service import DocumentService from .sync.file_change_scanner import FileChangeScanner -from .sync.document_sync_service import DocumentSyncService from .entity_service import EntityService from .file_service import FileService from .knowledge import KnowledgeService @@ -13,7 +12,6 @@ from .service import BaseService __all__ = [ "BaseService", "DocumentService", - "DocumentSyncService", "EntityService", "FileService", "ObservationService", diff --git a/src/basic_memory/services/sync/entity_sync_service.py b/src/basic_memory/services/sync/entity_sync_service.py new file mode 100644 index 00000000..071886b5 --- /dev/null +++ b/src/basic_memory/services/sync/entity_sync_service.py @@ -0,0 +1,118 @@ +"""Service for managing entities in the database.""" +from typing import Dict + +from loguru import logger + +from basic_memory.models import Entity as EntityModel, Observation, Relation +from basic_memory.repository.entity_repository import EntityRepository +from basic_memory.markdown.schemas import EntityMarkdown +from basic_memory.services import EntityService +from basic_memory.services.service import BaseService + + +def entity_model_from_markdown( + markdown: EntityMarkdown +) -> EntityModel: + """Convert markdown entity to model. + + Args: + markdown: Parsed markdown entity + include_relations: Whether to include relations. Set False for first sync pass. + """ + model = EntityModel( + name=markdown.content.title, + entity_type=markdown.frontmatter.type, + path_id=markdown.frontmatter.id, + file_path=markdown.frontmatter.id, + description=markdown.content.description, + observations=[Observation(content=obs.content) for obs in markdown.content.observations], + ) + return model + + +class EntitySyncService(EntityService): + """Service for managing entities in the database.""" + + def __init__(self, entity_repository: EntityRepository): + super().__init__(entity_repository) + + async def create_entity_without_relations(self, markdown: EntityMarkdown) -> EntityModel: + """First pass: Create entity and observations only. + + Creates the entity with null checksum to indicate sync not complete. + Relations will be added in second pass. + """ + logger.debug(f"Creating entity without relations: {markdown.frontmatter.id}") + model = entity_model_from_markdown(markdown) + model.checksum = None # Mark as incomplete sync + return await self.repository.add(model) + + async def update_entity_without_relations( + self, path_id: str, markdown: EntityMarkdown + ) -> EntityModel: + """First pass: Update entity fields and observations. + + Updates everything except relations and sets null checksum + to indicate sync not complete. + """ + logger.debug(f"Updating entity without relations: {path_id}") + db_entity = await self.get_by_path_id(path_id) + + # Update fields from markdown + db_entity.name = markdown.content.title + db_entity.entity_type = markdown.frontmatter.type + db_entity.description = markdown.content.description + + # Update observations + db_entity.observations = [ + Observation(content=obs.content) for obs in markdown.content.observations + ] + + # Mark as incomplete + db_entity.checksum = None + + return await self.repository.update( + db_entity.id, + { + "name": db_entity.name, + "entity_type": db_entity.entity_type, + "description": db_entity.description, + "observations": db_entity.observations, + "checksum": None, + }, + ) + + async def update_entity_relations(self, markdown: EntityMarkdown, checksum: str) -> EntityModel: + """Second pass: Update relations and set checksum. + + Args: + markdown: Parsed markdown entity with relations + checksum: Final checksum to set after relations are updated + """ + logger.debug(f"Updating relations for entity: {markdown.frontmatter.id}") + db_entity = await self.get_by_path_id(markdown.frontmatter.id) + + # get all entities from relations + target_entity_path_ids = [rel.target for rel in markdown.content.relations] + target_entities = await self.repository.find_by_path_ids(target_entity_path_ids) + + # zip dict by path + entity_by_path: Dict[str, EntityModel] = dict(zip(target_entity_path_ids, target_entities)) + + # Update relations from markdown + db_entity.to_relations = [ + Relation( + from_id=db_entity.id, + to_id=entity_by_path[rel.target].id, + relation_type=rel.type, + context=rel.context, + ) + for rel in markdown.content.relations + ] + + # Set final checksum to mark sync complete + db_entity.checksum = checksum + + return await self.repository.update( + db_entity.id, {"relations": db_entity.relations, "checksum": checksum} + ) diff --git a/src/basic_memory/services/sync/sync_service.py b/src/basic_memory/services/sync/sync_service.py index f4c8cc18..cbddd32a 100644 --- a/src/basic_memory/services/sync/sync_service.py +++ b/src/basic_memory/services/sync/sync_service.py @@ -5,6 +5,7 @@ from loguru import logger from basic_memory.services import FileChangeScanner, EntityService, DocumentService from basic_memory.markdown import KnowledgeParser +from basic_memory.services.sync.entity_sync_service import EntitySyncService class SyncService: @@ -19,7 +20,7 @@ class SyncService: self, scanner: FileChangeScanner, document_service: DocumentService, - entity_service: EntityService, + entity_service: EntitySyncService, knowledge_parser: KnowledgeParser, ): self.scanner = scanner diff --git a/tests/services/sync/test_entity_sync_service.py b/tests/services/sync/test_entity_sync_service.py new file mode 100644 index 00000000..c255338b --- /dev/null +++ b/tests/services/sync/test_entity_sync_service.py @@ -0,0 +1,211 @@ +"""Tests for EntitySyncService.""" + +import pytest +import pytest_asyncio +from datetime import datetime +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory.models import Entity as EntityModel +from basic_memory.repository.entity_repository import EntityRepository +from basic_memory.services.sync.entity_sync_service import EntitySyncService +from basic_memory.markdown.schemas import ( + EntityMarkdown, + EntityContent, + EntityFrontmatter, + EntityMetadata, + Observation as MarkdownObservation, + Relation as MarkdownRelation +) + + +@pytest_asyncio.fixture +async def entity_repository(session_maker: async_sessionmaker[AsyncSession]) -> EntityRepository: + """Create an EntityRepository instance.""" + return EntityRepository(session_maker) + + +@pytest_asyncio.fixture +async def entity_sync_service(entity_repository: EntityRepository) -> EntitySyncService: + """Create EntitySyncService with repository.""" + return EntitySyncService(entity_repository) + + +@pytest_asyncio.fixture +def test_frontmatter() -> EntityFrontmatter: + """Create test frontmatter.""" + return EntityFrontmatter( + type="concept", + id="concept/test_entity", + created=datetime.now(), + modified=datetime.now(), + tags=["test", "sync"] + ) + + +@pytest_asyncio.fixture +def test_content() -> EntityContent: + """Create test content with observations and relations.""" + return EntityContent( + title="Test Entity", + description="A test entity description", + observations=[ + MarkdownObservation(content="First observation"), + MarkdownObservation(content="Second observation") + ], + relations=[ + MarkdownRelation(type="depends_on", target="concept/other_entity"), + MarkdownRelation(type="related_to", target="concept/another_entity") + ] + ) + + +@pytest_asyncio.fixture +def test_markdown(test_frontmatter, test_content) -> EntityMarkdown: + """Create complete test markdown entity.""" + return EntityMarkdown( + frontmatter=test_frontmatter, + content=test_content, + entity_metadata=EntityMetadata() + ) + + +@pytest.mark.asyncio +async def test_create_entity_without_relations( + entity_sync_service: EntitySyncService, + test_markdown: EntityMarkdown +): + """Test first pass creation without relations.""" + # Create entity first pass + entity = await entity_sync_service.create_entity_without_relations(test_markdown) + + # Check basic fields + assert entity.name == "Test Entity" + assert entity.entity_type == "concept" + assert entity.path_id == "concept/test_entity" + assert entity.description == "A test entity description" + + # Check observations + assert len(entity.observations) == 2 + assert entity.observations[0].content == "First observation" + assert entity.observations[1].content == "Second observation" + + # Check no relations added + assert len(entity.relations) == 0 + + # Check checksum is None (indicating incomplete sync) + assert entity.checksum is None + + +@pytest.mark.asyncio +async def test_update_entity_without_relations( + entity_sync_service: EntitySyncService, + test_markdown: EntityMarkdown +): + """Test first pass update.""" + # First create entity + entity = await entity_sync_service.create_entity_without_relations(test_markdown) + + # Modify markdown content + test_markdown.content.title = "Updated Title" + test_markdown.content.description = "Updated description" + test_markdown.content.observations = [ + MarkdownObservation(content="Updated observation") + ] + + # Update entity + updated = await entity_sync_service.update_entity_without_relations( + entity.path_id, + test_markdown + ) + + # Check fields updated + assert updated.name == "Updated Title" + assert updated.description == "Updated description" + assert len(updated.observations) == 1 + assert updated.observations[0].content == "Updated observation" + + # Check checksum cleared + assert updated.checksum is None + + +@pytest.mark.asyncio +async def test_update_entity_relations( + entity_sync_service: EntitySyncService, + test_markdown: EntityMarkdown +): + """Test second pass relation updates.""" + # Create main entity first + entity = await entity_sync_service.create_entity_without_relations(test_markdown) + + # Create target entities that relations point to + other_entity = EntityModel( + name="Other Entity", + entity_type="concept", + path_id="concept/other_entity" + ) + another_entity = EntityModel( + name="Another Entity", + entity_type="concept", + path_id="concept/another_entity" + ) + await entity_sync_service.repository.add(other_entity) + await entity_sync_service.repository.add(another_entity) + + # Update relations and set checksum + test_checksum = "test-checksum-123" + updated = await entity_sync_service.update_entity_relations(test_markdown, test_checksum) + + # Check relations + assert len(updated.relations) == 2 + + # Check relation details + relations = sorted(updated.relations, key=lambda r: r.relation_type) + + assert relations[0].relation_type == "depends_on" + assert relations[0].from_id == entity.id + assert relations[0].to_id == other_entity.id + + assert relations[1].relation_type == "related_to" + assert relations[1].from_id == entity.id + assert relations[1].to_id == another_entity.id + + # Check checksum set + assert updated.checksum == test_checksum + + +@pytest.mark.asyncio +async def test_two_pass_sync_flow( + entity_sync_service: EntitySyncService, + test_markdown: EntityMarkdown +): + """Test complete two-pass sync flow.""" + # Create target entities first + other_entity = EntityModel( + name="Other Entity", + entity_type="concept", + path_id="concept/other_entity" + ) + another_entity = EntityModel( + name="Another Entity", + entity_type="concept", + path_id="concept/another_entity" + ) + await entity_sync_service.repository.add(other_entity) + await entity_sync_service.repository.add(another_entity) + + # First pass - create without relations + entity = await entity_sync_service.create_entity_without_relations(test_markdown) + assert len(entity.relations) == 0 + assert entity.checksum is None + + # Second pass - add relations + checksum = "final-checksum-456" + updated = await entity_sync_service.update_entity_relations(test_markdown, checksum) + + # Verify final state + assert len(updated.relations) == 2 + assert updated.checksum == checksum + + relations = sorted(updated.relations, key=lambda r: r.relation_type) + assert relations[0].to_id == other_entity.id + assert relations[1].to_id == another_entity.id