"""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 ' 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()