refactor schemas and services

src/basic_memory/fileio.py

@@ -4,7 +4,7 @@ Handles reading and writing entities and observations to the filesystem.
 """
 from pathlib import Path
 
-from basic_memory.schemas import Entity, Observation, Relation
+from basic_memory.schemas import EntityIn, ObservationIn, RelationIn
 
 
 class FileOperationError(Exception):
@@ -17,7 +17,7 @@ class EntityNotFoundError(Exception):
     pass
 
 
-async def write_entity_file(entities_path: Path, entity: Entity) -> bool:
+async def write_entity_file(entities_path: Path, entity: EntityIn) -> bool:
     """
     Write entity to filesystem in markdown format.
     
@@ -49,7 +49,10 @@ async def write_entity_file(entities_path: Path, entity: Entity) -> bool:
 
     # Add observations
     for obs in entity.observations:
-        content.append(f"- {obs.content}\n")
+        obs_line = f"- {obs.content}"
+        if obs.context:
+            obs_line += f" | {obs.context}"
+        content.append(f"{obs_line}\n")
         
     # Add relations section if we have relations
     if hasattr(entity, 'relations') and entity.relations:
@@ -80,7 +83,7 @@ async def write_entity_file(entities_path: Path, entity: Entity) -> bool:
     return True
 
 
-async def read_entity_file(entities_path: Path, entity_id: str) -> Entity:
+async def read_entity_file(entities_path: Path, entity_id: str) -> EntityIn:
     """
     Read entity data from filesystem.
     
@@ -131,7 +134,12 @@ async def read_entity_file(entities_path: Path, entity_id: str) -> Entity:
             in_observations = False
             in_relations = True
         elif in_observations and line.startswith("- "):
-            observations.append(Observation(content=line[2:]))
+            # Parse observation line: content | context
+            line = line[2:]  # Remove the "- "
+            parts = line.split(" | ", 1)
+            content = parts[0]
+            context = parts[1] if len(parts) > 1 else None
+            observations.append(ObservationIn(content=content, context=context))
         elif in_relations and line.startswith("- "):
             # Parse relation line: - [target_id] relation_type | context
             line = line[2:]  # Remove the bullet point
@@ -148,17 +156,18 @@ async def read_entity_file(entities_path: Path, entity_id: str) -> Entity:
             context = parts[1] if len(parts) > 1 else None
             
             # Create temporary entities for the relation
-            target_entity = Entity(id=target_id, name=target_id, entity_type="unknown")
-            source_entity = Entity(id=entity_id, name=name, entity_type=entity_type)
+            # TODO what is this for?
+            target_entity = EntityIn(id=target_id, name=target_id, entity_type="unknown")
+            source_entity = EntityIn(id=entity_id, name=name, entity_type=entity_type)
             
-            relations.append(Relation(
+            relations.append(RelationIn(
                 from_id=source_entity.id,
                 to_id=target_entity.id,
                 relation_type=relation_type,
                 context=context
             ))
     
-    return Entity(
+    return EntityIn(
         id=entity_id,
         name=name,
         entity_type=entity_type,

src/basic_memory/repository.py

@@ -179,28 +179,23 @@ class EntityRepository(Repository[Entity]):
     
     async def find_by_id(self, entity_id: str) -> Optional[Entity]:
         """
-        Find entity by ID with relations eagerly loaded.
+        Find entity by ID with all relationships eagerly loaded.
         
-        Uses selectinload to eagerly load outgoing and incoming relations in a single query.
-        This is necessary because:
+        Uses selectinload to eagerly load observations, 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
-        
+        2. Our service layer often needs the complete entity 
+        3. Using a single query with selectinload is more efficient
+
         :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}")
+        :return: Entity if found with everything loaded, None otherwise
         """
         try:
             result = await self.session.execute(
                 select(Entity)
                 .filter(Entity.id == entity_id)
                 .options(
+                    selectinload(Entity.observations),
                     selectinload(Entity.outgoing_relations),
                     selectinload(Entity.incoming_relations)
                 )
@@ -216,7 +211,15 @@ class EntityRepository(Repository[Entity]):
         :param name: Entity name to search for
         :return: Entity if found, None otherwise
         """
-        query = select(Entity).filter(Entity.name == name)
+        query = (
+            select(Entity)
+            .filter(Entity.name == name)
+            .options(
+                selectinload(Entity.observations),
+                selectinload(Entity.outgoing_relations),
+                selectinload(Entity.incoming_relations)
+            )
+        )
         return await self.find_one(query)
     
     async def search_by_type(self, entity_type: str, skip: int = 0, limit: int = 100) -> Sequence[Entity]:

src/basic_memory/schemas.py

@@ -29,62 +29,55 @@ class ObservationsOut(BaseModel):
     entity_id: str
     observations: List[ObservationOut]
 
-# Original schemas kept for now until we migrate everything
-class Observation(BaseModel):
-    """An atomic piece of information about an entity."""
-    id: Optional[int] = None  # Let the database handle ID generation
-    content: str
-    context: Optional[str] = None
 
-class Relation(BaseModel):
+class RelationIn(BaseModel):
     """
     Represents a directed edge between entities in the knowledge graph.
     Relations are always stored in active voice (e.g. "created", "teaches", etc.)
     """
-    id: Optional[int] = None
     from_id: str
     to_id: str              
     relation_type: str
     context: Optional[str] = None
 
-class Entity(BaseModel):
+class RelationOut(BaseModel):
+    id: int
+
+class EntityBase(BaseModel):
+    # id assigned at creation via model_validator
+    id: str
+    name: str
+    entity_type: str
+
+    @model_validator(mode='before')
+    @classmethod
+    def generate_id(cls, data: dict) -> dict:
+        """Generate an ID for this entity, eg `20240101-basic-memory`"""
+        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}"
+        return data
+
+    def file_name(self) -> str:
+        """Get the markdown file name for this entity."""
+        return f"{self.id}.md"
+
+class EntityIn(EntityBase):
     """
     Represents a node in our knowledge graph - could be a person, project,
     concept, etc. Each entity has a unique name, a type, and a list of
     associated observations.
     """
-    id: str                         # Text ID for filesystem references
-    name: str
-    entity_type: str
-    observations: List[Observation] = []
-    relations: List[Relation] = []
+    observations: List[ObservationIn] = []
+    relations: List[RelationIn] = []
 
-    @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"
+class EntityOut(EntityBase):
+    """
+    Represents a node in our knowledge graph - could be a person, project,
+    concept, etc. Each entity has a unique name, a type, and a list of
+    associated observations.
+    """
+    observations: List[ObservationOut] = []
+    relations: List[RelationOut] = []

src/basic_memory/services/entity_service.py

@@ -3,8 +3,8 @@ from datetime import datetime, UTC
 from pathlib import Path
 
 from basic_memory.repository import EntityRepository
-from basic_memory.schemas import Entity
-from basic_memory.models import Entity as EntityModel
+from basic_memory.schemas import EntityIn
+from basic_memory.models import Entity
 from . import ServiceError
 
 class EntityService:
@@ -17,7 +17,7 @@ class EntityService:
         self.project_path = project_path
         self.entity_repo = entity_repo
 
-    async def create_entity(self, entity: Entity) -> EntityModel:
+    async def create_entity(self, entity: EntityIn) -> Entity:
         """Create a new entity in the database."""
         # Create DB record
         db_data = {
@@ -26,7 +26,7 @@ class EntityService:
         }
         return await self.entity_repo.create(db_data)
 
-    async def get_entity(self, entity_id: str) -> EntityModel:
+    async def get_entity(self, entity_id: str) -> Entity:
         """Get entity by ID."""
         db_entity = await self.entity_repo.find_by_id(entity_id)
         if not db_entity:
@@ -35,7 +35,7 @@ class EntityService:
         return db_entity
 
     # TODO name is not uniaue
-    async def get_by_name(self, name: str) -> EntityModel:
+    async def get_by_name(self, name: str) -> Entity:
         """Get entity by name."""
         db_entity = await self.entity_repo.find_by_name(name)
         if not db_entity:

src/basic_memory/services/memory_service.py

@@ -3,9 +3,9 @@ import asyncio
 from typing import List, Dict, Any, Optional
 from pathlib import Path
 
+from basic_memory.models import Entity, Observation
 from basic_memory.schemas import (
-    Entity, Observation, Relation,
-    ObservationsIn, ObservationsOut, ObservationOut
+    ObservationsIn, ObservationsOut, ObservationOut, EntityIn, RelationIn, RelationOut
 )
 from basic_memory.fileio import write_entity_file, read_entity_file, delete_entity_file
 from basic_memory.services import EntityService, RelationService, ObservationService
@@ -28,24 +28,36 @@ class MemoryService:
 
     async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]:
         """Create multiple entities with their observations."""
-        entities = [Entity.model_validate(data) for data in entities_data]
+        entities_in = [EntityIn.model_validate(data) for data in entities_data]
+        print(f"\nCreating entities with observations:")
+        for e in entities_in:
+            print(f"Entity {e.name}: {len(e.observations)} observations")
 
         # Write files in parallel (filesystem is source of truth)
-        async def write_file(entity: Entity):
+        async def write_file(entity: EntityIn):
             await write_entity_file(self.entities_path, entity)
 
-        file_writes = [write_file(entity) for entity in entities]
+        file_writes = [write_file(entity) for entity in entities_in]
         await asyncio.gather(*file_writes)
 
-        # Update database index sequentially
-        for entity in entities:
-            await self.entity_service.create_entity(entity)
+        async def create_entity_in_db(entity_in: EntityIn):
+            print(f"\nCreating entity in DB: {entity_in.name}")
+            db_entity = await self.entity_service.create_entity(entity_in)
+            print(f"Adding {len(entity_in.observations)} observations to DB for {entity_in.name}")
+            await self.observation_service.add_observations(entity_in, entity_in.observations)
+            [await self.relation_service.create_relation(relation_in) for relation_in in entity_in.relations]
+            # query the entity again to return relations
+            final_entity = await self.entity_service.get_entity(entity_in.id)
+            print(f"Final entity {final_entity.name} has {len(final_entity.observations)} observations in DB")
+            return final_entity
 
+        # Update database index sequentially
+        entities = [await create_entity_in_db(entities_in) for entities_in in entities_in]
         return entities
 
-    async def create_relations(self, relations_data: List[Dict[str, Any]]) -> List[Relation]:
+    async def create_relations(self, relations_data: List[Dict[str, Any]]) -> List[RelationOut]:
         """Create multiple relations between entities."""
-        relations = [Relation.model_validate(data) for data in relations_data]
+        relations = [RelationIn.model_validate(data) for data in relations_data]
 
         for relation in relations:
             # First read complete entities from filesystem
@@ -68,46 +80,39 @@ class MemoryService:
 
         return relations
 
-    async def add_observations(self, observations_in: Dict[str, Any]) -> ObservationsOut:
+    async def add_observations(self, observations_in: Dict[str, Any]) -> List[Observation]:
         """Add observations to an existing entity.
         
         Args:
-            observations_in: input containing entity_name and observations
+            observations_in: input containing entity_id and observations
             
         Returns:
-            ObservationsOut containing the created observations with IDs
+            List[Observation] with the newly created observations
         """
         # Create new observations
         new_observations = ObservationsIn.model_validate(observations_in)
+        print(f"\nAdding new observations to entity {new_observations.entity_id}")
+        print(f"New observations to add: {len(new_observations.observations)}")
 
         # Read entity from filesystem
         entity = await read_entity_file(self.entities_path, new_observations.entity_id)
-
-        # Convert ObservationIn to Observation before adding to entity
-        entity_observations = [
-            Observation(content=obs.content, context=obs.context)
-            for obs in new_observations.observations
-        ]
-        entity.observations.extend(entity_observations)
+        print(f"Entity {entity.id} from file has {len(entity.observations)} observations")
+        
+        # Create new observations for the entity
+        for obs in new_observations.observations:
+            entity.observations.append(obs)
+        print(f"After appending, entity has {len(entity.observations)} observations")
 
         # Write updated entity file
         await write_entity_file(self.entities_path, entity)
         
         # Update database index
         added_observations = await self.observation_service.add_observations(entity, new_observations.observations)
+        print(f"Added {len(added_observations)} observations to DB")
         
-        # Create and return output model
-        return ObservationsOut(
-            entity_id=entity.id,
-            observations=[
-                ObservationOut(
-                    id=obs.id,
-                    content=obs.content,
-                    context=obs.context
-                )
-                for obs in added_observations
-            ]
-        )
+        db_entity = await self.entity_service.get_entity(entity.id)
+        print(f"Entity {entity.id} in DB now has {len(db_entity.observations)} observations")
+        return added_observations
 
     async def delete_entities(self, entity_names: List[str]) -> None:
         """Delete multiple entities and their associated data."""

src/basic_memory/services/observation_service.py

@@ -1,15 +1,13 @@
 """Service for managing observations in both filesystem and database."""
 from datetime import datetime, UTC
 from pathlib import Path
-from typing import Optional, List
-from uuid import uuid4
+from typing import List
 from sqlalchemy import select, delete
 
-from basic_memory.models import Observation as DbObservation
+from basic_memory.models import Observation
 from basic_memory.repository import ObservationRepository
-from basic_memory.schemas import Entity, Observation, ObservationIn
-from basic_memory.models import Observation as ObservationModel
-from . import ServiceError, DatabaseSyncError
+from basic_memory.schemas import EntityIn, ObservationIn
+from . import DatabaseSyncError
 
 
 class ObservationService:
@@ -22,38 +20,31 @@ class ObservationService:
         self.project_path = project_path
         self.observation_repo = observation_repo
         
-    async def add_observations(self, entity: Entity, observations: List[ObservationIn]) -> List[Observation]:
+    async def add_observations(self, entity: EntityIn, observations: List[ObservationIn]) -> List[Observation]:
         """
         Add multiple observations to an entity.
         Returns the created observations with IDs set.
         """
-        created_observations = []
+        print(f"\nObservationService.add_observations called for entity {entity.id}")
+        print(f"Adding {len(observations)} observations")
 
         async def add_observation(observation: ObservationIn) -> Observation:
             try:
-                db_observation = await self.observation_repo.create({
+                return await self.observation_repo.create({
                     'entity_id': entity.id,
                     'content': observation.content,
                     'context': observation.context,
                     'created_at': datetime.now(UTC)
                 })
-                # Convert db model to schema
-                return Observation(
-                    id=db_observation.id,
-                    content=db_observation.content,
-                    context=db_observation.context
-                )
             except Exception as e:
                 raise DatabaseSyncError(f"Failed to add observation to database: {str(e)}") from e
 
         # Add each observation and collect the results
         created_observations = [await add_observation(obs) for obs in observations]
-
-        # Update entity in memory with the created observations that have IDs
-        entity.observations.extend(created_observations)
+        print(f"Created {len(created_observations)} observations in DB")
         return created_observations
 
-    async def search_observations(self, query: str) -> List[ObservationModel]:
+    async def search_observations(self, query: str) -> List[Observation]:
         """
         Search for observations across all entities.
         
@@ -64,8 +55,8 @@ class ObservationService:
             List of matching observations with their entity contexts
         """
         result = await self.observation_repo.execute_query(
-            select(DbObservation).filter(
-                DbObservation.content.contains(query)
+            select(Observation).filter(
+                Observation.content.contains(query)
             )
         )
         return [
@@ -73,29 +64,10 @@ class ObservationService:
             for obs in result.scalars().all()
         ]
         
-    async def get_observations_by_context(self, context: str) -> List[ObservationModel]:
+    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
-        ]
-
-    async def rebuild_observation_index(self, entity: Entity) -> None:
-        """
-        Rebuild the observation database index for a specific entity.
-        Used for recovery or ensuring sync.
-        """
-        # Clear existing observations for this entity
-        await self.observation_repo.execute_query(
-            delete(DbObservation).where(DbObservation.entity_id == entity.id)
-        )
-        
-        # Rebuild from entity's observations
-        for obs in entity.observations:
-            await self.observation_repo.create({
-                'entity_id': entity.id,
-                'content': obs.content,
-                'context': obs.context,
-                'created_at': datetime.now(UTC)
-            })
+        ]

src/basic_memory/services/relation_service.py

@@ -4,9 +4,9 @@ from pathlib import Path
 from typing import Dict, Any
 from sqlalchemy import delete
 
-from basic_memory.models import Relation as DbRelation
+from basic_memory.models import Relation as DbRelation, Relation
 from basic_memory.repository import RelationRepository
-from basic_memory.schemas import Entity, Relation
+from basic_memory.schemas import EntityIn, RelationIn
 from . import ServiceError, DatabaseSyncError, RelationError
 
 
@@ -20,7 +20,7 @@ class RelationService:
         self.project_path = project_path
         self.relation_repo = relation_repo
 
-    async def create_relation(self, relation: Relation) -> Relation:
+    async def create_relation(self, relation: RelationIn) -> Relation:
         """Create a new relation in the database."""
         try:
             db_data = relation.model_dump()
@@ -30,7 +30,7 @@ class RelationService:
         except Exception as e:
             raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e
 
-    async def delete_relation(self, from_entity: Entity, to_entity: Entity, relation_type: str) -> bool:
+    async def delete_relation(self, from_entity: EntityIn, to_entity: EntityIn, relation_type: str) -> bool:
         """Delete a specific relation between entities."""
         # Find and remove the relation from the entity's relations
         if hasattr(from_entity, 'relations'):

tests/test_memory_service.py

@@ -2,8 +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
-from basic_memory.schemas import ObservationsIn, ObservationIn
+from basic_memory.models import Entity as EntityModel, Observation, Relation
+from basic_memory.schemas import EntityIn
 
 test_entities_data = [
     {
@@ -21,10 +21,10 @@ test_entities_data = [
 @pytest.mark.asyncio
 async def test_create_entities(memory_service: MemoryService):
     """Should create multiple entities in parallel with their observations."""
-    # Create entities
+    # Create entities - returns List[models.Entity]
     entities = await memory_service.create_entities(test_entities_data)
 
-    # Verify the entities were created
+    # Verify the SQLAlchemy models were created
     assert len(entities) == 2
 
     # Check first entity
@@ -41,16 +41,17 @@ async def test_create_entities(memory_service: MemoryService):
     assert entities[1].observations[0].content == "Observation 2.1"
     assert entities[1].observations[1].content == "Observation 2.2"
 
-    # Verify files were created
-    entity1_path = memory_service.entities_path / entities[0].file_name()
-    entity2_path = memory_service.entities_path / entities[1].file_name()
-    assert entity1_path.exists()
-    assert entity2_path.exists()
+    # Verify files were created (returns Pydantic Entity)
+    file_entity1 = await read_entity_file(memory_service.entities_path, entities[0].id)
+    file_entity2 = await read_entity_file(memory_service.entities_path, entities[1].id)
+    
+    assert file_entity1.name == "Test_Entity_1"
+    assert file_entity2.name == "Test_Entity_2"
 
 @pytest.mark.asyncio
 async def test_add_observations(memory_service: MemoryService):
     """Should add observations to an existing entity."""
-    # First create an entity
+    # First create an entity - returns SQLAlchemy Entity
     entities = await memory_service.create_entities([test_entities_data[0]])
     entity = entities[0]
 
@@ -63,25 +64,24 @@ async def test_add_observations(memory_service: MemoryService):
         ]
     }
 
-    # Add observations
-    result = await memory_service.add_observations(observations_data)
+    # Add observations - returns List[models.Observation]
+    added_observations = await memory_service.add_observations(observations_data)
 
-    # Check the result
-    assert result.entity_id == entity.id
-    assert len(result.observations) == 2
-    assert result.observations[0].content == "New observation 1"
-    assert result.observations[0].context is None
-    assert result.observations[1].content == "New observation 2"
-    assert result.observations[1].context == "test context"
+    # Check the SQLAlchemy model results
+    assert len(added_observations) == 2
+    assert added_observations[0].content == "New observation 1"
+    assert added_observations[0].context is None
+    assert added_observations[1].content == "New observation 2" 
+    assert added_observations[1].context == "test context"
 
-    # Verify file was updated
+    # Verify file was updated - returns Pydantic Entity
     updated_entity = await read_entity_file(memory_service.entities_path, entity.id)
     assert len(updated_entity.observations) == 4  # 2 original + 2 new
     assert updated_entity.observations[2].content == "New observation 1"
     assert updated_entity.observations[3].content == "New observation 2"
     assert updated_entity.observations[3].context == "test context"
 
-    # Verify database was updated via observation service
+    # Verify database - returns SQLAlchemy Entity
     db_entity = await memory_service.entity_service.get_entity(entity.id)
     assert len(db_entity.observations) == 4
 
@@ -99,7 +99,7 @@ async def test_add_observations_nonexistent_entity(memory_service: MemoryService
 @pytest.mark.asyncio
 async def test_create_relations(memory_service: MemoryService):
     """Should create relations between entities and update both filesystem and database."""
-    # First create the entities
+    # First create entities - returns List[models.Entity]
     entities = await memory_service.create_entities(test_entities_data)
     entity1, entity2 = entities
 
@@ -118,10 +118,10 @@ async def test_create_relations(memory_service: MemoryService):
         }
     ]
 
-    # Create relations
+    # Create relations - returns List[models.Relation]
     relations = await memory_service.create_relations(test_relations_data)
 
-    # Verify relations were created
+    # Verify SQLAlchemy Relation models were created
     assert len(relations) == 2
 
     # Check first relation
@@ -136,7 +136,7 @@ async def test_create_relations(memory_service: MemoryService):
     assert relations[1].relation_type == "references"
     assert relations[1].context == "test context"
 
-    # Read updated entities from filesystem to verify relations
+    # Read updated entities from filesystem - returns Pydantic Entities
     updated_entity1 = await read_entity_file(memory_service.entities_path, entity1.id)
     updated_entity2 = await read_entity_file(memory_service.entities_path, entity2.id)
 
@@ -153,10 +153,9 @@ async def test_create_relations(memory_service: MemoryService):
     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)
+    # Now verify database state - get SQLAlchemy Entity models
+    db_entity1 = await memory_service.entity_service.get_entity(entity1.id)
+    db_entity2 = 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
@@ -181,7 +180,7 @@ async def test_create_relations(memory_service: MemoryService):
 @pytest.mark.asyncio
 async def test_create_relations_with_invalid_entity_id(memory_service: MemoryService):
     """Should raise an appropriate error when trying to create relations with non-existent entity IDs."""
-    # Create only one entity
+    # Create one entity - returns SQLAlchemy Entity
     entities = await memory_service.create_entities([test_entities_data[0]])
     entity1 = entities[0]