rar-autopass/app/embedding_cache.py

"""Cache persistent de vectori embeddings in SQLite (tabela `embedding_cache`).

Design:
- Cheie (model, text_hash): schimbarea modelului nu foloseste vectori vechi.
- Vectorii raman `array('f')` (float32) end-to-end -- fara conversie la list[float].
- Scrieri/stergeri in tranzactii scurte, chunk-uite (conexiunile sunt autocommit,
  vezi app/db.py -- BEGIN/COMMIT explicit per chunk, altfel fiecare INSERT/DELETE
  e propria tranzactie).
- Degradare gratioasa: orice eroare SQLite -> log.warning, fara exceptie propagata
  din load/save/purge (caller-ul, ensure_embeddings_corpus, ramane neschimbat).
"""
from __future__ import annotations

import hashlib
import logging
import sqlite3
import sys
from array import array
from typing import Callable, Iterable, Sequence

log = logging.getLogger(__name__)

EMB_DIM = 384  # paraphrase-multilingual-MiniLM-L12-v2 (app/embeddings.py::FASTEMBED_MODEL)
_CHUNK_SIZE = 500  # A2/A3: tranzactii scurte vs worker BEGIN IMMEDIATE pe submissions


def text_hash(text: str) -> str:
    """SHA-256 hex al textului dat (apelantul normalizeaza inainte de hash)."""
    return hashlib.sha256(str(text).encode("utf-8")).hexdigest()


def vector_to_blob(vector: Sequence[float]) -> bytes:
    """Serializeaza un vector la float32 little-endian (independent de platforma)."""
    arr = array("f", vector)
    if sys.byteorder != "little":
        arr = array("f", arr)
        arr.byteswap()
    return arr.tobytes()


def blob_to_vector(blob: bytes) -> array:
    """Deserializeaza un blob float32 little-endian la `array('f')`."""
    arr = array("f")
    arr.frombytes(blob)
    if sys.byteorder != "little":
        arr.byteswap()
    return arr


def _chunks(seq: Sequence, size: int = _CHUNK_SIZE) -> Iterable[Sequence]:
    for i in range(0, len(seq), size):
        yield seq[i : i + size]


def load_cached_vectors(conn: sqlite3.Connection, model: str, hashes: Sequence[str]) -> dict[str, array]:
    """Citeste vectorii existenti pentru `model` + `hashes`. Blob corupt (lungime
    gresita) = tratat ca miss (nu apare in rezultat), cu log.warning.

    Degradare gratioasa: orice eroare SQLite -> dict gol/partial, fara exceptie.
    """
    out: dict[str, array] = {}
    if not hashes:
        return out
    expected_bytes = EMB_DIM * 4
    unique_hashes = list(dict.fromkeys(hashes))
    try:
        for chunk in _chunks(unique_hashes):
            placeholders = ",".join("?" for _ in chunk)
            rows = conn.execute(
                f"SELECT text_hash, vector FROM embedding_cache "
                f"WHERE model = ? AND text_hash IN ({placeholders})",
                (model, *chunk),
            ).fetchall()
            for row in rows:
                blob = row["vector"]
                if len(blob) != expected_bytes:
                    log.warning(
                        "embedding_cache: blob lungime gresita pentru hash=%s (asteptat %d, primit %d) -- tratat ca miss",
                        row["text_hash"], expected_bytes, len(blob),
                    )
                    continue
                out[row["text_hash"]] = blob_to_vector(blob)
    except sqlite3.OperationalError as exc:
        log.warning("embedding_cache: load_cached_vectors esuat (%s) -- fallback embed complet", exc)
    return out


def save_vectors(conn: sqlite3.Connection, model: str, items: Sequence[tuple[str, Sequence[float]]]) -> None:
    """INSERT OR REPLACE in chunk-uri de 500 randuri, BEGIN/COMMIT explicit per chunk
    (conexiunile sunt autocommit -- conn.commit() singur e no-op).

    Valideaza `len(vector) == EMB_DIM` la scriere: vector gresit = respins (log.warning),
    NU scris. Degradare gratioasa: eroare SQLite pe un chunk -> log.warning, chunk-urile
    ramase continua (cache partial e idempotent, se completeaza la urmatorul warmup).

    Invarianta: `conn` trebuie sa fie in autocommit (fara tranzactie deschisa de
    apelant) -- ROLLBACK-ul din except ar anula altfel tranzactia apelantului.
    """
    valid = []
    for h, vec in items:
        if len(vec) != EMB_DIM:
            log.warning(
                "embedding_cache: vector dimensiune gresita pentru hash=%s (asteptat %d, primit %d) -- respins",
                h, EMB_DIM, len(vec),
            )
            continue
        valid.append((h, model, vector_to_blob(vec)))

    for chunk in _chunks(valid):
        try:
            conn.execute("BEGIN")
            conn.executemany(
                "INSERT OR REPLACE INTO embedding_cache (text_hash, model, vector) VALUES (?, ?, ?)",
                chunk,
            )
            conn.execute("COMMIT")
        except sqlite3.OperationalError as exc:
            try:
                conn.execute("ROLLBACK")
            except sqlite3.OperationalError:
                pass
            log.warning("embedding_cache: save_vectors esuat pe un chunk (%s) -- cache ramane partial", exc)


def purge_stale(conn: sqlite3.Connection, model: str, corpus_hashes: Iterable[str]) -> None:
    """Sterge intrarile care nu mai apartin corpusului curent: orfane ale
    modelului curent (text_hash absent din `corpus_hashes`) SI toate intrarile
    modelelor VECHI (model != curent). Diff calculat in Python, DELETE chunk-uit
    pe PK (model, text_hash) -- tranzactii scurte.

    Degradare gratioasa: eroare SQLite -> log.warning, orfanele raman pana la
    urmatoarea trecere.

    Invarianta: `conn` trebuie sa fie in autocommit (fara tranzactie deschisa de
    apelant) -- ROLLBACK-ul din except ar anula altfel tranzactia apelantului.
    Apelantul trebuie sa cheme aceasta functie DOAR dupa o indexare reusita
    (vezi `ensure_embeddings_corpus`) -- un esec de indexare nu trebuie sa goleasca
    cache-ul de randuri inca valide.
    """
    keep = set(corpus_hashes)
    try:
        rows = conn.execute("SELECT model, text_hash FROM embedding_cache").fetchall()
    except sqlite3.OperationalError as exc:
        log.warning("embedding_cache: purge_stale citire esuata (%s)", exc)
        return

    to_delete = [
        (r["model"], r["text_hash"])
        for r in rows
        if r["model"] != model or r["text_hash"] not in keep
    ]
    if not to_delete:
        return

    for chunk in _chunks(to_delete):
        try:
            conn.execute("BEGIN")
            conn.executemany(
                "DELETE FROM embedding_cache WHERE model = ? AND text_hash = ?",
                chunk,
            )
            conn.execute("COMMIT")
        except sqlite3.OperationalError as exc:
            try:
                conn.execute("ROLLBACK")
            except sqlite3.OperationalError:
                pass
            log.warning("embedding_cache: purge_stale esuat pe un chunk (%s) -- orfanele raman", exc)


def sync_corpus_vectors(
    conn: sqlite3.Connection,
    model: str,
    texts: Sequence[str],
    embed_fn: Callable[[list[str]], Sequence[Sequence[float]]],
) -> list[array]:
    """Orchestreaza hash -> load -> embed(doar miss-uri) -> save -> vectori aliniati.

    Returneaza o lista de `array('f')` aliniata pozitional cu `texts`.

    `embed_fn` primeste lista textelor lipsa din cache si intoarce vectorii lor
    (aceeasi ordine). Daca `embed_fn` esueaza (arunca), exceptia se propaga;
    apelantul (ensure_embeddings_corpus) are deja degradare gratioasa (except -> pass).

    Un save partial esuat (lock SQLite) NU opreste intoarcerea vectorilor din RAM
    (A13c): vectorii noi raman in `cached` indiferent de rezultatul persistarii.

    NU purjeaza orfanele: purjarea e responsabilitatea apelantului, DUPA ce corpusul
    a fost indexat cu succes (`index_corpus`) -- un esec de indexare nu trebuie sa
    goleasca din greseala cache-ul de randuri inca valide.
    """
    hashes = [text_hash(t) for t in texts]
    cached = load_cached_vectors(conn, model, hashes)

    missing_positions = [i for i, h in enumerate(hashes) if h not in cached]
    if missing_positions:
        new_texts = [texts[i] for i in missing_positions]
        new_vecs = embed_fn(new_texts)
        if len(new_vecs) != len(new_texts):
            raise ValueError(
                f"embed_fn a intors {len(new_vecs)} vectori pentru {len(new_texts)} texte"
            )
        to_save = []
        for pos, vec in zip(missing_positions, new_vecs):
            arr = array("f", vec)
            cached[hashes[pos]] = arr
            to_save.append((hashes[pos], arr))
        save_vectors(conn, model, to_save)

    return [cached[h] for h in hashes]