all tests passing

src/basic_memory/services/entity_service.py

@@ -3,10 +3,12 @@ from datetime import datetime, UTC
 from pathlib import Path
 
 from basic_memory.repository import EntityRepository
-from basic_memory.schemas import EntityIn
-from basic_memory.models import Entity
+from basic_memory.schemas import EntityIn, ObservationIn
+from basic_memory.models import Entity, Observation
+from basic_memory.fileio import EntityNotFoundError
 from . import ServiceError
 
+
 class EntityService:
     """
     Service for managing entities in the database.
@@ -18,28 +20,35 @@ class EntityService:
         self.entity_repo = entity_repo
 
     async def create_entity(self, entity: EntityIn) -> Entity:
-        """Create a new entity in the database."""
-        # Create DB record
-        db_data = {
-            **entity.model_dump(),
+        """Create a new entity in the database.
+        
+        Note: ID is generated by the EntityIn validator before reaching this method.
+        """
+        # Create base entity first
+        base_data = {
+            "id": entity.id,  # Include the generated ID
+            "name": entity.name,
+            "entity_type": entity.entity_type,
             "created_at": datetime.now(UTC),
         }
-        return await self.entity_repo.create(db_data)
+        created_entity = await self.entity_repo.create(base_data)
+        await self.entity_repo.refresh(created_entity, ['observations', 'outgoing_relations', 'incoming_relations'])
+        return created_entity
 
     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:
-            raise ServiceError(f"Entity not found: {entity_id}")
+            raise EntityNotFoundError(f"Entity not found: {entity_id}")
             
         return db_entity
 
-    # TODO name is not uniaue
+    # TODO name is not unique
     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:
-            raise ServiceError(f"Entity not found: {name}")
+            raise EntityNotFoundError(f"Entity not found: {name}")
             
         return db_entity
 

src/basic_memory/services/memory_service.py

@@ -29,9 +29,6 @@ class MemoryService:
     async def create_entities(self, entities_data: List[Dict[str, Any]]) -> List[Entity]:
         """Create multiple entities with their observations."""
         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: EntityIn):
@@ -41,14 +38,11 @@ class MemoryService:
         await asyncio.gather(*file_writes)
 
         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.entity_service.create_entity(entity_in)
             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
@@ -91,27 +85,21 @@ class MemoryService:
         """
         # 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)
-        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")
-        
+
         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:

src/basic_memory/services/observation_service.py

@@ -25,23 +25,31 @@ class ObservationService:
         Add multiple observations to an entity.
         Returns the created observations with IDs set.
         """
-        print(f"\nObservationService.add_observations called for entity {entity.id}")
-        print(f"Adding {len(observations)} observations")
-
         async def add_observation(observation: ObservationIn) -> Observation:
             try:
-                return await self.observation_repo.create({
+                obs = await self.observation_repo.create({
                     'entity_id': entity.id,
                     'content': observation.content,
                     'context': observation.context,
                     'created_at': datetime.now(UTC)
                 })
+                # Ensure each observation is flushed
+                await self.observation_repo.session.flush()
+                # Refresh to get latest state
+                await self.observation_repo.session.refresh(obs)
+                return obs
             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]
-        print(f"Created {len(created_observations)} observations in DB")
+
+        # Make sure observations are in sync before returning
+        # This helps ensure related entities see the new observations
+        await self.observation_repo.session.flush()
+        for obs in created_observations:
+            await self.observation_repo.session.refresh(obs)
+            
         return created_observations
 
     async def search_observations(self, query: str) -> List[Observation]:

src/basic_memory/services/relation_service.py

@@ -25,8 +25,7 @@ class RelationService:
         try:
             db_data = relation.model_dump()
             db_data['created_at'] = datetime.now(UTC)
-            await self.relation_repo.create(db_data)
-            return relation
+            return await self.relation_repo.create(db_data)
         except Exception as e:
             raise DatabaseSyncError(f"Failed to sync relation to database: {str(e)}") from e