entity repository tests

src/basic_memory/models/knowledge.py

@@ -1,21 +1,25 @@
 """Knowledge graph models."""
+
 from datetime import datetime
-from typing import Optional, List
+from typing import Optional
+
 from sqlalchemy import Integer, String, Text, ForeignKey, UniqueConstraint, text, DateTime, Index
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from basic_memory.models.base import Base
 
+
 class Entity(Base):
     """
     Core entity in the knowledge graph.
-    
+
     Entities represent semantic nodes maintained by the AI layer. Each entity:
     - Has a unique numeric ID (database-generated)
-    - Maps to a document file on disk (optional) 
+    - Maps to a document file on disk (optional)
     - Maintains a checksum for change detection
     - Tracks both source document and semantic properties
     """
+
     __tablename__ = "entity"
     __table_args__ = (
         UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),
@@ -27,32 +31,27 @@ class Entity(Base):
     id: Mapped[int] = mapped_column(Integer, primary_key=True)
     name: Mapped[str] = mapped_column(String)
     entity_type: Mapped[str] = mapped_column(String)
-    
+
     # Content and validation
     description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
     path: Mapped[Optional[str]] = mapped_column(String, nullable=True)
     checksum: Mapped[Optional[str]] = mapped_column(String, nullable=True)
-    
+
     # Metadata and tracking
-    created_at: Mapped[datetime] = mapped_column(
-        DateTime, 
-        server_default=text("CURRENT_TIMESTAMP")
-    )
+    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
     updated_at: Mapped[datetime] = mapped_column(
-        DateTime, 
-        server_default=text("CURRENT_TIMESTAMP"), 
-        onupdate=text("CURRENT_TIMESTAMP")
+        DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP")
     )
-    
+
     # Relations
     doc_id: Mapped[Optional[int]] = mapped_column(
-        Integer, 
-        ForeignKey("documents.id", ondelete="SET NULL"), 
-        nullable=True
+        Integer, ForeignKey("documents.id", ondelete="SET NULL"), nullable=True
     )
 
     # Relationships
-    observations = relationship("Observation", back_populates="entity", cascade="all, delete-orphan")
+    observations = relationship(
+        "Observation", back_populates="entity", cascade="all, delete-orphan"
+    )
     from_relations = relationship(
         "Relation",
         back_populates="from_entity",
@@ -66,6 +65,10 @@ class Entity(Base):
         cascade="all, delete-orphan",
     )
 
+    @property
+    def relations(self):
+        return self.to_relations + self.from_relations
+
     def __repr__(self) -> str:
         return f"Entity(id={self.id}, name='{self.name}', type='{self.entity_type}')"
 
@@ -73,18 +76,16 @@ class Entity(Base):
 class Observation(Base):
     """
     An observation about an entity.
-    
+
     Observations are atomic facts or notes about an entity.
     """
+
     __tablename__ = "observations"
 
     id: Mapped[int] = mapped_column(Integer, primary_key=True)
     entity_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
     content: Mapped[str] = mapped_column(Text)
-    created_at: Mapped[datetime] = mapped_column(
-        DateTime,
-        server_default=text("CURRENT_TIMESTAMP")
-    )
+    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
 
     # Relationships
     entity = relationship("Entity", back_populates="observations")
@@ -97,6 +98,7 @@ class Relation(Base):
     """
     A directed relation between two entities.
     """
+
     __tablename__ = "relations"
     __table_args__ = (
         UniqueConstraint("from_id", "to_id", "relation_type", name="uix_relation"),
@@ -107,14 +109,11 @@ class Relation(Base):
     from_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
     to_id: Mapped[int] = mapped_column(Integer, ForeignKey("entity.id"))
     relation_type: Mapped[str] = mapped_column(String)
-    created_at: Mapped[datetime] = mapped_column(
-        DateTime,
-        server_default=text("CURRENT_TIMESTAMP")
-    )
+    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"))
 
     # Relationships
     from_entity = relationship("Entity", foreign_keys=[from_id], back_populates="from_relations")
     to_entity = relationship("Entity", foreign_keys=[to_id], back_populates="to_relations")
 
     def __repr__(self) -> str:
-        return f"Relation(from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')"
+        return f"Relation(from_id={self.from_id}, to_id={self.to_id}, type='{self.relation_type}')"