refactor(vector): define store contract
This commit is contained in:
@@ -0,0 +1,93 @@
|
||||
"""Thoth vector HTTP adapter using distinct read and write clients."""
|
||||
|
||||
from tht.ports.vector import (
|
||||
VectorCapabilities,
|
||||
VectorHealth,
|
||||
VectorHit,
|
||||
VectorRecord,
|
||||
VectorStoreError,
|
||||
VectorWriteUnavailable,
|
||||
)
|
||||
from tht.vectorstore.rest_client import VectorRestClient
|
||||
from tht.vectorstore.store import content_hash, 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, writer: VectorRestClient | None):
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
|
||||
@property
|
||||
def capabilities(self) -> VectorCapabilities:
|
||||
writable = self._writer is not None
|
||||
return VectorCapabilities(search=True, existing_hashes=writable, upsert=writable)
|
||||
|
||||
def health(self) -> VectorHealth:
|
||||
try:
|
||||
self._reader.list_tables()
|
||||
except Exception as exc:
|
||||
return VectorHealth(ok=False, detail=str(exc))
|
||||
return VectorHealth(ok=True)
|
||||
|
||||
def search(
|
||||
self,
|
||||
collections: list[str],
|
||||
embedding: list[float],
|
||||
*,
|
||||
limit: int,
|
||||
kinds: list[str] | None = None,
|
||||
) -> list[VectorHit]:
|
||||
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[VectorRecord]) -> int:
|
||||
writer = self._require_writer()
|
||||
rows = [self._row(record) for record in records]
|
||||
return writer.upsert_records(collection, rows)
|
||||
|
||||
@staticmethod
|
||||
def _row(record: VectorRecord) -> dict:
|
||||
extra = dict(record.metadata)
|
||||
try:
|
||||
embedding = extra.pop("embedding")
|
||||
except KeyError as exc:
|
||||
raise VectorStoreError(f"Vector record {record.id!r} has no embedding") from exc
|
||||
digest = extra.pop("content_hash", content_hash(record.content))
|
||||
metadata = {
|
||||
"kind": record.kind,
|
||||
"ref": record.ref,
|
||||
"record_key": record.id,
|
||||
"title": record.title,
|
||||
"content": record.content,
|
||||
**extra,
|
||||
}
|
||||
return {
|
||||
"record_key": record.id,
|
||||
"kind": record.kind,
|
||||
"content_hash": digest,
|
||||
"metadata": metadata,
|
||||
"embedding": embedding,
|
||||
}
|
||||
Reference in New Issue
Block a user