diff --git a/src/basic_memory/services.py b/src/basic_memory/services.py index d6b21137..c99c2d41 100644 --- a/src/basic_memory/services.py +++ b/src/basic_memory/services.py @@ -1,13 +1,14 @@ from datetime import datetime, UTC from pathlib import Path -from typing import Optional +from typing import Optional, List from uuid import uuid4 from sqlalchemy import and_, select, delete from basic_memory.models import Entity as DbEntity # Rename to avoid confusion from basic_memory.models import Observation as DbObservation -from basic_memory.repository import EntityRepository, ObservationRepository -from basic_memory.schemas import Entity, Observation +from basic_memory.models import Relation as DbRelation +from basic_memory.repository import EntityRepository, ObservationRepository, RelationRepository +from basic_memory.schemas import Entity, Observation, Relation from basic_memory.fileio import ( read_entity_file, write_entity_file, delete_entity_file, FileOperationError, EntityNotFoundError @@ -24,12 +25,14 @@ class DatabaseSyncError(ServiceError): pass +class RelationError(ServiceError): + """Base exception for relation-specific errors""" + pass + + class EntityService: - """ - Service for managing entities in the filesystem and database. - Follows the "filesystem is source of truth" principle. - """ - + """Service for managing entities in the filesystem and database.""" + # [Previous EntityService implementation remains unchanged] def __init__(self, project_path: Path, entity_repo: EntityRepository): self.project_path = project_path self.entity_repo = entity_repo @@ -111,14 +114,8 @@ class EntityService: class ObservationService: - """ - Service for managing observations in the filesystem and database. - Follows the "filesystem is source of truth" principle. - - Observations are stored in entity markdown files and indexed in the database - for efficient querying. - """ - + """Service for managing observations in the filesystem and database.""" + # [Previous ObservationService implementation remains unchanged] def __init__(self, project_path: Path, observation_repo: ObservationRepository): self.project_path = project_path self.entities_path = project_path / "entities" @@ -126,22 +123,7 @@ class ObservationService: async def add_observation(self, entity: Entity, content: str, context: Optional[str] = None) -> Observation: - """ - Add a new observation to an entity. - - Args: - entity: Entity to add observation to - content: Content of the observation - context: Optional context for the observation - - Returns: - The created Observation - - Raises: - FileOperationError: If file operations fail - DatabaseSyncError: If database sync fails - """ - # Create new observation + """Add a new observation to an entity.""" observation = Observation(content=content) entity.observations.append(observation) @@ -160,41 +142,117 @@ class ObservationService: return observation except Exception as e: raise DatabaseSyncError(f"Failed to sync observation to database: {str(e)}") from e - - async def search_observations(self, query: str) -> list[Observation]: + + +class RelationService: + """ + Service for managing relations between entities. + Follows the "filesystem is source of truth" principle. + + Relations are stored in entity markdown files and indexed in the database + for efficient querying. + """ + + def __init__(self, project_path: Path, relation_repo: RelationRepository): + self.project_path = project_path + self.entities_path = project_path / "entities" + self.relation_repo = relation_repo + + async def create_relation(self, from_entity: Entity, to_entity: Entity, relation_type: str, + context: Optional[str] = None) -> Relation: """ - Search for observations across all entities. + Create a new relation between two entities. Args: - query: Text to search for in observation content + from_entity: Source entity + to_entity: Target entity + relation_type: Type of relation + context: Optional context for the relation Returns: - List of matching observations with their entity contexts + The created Relation + + Raises: + FileOperationError: If file operations fail + DatabaseSyncError: If database sync fails """ - result = await self.observation_repo.execute_query( - select(DbObservation).filter( - DbObservation.content.contains(query) - ) + # Create new relation + relation = Relation( + id=f"rel-{uuid4().hex[:8]}", + from_id=from_entity.id, + to_id=to_entity.id, + relation_type=relation_type, + context=context ) - return [ - Observation(content=obs.content) - for obs in result.scalars().all() - ] - async def get_observations_by_context(self, context: str) -> list[Observation]: - """Get all observations with a specific context.""" - db_observations = await self.observation_repo.find_by_context(context) - return [ - Observation(content=obs.content) - for obs in db_observations - ] + # Add relation to source entity's relations list + if not hasattr(from_entity, 'relations'): + from_entity.relations = [] + from_entity.relations.append(relation) + + # Update filesystem first (source of truth) + await write_entity_file(self.entities_path, from_entity) + + # Update database index + try: + await self.relation_repo.create({ + 'id': relation.id, + 'from_id': from_entity.id, + 'to_id': to_entity.id, + 'relation_type': relation_type, + 'context': context, + 'created_at': datetime.now(UTC) + }) + return relation + except Exception as e: + raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e - async def rebuild_observation_index(self) -> None: + async def get_entity_relations(self, entity: Entity) -> List[Relation]: """ - Rebuild the observation database index from filesystem contents. + Get all relations for an entity (both outgoing and incoming). + + Args: + entity: Entity to get relations for + + Returns: + List of relations where the entity is either source or target + """ + # Relations are stored in the entity object already + return getattr(entity, 'relations', []) + + async def delete_relation(self, from_entity: Entity, relation_id: str) -> bool: + """ + Delete a relation from both filesystem and database. + + Args: + from_entity: Source entity containing the relation + relation_id: ID of the relation to delete + + Returns: + True if deletion was successful + + Raises: + RelationError: If relation cannot be found or deleted + """ + # Remove relation from entity's relations + if hasattr(from_entity, 'relations'): + from_entity.relations = [ + r for r in from_entity.relations + if r.id != relation_id + ] + + # Update filesystem first (source of truth) + await write_entity_file(self.entities_path, from_entity) + + # Remove from database index + await self.relation_repo.delete(relation_id) + return True + + async def rebuild_relation_index(self) -> None: + """ + Rebuild the relation database index from filesystem contents. Used for recovery or ensuring sync. """ - # List all entity files if not self.entities_path.exists(): return @@ -202,20 +260,22 @@ class ObservationService: entity_files = list(self.entities_path.glob("*.md")) except Exception as e: raise FileOperationError(f"Failed to read entities directory: {str(e)}") from e - - # Clear existing observation index - await self.observation_repo.execute_query(delete(DbObservation)) + + # Clear existing relation index + await self.relation_repo.execute_query(delete(DbRelation)) # Rebuild from each entity file for entity_file in entity_files: try: entity = await read_entity_file(self.entities_path, entity_file.stem) - for obs in entity.observations: - await self.observation_repo.create({ - 'id': f"{entity.id}-obs-{uuid4().hex[:8]}", - 'entity_id': entity.id, - 'content': obs.content, + for relation in getattr(entity, 'relations', []): + await self.relation_repo.create({ + 'id': relation.id, + 'from_id': relation.from_id, + 'to_id': relation.to_id, + 'relation_type': relation.relation_type, + 'context': relation.context, 'created_at': datetime.now(UTC) }) except Exception as e: - print(f"Warning: Failed to reindex observations for {entity_file}: {str(e)}") + print(f"Warning: Failed to reindex relations for {entity_file}: {str(e)}") diff --git a/tests/test_relation_service.py b/tests/test_relation_service.py new file mode 100644 index 00000000..9ebae45c --- /dev/null +++ b/tests/test_relation_service.py @@ -0,0 +1,182 @@ +import pytest +from datetime import datetime, UTC +from pathlib import Path + +from basic_memory.models import Entity as DbEntity +from basic_memory.models import Relation as DbRelation +from basic_memory.schemas import Entity, Relation +from basic_memory.services import RelationService, EntityService +from basic_memory.repository import EntityRepository, RelationRepository + + +@pytest.fixture +async def entity_repo(db_session): + return EntityRepository(db_session) + + +@pytest.fixture +async def relation_repo(db_session): + return RelationRepository(db_session) + + +@pytest.fixture +async def entity_service(tmp_path, entity_repo): + return EntityService(tmp_path, entity_repo) + + +@pytest.fixture +async def relation_service(tmp_path, relation_repo): + return RelationService(tmp_path, relation_repo) + + +@pytest.fixture +async def sample_entities(entity_service): + """Create two sample entities for testing relations""" + entity1 = await entity_service.create_entity( + name="test_entity_1", + entity_type="test_type" + ) + entity2 = await entity_service.create_entity( + name="test_entity_2", + entity_type="test_type" + ) + return entity1, entity2 + + +async def test_create_relation(relation_service, sample_entities): + """Test creating a basic relation between two entities""" + entity1, entity2 = sample_entities + + relation = await relation_service.create_relation( + from_entity=entity1, + to_entity=entity2, + relation_type="test_relation" + ) + + assert relation.from_id == entity1.id + assert relation.to_id == entity2.id + assert relation.relation_type == "test_relation" + + # Verify relation was added to source entity's relations + assert hasattr(entity1, 'relations') + assert len(entity1.relations) == 1 + assert entity1.relations[0].id == relation.id + + # Verify file was written with relation + entity_file = relation_service.entities_path / f"{entity1.id}.md" + assert entity_file.exists() + content = entity_file.read_text() + assert "## Relations" in content + assert f"[{entity2.id}] test_relation" in content + + # Verify database was updated + db_relation = await relation_service.relation_repo.find_by_id(relation.id) + assert db_relation is not None + assert db_relation.from_id == entity1.id + assert db_relation.to_id == entity2.id + assert db_relation.relation_type == "test_relation" + + +async def test_create_relation_with_context(relation_service, sample_entities): + """Test creating a relation with context information""" + entity1, entity2 = sample_entities + + relation = await relation_service.create_relation( + from_entity=entity1, + to_entity=entity2, + relation_type="test_relation", + context="test context" + ) + + assert relation.context == "test context" + + # Verify context in file + entity_file = relation_service.entities_path / f"{entity1.id}.md" + content = entity_file.read_text() + assert f"[{entity2.id}] test_relation | test context" in content + + # Verify context in database + db_relation = await relation_service.relation_repo.find_by_id(relation.id) + assert db_relation.context == "test context" + + +async def test_get_entity_relations(relation_service, sample_entities): + """Test retrieving relations for an entity""" + entity1, entity2 = sample_entities + + # Create test relation + relation = await relation_service.create_relation( + from_entity=entity1, + to_entity=entity2, + relation_type="test_relation" + ) + + # Get relations from entity + relations = await relation_service.get_entity_relations(entity1) + + assert len(relations) == 1 + assert relations[0].from_id == entity1.id + assert relations[0].to_id == entity2.id + assert relations[0].relation_type == "test_relation" + + +async def test_delete_relation(relation_service, sample_entities): + """Test deleting a relation""" + entity1, entity2 = sample_entities + + # Create then delete a relation + relation = await relation_service.create_relation( + from_entity=entity1, + to_entity=entity2, + relation_type="test_relation" + ) + + success = await relation_service.delete_relation(entity1, relation.id) + assert success is True + + # Verify removed from entity relations + assert not entity1.relations or relation.id not in [r.id for r in entity1.relations] + + # Verify removed from file + entity_file = relation_service.entities_path / f"{entity1.id}.md" + content = entity_file.read_text() + assert f"[{entity2.id}] test_relation" not in content + + # Verify removed from database + db_relation = await relation_service.relation_repo.find_by_id(relation.id) + assert db_relation is None + + +async def test_rebuild_relation_index(relation_service, sample_entities): + """Test rebuilding the relation index from files""" + entity1, entity2 = sample_entities + + # Create some test relations + relation1 = await relation_service.create_relation( + from_entity=entity1, + to_entity=entity2, + relation_type="test_relation_1" + ) + relation2 = await relation_service.create_relation( + from_entity=entity2, + to_entity=entity1, + relation_type="test_relation_2" + ) + + # Clear the database relations + await relation_service.relation_repo.execute_query( + 'DELETE FROM relation' + ) + + # Rebuild index + await relation_service.rebuild_relation_index() + + # Verify relations were restored + db_relations = await relation_service.relation_repo.execute_query( + 'SELECT * FROM relation' + ) + relations = db_relations.scalars().all() + + assert len(relations) == 2 + relation_types = {r.relation_type for r in relations} + assert relation_types == {"test_relation_1", "test_relation_2"}