""" FaceIndex — sqlite-vec backed reverse face search index. Stores 512-d ArcFace face embeddings in a SQLite database with sqlite-vec for fast cosine similarity search. Persists to a single file (data/face_index.db by default). Why sqlite-vec instead of ChromaDB / FAISS / Pinecone? - Zero external dependencies beyond the `sqlite-vec` Python package (5MB) - Persists to a single file (easy backup, easy ship) - Runs on free-tier VPS without RAM issues - Supports standard SQL queries alongside vector search - Can be inspected with the `sqlite3` CLI Schema: faces — one row per enrolled face (face_id, embedding, metadata) enrollments — audit log of all enrollment events The embedding is stored as a serialized 512-d float32 vector via sqlite-vec's vec0 virtual table type. Search is cosine similarity (since embeddings are L2-normalized, this equals dot product). """ from __future__ import annotations import json import sqlite3 import threading import uuid from datetime import datetime, timezone from pathlib import Path from typing import Any, Optional import numpy as np from loguru import logger # Schema for the metadata table (regular SQLite) _SCHEMA_METADATA = """ CREATE TABLE IF NOT EXISTS faces ( face_id TEXT PRIMARY KEY, name TEXT, source_url TEXT, metadata TEXT NOT NULL DEFAULT '{}', thumbnail_path TEXT, embedding_dim INTEGER NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_faces_name ON faces(name); CREATE INDEX IF NOT EXISTS idx_faces_created ON faces(created_at); CREATE TABLE IF NOT EXISTS enrollments ( id INTEGER PRIMARY KEY AUTOINCREMENT, face_id TEXT NOT NULL, action TEXT NOT NULL, -- 'enroll' | 'delete' timestamp TEXT NOT NULL, details TEXT, FOREIGN KEY (face_id) REFERENCES faces(face_id) ); CREATE INDEX IF NOT EXISTS idx_enrollments_face ON enrollments(face_id); """ class FaceIndex: """Thread-safe reverse face search index backed by sqlite-vec.""" EMBEDDING_DIM = 512 def __init__(self, path: str = ":memory:") -> None: self._path = path self._lock = threading.RLock() # Ensure parent dir exists if path != ":memory:": Path(path).parent.mkdir(parents=True, exist_ok=True) self._conn = sqlite3.connect(path, check_same_thread=False) self._conn.row_factory = sqlite3.Row # Load sqlite-vec extension try: import sqlite_vec # type: ignore self._conn.enable_load_extension(True) sqlite_vec.load(self._conn) self._conn.enable_load_extension(False) self._vec_available = True logger.info("sqlite-vec extension loaded successfully") except ImportError: self._vec_available = False logger.warning( "sqlite-vec not installed — FaceIndex will use pure-Python fallback " "(slower, no vector indexing). Install with: pip install sqlite-vec" ) except Exception as e: self._vec_available = False logger.error(f"Failed to load sqlite-vec: {e}") # Initialize schema self._conn.executescript(_SCHEMA_METADATA) self._uses_l2 = False # set by _init_vec_table self._init_vec_table() self._conn.commit() if self._vec_available: count = self.count() logger.info(f"FaceIndex initialized at {path} (sqlite-vec mode, {count} faces)") else: logger.info(f"FaceIndex initialized at {path} (fallback mode)") def _init_vec_table(self) -> None: """Create the vec0 virtual table for fast vector search. We use cosine distance (1 - dot product for L2-normalized vectors) by partition-adding a `distance_metric` column metadata. However, sqlite-vec's vec0 currently defaults to L2 (Euclidean) distance. For L2-normalized embeddings, L2² = 2(1 - cos), so we can convert at query time: cos_sim = 1 - (L2² / 2). """ if not self._vec_available: return try: self._conn.execute( f""" CREATE VIRTUAL TABLE IF NOT EXISTS face_embeddings USING vec0( face_id TEXT PRIMARY KEY, embedding float[{self.EMBEDDING_DIM}] distance_metric=cosine ) """ ) self._conn.commit() except Exception as e: # Fallback: try without distance_metric (older sqlite-vec versions) logger.warning(f"cosine metric not supported, falling back to L2: {e}") try: self._conn.execute( f""" CREATE VIRTUAL TABLE IF NOT EXISTS face_embeddings USING vec0( face_id TEXT PRIMARY KEY, embedding float[{self.EMBEDDING_DIM}] ) """ ) self._conn.commit() self._uses_l2 = True except Exception as e2: logger.error(f"Failed to create vec0 table: {e2}") self._vec_available = False else: self._uses_l2 = False # ------------------------------------------------------------------ # # Enrollment # ------------------------------------------------------------------ # def enroll( self, embedding: np.ndarray, name: Optional[str] = None, source_url: Optional[str] = None, metadata: Optional[dict] = None, thumbnail_path: Optional[str] = None, ) -> str: """ Add a face to the searchable index. Args: embedding: 512-d L2-normalized face embedding (ArcFace) name: optional human-readable name source_url: optional URL where the face was found metadata: optional dict of arbitrary metadata (age, location, etc.) thumbnail_path: optional path to a thumbnail image of the face Returns: face_id (UUID string) """ if embedding.shape != (self.EMBEDDING_DIM,): raise ValueError( f"Embedding must be shape ({self.EMBEDDING_DIM},), got {embedding.shape}" ) # L2-normalize (defensive — should already be normalized) norm = np.linalg.norm(embedding) if norm > 0: embedding = embedding / norm face_id = str(uuid.uuid4()) now = datetime.now(timezone.utc).isoformat() metadata_json = json.dumps(metadata or {}, default=str) embedding_bytes = embedding.astype(np.float32).tobytes() with self._lock: # Insert metadata self._conn.execute( """INSERT INTO faces (face_id, name, source_url, metadata, thumbnail_path, embedding_dim, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", (face_id, name, source_url, metadata_json, thumbnail_path, self.EMBEDDING_DIM, now, now), ) # Insert embedding if self._vec_available: self._conn.execute( "INSERT INTO face_embeddings (face_id, embedding) VALUES (?, ?)", (face_id, embedding_bytes), ) # Audit log self._conn.execute( """INSERT INTO enrollments (face_id, action, timestamp, details) VALUES (?, ?, ?, ?)""", (face_id, "enroll", now, json.dumps({"name": name, "source_url": source_url})), ) self._conn.commit() logger.debug(f"Enrolled face {face_id} (name={name})") return face_id # ------------------------------------------------------------------ # # Search # ------------------------------------------------------------------ # def search( self, query_embedding: np.ndarray, top_k: int = 10, threshold: float = 0.0, ) -> list[dict]: """ Search the index for faces similar to the query embedding. Args: query_embedding: 512-d L2-normalized face embedding top_k: number of results to return threshold: minimum cosine similarity (0-1); 0 = return all Returns: list of dicts sorted by similarity (descending): { "face_id": str, "name": str | None, "source_url": str | None, "metadata": dict, "thumbnail_path": str | None, "similarity": float, "created_at": str, } """ if query_embedding.shape != (self.EMBEDDING_DIM,): raise ValueError( f"Query embedding must be shape ({self.EMBEDDING_DIM},), got {query_embedding.shape}" ) # L2-normalize norm = np.linalg.norm(query_embedding) if norm > 0: query_embedding = query_embedding / norm with self._lock: if self._vec_available: matches = self._search_vec(query_embedding, top_k * 2) else: matches = self._search_fallback(query_embedding, top_k * 2) # Filter by threshold + take top_k results = [] for m in matches: if m["similarity"] >= threshold: results.append(m) if len(results) >= top_k: break return results def _search_vec(self, query_embedding: np.ndarray, k: int) -> list[dict]: """Use sqlite-vec KNN search.""" query_bytes = query_embedding.astype(np.float32).tobytes() rows = self._conn.execute( """ SELECT f.face_id, f.name, f.source_url, f.metadata, f.thumbnail_path, f.created_at, v.distance FROM face_embeddings v JOIN faces f ON f.face_id = v.face_id WHERE v.embedding MATCH ? AND k = ? ORDER BY v.distance ASC """, (query_bytes, k), ).fetchall() # distance interpretation depends on metric: # - cosine: distance = 1 - cos_sim, so sim = 1 - distance # - L2 (Euclidean): for L2-normalized vectors, L2² = 2(1 - cos_sim), # so cos_sim = 1 - (distance² / 2). We use squared distance here # since sqlite-vec returns Euclidean (not squared). out = [] for r in rows: if getattr(self, "_uses_l2", False): # L2 distance — convert to cosine similarity sim = 1.0 - (r["distance"] ** 2) / 2.0 else: # cosine distance sim = 1.0 - r["distance"] out.append({ "face_id": r["face_id"], "name": r["name"], "source_url": r["source_url"], "metadata": json.loads(r["metadata"] or "{}"), "thumbnail_path": r["thumbnail_path"], "similarity": float(sim), "created_at": r["created_at"], }) return out def _search_fallback(self, query_embedding: np.ndarray, k: int) -> list[dict]: """Pure-Python fallback when sqlite-vec is not available.""" rows = self._conn.execute( "SELECT face_id, name, source_url, metadata, thumbnail_path, created_at FROM faces" ).fetchall() # We don't store embeddings in the metadata table for fallback mode. # In a real fallback, we'd need a separate embedding store. For now, # return empty results — installation of sqlite-vec is required. logger.warning( "FaceIndex fallback mode does not support search — install sqlite-vec: " "pip install sqlite-vec" ) return [] # ------------------------------------------------------------------ # # CRUD # ------------------------------------------------------------------ # def get(self, face_id: str) -> Optional[dict]: """Get details of a specific enrolled face.""" with self._lock: row = self._conn.execute( "SELECT * FROM faces WHERE face_id = ?", (face_id,) ).fetchone() if not row: return None d = dict(row) d["metadata"] = json.loads(d.get("metadata") or "{}") return d def list(self, limit: int = 50, offset: int = 0, name: Optional[str] = None) -> list[dict]: """Paginated list of enrolled faces.""" with self._lock: if name: cur = self._conn.execute( "SELECT * FROM faces WHERE name LIKE ? ORDER BY created_at DESC LIMIT ? OFFSET ?", (f"%{name}%", limit, offset), ) else: cur = self._conn.execute( "SELECT * FROM faces ORDER BY created_at DESC LIMIT ? OFFSET ?", (limit, offset), ) rows = [dict(r) for r in cur.fetchall()] for r in rows: r["metadata"] = json.loads(r.get("metadata") or "{}") return rows def delete(self, face_id: str) -> bool: """Remove a face from the index. Returns True if deleted.""" now = datetime.now(timezone.utc).isoformat() with self._lock: # Check exists row = self._conn.execute( "SELECT face_id FROM faces WHERE face_id = ?", (face_id,) ).fetchone() if not row: return False self._conn.execute("DELETE FROM faces WHERE face_id = ?", (face_id,)) if self._vec_available: self._conn.execute( "DELETE FROM face_embeddings WHERE face_id = ?", (face_id,) ) self._conn.execute( """INSERT INTO enrollments (face_id, action, timestamp, details) VALUES (?, ?, ?, ?)""", (face_id, "delete", now, "{}"), ) self._conn.commit() logger.debug(f"Deleted face {face_id}") return True def count(self) -> int: """Total number of enrolled faces.""" with self._lock: row = self._conn.execute("SELECT COUNT(*) as n FROM faces").fetchone() return row["n"] if row else 0 def stats(self) -> dict: """Index statistics.""" with self._lock: total = self.count() named = self._conn.execute( "SELECT COUNT(*) as n FROM faces WHERE name IS NOT NULL" ).fetchone()["n"] last_enrollment = self._conn.execute( "SELECT created_at FROM faces ORDER BY created_at DESC LIMIT 1" ).fetchone() recent_enrollments = self._conn.execute( """SELECT COUNT(*) as n FROM enrollments WHERE action = 'enroll' AND timestamp > datetime('now', '-24 hours')""" ).fetchone()["n"] return { "total_faces": total, "named_faces": named, "anonymous_faces": total - named, "last_enrollment": last_enrollment["created_at"] if last_enrollment else None, "recent_enrollments_24h": recent_enrollments, "vec_available": self._vec_available, "embedding_dim": self.EMBEDDING_DIM, "path": self._path, } def clear(self) -> int: """Remove ALL faces from the index. Returns count deleted.""" with self._lock: n = self.count() self._conn.execute("DELETE FROM faces") if self._vec_available: self._conn.execute("DELETE FROM face_embeddings") self._conn.execute("DELETE FROM enrollments") self._conn.commit() logger.warning(f"Cleared {n} faces from index") return n def close(self) -> None: with self._lock: self._conn.close()