src/agent_harness/vector_backends.py

"""Pinned dense-index backends used by E01 and E05."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import sqlite3
import time
from typing import Sequence

import faiss
import numpy as np
import sqlite_vec

from .components import Candidate
from .retrieval import DenseRetriever, _candidate


@dataclass(frozen=True, slots=True)
class VectorBackendStats:
    backend: str
    build_seconds: float
    index_disk_bytes: int
    parameters: dict[str, int | str]


def _matrix(dense: DenseRetriever) -> np.ndarray:
    return np.asarray([list(item) for item in dense.vectors], dtype="float32")


class FaissFlatRetriever:
    def __init__(self, dense: DenseRetriever):
        started = time.monotonic()
        self.dense = dense
        matrix = _matrix(dense)
        self.index = faiss.IndexFlatIP(dense.spec.vector_dimension)
        self.index.add(matrix)
        self.stats = VectorBackendStats(
            backend="faiss_index_flat_ip",
            build_seconds=time.monotonic() - started,
            index_disk_bytes=0,
            parameters={"dimension": dense.spec.vector_dimension, "metric": "inner_product"},
        )

    def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
        vector = np.asarray([list(self.dense.query_vector(query))], dtype="float32")
        scores, indices = self.index.search(vector, min(limit, len(self.dense.chunks)))
        return tuple(
            _candidate(self.dense.chunks[int(index)], "dense_faiss_flat", float(score))
            for score, index in zip(scores[0], indices[0])
            if index >= 0
        )


class FaissHNSWRetriever:
    def __init__(
        self,
        dense: DenseRetriever,
        neighbors: int = 32,
        ef_construction: int = 80,
        ef_search: int = 64,
    ):
        started = time.monotonic()
        self.dense = dense
        matrix = _matrix(dense)
        self.index = faiss.IndexHNSWFlat(
            dense.spec.vector_dimension,
            neighbors,
            faiss.METRIC_INNER_PRODUCT,
        )
        self.index.hnsw.efConstruction = ef_construction
        self.index.hnsw.efSearch = ef_search
        self.index.add(matrix)
        self.stats = VectorBackendStats(
            backend="faiss_hnsw_ip",
            build_seconds=time.monotonic() - started,
            index_disk_bytes=0,
            parameters={
                "dimension": dense.spec.vector_dimension,
                "metric": "inner_product",
                "neighbors": neighbors,
                "ef_construction": ef_construction,
                "ef_search": ef_search,
            },
        )

    def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
        vector = np.asarray([list(self.dense.query_vector(query))], dtype="float32")
        scores, indices = self.index.search(vector, min(limit, len(self.dense.chunks)))
        return tuple(
            _candidate(self.dense.chunks[int(index)], "dense_faiss_hnsw", float(score))
            for score, index in zip(scores[0], indices[0])
            if index >= 0
        )


class SQLiteVecRetriever:
    def __init__(self, dense: DenseRetriever, path: Path):
        started = time.monotonic()
        self.dense = dense
        self.path = path
        path.parent.mkdir(parents=True, exist_ok=True)
        self.connection = sqlite3.connect(path)
        self.connection.enable_load_extension(True)
        sqlite_vec.load(self.connection)
        self.connection.enable_load_extension(False)
        self.connection.execute("DROP TABLE IF EXISTS vec_items")
        self.connection.execute(
            f"CREATE VIRTUAL TABLE vec_items USING vec0("
            f"embedding float[{dense.spec.vector_dimension}] distance_metric=cosine)"
        )
        self.connection.executemany(
            "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)",
            ((index + 1, vector.tobytes()) for index, vector in enumerate(dense.vectors)),
        )
        self.connection.commit()
        self.stats = VectorBackendStats(
            backend="sqlite_vec_exact",
            build_seconds=time.monotonic() - started,
            index_disk_bytes=path.stat().st_size,
            parameters={"dimension": dense.spec.vector_dimension, "metric": "cosine"},
        )

    def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
        vector = self.dense.query_vector(query)
        rows = self.connection.execute(
            "SELECT rowid, distance FROM vec_items WHERE embedding MATCH ? AND k = ? ORDER BY distance",
            (vector.tobytes(), min(limit, len(self.dense.chunks))),
        ).fetchall()
        return tuple(
            _candidate(
                self.dense.chunks[int(rowid) - 1],
                "dense_sqlite_vec",
                1.0 - float(distance),
            )
            for rowid, distance in rows
        )

    def close(self) -> None:
        self.connection.close()