basicmachines-co-basic-memory/src/basic_memory/models/knowledge.py

"""Knowledge graph models for basic-memory."""

from datetime import datetime
from typing import List, Optional

from sqlalchemy import (
    String, DateTime, ForeignKey, Text, Integer, 
    text, UniqueConstraint
)
from sqlalchemy.orm import Mapped, mapped_column, relationship

from basic_memory.models.base import Base
from basic_memory.utils import sanitize_name


class Entity(Base):
    """
    Core entity in the knowledge graph.

    Entities are the primary nodes in the knowledge graph. Each entity has:
    - A unique identifier (text, based on type/name path)
    - A name
    - An entity type (e.g., "person", "organization", "event")
    - A description (optional)
    - A list of observations
    """
    __tablename__ = "entity"
    __table_args__ = (
        UniqueConstraint("entity_type", "name", name="uix_entity_type_name"),
    )

    id: Mapped[str] = mapped_column(String, primary_key=True)
    name: Mapped[str] = mapped_column(String)
    entity_type: Mapped[str] = mapped_column(String)
    description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    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")
    )

    # Link to source document
    doc_id: Mapped[Optional[int]] = mapped_column(
        Integer,
        ForeignKey("documents.id", ondelete="SET NULL"),
        nullable=True
    )

    # Relationships
    document: Mapped["Document"] = relationship(
        "Document",
        back_populates="entities"
    )
    observations: Mapped[List["Observation"]] = relationship(
        "Observation",
        back_populates="entity",
        cascade="all, delete-orphan"
    )
    outgoing_relations: Mapped[List["Relation"]] = relationship(
        "Relation",
        foreign_keys="[Relation.from_id]",
        back_populates="from_entity",
        cascade="all, delete-orphan"
    )
    incoming_relations: Mapped[List["Relation"]] = relationship(
        "Relation",
        foreign_keys="[Relation.to_id]",
        back_populates="to_entity",
        cascade="all, delete-orphan"
    )

    @property
    def relations(self):
        return self.outgoing_relations + self.incoming_relations

    @classmethod
    def generate_id(cls, entity_type: str, name: str) -> str:
        """Generate a filesystem path-based ID for this entity."""
        # Use common normalization for filesystem safety
        safe_name = sanitize_name(name)
        return f"{entity_type}/{safe_name}"

    def get_file_path(self) -> str:
        """Get the filesystem path for this entity."""
        return f"{self.id}.md"  # id is already in path format

    def __repr__(self) -> str:
        return f"Entity(id='{self.id}', name='{self.name}', type='{self.entity_type}')"


class Observation(Base):
    """
    Observations are discrete pieces of information about an entity. They are:
    - Stored as strings
    - Attached to specific entities
    - Can be added or removed independently
    - Should be atomic (one fact per observation)
    """
    __tablename__ = "observation"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    entity_id: Mapped[str] = mapped_column(
        String,
        ForeignKey("entity.id", ondelete="CASCADE"),
        index=True
    )
    content: Mapped[str] = mapped_column(String)
    created_at: Mapped[datetime] = mapped_column(
        DateTime,
        server_default=text("CURRENT_TIMESTAMP")
    )
    context: Mapped[Optional[str]] = mapped_column(String, nullable=True)

    # Relationships
    entity: Mapped[Entity] = relationship("Entity", back_populates="observations")

    def __repr__(self) -> str:
        content = self.content[:50] + "..." if len(self.content) > 50 else self.content
        return f"Observation(id={self.id}, entity='{self.entity_id}', content='{content}')"


class Relation(Base):
    """
    Relations define directed connections between entities.
    They are always stored in active voice and describe how entities 
    interact or relate to each other.
    """
    __tablename__ = "relation"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    from_id: Mapped[str] = mapped_column(
        String,
        ForeignKey("entity.id", ondelete="CASCADE"),
        index=True
    )
    to_id: Mapped[str] = mapped_column(
        String,
        ForeignKey("entity.id", ondelete="CASCADE"),
        index=True
    )
    relation_type: Mapped[str] = mapped_column(String)
    created_at: Mapped[datetime] = mapped_column(
        DateTime,
        server_default=text("CURRENT_TIMESTAMP")
    )
    context: Mapped[Optional[str]] = mapped_column(String, nullable=True)

    # Relationships
    from_entity: Mapped[Entity] = relationship(
        "Entity",
        foreign_keys=[from_id],
        back_populates="outgoing_relations"
    )
    to_entity: Mapped[Entity] = relationship(
        "Entity",
        foreign_keys=[to_id],
        back_populates="incoming_relations"
    )

    def __repr__(self) -> str:
        return f"Relation(id={self.id}, from='{self.from_id}', type='{self.relation_type}', to='{self.to_id}')"