fix(evidence): enforce active vector generations
This commit is contained in:
@@ -5,6 +5,7 @@ from sqlalchemy import Engine
|
||||
from tht.ports.vector import (
|
||||
VectorCapabilities,
|
||||
VectorHealth,
|
||||
VectorStoreError,
|
||||
VectorWriteRecord,
|
||||
VectorWriteUnavailable,
|
||||
require_positive_limit,
|
||||
@@ -49,8 +50,11 @@ class LegacyDirectVectorStore:
|
||||
*,
|
||||
limit: int,
|
||||
kinds: list[str] | None = None,
|
||||
metadata_filter: dict[str, object] | None = None,
|
||||
) -> list[VectorHit]:
|
||||
require_positive_limit(limit)
|
||||
if metadata_filter is not None:
|
||||
raise VectorStoreError("Legacy vector store cannot enforce metadata filtering")
|
||||
hits: list[VectorHit] = []
|
||||
for collection in collections:
|
||||
table = TableVectorStore(
|
||||
|
||||
@@ -85,6 +85,8 @@ class PgVectorStore:
|
||||
search=self._reader is not None,
|
||||
existing_hashes=writable,
|
||||
upsert=writable,
|
||||
metadata_filter=self._reader is not None,
|
||||
delete_generation=writable,
|
||||
)
|
||||
|
||||
def _probe(
|
||||
@@ -242,6 +244,7 @@ class PgVectorStore:
|
||||
*,
|
||||
limit: int,
|
||||
kinds: list[str] | None = None,
|
||||
metadata_filter: dict[str, object] | None = None,
|
||||
) -> list[VectorHit]:
|
||||
require_positive_limit(limit)
|
||||
if self._reader is None:
|
||||
@@ -262,7 +265,27 @@ class PgVectorStore:
|
||||
)
|
||||
if kinds and not collection_kinds:
|
||||
continue
|
||||
where = sql.SQL(" WHERE kind = ANY(%s)") if collection_kinds else sql.SQL("")
|
||||
clauses = []
|
||||
filter_params = []
|
||||
if collection_kinds:
|
||||
clauses.append(sql.SQL("kind = ANY(%s)"))
|
||||
filter_params.append(collection_kinds)
|
||||
if metadata_filter is not None:
|
||||
if collection != "evidence" or set(metadata_filter) != {
|
||||
"vector_generation", "document_ids"
|
||||
}:
|
||||
raise VectorStoreError("Unsupported vector metadata filter")
|
||||
generation = metadata_filter["vector_generation"]
|
||||
document_ids = metadata_filter["document_ids"]
|
||||
if not isinstance(generation, str) or not isinstance(document_ids, list):
|
||||
raise VectorStoreError("Invalid vector metadata filter")
|
||||
clauses.append(sql.SQL("metadata->>'vector_generation' = %s"))
|
||||
clauses.append(sql.SQL("metadata->>'document_id' = ANY(%s)"))
|
||||
filter_params.extend((generation, document_ids))
|
||||
where = (
|
||||
sql.SQL(" WHERE ") + sql.SQL(" AND ").join(clauses)
|
||||
if clauses else sql.SQL("")
|
||||
)
|
||||
query = sql.SQL(
|
||||
"SELECT metadata, 1 - (embedding {} %s::{}) AS similarity "
|
||||
"FROM {}{} ORDER BY embedding {} %s::{}, record_key LIMIT %s"
|
||||
@@ -275,8 +298,7 @@ class PgVectorStore:
|
||||
_vector_type(self._schema),
|
||||
)
|
||||
params = [_vector_literal(embedding)]
|
||||
if collection_kinds:
|
||||
params.append(collection_kinds)
|
||||
params.extend(filter_params)
|
||||
params.extend((_vector_literal(embedding), limit))
|
||||
cursor.execute(query, params)
|
||||
hits.extend(hit_from_metadata(row[1], row[0]) for row in cursor.fetchall())
|
||||
@@ -381,5 +403,30 @@ class PgVectorStore:
|
||||
raw.close()
|
||||
return len(records)
|
||||
|
||||
def delete_generation(self, collection: str, generation: str) -> int:
|
||||
if collection != "evidence" or not generation.startswith("gen:"):
|
||||
raise VectorStoreError("Only exact Evidence generations may be deleted")
|
||||
raw = None
|
||||
try:
|
||||
raw = self._require_writer().raw_connection()
|
||||
with raw.cursor() as cursor:
|
||||
cursor.execute(
|
||||
sql.SQL(
|
||||
"DELETE FROM {} WHERE kind = 'evidence' "
|
||||
"AND metadata->>'vector_generation' = %s"
|
||||
).format(_collection(self._schema, collection)),
|
||||
(generation,),
|
||||
)
|
||||
count = cursor.rowcount
|
||||
raw.commit()
|
||||
return count
|
||||
except Exception as exc:
|
||||
if raw is not None:
|
||||
raw.rollback()
|
||||
raise VectorWriteUnavailable("Vector generation cleanup unavailable") from exc
|
||||
finally:
|
||||
if raw is not None:
|
||||
raw.close()
|
||||
|
||||
|
||||
__all__ = ["ALLOWED_COLLECTIONS", "PgVectorStore"]
|
||||
|
||||
@@ -40,7 +40,8 @@ class ThothHttpVectorStore:
|
||||
def capabilities(self) -> VectorCapabilities:
|
||||
writable = self._writer is not None
|
||||
return VectorCapabilities(
|
||||
search=self._reader is not None, existing_hashes=writable, upsert=writable
|
||||
search=self._reader is not None, existing_hashes=writable, upsert=writable,
|
||||
metadata_filter=self._reader is not None, delete_generation=writable,
|
||||
)
|
||||
|
||||
def health(self) -> VectorHealth:
|
||||
@@ -91,6 +92,7 @@ class ThothHttpVectorStore:
|
||||
*,
|
||||
limit: int,
|
||||
kinds: list[str] | None = None,
|
||||
metadata_filter: dict[str, object] | None = None,
|
||||
) -> list[VectorHit]:
|
||||
require_positive_limit(limit)
|
||||
if self._reader is None:
|
||||
@@ -103,7 +105,13 @@ class ThothHttpVectorStore:
|
||||
for collection in collections:
|
||||
_collection("vectors", collection)
|
||||
try:
|
||||
rows = self._reader.search_similar(collection, embedding, limit, kinds=kinds)
|
||||
if metadata_filter is None:
|
||||
rows = self._reader.search_similar(collection, embedding, limit, kinds=kinds)
|
||||
else:
|
||||
rows = self._reader.search_similar(
|
||||
collection, embedding, limit, kinds=kinds,
|
||||
metadata_filter=metadata_filter,
|
||||
)
|
||||
except VectorRestError as exc:
|
||||
raise VectorStoreError(str(exc)) from exc
|
||||
hits.extend(
|
||||
@@ -144,6 +152,14 @@ class ThothHttpVectorStore:
|
||||
except VectorRestError as exc:
|
||||
raise VectorStoreError(str(exc)) from exc
|
||||
|
||||
def delete_generation(self, collection: str, generation: str) -> int:
|
||||
if collection != "evidence" or not generation.startswith("gen:"):
|
||||
raise VectorStoreError("Only exact Evidence generations may be deleted")
|
||||
try:
|
||||
return self._require_writer().delete_generation(collection, generation)
|
||||
except VectorRestError as exc:
|
||||
raise VectorStoreError(str(exc)) from exc
|
||||
|
||||
@staticmethod
|
||||
def _row(write_record: VectorWriteRecord) -> dict:
|
||||
record = write_record.record
|
||||
|
||||
Reference in New Issue
Block a user