527 lines
22 KiB
Python
527 lines
22 KiB
Python
"""Qdrant-backed vector store for one workspace-owned semantic collection."""
|
|
|
|
import re
|
|
from collections.abc import Callable
|
|
from uuid import NAMESPACE_URL, uuid5
|
|
|
|
import requests
|
|
|
|
from tht.adapters.vector._shared import (
|
|
COLLECTION_KINDS,
|
|
validate_collection,
|
|
validate_collection_kinds,
|
|
validate_known_kinds,
|
|
)
|
|
from tht.ports.vector import (
|
|
VectorCapabilities,
|
|
VectorHealth,
|
|
VectorStoreError,
|
|
VectorWriteRecord,
|
|
require_positive_limit,
|
|
)
|
|
from tht.vectorstore.records import qdrant_payload, qdrant_semantic_kind
|
|
from tht.vectorstore.store import VectorHit, hit_from_metadata
|
|
|
|
_GENERATION = re.compile(r"gen:[0-9a-f]{32}")
|
|
_WORKSPACE = re.compile(r"[a-z][a-z0-9_-]{0,63}")
|
|
_BM25_LANGUAGES = frozenset({"english", "italian"})
|
|
_KEYWORD_INDEXES = (
|
|
|
|
"content_hash",
|
|
"document_id",
|
|
"kind",
|
|
"record_key",
|
|
"record_kind",
|
|
"vector_generation",
|
|
"workspace_id",
|
|
"workspace_revision",
|
|
)
|
|
|
|
UPSERT_BATCH_SIZE = 256
|
|
|
|
|
|
def point_id(workspace_id: str, kind: str, record_key: str, workspace_revision: str | None = None) -> str:
|
|
# P3: schema/Evidence points are revision-scoped; memory/solved remain workspace-wide.
|
|
if workspace_revision is not None:
|
|
return str(uuid5(NAMESPACE_URL, f"thothii:{workspace_id}:{workspace_revision}:{kind}:{record_key}"))
|
|
return str(uuid5(NAMESPACE_URL, f"thothii:{workspace_id}:{kind}:{record_key}"))
|
|
|
|
|
|
def _sanitize_exception(exc: Exception) -> str:
|
|
if isinstance(exc, requests.Timeout):
|
|
return "Qdrant request failed: timeout"
|
|
return f"Qdrant request failed: {type(exc).__name__}"
|
|
|
|
|
|
class QdrantVectorStore:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
base_url: str,
|
|
collection: str,
|
|
workspace_id: str,
|
|
workspace_revision: str | None = None,
|
|
expected_dimension: int | None = None,
|
|
collection_lifecycle: str = "self_heal",
|
|
request: Callable[..., object] | None = None,
|
|
connect_timeout: float = 2.0,
|
|
read_timeout: float = 10.0,
|
|
):
|
|
self._base_url = base_url.rstrip("/")
|
|
self._collection = collection
|
|
self._workspace_id = workspace_id
|
|
self._workspace_revision = None
|
|
self._workspace_revision = workspace_revision
|
|
self._expected_dimension = expected_dimension
|
|
self._collection_lifecycle = collection_lifecycle
|
|
self._request = request or requests.request
|
|
self._timeout = (connect_timeout, read_timeout)
|
|
|
|
@property
|
|
def capabilities(self) -> VectorCapabilities:
|
|
return VectorCapabilities(
|
|
search=True,
|
|
existing_hashes=True,
|
|
upsert=True,
|
|
metadata_filter=True,
|
|
delete_generation=True,
|
|
list_evidence_generations=True,
|
|
)
|
|
|
|
def health(self) -> VectorHealth:
|
|
try:
|
|
info = self._ensure_collection(strict=False)
|
|
except VectorStoreError as exc:
|
|
return VectorHealth(
|
|
ok=False,
|
|
detail=str(exc),
|
|
read_configured=True,
|
|
read_reachable=False,
|
|
read_detail=str(exc),
|
|
write_configured=True,
|
|
write_reachable=False,
|
|
write_detail=str(exc),
|
|
expected_dimension=self._expected_dimension,
|
|
bm25_compatible=None,
|
|
)
|
|
|
|
dimension = info["config"]["params"]["vectors"]["size"]
|
|
dimensions = (dimension,)
|
|
compatible = (
|
|
None if self._expected_dimension is None else dimensions == (self._expected_dimension,)
|
|
)
|
|
return VectorHealth(
|
|
ok=compatible is not False,
|
|
read_configured=True,
|
|
read_reachable=True,
|
|
write_configured=True,
|
|
write_reachable=True,
|
|
expected_dimension=self._expected_dimension,
|
|
observed_dimensions=dimensions,
|
|
dimension_compatible=compatible,
|
|
bm25_compatible=self._bm25_compatible(info),
|
|
)
|
|
|
|
def search(
|
|
self,
|
|
collections: list[str],
|
|
embedding: list[float],
|
|
*,
|
|
limit: int,
|
|
kinds: list[str] | None = None,
|
|
metadata_filter: dict[str, object] | None = None,
|
|
query_text: str | None = None,
|
|
query_language: str | None = None,
|
|
) -> list[VectorHit]:
|
|
require_positive_limit(limit)
|
|
self._validate_embedding(embedding, query=True)
|
|
allowed_record_kinds = self._allowed_record_kinds(collections, kinds)
|
|
if not allowed_record_kinds:
|
|
return []
|
|
filter_must = self._workspace_filter()
|
|
filter_must.extend(self._revision_filter(allowed_record_kinds))
|
|
filter_must.append(self._semantic_kind_filter(allowed_record_kinds))
|
|
filter_must.append({"key": "record_kind", "match": {"any": allowed_record_kinds}})
|
|
if metadata_filter is not None:
|
|
allowed_filters = {
|
|
"vector_generation", "document_ids", "workspace_id", "purpose",
|
|
"required_kinds", "required_concepts", "required_tables", "required_columns",
|
|
}
|
|
if not {"vector_generation", "document_ids", "workspace_id"} <= set(metadata_filter) or (
|
|
set(metadata_filter) - allowed_filters
|
|
):
|
|
raise VectorStoreError("Unsupported vector metadata filter")
|
|
generation = metadata_filter["vector_generation"]
|
|
document_ids = metadata_filter["document_ids"]
|
|
workspace_id = metadata_filter["workspace_id"]
|
|
if (
|
|
not isinstance(generation, str)
|
|
or not isinstance(document_ids, list)
|
|
or not all(isinstance(item, str) for item in document_ids)
|
|
or not isinstance(workspace_id, str)
|
|
):
|
|
raise VectorStoreError("Invalid vector metadata filter")
|
|
self._require_bound_workspace(workspace_id)
|
|
filter_must.extend([
|
|
{"key": "vector_generation", "match": {"value": generation}},
|
|
{"key": "document_id", "match": {"any": document_ids}},
|
|
])
|
|
purpose = metadata_filter.get("purpose")
|
|
if purpose is not None:
|
|
if not isinstance(purpose, str):
|
|
raise VectorStoreError("Invalid vector metadata filter")
|
|
filter_must.append({"key": "purposes", "match": {"value": purpose}})
|
|
required_kinds = metadata_filter.get("required_kinds", [])
|
|
if not isinstance(required_kinds, list) or not all(isinstance(item, str) for item in required_kinds):
|
|
raise VectorStoreError("Invalid vector metadata filter")
|
|
if required_kinds:
|
|
filter_must.append({"key": "evidence_kind", "match": {"any": required_kinds}})
|
|
for filter_key, payload_key in (
|
|
("required_concepts", "scope.concepts"),
|
|
("required_tables", "scope.tables"),
|
|
("required_columns", "scope.columns"),
|
|
):
|
|
values = metadata_filter.get(filter_key, [])
|
|
if not isinstance(values, list) or not all(isinstance(item, str) for item in values):
|
|
raise VectorStoreError("Invalid vector metadata filter")
|
|
filter_must.extend({"key": payload_key, "match": {"value": item}} for item in values)
|
|
if query_text is None:
|
|
if allowed_record_kinds == ["evidence"]:
|
|
raise VectorStoreError("Evidence hybrid query text is required")
|
|
response = self._call(
|
|
"POST",
|
|
f"/collections/{self._collection}/points/query",
|
|
{
|
|
"vector": embedding,
|
|
"limit": limit,
|
|
"with_payload": True,
|
|
"filter": {"must": filter_must},
|
|
},
|
|
)
|
|
else:
|
|
if allowed_record_kinds != ["evidence"]:
|
|
raise VectorStoreError("Hybrid BM25 is only available for Evidence")
|
|
if query_text.strip() == "" or query_language not in _BM25_LANGUAGES:
|
|
raise VectorStoreError("Evidence BM25 query is invalid")
|
|
self._ensure_collection(strict=False, require_bm25=True)
|
|
shared_filter = {"must": filter_must}
|
|
response = self._call(
|
|
"POST",
|
|
f"/collections/{self._collection}/points/query",
|
|
{
|
|
"prefetch": [
|
|
{"query": embedding, "limit": limit * 2, "filter": shared_filter},
|
|
{
|
|
"query": self._bm25_document(query_text, query_language),
|
|
"using": "bm25",
|
|
"limit": limit * 2,
|
|
"filter": shared_filter,
|
|
},
|
|
],
|
|
"query": {"rrf": {}},
|
|
"limit": limit,
|
|
"with_payload": True,
|
|
},
|
|
)
|
|
points = response.get("result", {}).get("points")
|
|
if not isinstance(points, list):
|
|
raise VectorStoreError("Qdrant returned malformed query response")
|
|
hits = [self._hit_from_point(point) for point in points]
|
|
return sorted(hits, key=lambda hit: (-hit.similarity, hit.id))[:limit]
|
|
|
|
def existing_hashes(self, collection: str, kinds: list[str]) -> dict[str, str]:
|
|
validate_collection(collection)
|
|
validate_collection_kinds(collection, kinds)
|
|
points = self._scroll(
|
|
[
|
|
*self._workspace_filter(),
|
|
self._semantic_kind_filter(kinds),
|
|
{"key": "record_kind", "match": {"any": sorted(kinds)}},
|
|
]
|
|
)
|
|
hashes: dict[str, str] = {}
|
|
for point in points:
|
|
payload = point.get("payload")
|
|
if not isinstance(payload, dict):
|
|
raise VectorStoreError("Qdrant returned malformed scroll response")
|
|
record_key = payload.get("record_key")
|
|
content_hash = payload.get("content_hash")
|
|
if not isinstance(record_key, str) or not isinstance(content_hash, str):
|
|
raise VectorStoreError("Qdrant returned malformed scroll response")
|
|
hashes[record_key] = content_hash
|
|
return hashes
|
|
|
|
def upsert(self, collection: str, records: list[VectorWriteRecord]) -> int:
|
|
validate_collection(collection)
|
|
self._ensure_collection(
|
|
strict=True,
|
|
require_bm25=any(record.sparse_text is not None for record in records),
|
|
)
|
|
points = []
|
|
for write_record in records:
|
|
validate_collection_kinds(collection, [write_record.record.kind])
|
|
self._validate_embedding(write_record.embedding, query=False)
|
|
semantic_kind = qdrant_semantic_kind(write_record.record.kind)
|
|
vector: list[float] | dict = write_record.embedding
|
|
if write_record.sparse_text is not None:
|
|
if semantic_kind != "evidence" or write_record.sparse_language not in _BM25_LANGUAGES:
|
|
raise VectorStoreError("Evidence BM25 document is invalid")
|
|
vector = {
|
|
"": write_record.embedding,
|
|
"bm25": self._bm25_document(write_record.sparse_text, write_record.sparse_language),
|
|
}
|
|
points.append(
|
|
{
|
|
"id": point_id(
|
|
self._workspace_id,
|
|
semantic_kind,
|
|
write_record.record.id,
|
|
self._workspace_revision if semantic_kind in ("schema_table", "schema_column", "evidence") else None,
|
|
),
|
|
"vector": vector,
|
|
"payload": qdrant_payload(
|
|
write_record.record,
|
|
content_hash=write_record.content_hash,
|
|
workspace_id=self._workspace_id,
|
|
workspace_revision=self._workspace_revision,
|
|
),
|
|
}
|
|
)
|
|
# Qdrant rejects request bodies larger than its JSON limit (32 MiB by default).
|
|
# A large schema/Evidence corpus therefore must be upserted in bounded chunks.
|
|
for start in range(0, len(points), UPSERT_BATCH_SIZE):
|
|
self._call(
|
|
"PUT",
|
|
f"/collections/{self._collection}/points?wait=true",
|
|
{"points": points[start:start + UPSERT_BATCH_SIZE]},
|
|
)
|
|
return len(records)
|
|
|
|
def delete_kinds(self, collection: str, kinds: list[str]) -> int:
|
|
validate_collection(collection)
|
|
validate_collection_kinds(collection, kinds)
|
|
must = [
|
|
*self._workspace_filter(),
|
|
self._semantic_kind_filter(kinds),
|
|
{"key": "record_kind", "match": {"any": sorted(kinds)}},
|
|
]
|
|
before = len(self._scroll(must))
|
|
self._call(
|
|
"POST",
|
|
f"/collections/{self._collection}/points/delete?wait=true",
|
|
{"filter": {"must": must}},
|
|
)
|
|
return before
|
|
|
|
def delete_generation(self, collection: str, generation: str, workspace_id: str) -> int:
|
|
if collection != "evidence" or _GENERATION.fullmatch(generation) is None:
|
|
raise VectorStoreError("Only exact Evidence generations may be deleted")
|
|
if _WORKSPACE.fullmatch(workspace_id) is None:
|
|
raise VectorStoreError("Invalid Evidence workspace namespace")
|
|
self._require_bound_workspace(workspace_id)
|
|
must = [
|
|
*self._workspace_filter(),
|
|
{"key": "kind", "match": {"value": "evidence"}},
|
|
{"key": "record_kind", "match": {"any": ["evidence"]}},
|
|
{"key": "vector_generation", "match": {"value": generation}},
|
|
]
|
|
before = len(
|
|
self._scroll(must)
|
|
)
|
|
self._call(
|
|
"POST",
|
|
f"/collections/{self._collection}/points/delete?wait=true",
|
|
{"filter": {"must": must}},
|
|
)
|
|
return before
|
|
|
|
def list_evidence_generations(self, collection: str, workspace_id: str) -> list[str]:
|
|
if collection != "evidence":
|
|
raise VectorStoreError("Only exact Evidence generations may be listed")
|
|
if _WORKSPACE.fullmatch(workspace_id) is None:
|
|
raise VectorStoreError("Invalid Evidence workspace namespace")
|
|
self._require_bound_workspace(workspace_id)
|
|
points = self._scroll(
|
|
[
|
|
*self._workspace_filter(),
|
|
{"key": "kind", "match": {"value": "evidence"}},
|
|
{"key": "record_kind", "match": {"any": ["evidence"]}},
|
|
]
|
|
)
|
|
generations = {
|
|
payload["vector_generation"]
|
|
for point in points
|
|
if isinstance((payload := point.get("payload")), dict)
|
|
and isinstance(payload.get("vector_generation"), str)
|
|
and _GENERATION.fullmatch(payload["vector_generation"]) is not None
|
|
}
|
|
return sorted(generations)
|
|
|
|
def _workspace_filter(self) -> list[dict]:
|
|
return [{"key": "workspace_id", "match": {"value": self._workspace_id}}]
|
|
|
|
@staticmethod
|
|
def _bm25_document(text: str, language: str) -> dict:
|
|
return {
|
|
"text": text,
|
|
"model": "qdrant/bm25",
|
|
"options": {"language": language},
|
|
}
|
|
|
|
def _revision_filter(self, kinds: list[str]) -> list[dict]:
|
|
if self._workspace_revision is None:
|
|
return []
|
|
if not any(kind in ("schema_table", "schema_column", "evidence") for kind in kinds):
|
|
return []
|
|
return [{"key": "workspace_revision", "match": {"value": self._workspace_revision}}]
|
|
|
|
def _semantic_kind_filter(self, record_kinds: list[str]) -> dict:
|
|
semantic_kinds = sorted({qdrant_semantic_kind(kind) for kind in record_kinds})
|
|
return {"key": "kind", "match": {"any": semantic_kinds}}
|
|
|
|
def bind_workspace_revision(self, workspace_revision: str) -> None:
|
|
if not re.fullmatch(r"[0-9a-f]{40}", workspace_revision):
|
|
raise VectorStoreError("workspace revision is invalid")
|
|
self._workspace_revision = workspace_revision
|
|
|
|
def _require_bound_workspace(self, workspace_id: str) -> None:
|
|
if workspace_id != self._workspace_id:
|
|
raise VectorStoreError("Evidence workspace namespace does not match bound workspace")
|
|
|
|
def _allowed_record_kinds(
|
|
self, collections: list[str], kinds: list[str] | None
|
|
) -> list[str]:
|
|
selected: set[str] = set()
|
|
for collection in collections:
|
|
validate_collection(collection)
|
|
selected.update(COLLECTION_KINDS[collection])
|
|
if kinds:
|
|
validate_known_kinds(kinds)
|
|
selected &= set(kinds)
|
|
return sorted(selected)
|
|
|
|
def _validate_embedding(self, embedding: list[float], *, query: bool) -> None:
|
|
if self._expected_dimension is not None and len(embedding) != self._expected_dimension:
|
|
raise VectorStoreError(
|
|
"Query embedding dimension does not match configured dimension"
|
|
if query
|
|
else "Embedding dimension does not match configured dimension"
|
|
)
|
|
|
|
@staticmethod
|
|
def _bm25_compatible(info: dict) -> bool:
|
|
sparse_vectors = info.get("config", {}).get("params", {}).get("sparse_vectors")
|
|
if not isinstance(sparse_vectors, dict):
|
|
return False
|
|
bm25 = sparse_vectors.get("bm25")
|
|
return isinstance(bm25, dict) and bm25.get("modifier") == "idf"
|
|
|
|
def _ensure_collection(self, *, strict: bool, require_bm25: bool = False) -> dict | None:
|
|
response = self._call("GET", f"/collections/{self._collection}", None, allow_missing=True)
|
|
if response is None:
|
|
if not strict:
|
|
raise VectorStoreError("Qdrant collection is missing")
|
|
if self._collection_lifecycle == "require_existing":
|
|
raise VectorStoreError("semantic_index_incompatible")
|
|
self._call(
|
|
"PUT",
|
|
f"/collections/{self._collection}",
|
|
{"vectors": {"size": self._expected_dimension or 1024, "distance": "Cosine"}},
|
|
)
|
|
for field_name in _KEYWORD_INDEXES:
|
|
self._call(
|
|
"PUT",
|
|
f"/collections/{self._collection}/index",
|
|
{"field_name": field_name, "field_schema": "keyword"},
|
|
)
|
|
response = self._call("GET", f"/collections/{self._collection}", None)
|
|
result = response.get("result") if isinstance(response, dict) else None
|
|
config = result.get("config", {}).get("params", {}).get("vectors") if isinstance(result, dict) else None
|
|
if not isinstance(config, dict):
|
|
raise VectorStoreError("Qdrant returned malformed collection response")
|
|
size = config.get("size")
|
|
distance = config.get("distance")
|
|
if (
|
|
self._expected_dimension is not None
|
|
and (size != self._expected_dimension or distance != "Cosine")
|
|
):
|
|
if strict and self._collection_lifecycle == "require_existing":
|
|
raise VectorStoreError("semantic_index_incompatible")
|
|
raise VectorStoreError("Qdrant collection configuration mismatch")
|
|
for field_name in _KEYWORD_INDEXES:
|
|
if field_name not in result.get("payload_schema", {}):
|
|
if strict and self._collection_lifecycle == "require_existing":
|
|
raise VectorStoreError("semantic_index_incompatible")
|
|
if not strict:
|
|
raise VectorStoreError("Qdrant collection payload indexes mismatch")
|
|
self._call(
|
|
"PUT",
|
|
f"/collections/{self._collection}/index",
|
|
{"field_name": field_name, "field_schema": "keyword"},
|
|
)
|
|
if require_bm25 and not self._bm25_compatible(result):
|
|
raise VectorStoreError("Evidence BM25 collection configuration mismatch")
|
|
return result
|
|
|
|
def _scroll(self, must: list[dict]) -> list[dict]:
|
|
points: list[dict] = []
|
|
offset = None
|
|
seen_offsets = set()
|
|
while True:
|
|
response = self._call(
|
|
"POST",
|
|
f"/collections/{self._collection}/points/scroll",
|
|
{
|
|
"with_payload": True,
|
|
"limit": 10000,
|
|
"filter": {"must": must},
|
|
"offset": offset,
|
|
},
|
|
)
|
|
result = response.get("result", {})
|
|
page = result.get("points")
|
|
if not isinstance(page, list):
|
|
raise VectorStoreError("Qdrant returned malformed scroll response")
|
|
points.extend(page)
|
|
next_page_offset = result.get("next_page_offset")
|
|
if next_page_offset is None:
|
|
return points
|
|
if next_page_offset in seen_offsets:
|
|
raise VectorStoreError("Qdrant returned malformed scroll response")
|
|
seen_offsets.add(next_page_offset)
|
|
offset = next_page_offset
|
|
|
|
def _hit_from_point(self, point: dict) -> VectorHit:
|
|
payload = point.get("payload")
|
|
score = point.get("score")
|
|
if not isinstance(payload, dict) or not isinstance(score, (int, float)):
|
|
raise VectorStoreError("Qdrant returned malformed query response")
|
|
return hit_from_metadata(float(score), payload)
|
|
|
|
def _call(self, method: str, path: str, payload: dict | None, allow_missing: bool = False) -> dict | None:
|
|
try:
|
|
response = self._request(
|
|
method,
|
|
f"{self._base_url}{path}",
|
|
json=payload,
|
|
timeout=self._timeout,
|
|
)
|
|
except requests.RequestException as exc:
|
|
raise VectorStoreError(_sanitize_exception(exc)) from exc
|
|
if response.status_code == 404 and allow_missing:
|
|
return None
|
|
if not response.ok:
|
|
raise VectorStoreError(f"Qdrant request failed: HTTP {response.status_code}")
|
|
if response.status_code == 204 or not getattr(response, "text", ""):
|
|
return {}
|
|
try:
|
|
data = response.json()
|
|
except Exception as exc:
|
|
raise VectorStoreError("Qdrant returned malformed JSON response") from exc
|
|
if not isinstance(data, dict):
|
|
raise VectorStoreError("Qdrant returned malformed JSON response")
|
|
return data
|
|
|
|
|
|
__all__ = ["QdrantVectorStore", "point_id"]
|