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 search runs
# as an indexed `ORDER BY embedding <=> :q LIMIT k` in Postgres rather than a
# full-table Python cosine scan.
EMBEDDING_DIM = 384


class NoteEmbedding(Base):
    """One embedding vector per CHUNK of a note (#280, 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
    # (#4132). 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 (milestone 307, note 3026).

    A SIBLING of NoteEmbedding rather than a generalisation of it, decided
    deliberately. The reasoning below turned out to be the only written
    statement of a rule Scribe applies to every record type, so it is now also
    NOTE 3163 — "When a record type earns its own table" — with the worked
    cases and the cost of leaving `notes`. Read that before splitting a record
    type off; this stays here because it is where the decision was made.

    - 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 dated incident narrative and would drag every rule toward one
    # centroid (measured in note 2485).
    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 (milestone 415).

    The third sibling, for note 3163'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. Before this, a roadmap written as milestones was invisible to recall,
    and every later plan was opened as a new milestone beside the one that
    already described 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 (#4251).

    The fourth sibling, for the reason note 3163 gives about the third: 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".
    Before this there was 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 could 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),
    )