135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
"""Thoth vector HTTP adapter using distinct read and write clients."""
|
|
|
|
from tht.ports.vector import (
|
|
VectorCapabilities,
|
|
VectorHealth,
|
|
VectorHit,
|
|
VectorReadUnavailable,
|
|
VectorWriteRecord,
|
|
VectorWriteUnavailable,
|
|
require_positive_limit,
|
|
)
|
|
from tht.vectorstore.rest_client import VectorRestClient
|
|
from tht.vectorstore.store import hit_from_metadata
|
|
|
|
|
|
def _merge(hits: list[VectorHit], limit: int) -> list[VectorHit]:
|
|
return sorted(hits, key=lambda hit: hit.similarity, reverse=True)[: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")
|
|
hits: list[VectorHit] = []
|
|
for collection in collections:
|
|
rows = self._reader.search_similar(collection, embedding, limit, kinds=kinds)
|
|
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]:
|
|
return self._require_writer().existing_hashes(collection, kinds)
|
|
|
|
def upsert(self, collection: str, records: list[VectorWriteRecord]) -> int:
|
|
writer = self._require_writer()
|
|
rows = [self._row(record) for record in records]
|
|
return writer.upsert_records(collection, rows)
|
|
|
|
@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,
|
|
}
|