FabledScribe/src/scribe/models/embedding.py

from datetime import datetime, timezone

from pgvector.sqlalchemy import Vector
from sqlalchemy import BigInteger, DateTime, ForeignKey, Integer, Text
from sqlalchemy.orm import Mapped, mapped_column

from scribe.models import Base

# bge-small-en-v1.5 produces 384-dim unit-normalized vectors. The column is a
# native pgvector `vector(384)` (see migration 0067) so similarity is ranked in
# Postgres by `<=>` rather than by a full-table Python cosine scan. The semantic
# searches rank the CALLER'S in-scope chunks exactly rather than through the
# HNSW index, which ranks before it filters (see embeddings._rank_scoped).
EMBEDDING_DIM = 384


class NoteEmbedding(Base):
    """One embedding vector per CHUNK of a note (migration 0077).

    The model reads at most 512 tokens, so a single whole-document vector
    permanently lost everything past ~400 words. A note now stores one row per
    chunk of `embeddings.chunk_document`, and a query matches the note if it
    matches ANY chunk — retrieval collapses rows to best-chunk-per-note.
    A short note has exactly one row (chunk_index 0) whose text is the
    historical `title\\nbody` shape.
    """

    __tablename__ = "note_embeddings"

    note_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("notes.id", ondelete="CASCADE"),
        primary_key=True,
    )
    chunk_index: Mapped[int] = mapped_column(Integer, primary_key=True)
    user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
    # Exactly what this vector encodes — inspectable when a ranking surprises,
    # and the hook for surfacing WHICH section matched, later.
    chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
    # embeddings.CHUNKER_VERSION at write time. The startup backfill re-embeds
    # any note whose rows carry a stale version — shape changes become a
    # version bump instead of a table wipe.
    chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
    # embeddings.EMBEDDING_MODEL at write time — the SPACE the vector lives in.
    # The column is `vector(384)`, a width and not an identity, so a
    # swap to another 384-dim model writes a second geometry beside the first
    # with no error, and cosine across the two is a number that means nothing.
    # The version above says what text was embedded; this says in whose
    # geometry. Either one moving makes the row stale, and the backfill
    # re-embeds on both.
    embedding_model: Mapped[str] = mapped_column(Text, nullable=False)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
    )


class RuleEmbedding(Base):
    """One embedding vector per CHUNK of a rule.

    A SIBLING of NoteEmbedding rather than a generalisation of it, decided
    deliberately. The reasoning below is the general test for when a record
    type earns its own embedding table; the later siblings apply it too.

    - The embedding ROW could have been made polymorphic. The SEARCH could not.
      `semantic_search_notes` is a long function of Note-specific scoping —
      the visibility clause, the supersession penalty, note_type/task_kind and
      system filters — and a rule shares none of it. Rules scope by rulebook
      ownership and project applicability instead.
    - Generalising the row while still needing two searches is the worst of
      both: a polymorphic key with referential integrity to neither table, on
      the path every session start runs, to share four columns.
    - What is genuinely common is BEHAVIOUR, not storage — get_embedding,
      chunk_document, embedding_text and CHUNKER_VERSION are already free
      functions and are reused as-is. Sharing those is the DRY win; sharing
      the table would have been the DRY costume.

    No `user_id`: NoteEmbedding carries one and its own search deliberately
    ignores it (scoping on the note instead, or shared records become
    unreachable). Rather than repeat a column that exists to be ignored, a
    rule's reach is resolved by joining the rule.
    """

    __tablename__ = "rule_embeddings"

    rule_id: Mapped[int] = mapped_column(
        BigInteger,
        ForeignKey("rules.id", ondelete="CASCADE"),
        primary_key=True,
    )
    chunk_index: Mapped[int] = mapped_column(Integer, primary_key=True)
    embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
    # Exactly what this vector encodes — inspectable when a ranking surprises.
    # For a rule this is the trigger-first document, NOT the rule's `why`:
    # `why` is incident narrative and drags every rule toward one centroid.
    chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
    chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
    embedding_model: Mapped[str] = mapped_column(Text, nullable=False)  # see NoteEmbedding
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
    )


class MilestoneEmbedding(Base):
    """One embedding vector per CHUNK of a milestone.

    The third sibling, for RuleEmbedding's reason: the row could be shared, the
    search cannot. A milestone is scoped by its project, has no share of its
    own, and is searched to answer one question — "is there already a plan for
    this?" — which no note or rule search can answer, because a plan is not a
    note. Without it, a roadmap written as milestones is invisible to recall,
    and a later plan gets opened as a new milestone beside the one that
    already describes it.

    The document is the title, the one-line description and the plan body, the
    parts a reader uses to recognise a plan. Derived data: the startup backfill
    regenerates it, which is also how a chunker-version bump is handled.
    """

    __tablename__ = "milestone_embeddings"

    milestone_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("milestones.id", ondelete="CASCADE"),
        primary_key=True,
    )
    chunk_index: Mapped[int] = mapped_column(Integer, primary_key=True)
    embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
    chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
    chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
    embedding_model: Mapped[str] = mapped_column(Text, nullable=False)  # see NoteEmbedding
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
    )


class SystemEmbedding(Base):
    """One embedding vector per CHUNK of a System's charter.

    The fourth sibling, for the same reason as the others: the row could be
    shared, the search cannot. A System's `description` is a charter —
    several hundred words saying what belongs in that area and what does not —
    and it is the answer to "which part of this codebase does X live in".
    Without it there is no semantic path to one: `list_systems` enumerates,
    and `search(system_id=…)` uses a System as a FILTER over notes, so a System
    could narrow a search and never be the answer to one.

    "Where does this belong?" is a different question from "what prior art is
    there?", which is why this is its own search rather than a note_type: a
    charter competing with two thousand notes for the same top-k would be
    outranked by the records filed under it, and the right answer would be
    crowded out by its own contents.

    The document is the name and the charter. Derived data: the startup
    backfill regenerates it, which is also how a chunker-version bump is
    handled.
    """

    __tablename__ = "system_embeddings"

    system_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("systems.id", ondelete="CASCADE"),
        primary_key=True,
    )
    chunk_index: Mapped[int] = mapped_column(Integer, primary_key=True)
    embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
    chunk_text: Mapped[str] = mapped_column(Text, nullable=False)
    chunker_version: Mapped[int] = mapped_column(Integer, nullable=False)
    embedding_model: Mapped[str] = mapped_column(Text, nullable=False)  # see NoteEmbedding
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
    )