second-brain/src/second_brain/models.py

"""
SQLAlchemy ORM models for second-brain.

Tables:
- Source      — a queued/processed URL (video or article)
- Extraction  — structured LLM output for a source
- WikiPage    — track which vault pages have been written / updated
"""

from __future__ import annotations

import enum
from datetime import datetime, timezone
from typing import Optional


def utcnow() -> datetime:
    """Naive UTC timestamp. Replaces `datetime.utcnow()` (deprecated in 3.12)."""
    return datetime.now(timezone.utc).replace(tzinfo=None)

from sqlalchemy import (
    DateTime,
    Enum,
    ForeignKey,
    Integer,
    JSON,
    String,
    Text,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class Base(DeclarativeBase):
    """Shared declarative base."""


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------


class SourceType(enum.Enum):
    VIDEO = "video"
    ARTICLE = "article"


class SourceStatus(enum.Enum):
    PENDING = "pending"
    PULLED = "pulled"
    TRANSCRIBED = "transcribed"
    ANALYZED = "analyzed"
    ACCEPTED = "accepted"
    PUBLISHED = "published"
    FAILED = "failed"


# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------


class Source(Base):
    """A URL queued for processing."""

    __tablename__ = "sources"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

    url: Mapped[str] = mapped_column(String(2000), unique=True, nullable=False)
    title: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
    source_type: Mapped[SourceType] = mapped_column(
        Enum(SourceType), nullable=False, default=SourceType.VIDEO
    )

    # Domain and focus
    domain: Mapped[str] = mapped_column(String(50), nullable=False, default="development")
    focus: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Pipeline status
    status: Mapped[SourceStatus] = mapped_column(
        Enum(SourceStatus), nullable=False, default=SourceStatus.PENDING
    )

    # File paths (set after pull/transcribe steps)
    media_path: Mapped[Optional[str]] = mapped_column(String(1000), nullable=True)
    transcript_path: Mapped[Optional[str]] = mapped_column(String(1000), nullable=True)
    transcript_text: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Metadata from the source
    published_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    duration_seconds: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    author: Mapped[Optional[str]] = mapped_column(String(200), nullable=True)

    # Timestamps
    ingested_at: Mapped[datetime] = mapped_column(DateTime, default=utcnow)
    updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
    error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Relationships
    extraction: Mapped[Optional["Extraction"]] = relationship(
        "Extraction", back_populates="source", uselist=False, cascade="all, delete-orphan"
    )

    def __repr__(self) -> str:
        return (
            f"Source(id={self.id}, domain={self.domain!r}, "
            f"status={self.status.value}, url={self.url!r})"
        )


class Extraction(Base):
    """Structured LLM extraction output for a source."""

    __tablename__ = "extractions"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    source_id: Mapped[int] = mapped_column(
        ForeignKey("sources.id", ondelete="CASCADE"), nullable=False, unique=True
    )

    # Core fields (mirrors the extraction schema contract)
    summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    key_points: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    entities: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    claims: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    open_questions: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    contradictions: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)

    # Raw LLM response (for debugging / re-parsing)
    raw_response: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Token usage
    input_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    output_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)

    created_at: Mapped[datetime] = mapped_column(DateTime, default=utcnow)

    # Relationships
    source: Mapped["Source"] = relationship("Source", back_populates="extraction")

    def __repr__(self) -> str:
        return f"Extraction(id={self.id}, source_id={self.source_id})"


class WikiPage(Base):
    """Tracks vault pages written or updated by the wiki compiler."""

    __tablename__ = "wiki_pages"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)

    # Relative path inside the vault
    vault_path: Mapped[str] = mapped_column(String(500), unique=True, nullable=False)
    domain: Mapped[str] = mapped_column(String(50), nullable=False)
    title: Mapped[str] = mapped_column(String(300), nullable=False)

    # Git commit shas before/after last compile
    git_sha_before: Mapped[Optional[str]] = mapped_column(String(40), nullable=True)
    git_sha_after: Mapped[Optional[str]] = mapped_column(String(40), nullable=True)

    created_at: Mapped[datetime] = mapped_column(DateTime, default=utcnow)
    updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)

    def __repr__(self) -> str:
        return f"WikiPage(id={self.id}, vault_path={self.vault_path!r})"


__all__ = [
    "Base",
    "Source",
    "SourceStatus",
    "SourceType",
    "Extraction",
    "WikiPage",
    "utcnow",
]