second-brain/src/second_brain/models.py

"""
SQLAlchemy ORM models for second-brain.

All tables live in the Postgres `second_brain` schema (in the petalbrain DB).
PK style is int autoincrement (SERIAL in Postgres); timestamps are naive UTC
(`timestamp without time zone`) per the locked design decisions in CLAUDE.md.

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, time, timezone
from typing import Optional

from sqlalchemy import (
    JSON,
    Boolean,
    DateTime,
    Enum,
    ForeignKey,
    Integer,
    MetaData,
    String,
    Text,
    Time,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


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


SCHEMA = "second_brain"


class Base(DeclarativeBase):
    """Shared declarative base. All tables live in the `second_brain` schema."""

    metadata = MetaData(schema=SCHEMA)


# ---------------------------------------------------------------------------
# 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"


# Native Postgres enum types, scoped to the second_brain schema so they
# don't collide with anything else in the cluster.
_SourceTypeEnum = Enum(
    SourceType,
    name="source_type",
    schema=SCHEMA,
    values_callable=lambda et: [e.value for e in et],
)
_SourceStatusEnum = Enum(
    SourceStatus,
    name="source_status",
    schema=SCHEMA,
    values_callable=lambda et: [e.value for e in et],
)


# ---------------------------------------------------------------------------
# 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(
        _SourceTypeEnum, 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(
        _SourceStatusEnum, 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)
    # Segment-level output from faster-whisper. List of dicts
    # `{id, start, end, text, words: [{start, end, word, probability}, …]}`.
    # transcript_text remains the canonical full text the extractor reads;
    # this is an additive index for future timestamp-anchored use cases.
    transcript_segments: Mapped[Optional[list]] = mapped_column(JSONB, 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)

    # Claim mechanism — see `second_brain.claim` for the queue dance.
    claimed_by: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    claimed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, 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(f"{SCHEMA}.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})"


class PipelineSettings(Base):
    """Single-row config the web dashboard edits and workers read.

    The transcription_* fields persist now and will be honored by the
    tower-side transcribe worker once it exists; the extraction_* fields
    are already wired into the dev-side scheduler. The row is pinned to
    `id=1` by a DB-level CHECK so upserts-by-PK keep the table singleton.
    """

    __tablename__ = "pipeline_settings"

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

    # Transcription (consumed by the tower-side worker — not yet built).
    transcription_enabled: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=True
    )
    transcription_active_hours_start: Mapped[Optional[time]] = mapped_column(
        Time, nullable=True
    )
    transcription_active_hours_end: Mapped[Optional[time]] = mapped_column(
        Time, nullable=True
    )
    transcription_max_concurrent_gpu_jobs: Mapped[int] = mapped_column(
        Integer, nullable=False, default=1
    )
    transcription_max_video_length_seconds: Mapped[Optional[int]] = mapped_column(
        Integer, nullable=True
    )
    transcription_max_items_per_run: Mapped[int] = mapped_column(
        Integer, nullable=False, default=10
    )

    # Extraction (consumed by the dev-side scheduler today).
    extraction_enabled: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=True
    )
    extraction_active_hours_start: Mapped[Optional[time]] = mapped_column(
        Time, nullable=True
    )
    extraction_active_hours_end: Mapped[Optional[time]] = mapped_column(
        Time, nullable=True
    )
    extraction_max_items_per_run: Mapped[int] = mapped_column(
        Integer, nullable=False, default=20
    )

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

    def __repr__(self) -> str:
        return (
            f"PipelineSettings(transcription_enabled={self.transcription_enabled}, "
            f"extraction_enabled={self.extraction_enabled})"
        )


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