Behavision/behavision/gallery/store.py

"""SQLite persistence for identities, embeddings and sightings.

Single writer class with an internal lock; WAL mode so the API can read
while the pipeline writes. Embeddings are stored as float32 BLOBs — SQLite
is the source of truth and the vector index is rebuilt from here at boot.
"""
from __future__ import annotations

import json
import sqlite3
import threading
import time
from pathlib import Path

import numpy as np

_SCHEMA = """
CREATE TABLE IF NOT EXISTS identities (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    label TEXT NOT NULL,
    kind TEXT NOT NULL DEFAULT 'auto',
    created_at REAL NOT NULL,
    last_seen_at REAL,
    sighting_count INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS embeddings (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    identity_id INTEGER NOT NULL REFERENCES identities(id) ON DELETE CASCADE,
    vector BLOB NOT NULL,
    model TEXT NOT NULL DEFAULT '',
    quality REAL NOT NULL DEFAULT 0,
    created_at REAL NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_embeddings_identity ON embeddings(identity_id);
CREATE TABLE IF NOT EXISTS sightings (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    identity_id INTEGER NOT NULL REFERENCES identities(id) ON DELETE CASCADE,
    camera_id TEXT NOT NULL,
    ts REAL NOT NULL,
    similarity REAL NOT NULL DEFAULT 0,
    quality REAL NOT NULL DEFAULT 0,
    attributes TEXT
);
CREATE INDEX IF NOT EXISTS idx_sightings_identity ON sightings(identity_id);
CREATE INDEX IF NOT EXISTS idx_sightings_ts ON sightings(ts);
"""


class IdentityStore:
    def __init__(self, db_path: "Path | str"):
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._db = sqlite3.connect(str(db_path), check_same_thread=False)
        self._db.row_factory = sqlite3.Row
        with self._lock:
            self._db.execute("PRAGMA journal_mode=WAL")
            self._db.execute("PRAGMA foreign_keys=ON")
            self._db.executescript(_SCHEMA)
            self._db.commit()

    # -- identities -----------------------------------------------------
    def create_identity(self, label: str, kind: str = "auto") -> int:
        with self._lock:
            cur = self._db.execute(
                "INSERT INTO identities(label, kind, created_at) VALUES(?,?,?)",
                (label, kind, time.time()))
            self._db.commit()
            return int(cur.lastrowid)

    def create_auto_identity(self) -> "tuple[int, str]":
        """Create an auto-enrolled identity labelled 'Visitor <id>' in one
        transaction; returns (id, label)."""
        with self._lock:
            cur = self._db.execute(
                "INSERT INTO identities(label, kind, created_at) VALUES(?,?,?)",
                ("pending", "auto", time.time()))
            identity_id = int(cur.lastrowid)
            label = f"Visitor {identity_id}"
            self._db.execute(
                "UPDATE identities SET label=? WHERE id=?", (label, identity_id))
            self._db.commit()
            return identity_id, label

    def rename_identity(self, identity_id: int, label: str) -> bool:
        with self._lock:
            cur = self._db.execute(
                "UPDATE identities SET label=?, kind='enrolled' WHERE id=?",
                (label, identity_id))
            self._db.commit()
            return cur.rowcount > 0

    def delete_identity(self, identity_id: int) -> "list[int]":
        """Delete an identity; returns removed embedding ids (for the index)."""
        with self._lock:
            rows = self._db.execute(
                "SELECT id FROM embeddings WHERE identity_id=?",
                (identity_id,)).fetchall()
            self._db.execute("DELETE FROM identities WHERE id=?", (identity_id,))
            self._db.commit()
            return [int(r["id"]) for r in rows]

    def merge_identities(self, source_id: int, target_id: int,
                         max_embeddings: int = 5) -> "dict | None":
        """Fold `source_id` into `target_id`; returns a summary, or None if
        either identity is missing.

        Embeddings and sightings are re-pointed rather than copied, which is
        what keeps this cheap AND keeps the vector index valid: the index maps
        *embedding* id to vector, and those ids do not change here, so a merge
        needs no reindex. Only trimmed embeddings have to be dropped from it,
        which is why they are returned.

        Everything happens in one transaction. A half-merge — sightings moved,
        embeddings not — would leave two identities each holding part of one
        person, which is strictly worse than the duplicate we started with.
        """
        with self._lock:
            src = self._db.execute("SELECT * FROM identities WHERE id=?",
                                   (source_id,)).fetchone()
            dst = self._db.execute("SELECT * FROM identities WHERE id=?",
                                   (target_id,)).fetchone()
            if src is None or dst is None or source_id == target_id:
                return None
            try:
                emb = self._db.execute(
                    "UPDATE embeddings SET identity_id=? WHERE identity_id=?",
                    (target_id, source_id)).rowcount
                sig = self._db.execute(
                    "UPDATE sightings SET identity_id=? WHERE identity_id=?",
                    (target_id, source_id)).rowcount

                # A human-assigned name outranks an auto "Visitor N" whichever
                # direction the operator merged in — silently turning "Alice"
                # back into "Visitor 3" would be a data-loss bug, not a policy.
                label, kind = dst["label"], dst["kind"]
                if dst["kind"] == "auto" and src["kind"] != "auto":
                    label, kind = src["label"], src["kind"]

                # The merged identity's history starts at the earlier of the
                # two first-sightings; it is one person and always was.
                created = min(float(src["created_at"]), float(dst["created_at"]))

                # Trim to the highest-quality views. Merging two identities
                # that each held the cap would otherwise leave one holding
                # double, quietly overweighting that person in every search.
                dropped = [int(r["id"]) for r in self._db.execute(
                    "SELECT id FROM embeddings WHERE identity_id=? "
                    "ORDER BY quality DESC, id ASC LIMIT -1 OFFSET ?",
                    (target_id, max_embeddings)).fetchall()]
                if dropped:
                    self._db.execute(
                        "DELETE FROM embeddings WHERE id IN (%s)"
                        % ",".join("?" * len(dropped)), dropped)

                # Recomputed, never summed: sighting_count on the source may
                # itself be stale, and COUNT(*) is the only figure that cannot
                # drift away from the rows actually present.
                agg = self._db.execute(
                    "SELECT COUNT(*) AS n, MAX(ts) AS last FROM sightings "
                    "WHERE identity_id=?", (target_id,)).fetchone()
                self._db.execute(
                    "UPDATE identities SET label=?, kind=?, created_at=?, "
                    "sighting_count=?, last_seen_at=? WHERE id=?",
                    (label, kind, created, int(agg["n"]), agg["last"],
                     target_id))
                self._db.execute("DELETE FROM identities WHERE id=?",
                                 (source_id,))
                self._db.commit()
            except Exception:
                self._db.rollback()
                raise
            return {"source": source_id, "target": target_id, "label": label,
                    "embeddings_moved": int(emb), "sightings_moved": int(sig),
                    "dropped_embeddings": dropped,
                    "sighting_count": int(agg["n"])}

    def identity_embeddings(self, identity_id: int, dim: int,
                            model: "str | None" = None
                            ) -> "tuple[list[int], np.ndarray]":
        """One identity's stored vectors, for comparing two identities to each
        other. Model-filtered for the same reason the index is."""
        with self._lock:
            if model is None:
                rows = self._db.execute(
                    "SELECT id, vector FROM embeddings WHERE identity_id=? "
                    "ORDER BY id", (identity_id,)).fetchall()
            else:
                rows = self._db.execute(
                    "SELECT id, vector FROM embeddings WHERE identity_id=? "
                    "AND model=? ORDER BY id",
                    (identity_id, model)).fetchall()
        ids = [int(r["id"]) for r in rows]
        if not ids:
            return [], np.empty((0, dim), dtype=np.float32)
        return ids, np.vstack([
            np.frombuffer(r["vector"], dtype=np.float32) for r in rows])

    def get_identity(self, identity_id: int) -> "dict | None":
        with self._lock:
            row = self._db.execute(
                "SELECT * FROM identities WHERE id=?", (identity_id,)).fetchone()
        return dict(row) if row else None

    def list_identities(self, limit: int = 200) -> "list[dict]":
        with self._lock:
            rows = self._db.execute(
                "SELECT i.*, COUNT(e.id) AS embedding_count FROM identities i "
                "LEFT JOIN embeddings e ON e.identity_id = i.id "
                "GROUP BY i.id ORDER BY i.last_seen_at DESC LIMIT ?",
                (limit,)).fetchall()
        return [dict(r) for r in rows]

    # -- embeddings -----------------------------------------------------
    def add_embedding(self, identity_id: int, vector: np.ndarray,
                      quality: float, model: str = "") -> int:
        blob = np.asarray(vector, dtype=np.float32).tobytes()
        with self._lock:
            cur = self._db.execute(
                "INSERT INTO embeddings(identity_id, vector, model, quality,"
                " created_at) VALUES(?,?,?,?,?)",
                (identity_id, blob, model, quality, time.time()))
            self._db.commit()
            return int(cur.lastrowid)

    def embedding_count(self, identity_id: int) -> int:
        with self._lock:
            row = self._db.execute(
                "SELECT COUNT(*) AS n FROM embeddings WHERE identity_id=?",
                (identity_id,)).fetchone()
        return int(row["n"])

    def identity_for_embedding(self, embedding_id: int) -> "dict | None":
        with self._lock:
            row = self._db.execute(
                "SELECT i.* FROM identities i JOIN embeddings e "
                "ON e.identity_id = i.id WHERE e.id=?",
                (embedding_id,)).fetchone()
        return dict(row) if row else None

    def all_embeddings(self, dim: int, model: "str | None" = None
                       ) -> "tuple[list[int], np.ndarray]":
        """Embeddings for the vector index. Filtering by `model` is what
        keeps vectors from different encoders out of the same search space —
        they are numerically incompatible."""
        with self._lock:
            if model is None:
                rows = self._db.execute(
                    "SELECT id, vector FROM embeddings ORDER BY id").fetchall()
            else:
                rows = self._db.execute(
                    "SELECT id, vector FROM embeddings WHERE model=? "
                    "ORDER BY id", (model,)).fetchall()
        ids = [int(r["id"]) for r in rows]
        if not ids:
            return [], np.empty((0, dim), dtype=np.float32)
        vecs = np.vstack([
            np.frombuffer(r["vector"], dtype=np.float32) for r in rows])
        return ids, vecs

    def best_embedding(self, identity_id: int, model: "str | None" = None
                       ) -> "tuple[np.ndarray, float] | None":
        """The highest-quality stored view of one identity.

        For handing an identity to the server: sending the best view rather
        than the mean because a mean of two disagreeing views is a vector that
        matches neither, which is precisely how one person becomes two
        identities.
        """
        with self._lock:
            if model is None:
                row = self._db.execute(
                    "SELECT vector, quality FROM embeddings WHERE identity_id=? "
                    "ORDER BY quality DESC, id ASC LIMIT 1",
                    (identity_id,)).fetchone()
            else:
                row = self._db.execute(
                    "SELECT vector, quality FROM embeddings WHERE identity_id=? "
                    "AND model=? ORDER BY quality DESC, id ASC LIMIT 1",
                    (identity_id, model)).fetchone()
        if row is None:
            return None
        return np.frombuffer(row["vector"], dtype=np.float32), float(row["quality"])

    def model_counts(self) -> "dict[str, int]":
        """How many stored embeddings each encoder produced.

        The gallery only ever searches vectors tagged with the *running*
        encoder, so this is what says whether the rest of the gallery is
        reachable at all. See `Gallery.health` for why that matters.
        """
        with self._lock:
            rows = self._db.execute(
                "SELECT model, COUNT(*) AS n FROM embeddings "
                "GROUP BY model").fetchall()
        return {str(r["model"]): int(r["n"]) for r in rows}

    def identities_with_model(self, model: str) -> int:
        """Identities holding at least one embedding from this encoder.

        Not the same as the identity count: an identity whose only vectors
        came from a previous encoder still exists, and is unrecognisable.
        """
        with self._lock:
            row = self._db.execute(
                "SELECT COUNT(DISTINCT identity_id) AS n FROM embeddings "
                "WHERE model=?", (model,)).fetchone()
        return int(row["n"]) if row else 0

    def embedding_owners(self, model: "str | None" = None) -> "dict[int, int]":
        """embedding_id -> identity_id, for turning index hits into identity
        pairs without a round trip to SQLite per hit."""
        with self._lock:
            if model is None:
                rows = self._db.execute(
                    "SELECT id, identity_id FROM embeddings").fetchall()
            else:
                rows = self._db.execute(
                    "SELECT id, identity_id FROM embeddings WHERE model=?",
                    (model,)).fetchall()
        return {int(r["id"]): int(r["identity_id"]) for r in rows}

    # -- sightings ------------------------------------------------------
    def record_sighting(self, identity_id: int, camera_id: str, ts: float,
                        similarity: float, quality: float,
                        attributes: "dict | None" = None) -> None:
        with self._lock:
            self._db.execute(
                "INSERT INTO sightings(identity_id, camera_id, ts, similarity,"
                " quality, attributes) VALUES(?,?,?,?,?,?)",
                (identity_id, camera_id, ts, similarity, quality,
                 json.dumps(attributes) if attributes else None))
            self._db.execute(
                "UPDATE identities SET last_seen_at=?, "
                "sighting_count=sighting_count+1 WHERE id=?", (ts, identity_id))
            self._db.commit()

    def recent_sightings(self, limit: int = 100) -> "list[dict]":
        with self._lock:
            rows = self._db.execute(
                "SELECT s.*, i.label FROM sightings s JOIN identities i "
                "ON i.id = s.identity_id ORDER BY s.ts DESC LIMIT ?",
                (limit,)).fetchall()
        out = []
        for r in rows:
            d = dict(r)
            if d.get("attributes"):
                d["attributes"] = json.loads(d["attributes"])
            out.append(d)
        return out

    def stats(self) -> dict:
        with self._lock:
            n_id = self._db.execute(
                "SELECT COUNT(*) AS n FROM identities").fetchone()["n"]
            n_emb = self._db.execute(
                "SELECT COUNT(*) AS n FROM embeddings").fetchone()["n"]
            n_sight = self._db.execute(
                "SELECT COUNT(*) AS n FROM sightings").fetchone()["n"]
        return {"identities": n_id, "embeddings": n_emb, "sightings": n_sight}

    def close(self) -> None:
        with self._lock:
            self._db.close()