hindsight-api 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

hindsight_api/models.py

@@ -1,49 +1,47 @@
 """
 SQLAlchemy models for the memory system.
 """
+
 from datetime import datetime
-from typing import Optional
-from uuid import UUID as PyUUID, uuid4
+from uuid import UUID as PyUUID
 
+from pgvector.sqlalchemy import Vector
 from sqlalchemy import (
     CheckConstraint,
-    Column,
     Float,
     ForeignKey,
     ForeignKeyConstraint,
     Index,
     Integer,
-    PrimaryKeyConstraint,
     Text,
     func,
+)
+from sqlalchemy import (
     text as sql_text,
 )
 from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP, UUID
 from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
-from pgvector.sqlalchemy import Vector
 
 
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models."""
+
     pass
 
 
 class Document(Base):
     """Source documents for memory units."""
+
     __tablename__ = "documents"
 
     id: Mapped[str] = mapped_column(Text, primary_key=True)
     bank_id: Mapped[str] = mapped_column(Text, primary_key=True)
-    original_text: Mapped[Optional[str]] = mapped_column(Text)
-    content_hash: Mapped[Optional[str]] = mapped_column(Text)
+    original_text: Mapped[str | None] = mapped_column(Text)
+    content_hash: Mapped[str | None] = mapped_column(Text)
     doc_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
-    created_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
-    updated_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
 
     # Relationships
     memory_units = relationship("MemoryUnit", back_populates="document", cascade="all, delete-orphan")
@@ -56,45 +54,42 @@ class Document(Base):
 
 class MemoryUnit(Base):
     """Individual sentence-level memories."""
+
     __tablename__ = "memory_units"
 
     id: Mapped[PyUUID] = mapped_column(
         UUID(as_uuid=True), primary_key=True, server_default=sql_text("gen_random_uuid()")
     )
     bank_id: Mapped[str] = mapped_column(Text, nullable=False)
-    document_id: Mapped[Optional[str]] = mapped_column(Text)
+    document_id: Mapped[str | None] = mapped_column(Text)
     text: Mapped[str] = mapped_column(Text, nullable=False)
     embedding = mapped_column(Vector(384))  # pgvector type
-    context: Mapped[Optional[str]] = mapped_column(Text)
-    event_date: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False)  # Kept for backward compatibility
-    occurred_start: Mapped[Optional[datetime]] = mapped_column(TIMESTAMP(timezone=True))  # When fact occurred (range start)
-    occurred_end: Mapped[Optional[datetime]] = mapped_column(TIMESTAMP(timezone=True))    # When fact occurred (range end)
-    mentioned_at: Mapped[Optional[datetime]] = mapped_column(TIMESTAMP(timezone=True))    # When fact was mentioned
+    context: Mapped[str | None] = mapped_column(Text)
+    event_date: Mapped[datetime] = mapped_column(
+        TIMESTAMP(timezone=True), nullable=False
+    )  # Kept for backward compatibility
+    occurred_start: Mapped[datetime | None] = mapped_column(
+        TIMESTAMP(timezone=True)
+    )  # When fact occurred (range start)
+    occurred_end: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True))  # When fact occurred (range end)
+    mentioned_at: Mapped[datetime | None] = mapped_column(TIMESTAMP(timezone=True))  # When fact was mentioned
     fact_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="world")
-    confidence_score: Mapped[Optional[float]] = mapped_column(Float)
+    confidence_score: Mapped[float | None] = mapped_column(Float)
     access_count: Mapped[int] = mapped_column(Integer, server_default="0")
-    unit_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))  # User-defined metadata (str->str)
-    created_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
-    updated_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
+    unit_metadata: Mapped[dict] = mapped_column(
+        "metadata", JSONB, server_default=sql_text("'{}'::jsonb")
+    )  # User-defined metadata (str->str)
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
 
     # Relationships
     document = relationship("Document", back_populates="memory_units")
     unit_entities = relationship("UnitEntity", back_populates="memory_unit", cascade="all, delete-orphan")
     outgoing_links = relationship(
-        "MemoryLink",
-        foreign_keys="MemoryLink.from_unit_id",
-        back_populates="from_unit",
-        cascade="all, delete-orphan"
+        "MemoryLink", foreign_keys="MemoryLink.from_unit_id", back_populates="from_unit", cascade="all, delete-orphan"
     )
     incoming_links = relationship(
-        "MemoryLink",
-        foreign_keys="MemoryLink.to_unit_id",
-        back_populates="to_unit",
-        cascade="all, delete-orphan"
+        "MemoryLink", foreign_keys="MemoryLink.to_unit_id", back_populates="to_unit", cascade="all, delete-orphan"
     )
 
     __table_args__ = (
@@ -110,7 +105,7 @@ class MemoryUnit(Base):
             "(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
             "(fact_type = 'observation') OR "
             "(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)",
-            name="confidence_score_fact_type_check"
+            name="confidence_score_fact_type_check",
         ),
         Index("idx_memory_units_bank_id", "bank_id"),
         Index("idx_memory_units_document_id", "document_id"),
@@ -119,39 +114,46 @@ class MemoryUnit(Base):
         Index("idx_memory_units_access_count", "access_count", postgresql_ops={"access_count": "DESC"}),
         Index("idx_memory_units_fact_type", "fact_type"),
         Index("idx_memory_units_bank_fact_type", "bank_id", "fact_type"),
-        Index("idx_memory_units_bank_type_date", "bank_id", "fact_type", "event_date", postgresql_ops={"event_date": "DESC"}),
+        Index(
+            "idx_memory_units_bank_type_date",
+            "bank_id",
+            "fact_type",
+            "event_date",
+            postgresql_ops={"event_date": "DESC"},
+        ),
         Index(
             "idx_memory_units_opinion_confidence",
             "bank_id",
             "confidence_score",
             postgresql_where=sql_text("fact_type = 'opinion'"),
-            postgresql_ops={"confidence_score": "DESC"}
+            postgresql_ops={"confidence_score": "DESC"},
         ),
         Index(
             "idx_memory_units_opinion_date",
             "bank_id",
             "event_date",
             postgresql_where=sql_text("fact_type = 'opinion'"),
-            postgresql_ops={"event_date": "DESC"}
+            postgresql_ops={"event_date": "DESC"},
         ),
         Index(
             "idx_memory_units_observation_date",
             "bank_id",
             "event_date",
             postgresql_where=sql_text("fact_type = 'observation'"),
-            postgresql_ops={"event_date": "DESC"}
+            postgresql_ops={"event_date": "DESC"},
         ),
         Index(
             "idx_memory_units_embedding",
             "embedding",
             postgresql_using="hnsw",
-            postgresql_ops={"embedding": "vector_cosine_ops"}
+            postgresql_ops={"embedding": "vector_cosine_ops"},
         ),
     )
 
 
 class Entity(Base):
     """Resolved entities (people, organizations, locations, etc.)."""
+
     __tablename__ = "entities"
 
     id: Mapped[PyUUID] = mapped_column(
@@ -160,12 +162,8 @@ class Entity(Base):
     canonical_name: Mapped[str] = mapped_column(Text, nullable=False)
     bank_id: Mapped[str] = mapped_column(Text, nullable=False)
     entity_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
-    first_seen: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
-    last_seen: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
+    first_seen: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
+    last_seen: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
     mention_count: Mapped[int] = mapped_column(Integer, server_default="1")
 
     # Relationships
@@ -175,13 +173,13 @@ class Entity(Base):
         "EntityCooccurrence",
         foreign_keys="EntityCooccurrence.entity_id_1",
         back_populates="entity_1",
-        cascade="all, delete-orphan"
+        cascade="all, delete-orphan",
     )
     cooccurrences_2 = relationship(
         "EntityCooccurrence",
         foreign_keys="EntityCooccurrence.entity_id_2",
         back_populates="entity_2",
-        cascade="all, delete-orphan"
+        cascade="all, delete-orphan",
     )
 
     __table_args__ = (
@@ -193,6 +191,7 @@ class Entity(Base):
 
 class UnitEntity(Base):
     """Association between memory units and entities."""
+
     __tablename__ = "unit_entities"
 
     unit_id: Mapped[PyUUID] = mapped_column(
@@ -214,6 +213,7 @@ class UnitEntity(Base):
 
 class EntityCooccurrence(Base):
     """Materialized cache of entity co-occurrences."""
+
     __tablename__ = "entity_cooccurrences"
 
     entity_id_1: Mapped[PyUUID] = mapped_column(
@@ -223,9 +223,7 @@ class EntityCooccurrence(Base):
         UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
     )
     cooccurrence_count: Mapped[int] = mapped_column(Integer, server_default="1")
-    last_cooccurred: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
+    last_cooccurred: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
 
     # Relationships
     entity_1 = relationship("Entity", foreign_keys=[entity_id_1], back_populates="cooccurrences_1")
@@ -241,6 +239,7 @@ class EntityCooccurrence(Base):
 
 class MemoryLink(Base):
     """Links between memory units (temporal, semantic, entity)."""
+
     __tablename__ = "memory_links"
 
     from_unit_id: Mapped[PyUUID] = mapped_column(
@@ -250,13 +249,11 @@ class MemoryLink(Base):
         UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
     )
     link_type: Mapped[str] = mapped_column(Text, primary_key=True)
-    entity_id: Mapped[Optional[PyUUID]] = mapped_column(
+    entity_id: Mapped[PyUUID | None] = mapped_column(
         UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
     )
     weight: Mapped[float] = mapped_column(Float, nullable=False, server_default="1.0")
-    created_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
 
     # Relationships
     from_unit = relationship("MemoryUnit", foreign_keys=[from_unit_id], back_populates="outgoing_links")
@@ -266,7 +263,7 @@ class MemoryLink(Base):
     __table_args__ = (
         CheckConstraint(
             "link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')",
-            name="memory_links_link_type_check"
+            name="memory_links_link_type_check",
         ),
         CheckConstraint("weight >= 0.0 AND weight <= 1.0", name="memory_links_weight_check"),
         Index("idx_memory_links_from", "from_unit_id"),
@@ -278,31 +275,22 @@ class MemoryLink(Base):
             "from_unit_id",
             "weight",
             postgresql_where=sql_text("weight >= 0.1"),
-            postgresql_ops={"weight": "DESC"}
+            postgresql_ops={"weight": "DESC"},
         ),
     )
 
 
 class Bank(Base):
     """Memory bank profiles with disposition traits and background."""
+
     __tablename__ = "banks"
 
     bank_id: Mapped[str] = mapped_column(Text, primary_key=True)
     disposition: Mapped[dict] = mapped_column(
-        JSONB,
-        nullable=False,
-        server_default=sql_text(
-            '\'{"skepticism": 3, "literalism": 3, "empathy": 3}\'::jsonb'
-        )
+        JSONB, nullable=False, server_default=sql_text('\'{"skepticism": 3, "literalism": 3, "empathy": 3}\'::jsonb')
     )
     background: Mapped[str] = mapped_column(Text, nullable=False, server_default="")
-    created_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
-    updated_at: Mapped[datetime] = mapped_column(
-        TIMESTAMP(timezone=True), server_default=func.now()
-    )
+    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
 
-    __table_args__ = (
-        Index("idx_banks_bank_id", "bank_id"),
-    )
+    __table_args__ = (Index("idx_banks_bank_id", "bank_id"),)