refactor schemas and services

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] = []