296 lines
12 KiB
Python
296 lines
12 KiB
Python
"""
|
|
SQLAlchemy models for the memory system.
|
|
"""
|
|
|
|
from datetime import datetime
|
|
from uuid import UUID as PyUUID
|
|
|
|
from pgvector.sqlalchemy import Vector
|
|
from sqlalchemy import (
|
|
CheckConstraint,
|
|
Float,
|
|
ForeignKey,
|
|
ForeignKeyConstraint,
|
|
Index,
|
|
Integer,
|
|
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
|
|
|
|
|
|
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[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())
|
|
|
|
# Relationships
|
|
memory_units = relationship("MemoryUnit", back_populates="document", cascade="all, delete-orphan")
|
|
|
|
__table_args__ = (
|
|
Index("idx_documents_bank_id", "bank_id"),
|
|
Index("idx_documents_content_hash", "content_hash"),
|
|
)
|
|
|
|
|
|
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[str | None] = mapped_column(Text)
|
|
text: Mapped[str] = mapped_column(Text, nullable=False)
|
|
embedding = mapped_column(Vector(384)) # pgvector type
|
|
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[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())
|
|
|
|
# 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"
|
|
)
|
|
incoming_links = relationship(
|
|
"MemoryLink", foreign_keys="MemoryLink.to_unit_id", back_populates="to_unit", cascade="all, delete-orphan"
|
|
)
|
|
|
|
__table_args__ = (
|
|
ForeignKeyConstraint(
|
|
["document_id", "bank_id"],
|
|
["documents.id", "documents.bank_id"],
|
|
name="memory_units_document_fkey",
|
|
ondelete="CASCADE",
|
|
),
|
|
CheckConstraint("fact_type IN ('world', 'experience', 'opinion', 'observation')"),
|
|
CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)"),
|
|
CheckConstraint(
|
|
"(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",
|
|
),
|
|
Index("idx_memory_units_bank_id", "bank_id"),
|
|
Index("idx_memory_units_document_id", "document_id"),
|
|
Index("idx_memory_units_event_date", "event_date", postgresql_ops={"event_date": "DESC"}),
|
|
Index("idx_memory_units_bank_date", "bank_id", "event_date", postgresql_ops={"event_date": "DESC"}),
|
|
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_opinion_confidence",
|
|
"bank_id",
|
|
"confidence_score",
|
|
postgresql_where=sql_text("fact_type = 'opinion'"),
|
|
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"},
|
|
),
|
|
Index(
|
|
"idx_memory_units_observation_date",
|
|
"bank_id",
|
|
"event_date",
|
|
postgresql_where=sql_text("fact_type = 'observation'"),
|
|
postgresql_ops={"event_date": "DESC"},
|
|
),
|
|
Index(
|
|
"idx_memory_units_embedding",
|
|
"embedding",
|
|
postgresql_using="hnsw",
|
|
postgresql_ops={"embedding": "vector_cosine_ops"},
|
|
),
|
|
)
|
|
|
|
|
|
class Entity(Base):
|
|
"""Resolved entities (people, organizations, locations, etc.)."""
|
|
|
|
__tablename__ = "entities"
|
|
|
|
id: Mapped[PyUUID] = mapped_column(
|
|
UUID(as_uuid=True), primary_key=True, server_default=sql_text("gen_random_uuid()")
|
|
)
|
|
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())
|
|
mention_count: Mapped[int] = mapped_column(Integer, server_default="1")
|
|
|
|
# Relationships
|
|
unit_entities = relationship("UnitEntity", back_populates="entity", cascade="all, delete-orphan")
|
|
memory_links = relationship("MemoryLink", back_populates="entity", cascade="all, delete-orphan")
|
|
cooccurrences_1 = relationship(
|
|
"EntityCooccurrence",
|
|
foreign_keys="EntityCooccurrence.entity_id_1",
|
|
back_populates="entity_1",
|
|
cascade="all, delete-orphan",
|
|
)
|
|
cooccurrences_2 = relationship(
|
|
"EntityCooccurrence",
|
|
foreign_keys="EntityCooccurrence.entity_id_2",
|
|
back_populates="entity_2",
|
|
cascade="all, delete-orphan",
|
|
)
|
|
|
|
__table_args__ = (
|
|
Index("idx_entities_bank_id", "bank_id"),
|
|
Index("idx_entities_canonical_name", "canonical_name"),
|
|
Index("idx_entities_bank_name", "bank_id", "canonical_name"),
|
|
)
|
|
|
|
|
|
class UnitEntity(Base):
|
|
"""Association between memory units and entities."""
|
|
|
|
__tablename__ = "unit_entities"
|
|
|
|
unit_id: Mapped[PyUUID] = mapped_column(
|
|
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
|
|
)
|
|
entity_id: Mapped[PyUUID] = mapped_column(
|
|
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
|
|
)
|
|
|
|
# Relationships
|
|
memory_unit = relationship("MemoryUnit", back_populates="unit_entities")
|
|
entity = relationship("Entity", back_populates="unit_entities")
|
|
|
|
__table_args__ = (
|
|
Index("idx_unit_entities_unit", "unit_id"),
|
|
Index("idx_unit_entities_entity", "entity_id"),
|
|
)
|
|
|
|
|
|
class EntityCooccurrence(Base):
|
|
"""Materialized cache of entity co-occurrences."""
|
|
|
|
__tablename__ = "entity_cooccurrences"
|
|
|
|
entity_id_1: Mapped[PyUUID] = mapped_column(
|
|
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
|
|
)
|
|
entity_id_2: Mapped[PyUUID] = mapped_column(
|
|
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())
|
|
|
|
# Relationships
|
|
entity_1 = relationship("Entity", foreign_keys=[entity_id_1], back_populates="cooccurrences_1")
|
|
entity_2 = relationship("Entity", foreign_keys=[entity_id_2], back_populates="cooccurrences_2")
|
|
|
|
__table_args__ = (
|
|
CheckConstraint("entity_id_1 < entity_id_2", name="entity_cooccurrence_order_check"),
|
|
Index("idx_entity_cooccurrences_entity1", "entity_id_1"),
|
|
Index("idx_entity_cooccurrences_entity2", "entity_id_2"),
|
|
Index("idx_entity_cooccurrences_count", "cooccurrence_count", postgresql_ops={"cooccurrence_count": "DESC"}),
|
|
)
|
|
|
|
|
|
class MemoryLink(Base):
|
|
"""Links between memory units (temporal, semantic, entity)."""
|
|
|
|
__tablename__ = "memory_links"
|
|
|
|
from_unit_id: Mapped[PyUUID] = mapped_column(
|
|
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
|
|
)
|
|
to_unit_id: Mapped[PyUUID] = mapped_column(
|
|
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[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())
|
|
|
|
# Relationships
|
|
from_unit = relationship("MemoryUnit", foreign_keys=[from_unit_id], back_populates="outgoing_links")
|
|
to_unit = relationship("MemoryUnit", foreign_keys=[to_unit_id], back_populates="incoming_links")
|
|
entity = relationship("Entity", back_populates="memory_links")
|
|
|
|
__table_args__ = (
|
|
CheckConstraint(
|
|
"link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')",
|
|
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"),
|
|
Index("idx_memory_links_to", "to_unit_id"),
|
|
Index("idx_memory_links_type", "link_type"),
|
|
Index("idx_memory_links_entity", "entity_id", postgresql_where=sql_text("entity_id IS NOT NULL")),
|
|
Index(
|
|
"idx_memory_links_from_weight",
|
|
"from_unit_id",
|
|
"weight",
|
|
postgresql_where=sql_text("weight >= 0.1"),
|
|
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')
|
|
)
|
|
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())
|
|
|
|
__table_args__ = (Index("idx_banks_bank_id", "bank_id"),)
|