From 122dbd7d31102f40ee3a01a9dfe3fdff579c3239 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sat, 7 Dec 2024 17:51:52 -0600 Subject: [PATCH] MemoryService.create_relation implemented --- src/basic_memory/fileio.py | 4 +- src/basic_memory/repository.py | 34 ++++++++++- src/basic_memory/schemas.py | 32 ++++++++++- src/basic_memory/services/entity_service.py | 31 ++++------- src/basic_memory/services/memory_service.py | 48 ++++++---------- tests/test_memory_service.py | 62 +++++++++++++++------ 6 files changed, 137 insertions(+), 74 deletions(-) diff --git a/src/basic_memory/fileio.py b/src/basic_memory/fileio.py index 39a434fe..85002e35 100644 --- a/src/basic_memory/fileio.py +++ b/src/basic_memory/fileio.py @@ -152,8 +152,8 @@ async def read_entity_file(entities_path: Path, entity_id: str) -> Entity: source_entity = Entity(id=entity_id, name=name, entity_type=entity_type) relations.append(Relation( - from_entity=source_entity, - to_entity=target_entity, + from_id=source_entity.id, + to_id=target_entity.id, relation_type=relation_type, context=context )) diff --git a/src/basic_memory/repository.py b/src/basic_memory/repository.py index bfc44da3..cef3307c 100644 --- a/src/basic_memory/repository.py +++ b/src/basic_memory/repository.py @@ -2,7 +2,7 @@ from typing import Type, Optional, Any, Sequence from sqlalchemy import select, func, Select, Executable, inspect, Result, Column, and_ from sqlalchemy.exc import NoResultFound from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import Mapped +from sqlalchemy.orm import Mapped, selectinload from basic_memory.models import Entity, Observation, Relation, Base @@ -177,6 +177,38 @@ class EntityRepository(Repository[Entity]): Repository for Entity model with memory-specific operations. """ + async def find_by_id(self, entity_id: str) -> Optional[Entity]: + """ + Find entity by ID with relations eagerly loaded. + + Uses selectinload to eagerly load outgoing and incoming relations in a single query. + This is necessary because: + 1. In async code, lazy loading relations after the session closes doesn't work + 2. Our service layer often needs the complete entity with its relations + 3. Using a single query with selectinload is more efficient than multiple lazy-loaded queries + + :param entity_id: Entity ID to search for + :return: Entity if found with all relations loaded, None otherwise + + Example: + entity = await repo.find_by_id('20240102-entity-123') + # Relations are already loaded - no additional queries needed + for relation in entity.outgoing_relations: + print(f"Related to {relation.to_id} via {relation.relation_type}") + """ + try: + result = await self.session.execute( + select(Entity) + .filter(Entity.id == entity_id) + .options( + selectinload(Entity.outgoing_relations), + selectinload(Entity.incoming_relations) + ) + ) + return result.scalars().one() + except NoResultFound: + return None + async def find_by_name(self, name: str) -> Optional[Entity]: """ Find an entity by its unique name. diff --git a/src/basic_memory/schemas.py b/src/basic_memory/schemas.py index 9b608828..580a9967 100644 --- a/src/basic_memory/schemas.py +++ b/src/basic_memory/schemas.py @@ -5,8 +5,10 @@ independent from storage/persistence concerns. """ from datetime import datetime, UTC -from typing import List, Optional -from pydantic import BaseModel +from typing import List, Optional, Dict, Any +from uuid import uuid4 + +from pydantic import BaseModel, model_validator class Observation(BaseModel): @@ -47,6 +49,32 @@ class Entity(BaseModel): observations: List[Observation] = [] relations: List[Relation] = [] + @model_validator(mode='before') + @classmethod + def generate_id(cls, data: dict) -> dict: + """Generate an ID if one wasn't provided during instantiation""" + if not data.get('id') and data.get('name'): + timestamp = datetime.now(UTC).strftime("%Y%m%d") + normalized_name = data['name'].lower().replace(" ", "-") + data['id'] = f"{timestamp}-{normalized_name}-{uuid4().hex[:8]}" + return data + + def model_dump(self, **kwargs) -> Dict[str, Any]: + """Serialize entity, handling relations to prevent circular references""" + # Get basic data without relations + exclude = kwargs.pop('exclude', set()) + exclude.add('relations') + basic_data = super().model_dump(exclude=exclude, **kwargs) + + # Add serialized relations if we have any + if 'relations' not in exclude and self.relations: + basic_data['relations'] = [ + relation.model_dump(**kwargs) + for relation in self.relations + ] + + return basic_data + def file_name(self) -> str: """Get the markdown file name for this entity.""" return f"{self.id}.md" \ No newline at end of file diff --git a/src/basic_memory/services/entity_service.py b/src/basic_memory/services/entity_service.py index 55d715fe..499065e0 100644 --- a/src/basic_memory/services/entity_service.py +++ b/src/basic_memory/services/entity_service.py @@ -1,12 +1,11 @@ """Service for managing entities in the database.""" from datetime import datetime, UTC from pathlib import Path -from typing import Optional -from basic_memory.models import Entity as DbEntity from basic_memory.repository import EntityRepository from basic_memory.schemas import Entity -from . import ServiceError, DatabaseSyncError +from basic_memory.models import Entity as EntityModel +from . import ServiceError class EntityService: """ @@ -18,42 +17,32 @@ class EntityService: self.project_path = project_path self.entity_repo = entity_repo - async def create_entity(self, entity: Entity) -> Entity: + async def create_entity(self, entity: Entity) -> EntityModel: """Create a new entity in the database.""" # Create DB record db_data = { **entity.model_dump(), "created_at": datetime.now(UTC), - "updated_at": datetime.now(UTC) } - await self.entity_repo.create(db_data) - return entity + return await self.entity_repo.create(db_data) - async def get_entity(self, entity_id: str) -> Entity: + async def get_entity(self, entity_id: str) -> EntityModel: """Get entity by ID.""" db_entity = await self.entity_repo.find_by_id(entity_id) if not db_entity: raise ServiceError(f"Entity not found: {entity_id}") - return Entity( - id=db_entity.id, - name=db_entity.name, - entity_type=db_entity.entity_type - ) + return db_entity - async def get_by_name(self, name: str) -> Entity: + # TODO name is not uniaue + async def get_by_name(self, name: str) -> EntityModel: """Get entity by name.""" db_entity = await self.entity_repo.find_by_name(name) if not db_entity: raise ServiceError(f"Entity not found: {name}") - return Entity( - id=db_entity.id, - name=db_entity.name, - entity_type=db_entity.entity_type - ) + return db_entity async def delete_entity(self, entity_id: str) -> bool: """Delete entity from database.""" - await self.entity_repo.delete(entity_id) - return True \ No newline at end of file + return await self.entity_repo.delete(entity_id) \ No newline at end of file diff --git a/src/basic_memory/services/memory_service.py b/src/basic_memory/services/memory_service.py index c1aed19b..3acef6a0 100644 --- a/src/basic_memory/services/memory_service.py +++ b/src/basic_memory/services/memory_service.py @@ -42,40 +42,28 @@ class MemoryService: async def create_relations(self, relations_data: List[Dict[str, Any]]) -> List[Relation]: """Create multiple relations between entities.""" - relations = [] - - for data in relations_data: - # Resolve entities by ID - from_entity, to_entity = await asyncio.gather( - self.entity_service.get_entity(data["from_id"]), - self.entity_service.get_entity(data["to_id"]) - ) - - # Create relation from schema data - relation = Relation( - from_id=from_entity.id, - to_id=to_entity.id, - relation_type=data["relation_type"], - context=data.get("context") - ) - - # Create in database - stored_relation = await self.relation_service.create_relation(relation) - - # Add to source entity's relations list + relations = [Relation.model_validate(data) for data in relations_data] + + for relation in relations: + # First read complete entities from filesystem + from_entity = await read_entity_file(self.entities_path, relation.from_id) + to_entity = await read_entity_file(self.entities_path, relation.to_id) + + # Add the new relation to the source entity if not hasattr(from_entity, 'relations'): from_entity.relations = [] - from_entity.relations.append(stored_relation) - relations.append((stored_relation, from_entity)) + from_entity.relations.append(relation) - # Write updated entities in parallel - async def write_file(entity: Entity): - await write_entity_file(self.entities_path, entity) + # Write updated entity files (filesystem is source of truth) + await asyncio.gather( + write_entity_file(self.entities_path, from_entity), + write_entity_file(self.entities_path, to_entity) + ) - file_writes = [write_file(entity) for _, entity in relations] - await asyncio.gather(*file_writes) + # Now update the database index + await self.relation_service.create_relation(relation) - return [relation for relation, _ in relations] + return relations async def add_observations(self, observations_data: List[Dict[str, Any]]) -> None: """Add observations to existing entities.""" @@ -129,7 +117,7 @@ class MemoryService: ] entity_updates.append(entity) - # Write updated files in parallel + # Write updated entities in parallel async def write_file(entity: Entity): await write_entity_file(self.entities_path, entity) diff --git a/tests/test_memory_service.py b/tests/test_memory_service.py index 24286272..9f68ccc2 100644 --- a/tests/test_memory_service.py +++ b/tests/test_memory_service.py @@ -2,6 +2,8 @@ import pytest from basic_memory.services import MemoryService +from basic_memory.fileio import read_entity_file +from basic_memory.models import Entity as EntityModel test_entities_data = [ { @@ -47,7 +49,7 @@ async def test_create_entities(memory_service: MemoryService): @pytest.mark.asyncio async def test_create_relations(memory_service: MemoryService): - """Should create relations between entities and update the filesystem.""" + """Should create relations between entities and update both filesystem and database.""" # First create the entities entities = await memory_service.create_entities(test_entities_data) entity1, entity2 = entities @@ -85,23 +87,47 @@ async def test_create_relations(memory_service: MemoryService): assert relations[1].relation_type == "references" assert relations[1].context == "test context" - # Verify relations were added to the entities - # Read updated entity1 from filesystem - entity1_path = memory_service.entities_path / entity1.file_name() - assert entity1_path.exists() - assert len(entity1.relations) == 1 - assert entity1.relations[0].from_id == entity1.id - assert entity1.relations[0].to_id == entity2.id - assert entity1.relations[0].relation_type == "connects_to" + # Read updated entities from filesystem to verify relations + updated_entity1 = await read_entity_file(memory_service.entities_path, entity1.id) + updated_entity2 = await read_entity_file(memory_service.entities_path, entity2.id) - # Read updated entity2 from filesystem - entity2_path = memory_service.entities_path / entity2.file_name() - assert entity2_path.exists() - assert len(entity2.relations) == 1 - assert entity2.relations[0].from_id == entity2.id - assert entity2.relations[0].to_id == entity1.id - assert entity2.relations[0].relation_type == "references" - assert entity2.relations[0].context == "test context" + # Verify relations were added to entity1 in filesystem + assert len(updated_entity1.relations) == 1 + assert updated_entity1.relations[0].from_id == entity1.id + assert updated_entity1.relations[0].to_id == entity2.id + assert updated_entity1.relations[0].relation_type == "connects_to" + + # Verify relations were added to entity2 in filesystem + assert len(updated_entity2.relations) == 1 + assert updated_entity2.relations[0].from_id == entity2.id + assert updated_entity2.relations[0].to_id == entity1.id + assert updated_entity2.relations[0].relation_type == "references" + assert updated_entity2.relations[0].context == "test context" + + # Now verify database state + # Get entities from database + db_entity1: EntityModel = await memory_service.entity_service.get_entity(entity1.id) + db_entity2: EntityModel = await memory_service.entity_service.get_entity(entity2.id) + + # Entity 1 should have one outgoing relation to entity 2 + assert len(db_entity1.outgoing_relations) == 1 + outgoing = db_entity1.outgoing_relations[0] + assert outgoing.from_id == entity1.id + assert outgoing.to_id == entity2.id + assert outgoing.relation_type == "connects_to" + assert outgoing.context is None + + # Entity 1 should have one incoming relation from entity 2 + assert len(db_entity1.incoming_relations) == 1 + incoming = db_entity1.incoming_relations[0] + assert incoming.from_id == entity2.id + assert incoming.to_id == entity1.id + assert incoming.relation_type == "references" + assert incoming.context == "test context" + + # Verify the same for entity 2 (reversed) + assert len(db_entity2.outgoing_relations) == 1 + assert len(db_entity2.incoming_relations) == 1 @pytest.mark.asyncio async def test_create_relations_with_invalid_entity_id(memory_service: MemoryService): @@ -118,4 +144,4 @@ async def test_create_relations_with_invalid_entity_id(memory_service: MemorySer } with pytest.raises(Exception) as exc: # We might want to define a specific error type - await memory_service.create_relations([bad_relation]) + await memory_service.create_relations([bad_relation]) \ No newline at end of file