backend/app/models/embedding.py

from __future__ import annotations

from datetime import datetime, timezone

import numpy as np
from sqlalchemy import (
    DateTime,
    ForeignKey,
    Integer,
    LargeBinary,
    String,
    UniqueConstraint,
    func,
)
from sqlalchemy.orm import Mapped, mapped_column, relationship

from .base import Base


class Embedding(Base):
    """An embedding vector for one picture under one model (spec §7.6).

    The vector is stored as L2-normalized float32 bytes so cosine similarity == dot product.
    Keeping embeddings in their own table lets us recompute across model versions without
    losing history.
    """

    __tablename__ = "embeddings"
    __table_args__ = (
        UniqueConstraint(
            "picture_id", "model_name", "model_version", name="uq_embedding_picture_model"
        ),
    )

    id: Mapped[int] = mapped_column(primary_key=True)
    picture_id: Mapped[int] = mapped_column(ForeignKey("pictures.id"), index=True)
    model_name: Mapped[str] = mapped_column(String(80))
    model_version: Mapped[str] = mapped_column(String(40))
    dim: Mapped[int] = mapped_column(Integer)
    vector: Mapped[bytes] = mapped_column(LargeBinary)  # float32, L2-normalized
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        default=lambda: datetime.now(timezone.utc),
        server_default=func.now(),
    )

    picture: Mapped["Picture"] = relationship(back_populates="embeddings")  # noqa: F821

    def as_array(self) -> np.ndarray:
        return np.frombuffer(self.vector, dtype=np.float32)

    @staticmethod
    def to_bytes(vec: np.ndarray) -> bytes:
        return np.asarray(vec, dtype=np.float32).tobytes()