vecspan 0.1.0__py3-none-any.whl

vecspan/backends/qdrant.py

@@ -0,0 +1,532 @@
+"""Qdrant adapter built on `qdrant_client.QdrantClient`.
+
+Collection specs are stored in a private `_vecspan_schemas` collection because
+Qdrant does not keep the id type or unindexed payload fields.
+"""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any, Mapping, Sequence
+from urllib.parse import urlsplit
+
+from vecspan.backends.base import VectorBackend
+from vecspan.errors import (
+    BackendError,
+    CollectionAlreadyExistsError,
+    CollectionNotFoundError,
+    ValidationError,
+    VecspanError,
+)
+from vecspan.filters import And, Cmp, Eq, Expr, Filter, In, IsNull, Not, Or
+from vecspan.naming import META_COLLECTION
+from vecspan.types import (
+    CollectionSpec,
+    PayloadType,
+    Point,
+    ScrollPage,
+    SearchHit,
+    as_plain,
+    as_vector,
+    spec_from_dict,
+    spec_to_dict,
+)
+from vecspan.validate import coerce_stored_id, prepare_search_options
+
+_META_NAMESPACE = uuid.UUID("8c1c6b2e-9a1d-4f0e-9b7a-6e5d4c3b2a19")
+_OPTIONS = {"api_key", "timeout", "prefer_grpc"}
+_SCHEMA = {
+    PayloadType.KEYWORD: "keyword",
+    PayloadType.INT: "integer",
+    PayloadType.FLOAT: "float",
+    PayloadType.BOOL: "bool",
+}
+
+
+class QdrantBackend(VectorBackend):
+    """Qdrant collections with named dense vectors."""
+
+    name = "qdrant"
+
+    def __init__(self, client: Any) -> None:
+        self._client = client
+        self._specs: dict[str, CollectionSpec] = {}
+        self._closed = False
+
+    @classmethod
+    def from_url(cls, url: str, **kwargs: Any) -> QdrantBackend:
+        """Open a client from a `qdrant://` or `qdrant+https://` URL."""
+        params = parse_qdrant_target(url, **kwargs)
+        try:
+            from qdrant_client import QdrantClient
+        except ImportError as exc:
+            raise BackendError("Qdrant support requires the qdrant extra: pip install 'vecspan[qdrant]'") from exc
+        try:
+            client = QdrantClient(**params)
+        except Exception as exc:
+            raise BackendError(f"could not connect to Qdrant: {exc}") from exc
+        return cls(client)
+
+    def create_collection(self, spec: CollectionSpec) -> None:
+        def run() -> None:
+            if self._exists(spec.name):
+                raise CollectionAlreadyExistsError(f"collection {spec.name!r} already exists")
+            plan = qdrant_plan(spec)
+            models = _models()
+            vectors = {
+                name: models.VectorParams(
+                    size=item["size"],
+                    distance=getattr(models.Distance, item["distance"]),
+                    hnsw_config=models.HnswConfigDiff(m=item["m"], ef_construct=item["ef_construct"]),
+                )
+                for name, item in plan["vectors"].items()
+            }
+            self._client.create_collection(collection_name=spec.name, vectors_config=vectors)
+            try:
+                for index in plan["payload_indexes"]:
+                    self._client.create_payload_index(
+                        collection_name=spec.name,
+                        field_name=index["name"],
+                        field_schema=getattr(models.PayloadSchemaType, index["schema"].upper()),
+                    )
+                self._write_meta(spec)
+            except Exception:
+                self._client.delete_collection(collection_name=spec.name)
+                raise
+            self._specs[spec.name] = spec
+
+        self._run(run)
+
+    def drop_collection(self, name: str) -> None:
+        def run() -> None:
+            if name == META_COLLECTION:
+                raise ValidationError(f"reserved collection name: {name!r}")
+            if not self._exists(name):
+                raise CollectionNotFoundError(f"collection {name!r} does not exist")
+            self._client.delete_collection(collection_name=name)
+            self._delete_meta(name)
+            self._specs.pop(name, None)
+
+        self._run(run)
+
+    def has_collection(self, name: str) -> bool:
+        return self._run(lambda: self._exists(name) and name != META_COLLECTION)
+
+    def describe_collection(self, name: str) -> CollectionSpec:
+        return self._run(lambda: self._describe(name))
+
+    def list_collections(self) -> list[str]:
+        def run() -> list[str]:
+            listed = self._client.get_collections().collections
+            return [item.name for item in listed if item.name != META_COLLECTION]
+
+        return self._run(run)
+
+    def upsert(self, name: str, points: Sequence[Point]) -> None:
+        def run() -> None:
+            self._cached(name)
+            if not points:
+                return
+            models = _models()
+            structs = [
+                models.PointStruct(id=point.id, vector=dict(point.vectors), payload=dict(point.payload))
+                for point in points
+            ]
+            self._client.upsert(collection_name=name, points=structs, wait=True)
+
+        self._run(run)
+
+    def update_payload(self, name: str, ids: Sequence[int | str], payload: dict) -> None:
+        def run() -> None:
+            spec = self._cached(name)
+            self._require_existing(name, spec, ids)
+            self._client.set_payload(collection_name=name, payload=dict(payload), points=list(ids), wait=True)
+
+        self._run(run)
+
+    def update_vectors(
+        self,
+        name: str,
+        ids: Sequence[int | str],
+        vectors: Sequence[dict[str, list[float]]],
+    ) -> None:
+        def run() -> None:
+            spec = self._cached(name)
+            self._require_existing(name, spec, ids)
+            models = _models()
+            points = [
+                models.PointVectors(id=point_id, vector=dict(vector))
+                for point_id, vector in zip(ids, vectors)
+            ]
+            self._client.update_vectors(collection_name=name, points=points, wait=True)
+
+        self._run(run)
+
+    def delete(
+        self,
+        name: str,
+        *,
+        ids: Sequence[int | str] | None = None,
+        filter: Filter | None = None,
+    ) -> None:
+        def run() -> None:
+            self._cached(name)
+            models = _models()
+            if ids is not None:
+                selector = models.PointIdsList(points=list(ids))
+            elif filter is not None:
+                selector = models.FilterSelector(filter=compile_filter(filter))
+            else:
+                raise ValidationError("pass exactly one of ids or filter")
+            self._client.delete(collection_name=name, points_selector=selector, wait=True)
+
+        self._run(run)
+
+    def get(self, name: str, ids: Sequence[int | str], *, with_vectors: bool = False) -> list[Point]:
+        def run() -> list[Point]:
+            spec = self._cached(name)
+            records = self._client.retrieve(
+                collection_name=name,
+                ids=list(ids),
+                with_payload=True,
+                with_vectors=with_vectors,
+            )
+            by_id = {
+                point.id: point
+                for point in (
+                    record_to_point(record, spec, with_vectors=with_vectors) for record in records
+                )
+            }
+            return [by_id[point_id] for point_id in ids if point_id in by_id]
+
+        return self._run(run)
+
+    def count(self, name: str, *, filter: Filter | None = None) -> int:
+        def run() -> int:
+            self._cached(name)
+            result = self._client.count(
+                collection_name=name,
+                count_filter=None if filter is None else compile_filter(filter),
+                exact=True,
+            )
+            return int(result.count)
+
+        return self._run(run)
+
+    def scroll(
+        self,
+        name: str,
+        *,
+        filter: Filter | None = None,
+        limit: int = 100,
+        offset: int | str | None = None,
+        with_vectors: bool = False,
+    ) -> ScrollPage:
+        def run() -> ScrollPage:
+            spec = self._cached(name)
+            if offset is not None and not _valid_offset(offset):
+                raise ValidationError("Qdrant scroll offset must come from a previous page")
+            records, next_offset = self._client.scroll(
+                collection_name=name,
+                scroll_filter=None if filter is None else compile_filter(filter),
+                limit=limit,
+                offset=offset,
+                with_payload=True,
+                with_vectors=with_vectors,
+            )
+            points = [record_to_point(record, spec, with_vectors=with_vectors) for record in records]
+            return ScrollPage(points=points, next_offset=_offset_token(next_offset))
+
+        return self._run(run)
+
+    def search(
+        self,
+        name: str,
+        vector: Sequence[float],
+        *,
+        vector_name: str,
+        filter: Filter | None = None,
+        limit: int = 10,
+        with_vectors: bool = False,
+        backend_options: dict | None = None,
+    ) -> list[SearchHit]:
+        def run() -> list[SearchHit]:
+            spec = self._cached(name)
+            options = prepare_search_options(backend_options, allow_nprobe=False)
+            models = _models()
+            response = self._client.query_points(
+                collection_name=name,
+                query=list(vector),
+                using=vector_name,
+                query_filter=None if filter is None else compile_filter(filter),
+                limit=limit,
+                with_payload=True,
+                with_vectors=with_vectors,
+                search_params=models.SearchParams(hnsw_ef=options.get("ef", max(64, limit))),
+            )
+            return [
+                scored_to_hit(point, spec, vector_name, with_vectors=with_vectors)
+                for point in response.points
+            ]
+
+        return self._run(run)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        closer = getattr(self._client, "close", None)
+        if not callable(closer):
+            return
+        try:
+            closer()
+        except Exception as exc:
+            raise BackendError(f"qdrant close failed: {exc}") from exc
+
+    def _exists(self, name: str) -> bool:
+        exists = getattr(self._client, "collection_exists", None)
+        if callable(exists):
+            return bool(exists(collection_name=name) if _accepts_collection_name(exists) else exists(name))
+        listed = self._client.get_collections().collections
+        return any(item.name == name for item in listed)
+
+    def _describe(self, name: str) -> CollectionSpec:
+        if name == META_COLLECTION or not self._exists(name):
+            raise CollectionNotFoundError(f"collection {name!r} does not exist")
+        if not self._exists(META_COLLECTION):
+            raise BackendError(f"collection {name!r} was not created by vecspan")
+        records = self._client.retrieve(
+            collection_name=META_COLLECTION,
+            ids=[_meta_id(name)],
+            with_payload=True,
+            with_vectors=False,
+        )
+        if not records:
+            raise BackendError(f"collection {name!r} was not created by vecspan")
+        payload = _field(records[0], "payload") or {}
+        spec = spec_from_dict(payload["spec"])
+        self._specs[name] = spec
+        return spec
+
+    def _cached(self, name: str) -> CollectionSpec:
+        spec = self._specs.get(name)
+        if spec is None:
+            return self._describe(name)
+        return spec
+
+    def _require_existing(self, name: str, spec: CollectionSpec, ids: Sequence[int | str]) -> None:
+        records = self._client.retrieve(
+            collection_name=name,
+            ids=list(ids),
+            with_payload=False,
+            with_vectors=False,
+        )
+        found = {coerce_stored_id(spec, _field(record, "id")) for record in records}
+        for point_id in ids:
+            if point_id not in found:
+                raise ValidationError(f"unknown point id: {point_id!r}")
+
+    def _write_meta(self, spec: CollectionSpec) -> None:
+        self._ensure_meta()
+        models = _models()
+        self._client.upsert(
+            collection_name=META_COLLECTION,
+            points=[
+                models.PointStruct(
+                    id=_meta_id(spec.name),
+                    vector={"_": [0.0]},
+                    payload={"spec": spec_to_dict(spec)},
+                )
+            ],
+            wait=True,
+        )
+
+    def _ensure_meta(self) -> None:
+        if self._exists(META_COLLECTION):
+            return
+        models = _models()
+        self._client.create_collection(
+            collection_name=META_COLLECTION,
+            vectors_config={"_": models.VectorParams(size=1, distance=models.Distance.DOT)},
+        )
+
+    def _delete_meta(self, name: str) -> None:
+        if not self._exists(META_COLLECTION):
+            return
+        models = _models()
+        self._client.delete(
+            collection_name=META_COLLECTION,
+            points_selector=models.PointIdsList(points=[_meta_id(name)]),
+            wait=True,
+        )
+
+    def _run(self, action):
+        if self._closed:
+            raise BackendError("client is closed")
+        try:
+            return action()
+        except VecspanError:
+            raise
+        except Exception as exc:
+            raise BackendError(f"qdrant request failed: {exc}") from exc
+
+
+def parse_qdrant_target(url: str, **kwargs: Any) -> dict[str, Any]:
+    """Return `QdrantClient` keyword arguments for a qdrant URL."""
+    unknown = sorted(set(kwargs) - _OPTIONS)
+    if unknown:
+        raise ValidationError(f"unknown option: {unknown[0]}")
+    parsed = urlsplit(url)
+    scheme = parsed.scheme.lower()
+    if scheme not in {"qdrant", "qdrant+https"}:
+        raise ValidationError("Qdrant URL must start with qdrant:// or qdrant+https://")
+    if parsed.username or parsed.password:
+        raise ValidationError("pass api_key= instead of putting credentials in the Qdrant URL")
+    if parsed.query or parsed.fragment or parsed.path not in ("", "/"):
+        raise ValidationError("qdrant URL does not take a path or query string")
+    host = parsed.hostname
+    if not host:
+        raise ValidationError("qdrant URL requires a host")
+    https = scheme == "qdrant+https"
+    port = parsed.port or (443 if https else 6333)
+    params: dict[str, Any] = {"url": f"{'https' if https else 'http'}://{host}:{port}"}
+    api_key = kwargs.get("api_key")
+    if api_key is not None:
+        if not isinstance(api_key, str):
+            raise ValidationError("api_key must be a string")
+        params["api_key"] = api_key
+    timeout = kwargs.get("timeout")
+    if timeout is not None:
+        if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or timeout < 0:
+            raise ValidationError("timeout must be a non-negative number")
+        params["timeout"] = float(timeout)
+    if "prefer_grpc" in kwargs:
+        if not isinstance(kwargs["prefer_grpc"], bool):
+            raise ValidationError("prefer_grpc must be a bool")
+        params["prefer_grpc"] = kwargs["prefer_grpc"]
+    return params
+
+
+def qdrant_plan(spec: CollectionSpec) -> dict[str, Any]:
+    """Describe named vectors and payload indexes for `spec`."""
+    metric = {"cosine": "COSINE", "dot": "DOT", "l2": "EUCLID"}
+    return {
+        "vectors": {
+            item.name: {
+                "size": item.size,
+                "distance": metric[item.distance.value],
+                "m": item.m,
+                "ef_construct": item.ef_construct,
+            }
+            for item in spec.vectors
+        },
+        "payload_indexes": [
+            {"name": item.name, "schema": _SCHEMA[item.type]}
+            for item in spec.payload_fields
+            if item.indexed
+        ],
+    }
+
+
+def compile_filter(filt: Filter) -> Any:
+    """Compile a filter into a `qdrant_client.models.Filter` or condition."""
+    return _compile(filt.expr, _models())
+
+
+def record_to_point(record: Any, spec: CollectionSpec, *, with_vectors: bool) -> Point:
+    """Convert a Qdrant record into a point."""
+    payload = as_plain(_field(record, "payload") or {})
+    if not isinstance(payload, dict):
+        payload = {}
+    vectors: dict[str, list[float]] = {}
+    if with_vectors:
+        raw = _field(record, "vector")
+        if raw is None:
+            raw = _field(record, "vectors")
+        if isinstance(raw, Mapping):
+            for field in spec.vectors:
+                if field.name in raw:
+                    vectors[field.name] = as_vector(raw[field.name])
+        elif raw is not None and spec.vectors:
+            vectors[spec.vectors[0].name] = as_vector(raw)
+    return Point(
+        id=coerce_stored_id(spec, _field(record, "id")),
+        vectors=vectors,
+        payload=payload,
+    )
+
+
+def scored_to_hit(record: Any, spec: CollectionSpec, vector_name: str, *, with_vectors: bool) -> SearchHit:
+    """Convert a scored Qdrant point into a search hit."""
+    point = record_to_point(record, spec, with_vectors=with_vectors)
+    score = _field(record, "score")
+    if score is None:
+        raise BackendError("Qdrant search hit is missing a score")
+    return SearchHit(
+        id=point.id,
+        score=float(score),
+        distance=spec.vector(vector_name).distance,
+        payload=dict(point.payload),
+        vectors=dict(point.vectors) if with_vectors else None,
+    )
+
+
+def _compile(expr: Expr, models: Any) -> Any:
+    if isinstance(expr, Eq):
+        return models.FieldCondition(key=expr.field, match=models.MatchValue(value=expr.value))
+    if isinstance(expr, Cmp):
+        return models.FieldCondition(key=expr.field, range=models.Range(**{expr.op: expr.value}))
+    if isinstance(expr, In):
+        return models.FieldCondition(key=expr.field, match=models.MatchAny(any=list(expr.values)))
+    if isinstance(expr, IsNull):
+        return models.IsNullCondition(is_null=models.PayloadField(key=expr.field))
+    if isinstance(expr, Not):
+        return models.Filter(must_not=[_compile(expr.inner, models)])
+    if isinstance(expr, And):
+        return models.Filter(must=[_compile(child, models) for child in expr.children])
+    if isinstance(expr, Or):
+        return models.Filter(should=[_compile(child, models) for child in expr.children])
+    raise ValidationError(f"unsupported filter node: {type(expr).__name__}")
+
+
+def _models() -> Any:
+    try:
+        from qdrant_client import models
+    except ImportError as exc:
+        raise BackendError("Qdrant support requires the qdrant extra: pip install 'vecspan[qdrant]'") from exc
+    return models
+
+
+def _meta_id(name: str) -> str:
+    return str(uuid.uuid5(_META_NAMESPACE, name))
+
+
+def _field(record: Any, name: str) -> Any:
+    if isinstance(record, Mapping):
+        return record.get(name)
+    return getattr(record, name, None)
+
+
+def _valid_offset(offset: object) -> bool:
+    if isinstance(offset, bool):
+        return False
+    return isinstance(offset, (int, str))
+
+
+def _offset_token(value: object) -> int | str | None:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        raise BackendError("unexpected Qdrant scroll offset")
+    if isinstance(value, int):
+        return value
+    return str(value)
+
+
+def _accepts_collection_name(func: Any) -> bool:
+    try:
+        import inspect
+
+        return "collection_name" in inspect.signature(func).parameters
+    except (TypeError, ValueError):
+        return False

vecspan/capabilities.py

@@ -0,0 +1,16 @@
+"""Operations a backend may or may not implement."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class Capability(str, Enum):
+    """A portable operation that some backends might not support."""
+
+    DELETE_BY_FILTER = "delete_by_filter"
+    UPDATE_PAYLOAD = "update_payload"
+    UPDATE_VECTORS = "update_vectors"
+
+
+ALL_CAPABILITIES = frozenset(Capability)