From 644328f62301fb06b717b695fcfe87f21c689d57 Mon Sep 17 00:00:00 2001 From: phernandez Date: Fri, 27 Dec 2024 12:24:53 -0600 Subject: [PATCH] refactor knowledge_service --- .../services/knowledge/entity_operations.py | 30 +++-- .../services/knowledge/file_operations.py | 19 ++- .../services/knowledge/knowledge_service.py | 127 +++++++++++++++--- .../knowledge/observation_operations.py | 21 ++- .../services/knowledge/relation_operations.py | 26 ++-- tests/services/test_knowledge_service.py | 34 ++--- 6 files changed, 193 insertions(+), 64 deletions(-) diff --git a/src/basic_memory/services/knowledge/entity_operations.py b/src/basic_memory/services/knowledge/entity_operations.py index 6cbaf46e..0a3e1b99 100644 --- a/src/basic_memory/services/knowledge/entity_operations.py +++ b/src/basic_memory/services/knowledge/entity_operations.py @@ -6,25 +6,36 @@ from loguru import logger from basic_memory.models import Entity as EntityModel from basic_memory.schemas import Entity as EntitySchema +from basic_memory.services.entity_service import EntityService +from basic_memory.services.exceptions import EntityNotFoundError from .file_operations import FileOperations -from ..exceptions import EntityCreationError, EntityNotFoundError -class EntityOperations(FileOperations): - """Entity operations mixin for KnowledgeService.""" +class EntityOperations: + """Entity operations for knowledge service.""" + + def __init__(self, + entity_service: EntityService, + file_operations: FileOperations + ): + self.entity_service = entity_service + self.file_operations = file_operations + + async def get_by_path_id(self, path_id: str) -> EntityModel: + """Get entity by path ID.""" + return await self.entity_service.get_by_path_id(path_id) async def create_entity(self, entity: EntitySchema) -> EntityModel: """Create a new entity and write to filesystem.""" logger.debug(f"Creating entity: {entity}") db_entity = None - file_path = None try: # 1. Create entity in DB db_entity = await self.entity_service.create_entity(entity) # 2. Write file and get checksum - file_path, checksum = await self.write_entity_file(db_entity) + _, checksum = await self.file_operations.write_entity_file(db_entity) # 3. Update DB with checksum updated = await self.entity_service.update_entity( @@ -37,8 +48,7 @@ class EntityOperations(FileOperations): # Clean up on any failure if db_entity: await self.entity_service.delete_entity(db_entity.path_id) - if file_path: - await self.file_service.delete_file(file_path) + await self.file_operations.delete_entity_file(db_entity) logger.error(f"Failed to create entity: {e}") raise @@ -62,8 +72,7 @@ class EntityOperations(FileOperations): entity = await self.entity_service.get_by_path_id(path_id) # Delete file first (it's source of truth) - path = self.get_entity_path(entity) - await self.file_service.delete_file(path) + await self.file_operations.delete_entity_file(entity) # Delete from DB (this will cascade to observations/relations) return await self.entity_service.delete_entity(path_id) @@ -81,9 +90,8 @@ class EntityOperations(FileOperations): logger.debug(f"Deleting entities: {path_ids}") success = True - # Let errors bubble up for path_id in path_ids: await self.delete_entity(path_id) success = True - return success + return success \ No newline at end of file diff --git a/src/basic_memory/services/knowledge/file_operations.py b/src/basic_memory/services/knowledge/file_operations.py index 2b257c36..09c2d09f 100644 --- a/src/basic_memory/services/knowledge/file_operations.py +++ b/src/basic_memory/services/knowledge/file_operations.py @@ -13,7 +13,7 @@ from basic_memory.services.file_service import FileService class FileOperations: - """File operations mixin for KnowledgeService.""" + """File operations for knowledge entities.""" def __init__( self, @@ -27,12 +27,18 @@ class FileOperations: self.knowledge_writer = knowledge_writer self.base_path = base_path + async def file_exists(self, path: Path) -> bool: + return await self.file_service.exists(path) + + async def read_file(self, path: Path) -> Tuple[str, str]: + return await self.file_service.read_file(path) + def get_entity_path(self, entity: EntityModel) -> Path: """Generate filesystem path for entity.""" return self.base_path / entity.entity_type / f"{entity.name}.md" async def write_entity_file(self, entity: EntityModel) -> Tuple[Path, str]: - """Write entity to filesystem and return checksum.""" + """Write entity to filesystem and return path and checksum.""" try: # Ensure we have a fresh entity with all relations loaded entity = await self.entity_service.get_by_path_id(entity.path_id) @@ -53,3 +59,12 @@ class FileOperations: except Exception as e: logger.error(f"Failed to write entity file: {e}") raise FileOperationError(f"Failed to write entity file: {e}") + + async def delete_entity_file(self, entity: EntityModel) -> None: + """Delete entity file from filesystem.""" + try: + path = self.get_entity_path(entity) + await self.file_service.delete_file(path) + except Exception as e: + logger.error(f"Failed to delete entity file: {e}") + raise FileOperationError(f"Failed to delete entity file: {e}") \ No newline at end of file diff --git a/src/basic_memory/services/knowledge/knowledge_service.py b/src/basic_memory/services/knowledge/knowledge_service.py index 3491885c..7d9b687c 100644 --- a/src/basic_memory/services/knowledge/knowledge_service.py +++ b/src/basic_memory/services/knowledge/knowledge_service.py @@ -1,31 +1,38 @@ """Main knowledge service implementation.""" from pathlib import Path +from typing import List, Sequence, Tuple + +from loguru import logger from basic_memory.markdown.knowledge_writer import KnowledgeWriter +from basic_memory.models import Entity as EntityModel +from basic_memory.schemas import Entity as EntitySchema +from basic_memory.schemas import Relation as RelationSchema +from basic_memory.schemas.request import ObservationCreate from basic_memory.services.entity_service import EntityService from basic_memory.services.file_service import FileService from basic_memory.services.observation_service import ObservationService from basic_memory.services.relation_service import RelationService + +from .file_operations import FileOperations +from .entity_operations import EntityOperations +from .relation_operations import RelationOperations from .observation_operations import ObservationOperations -class KnowledgeService(ObservationOperations): +class KnowledgeService: """ Service for managing knowledge graph entities and their persistence. - Orchestrates operations between: - - EntityService for core entity operations - - ObservationService for atomic facts - - RelationService for entity connections - - FileService for persistence - - KnowledgeParser for file formatting - - Operations are split across mixins: - - FileOperations: Core file handling - - EntityOperations: Entity CRUD operations - - RelationOperations: Relation management - - ObservationOperations: Observation handling + Composes specialized operations for: + - File handling and persistence + - Entity CRUD operations + - Relations between entities + - Observations about entities + + Acts as the main coordinator for all knowledge operations, ensuring + consistency between database and filesystem. """ def __init__( @@ -37,11 +44,95 @@ class KnowledgeService(ObservationOperations): knowledge_writer: KnowledgeWriter, base_path: Path, ): - super().__init__( + + self.base_path = base_path + + # Initialize operations in dependency order + self.file_ops = FileOperations( entity_service=entity_service, - observation_service=observation_service, - relation_service=relation_service, file_service=file_service, knowledge_writer=knowledge_writer, - base_path=base_path, - ) \ No newline at end of file + base_path=base_path + ) + + self.entity_ops = EntityOperations( + entity_service=entity_service, + file_operations=self.file_ops + ) + + self.relation_ops = RelationOperations( + relation_service=relation_service, + entity_service=entity_service, + file_operations=self.file_ops + ) + + self.observation_ops = ObservationOperations( + observation_service=observation_service, + entity_service=entity_service, + file_operations=self.file_ops + ) + + # Entity operations + async def get_entity_by_path_id(self, path_id: str) -> EntityModel: + return await self.entity_ops.get_by_path_id(path_id) + + async def create_entity(self, entity: EntitySchema) -> EntityModel: + """Create a new entity.""" + return await self.entity_ops.create_entity(entity) + + async def create_entities(self, entities: List[EntitySchema]) -> Sequence[EntityModel]: + """Create multiple entities.""" + return await self.entity_ops.create_entities(entities) + + async def delete_entity(self, path_id: str) -> bool: + """Delete an entity and its file.""" + return await self.entity_ops.delete_entity(path_id) + + async def delete_entities(self, path_ids: List[str]) -> bool: + """Delete multiple entities and their files.""" + return await self.entity_ops.delete_entities(path_ids) + + async def file_exists(self, path: Path) -> bool: + """Check if entity file exists.""" + return await self.file_ops.file_exists(path) + + async def read_file(self, path: Path) -> Tuple[str, str]: + """Check if entity file exists.""" + return await self.file_ops.read_file(path) + + # Relation operations + async def create_relations(self, relations: List[RelationSchema]) -> Sequence[EntityModel]: + """Create relations between entities.""" + return await self.relation_ops.create_relations(relations) + + async def delete_relations(self, to_delete: List[RelationSchema]) -> Sequence[EntityModel]: + """Delete relations between entities.""" + return await self.relation_ops.delete_relations(to_delete) + + # Observation operations + async def add_observations( + self, + path_id: str, + observations: List[ObservationCreate], + context: str | None = None + ) -> EntityModel: + """Add observations to an entity.""" + return await self.observation_ops.add_observations(path_id, observations, context) + + async def delete_observations( + self, + path_id: str, + observations: List[str] + ) -> EntityModel: + """Delete observations from an entity.""" + return await self.observation_ops.delete_observations(path_id, observations) + + # File operations for direct access if needed + def get_entity_path(self, entity: EntityModel) -> Path: + """Get filesystem path for entity.""" + return self.file_ops.get_entity_path(entity) + + async def write_entity_file(self, entity: EntityModel) -> Path: + """Write entity to filesystem.""" + path, _ = await self.file_ops.write_entity_file(entity) + return path \ No newline at end of file diff --git a/src/basic_memory/services/knowledge/observation_operations.py b/src/basic_memory/services/knowledge/observation_operations.py index fd1a4b80..86117fdf 100644 --- a/src/basic_memory/services/knowledge/observation_operations.py +++ b/src/basic_memory/services/knowledge/observation_operations.py @@ -7,16 +7,23 @@ from loguru import logger from basic_memory.models import Entity as EntityModel from basic_memory.services.exceptions import EntityNotFoundError from basic_memory.services.observation_service import ObservationService +from basic_memory.services.entity_service import EntityService from basic_memory.schemas.request import ObservationCreate -from .relation_operations import RelationOperations +from .file_operations import FileOperations -class ObservationOperations(RelationOperations): - """Observation operations mixin for KnowledgeService.""" +class ObservationOperations: + """Observation operations for knowledge service.""" - def __init__(self, *args, observation_service: ObservationService, **kwargs): - super().__init__(*args, **kwargs) + def __init__( + self, + observation_service: ObservationService, + entity_service: EntityService, + file_operations: FileOperations + ): self.observation_service = observation_service + self.entity_service = entity_service + self.file_operations = file_operations async def add_observations( self, @@ -50,7 +57,7 @@ class ObservationOperations(RelationOperations): entity = await self.entity_service.get_by_path_id(path_id) # Write updated file and checksum - _, checksum = await self.write_entity_file(entity) + _, checksum = await self.file_operations.write_entity_file(entity) await self.entity_service.update_entity(path_id, {"checksum": checksum}) # Return final entity with all updates and relations @@ -83,7 +90,7 @@ class ObservationOperations(RelationOperations): await self.observation_service.delete_observations(entity.id, observations) # Write updated file - _, checksum = await self.write_entity_file(entity) + _, checksum = await self.file_operations.write_entity_file(entity) await self.entity_service.update_entity(path_id, {"checksum": checksum}) # Return final entity with all updates diff --git a/src/basic_memory/services/knowledge/relation_operations.py b/src/basic_memory/services/knowledge/relation_operations.py index 248e744d..aba8f765 100644 --- a/src/basic_memory/services/knowledge/relation_operations.py +++ b/src/basic_memory/services/knowledge/relation_operations.py @@ -1,6 +1,6 @@ """Relation operations for knowledge service.""" -from typing import Sequence, List, Dict, Any +from typing import Sequence, List from loguru import logger @@ -9,15 +9,23 @@ from basic_memory.models import Relation as RelationModel from basic_memory.schemas import Relation as RelationSchema from basic_memory.services.exceptions import EntityNotFoundError from basic_memory.services.relation_service import RelationService +from basic_memory.services.entity_service import EntityService from .entity_operations import EntityOperations +from .file_operations import FileOperations -class RelationOperations(EntityOperations): - """Relation operations mixin for KnowledgeService.""" +class RelationOperations: + """Relation operations for knowledge service.""" - def __init__(self, *args, relation_service: RelationService, **kwargs): - super().__init__(*args, **kwargs) + def __init__( + self, + relation_service: RelationService, + entity_service: EntityService, + file_operations: FileOperations + ): self.relation_service = relation_service + self.entity_service = entity_service + self.file_operations = file_operations async def create_relations(self, relations: List[RelationSchema]) -> Sequence[EntityModel]: """Create relations and return updated entities.""" @@ -44,7 +52,7 @@ class RelationOperations(EntityOperations): entities_to_update.add(rs.to_id) except Exception as e: - logger.error(f"Failed to create rs: {e}") + logger.error(f"Failed to create relation: {e}") continue # Get fresh copies of all updated entities @@ -54,7 +62,7 @@ class RelationOperations(EntityOperations): entity = await self.entity_service.get_by_path_id(path_id) # Write updated file - _, checksum = await self.write_entity_file(entity) + _, checksum = await self.file_operations.write_entity_file(entity) updated = await self.entity_service.update_entity(path_id, {"checksum": checksum}) updated_entities.append(updated) @@ -97,7 +105,7 @@ class RelationOperations(EntityOperations): raise EntityNotFoundError(f"Entity not found: {path_id}") # Write updated file - _, checksum = await self.write_entity_file(entity) + _, checksum = await self.file_operations.write_entity_file(entity) updated = await self.entity_service.update_entity( path_id, {"checksum": checksum} ) @@ -112,4 +120,4 @@ class RelationOperations(EntityOperations): except Exception as e: logger.error(f"Failed to delete relations: {e}") - raise + raise \ No newline at end of file diff --git a/tests/services/test_knowledge_service.py b/tests/services/test_knowledge_service.py index c338815f..cd7d351d 100644 --- a/tests/services/test_knowledge_service.py +++ b/tests/services/test_knowledge_service.py @@ -40,7 +40,7 @@ async def test_create_entity(knowledge_service: KnowledgeService): # Verify file was written file_path = knowledge_service.get_entity_path(created) - assert await knowledge_service.file_service.exists(file_path) + assert await knowledge_service.file_exists(file_path) @pytest.mark.asyncio @@ -57,7 +57,7 @@ async def test_create_multiple_entities(knowledge_service: KnowledgeService): for i, entity in enumerate(created): assert entity.name == f"entity-{i}" file_path = knowledge_service.get_entity_path(entity) - assert await knowledge_service.file_service.exists(file_path) + assert await knowledge_service.file_exists(file_path) @pytest.mark.asyncio @@ -88,7 +88,7 @@ async def test_create_relations(knowledge_service: KnowledgeService, entity_serv for entity in [entity1, entity2]: found = await entity_service.get_by_path_id(entity.path_id) file_path = knowledge_service.get_entity_path(found) - content, _ = await knowledge_service.file_service.read_file(file_path) + content, _ = await knowledge_service.read_file(file_path) assert "test_relation" in content @@ -121,7 +121,7 @@ async def test_add_observations_observation(knowledge_service: KnowledgeService) # Verify file was updated file_path = knowledge_service.get_entity_path(updated_entity) - content, _ = await knowledge_service.file_service.read_file(file_path) + content, _ = await knowledge_service.read_file(file_path) for obs in observations: expected_line = f"- [{obs.category.value}] {obs.content} ({context})" @@ -141,18 +141,18 @@ async def test_delete_entity(knowledge_service: KnowledgeService): file_path = knowledge_service.get_entity_path(entity) # Verify file exists - assert await knowledge_service.file_service.exists(file_path) + assert await knowledge_service.file_exists(file_path) # Delete entity success = await knowledge_service.delete_entity(entity.path_id) assert success # Verify file was deleted - assert not await knowledge_service.file_service.exists(file_path) + assert not await knowledge_service.file_exists(file_path) # Verify entity was deleted from DB with pytest.raises(EntityNotFoundError): - await knowledge_service.entity_service.get_by_path_id(entity.path_id) + await knowledge_service.get_entity_by_path_id(entity.path_id) @pytest.mark.asyncio @@ -173,9 +173,9 @@ async def test_delete_multiple_entities(knowledge_service: KnowledgeService): # Verify files were deleted for entity in entities: file_path = knowledge_service.get_entity_path(entity) - assert not await knowledge_service.file_service.exists(file_path) + assert not await knowledge_service.file_exists(file_path) with pytest.raises(EntityNotFoundError): - await knowledge_service.entity_service.get_by_path_id(entity.path_id) + await knowledge_service.get_entity_by_path_id(entity.path_id) @pytest.mark.asyncio @@ -185,7 +185,7 @@ async def test_handle_file_operation_errors(knowledge_service: KnowledgeService, async def mock_write_file(*args): raise FileOperationError("Test error") - monkeypatch.setattr(knowledge_service.file_service, "write_file", mock_write_file) + monkeypatch.setattr(knowledge_service.file_ops.file_service, "write_file", mock_write_file) with pytest.raises(FileOperationError): await knowledge_service.create_entity( @@ -206,7 +206,7 @@ async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService, entity_ids: List[str] = [] # Capture created entity ID - original_create = knowledge_service.entity_service.create_entity + original_create = knowledge_service.entity_ops.entity_service.create_entity async def mock_create_entity(*args, **kwargs): entity = await original_create(*args, **kwargs) @@ -217,8 +217,8 @@ async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService, async def mock_write_file(*args): raise FileOperationError("Test error") - monkeypatch.setattr(knowledge_service.entity_service, "create_entity", mock_create_entity) - monkeypatch.setattr(knowledge_service.file_service, "write_file", mock_write_file) + monkeypatch.setattr(knowledge_service.entity_ops.entity_service, "create_entity", mock_create_entity) + monkeypatch.setattr(knowledge_service.file_ops.file_service, "write_file", mock_write_file) # Attempt creation (should fail) with pytest.raises(FileOperationError): @@ -229,7 +229,7 @@ async def test_cleanup_on_creation_failure(knowledge_service: KnowledgeService, # Verify entity was cleaned up assert len(entity_ids) == 1 with pytest.raises(EntityNotFoundError): - await knowledge_service.entity_service.get_by_path_id(entity_ids[0]) + await knowledge_service.entity_ops.entity_service.get_by_path_id(entity_ids[0]) @pytest.mark.asyncio @@ -271,17 +271,17 @@ async def test_update_relations_in_files(knowledge_service: KnowledgeService): # Verify both files contain relation for entity in [entity1, entity2]: file_path = knowledge_service.get_entity_path(entity) - content, _ = await knowledge_service.file_service.read_file(file_path) + content, _ = await knowledge_service.read_file(file_path) assert "connects_to" in content # Source should show outgoing relation - content, _ = await knowledge_service.file_service.read_file( + content, _ = await knowledge_service.read_file( knowledge_service.get_entity_path(entity1) ) assert "target" in content # Target should show incoming relation - content, _ = await knowledge_service.file_service.read_file( + content, _ = await knowledge_service.read_file( knowledge_service.get_entity_path(entity2) ) assert "source" in content