ThothII/harness/tht/adapters/vector/thoth_http.py

"""Thoth vector HTTP adapter using distinct read and write clients."""

from tht.ports.vector import (
    VectorCapabilities,
    VectorHealth,
    VectorHit,
    VectorReadUnavailable,
    VectorStoreError,
    VectorWriteRecord,
    VectorWriteUnavailable,
    require_positive_limit,
)
from tht.vectorstore.rest_client import VectorRestClient, VectorRestError
from tht.vectorstore.store import hit_from_metadata
from tht.adapters.vector.pgvector import (
    _collection,
    _validate_collection_kinds,
    _validate_known_kinds,
)


def _merge(hits: list[VectorHit], limit: int) -> list[VectorHit]:
    return sorted(hits, key=lambda hit: (-hit.similarity, hit.id))[:limit]


class ThothHttpVectorStore:
    """Vector port backed by the existing allowlisted REST RPCs."""

    def __init__(
        self,
        reader: VectorRestClient | None,
        writer: VectorRestClient | None,
        expected_dimension: int | None = None,
    ):
        self._reader = reader
        self._writer = writer
        self._expected_dimension = expected_dimension

    @property
    def capabilities(self) -> VectorCapabilities:
        writable = self._writer is not None
        return VectorCapabilities(
            search=self._reader is not None, existing_hashes=writable, upsert=writable
        )

    def health(self) -> VectorHealth:
        read_reachable, read_detail, read_tables = self._probe(self._reader)
        write_reachable, write_detail, write_tables = self._probe(self._writer)
        dimensions = tuple(sorted({
            dimension
            for row in [*read_tables, *write_tables]
            if type(dimension := row.get("vector_dimensions")) is int
        }))
        compatible = (
            None
            if self._expected_dimension is None or not dimensions
            else dimensions == (self._expected_dimension,)
        )
        reachable = [
            status for status in (read_reachable, write_reachable) if status is not None
        ]
        ok = bool(reachable) and all(reachable) and compatible is not False
        details = [detail for detail in (read_detail, write_detail) if detail]
        return VectorHealth(
            ok=ok,
            detail="; ".join(details) or None,
            read_configured=self._reader is not None,
            read_reachable=read_reachable,
            read_detail=read_detail,
            write_configured=self._writer is not None,
            write_reachable=write_reachable,
            write_detail=write_detail,
            expected_dimension=self._expected_dimension,
            observed_dimensions=dimensions,
            dimension_compatible=compatible,
        )

    @staticmethod
    def _probe(client: VectorRestClient | None) -> tuple[bool | None, str | None, list[dict]]:
        if client is None:
            return None, None, []
        try:
            return True, None, client.list_tables()
        except Exception as exc:
            return False, str(exc), []

    def search(
        self,
        collections: list[str],
        embedding: list[float],
        *,
        limit: int,
        kinds: list[str] | None = None,
    ) -> list[VectorHit]:
        require_positive_limit(limit)
        if self._reader is None:
            raise VectorReadUnavailable("Vector reader credential is not configured")
        if self._expected_dimension is not None and len(embedding) != self._expected_dimension:
            raise VectorStoreError("Query embedding dimension does not match configured dimension")
        if kinds:
            _validate_known_kinds(kinds)
        hits: list[VectorHit] = []
        for collection in collections:
            _collection("vectors", collection)
            try:
                rows = self._reader.search_similar(collection, embedding, limit, kinds=kinds)
            except VectorRestError as exc:
                raise VectorStoreError(str(exc)) from exc
            hits.extend(
                hit_from_metadata(row.get("similarity", 0.0), row.get("metadata"))
                for row in rows
            )
        if kinds:
            allowed = set(kinds)
            hits = [hit for hit in hits if hit.kind in allowed]
        return _merge(hits, limit)

    def _require_writer(self) -> VectorRestClient:
        if self._writer is None:
            raise VectorWriteUnavailable("Vector writer credential is not configured")
        return self._writer

    def existing_hashes(self, collection: str, kinds: list[str]) -> dict[str, str]:
        _collection("vectors", collection)
        _validate_collection_kinds(collection, kinds)
        try:
            return self._require_writer().existing_hashes(collection, kinds)
        except VectorRestError as exc:
            raise VectorStoreError(str(exc)) from exc

    def upsert(self, collection: str, records: list[VectorWriteRecord]) -> int:
        writer = self._require_writer()
        _collection("vectors", collection)
        for record in records:
            _validate_collection_kinds(collection, [record.record.kind])
            if (
                self._expected_dimension is not None
                and len(record.embedding) != self._expected_dimension
            ):
                raise VectorStoreError("Embedding dimension does not match configured dimension")
        rows = [self._row(record) for record in records]
        try:
            return writer.upsert_records(collection, rows)
        except VectorRestError as exc:
            raise VectorStoreError(str(exc)) from exc

    @staticmethod
    def _row(write_record: VectorWriteRecord) -> dict:
        record = write_record.record
        metadata = {
            "kind": record.kind,
            "ref": record.ref,
            "record_key": record.id,
            "title": record.title,
            "content": record.content,
            **record.metadata,
        }
        return {
            "record_key": record.id,
            "kind": record.kind,
            "content_hash": write_record.content_hash,
            "metadata": metadata,
            "embedding": write_record.embedding,
        }