refactor schemas and services

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."""