vecspan 0.1.0__py3-none-any.whl

vecspan/client.py

@@ -0,0 +1,205 @@
+"""Facade that validates calls and delegates them to a backend."""
+
+from __future__ import annotations
+
+from typing import Any, Mapping, Sequence
+
+from vecspan.backends.base import VectorBackend
+from vecspan.capabilities import Capability
+from vecspan.errors import UnsupportedOperation, ValidationError
+from vecspan.filters import Filter
+from vecspan.naming import check_collection_name
+from vecspan.registry import create_backend
+from vecspan.types import CollectionSpec, Point, ScrollPage, SearchHit
+from vecspan.validate import (
+    check_limit,
+    check_with_vectors,
+    prepare_filter,
+    prepare_ids,
+    prepare_payload,
+    prepare_points,
+    prepare_query_vector,
+    prepare_vector_updates,
+)
+
+PointId = int | str
+
+
+class VectorClient:
+    """Synchronous client for one vector database.
+
+    Connect with a URL, or pass a backend you already built:
+
+        client = VectorClient.connect("memory://")
+        client = VectorClient.connect("milvus://localhost:19530")
+        client = VectorClient.connect("qdrant://localhost:6333", api_key="...")
+    """
+
+    def __init__(self, backend: VectorBackend) -> None:
+        self.backend = backend
+        self._specs: dict[str, CollectionSpec] = {}
+
+    @classmethod
+    def connect(cls, target: str | VectorBackend, **kwargs: Any) -> VectorClient:
+        """Connect to `target`. A string is a URL. A backend is used as-is."""
+        if isinstance(target, str):
+            return cls(create_backend(target, **kwargs))
+        if kwargs:
+            raise ValidationError("options are only accepted when connecting with a URL")
+        return cls(target)
+
+    def create_collection(self, spec: CollectionSpec) -> None:
+        """Create a collection. The id type and vector fields stay fixed afterwards."""
+        if not isinstance(spec, CollectionSpec):
+            raise ValidationError("spec must be a CollectionSpec")
+        self.backend.create_collection(spec)
+        self._specs[spec.name] = spec
+
+    def drop_collection(self, name: str) -> None:
+        """Drop a collection."""
+        checked = check_collection_name(name)
+        self.backend.drop_collection(checked)
+        self._specs.pop(checked, None)
+
+    def has_collection(self, name: str) -> bool:
+        """Return whether a collection exists."""
+        return self.backend.has_collection(check_collection_name(name))
+
+    def describe_collection(self, name: str) -> CollectionSpec:
+        """Return the schema of a collection."""
+        checked = check_collection_name(name)
+        spec = self.backend.describe_collection(checked)
+        self._specs[checked] = spec
+        return spec
+
+    def list_collections(self) -> list[str]:
+        """Return collection names in sorted order."""
+        return sorted(self.backend.list_collections())
+
+    def upsert(self, name: str, points: Point | Mapping[str, Any] | Sequence[object]) -> None:
+        """Insert points, or replace the vector and payload of an existing id."""
+        spec = self._spec(name)
+        prepared = prepare_points(spec, points)
+        if prepared:
+            self.backend.upsert(spec.name, prepared)
+
+    def update_payload(self, name: str, ids: object, payload: Mapping[str, Any]) -> None:
+        """Merge payload fields into existing points. Other fields stay as they are."""
+        self._need(Capability.UPDATE_PAYLOAD)
+        spec = self._spec(name)
+        prepared_ids = prepare_ids(spec, ids)
+        patch = prepare_payload(spec, payload, allow_empty=False)
+        self.backend.update_payload(spec.name, prepared_ids, patch)
+
+    def update_vectors(self, name: str, ids: object, vectors: Sequence[Mapping[str, Sequence[float]]]) -> None:
+        """Replace the named vectors and leave every other vector field untouched."""
+        self._need(Capability.UPDATE_VECTORS)
+        spec = self._spec(name)
+        updates = prepare_vector_updates(spec, ids, vectors)
+        self.backend.update_vectors(
+            spec.name,
+            [point_id for point_id, _ in updates],
+            [vector for _, vector in updates],
+        )
+
+    def delete(
+        self,
+        name: str,
+        *,
+        ids: object = None,
+        filter: Filter | None = None,
+    ) -> None:
+        """Delete points by id or by filter. Pass one of the two, not both."""
+        spec = self._spec(name)
+        if (ids is None) == (filter is None):
+            raise ValidationError("pass exactly one of ids or filter")
+        if filter is not None:
+            self._need(Capability.DELETE_BY_FILTER)
+            prepared = prepare_filter(spec, filter)
+            self.backend.delete(spec.name, filter=prepared)
+            return
+        prepared_ids = prepare_ids(spec, ids)
+        self.backend.delete(spec.name, ids=prepared_ids)
+
+    def get(self, name: str, ids: object, *, with_vectors: bool = False) -> list[Point]:
+        """Return existing points in the requested id order."""
+        spec = self._spec(name)
+        prepared_ids = prepare_ids(spec, ids)
+        include = check_with_vectors(with_vectors)
+        return self.backend.get(spec.name, prepared_ids, with_vectors=include)
+
+    def count(self, name: str, *, filter: Filter | None = None) -> int:
+        """Count points that match `filter`, or every point when it is omitted."""
+        spec = self._spec(name)
+        return self.backend.count(spec.name, filter=prepare_filter(spec, filter))
+
+    def scroll(
+        self,
+        name: str,
+        *,
+        filter: Filter | None = None,
+        limit: int = 100,
+        offset: int | str | None = None,
+        with_vectors: bool = False,
+    ) -> ScrollPage:
+        """Read one page of points. Pass `next_offset` back in as `offset`."""
+        spec = self._spec(name)
+        return self.backend.scroll(
+            spec.name,
+            filter=prepare_filter(spec, filter),
+            limit=check_limit(limit),
+            offset=offset,
+            with_vectors=check_with_vectors(with_vectors),
+        )
+
+    def search(
+        self,
+        name: str,
+        vector: Sequence[float],
+        *,
+        vector_name: str | None = None,
+        filter: Filter | None = None,
+        limit: int = 10,
+        with_vectors: bool = False,
+        backend_options: Mapping[str, Any] | None = None,
+    ) -> list[SearchHit]:
+        """Return the nearest points.
+
+        `score` is the backend's own value. Cosine and dot are larger when
+        closer. L2 is smaller when closer. Memory and Qdrant L2 values are
+        Euclidean distance. Milvus L2 values are squared Euclidean distance.
+        """
+        spec = self._spec(name)
+        field, query = prepare_query_vector(spec, vector, vector_name)
+        return self.backend.search(
+            spec.name,
+            query,
+            vector_name=field.name,
+            filter=prepare_filter(spec, filter),
+            limit=check_limit(limit),
+            with_vectors=check_with_vectors(with_vectors),
+            backend_options=None if backend_options is None else dict(backend_options),
+        )
+
+    def close(self) -> None:
+        """Close the underlying connection."""
+        self._specs.clear()
+        self.backend.close()
+
+    def __enter__(self) -> VectorClient:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def _spec(self, name: str) -> CollectionSpec:
+        checked = check_collection_name(name)
+        spec = self._specs.get(checked)
+        if spec is None:
+            spec = self.backend.describe_collection(checked)
+            self._specs[checked] = spec
+        return spec
+
+    def _need(self, capability: Capability) -> None:
+        if capability not in self.backend.capabilities:
+            raise UnsupportedOperation(f"{self.backend.name} does not support {capability.value}")

vecspan/errors.py

@@ -0,0 +1,27 @@
+"""Errors raised by vecspan."""
+
+from __future__ import annotations
+
+
+class VecspanError(Exception):
+    """Base class for errors raised by this library."""
+
+
+class ValidationError(VecspanError):
+    """The call does not match the collection schema or the public API."""
+
+
+class UnsupportedOperation(VecspanError):
+    """The connected backend does not implement this operation."""
+
+
+class CollectionNotFoundError(VecspanError):
+    """No collection with this name exists."""
+
+
+class CollectionAlreadyExistsError(VecspanError):
+    """A collection with this name already exists."""
+
+
+class BackendError(VecspanError):
+    """The underlying database rejected the call or could not be reached."""

vecspan/filters.py

@@ -0,0 +1,258 @@
+"""Filter expression tree, in-memory matcher, and Milvus compiler.
+
+Qdrant compilation lives next to that backend because it builds Qdrant models.
+The tree itself has no database dependency.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Mapping, Sequence
+
+from vecspan.errors import ValidationError
+from vecspan.naming import check_field_name
+
+Expr = "Eq | Cmp | In | IsNull | Not | And | Or"
+
+
+@dataclass(frozen=True)
+class Eq:
+    field: str
+    value: object
+
+
+@dataclass(frozen=True)
+class Cmp:
+    field: str
+    op: str
+    value: object
+
+
+@dataclass(frozen=True)
+class In:
+    field: str
+    values: tuple[object, ...]
+
+
+@dataclass(frozen=True)
+class IsNull:
+    field: str
+
+
+@dataclass(frozen=True)
+class Not:
+    inner: Expr
+
+
+@dataclass(frozen=True)
+class And:
+    children: tuple[Expr, ...]
+
+
+@dataclass(frozen=True)
+class Or:
+    children: tuple[Expr, ...]
+
+
+_CMP_OPS = {
+    "gt": ">",
+    "gte": ">=",
+    "lt": "<",
+    "lte": "<=",
+}
+
+
+class Filter:
+    """A boolean expression over payload fields.
+
+    Build it with the class methods and combine pieces with `&`, `|`, and `~`.
+    Field names are identifiers. Values are scalars, so a value cannot change
+    the shape of the compiled expression.
+    """
+
+    def __init__(self, expr: Expr) -> None:
+        self.expr = expr
+
+    @classmethod
+    def eq(cls, field: str, value: object) -> Filter:
+        """Match a payload field equal to `value`."""
+        return cls(Eq(check_field_name(field), _equality_value(value)))
+
+    @classmethod
+    def gt(cls, field: str, value: object) -> Filter:
+        return cls._cmp("gt", field, value)
+
+    @classmethod
+    def gte(cls, field: str, value: object) -> Filter:
+        return cls._cmp("gte", field, value)
+
+    @classmethod
+    def lt(cls, field: str, value: object) -> Filter:
+        return cls._cmp("lt", field, value)
+
+    @classmethod
+    def lte(cls, field: str, value: object) -> Filter:
+        return cls._cmp("lte", field, value)
+
+    @classmethod
+    def in_(cls, field: str, values: Sequence[object]) -> Filter:
+        """Match a field whose value is one of `values`.
+
+        The method is named `in_` because `in` is a Python keyword.
+        """
+        if isinstance(values, (str, bytes)) or not isinstance(values, Sequence):
+            raise ValidationError("in values must be a sequence")
+        items = tuple(_equality_value(item) for item in values)
+        if not items:
+            raise ValidationError("in values must not be empty")
+        kinds = {_value_kind(item) for item in items}
+        if len(kinds) > 1:
+            raise ValidationError("in values must share one type")
+        return cls(In(check_field_name(field), items))
+
+    @classmethod
+    def is_null(cls, field: str) -> Filter:
+        """Match a missing payload field or a field explicitly set to null."""
+        return cls(IsNull(check_field_name(field)))
+
+    @classmethod
+    def _cmp(cls, op: str, field: str, value: object) -> Filter:
+        return cls(Cmp(check_field_name(field), op, _number_value(value)))
+
+    def __and__(self, other: object) -> Filter:
+        if not isinstance(other, Filter):
+            return NotImplemented
+        return Filter(And((self.expr, other.expr)))
+
+    def __or__(self, other: object) -> Filter:
+        if not isinstance(other, Filter):
+            return NotImplemented
+        return Filter(Or((self.expr, other.expr)))
+
+    def __invert__(self) -> Filter:
+        return Filter(Not(self.expr))
+
+    def __repr__(self) -> str:
+        return f"Filter({self.expr!r})"
+
+
+def matches(payload: Mapping[str, object], filt: Filter) -> bool:
+    """Return whether `payload` satisfies `filt`."""
+    return _matches(payload, filt.expr)
+
+
+def compile_milvus(filt: Filter) -> str:
+    """Compile a filter to a Milvus boolean expression."""
+    return _compile_milvus(filt.expr)
+
+
+def _matches(payload: Mapping[str, object], expr: Expr) -> bool:
+    if isinstance(expr, Eq):
+        if expr.field not in payload:
+            return False
+        return _same(payload[expr.field], expr.value)
+    if isinstance(expr, Cmp):
+        if expr.field not in payload:
+            return False
+        left = payload[expr.field]
+        if not _is_number(left) or not _is_number(expr.value):
+            return False
+        if expr.op == "gt":
+            return left > expr.value
+        if expr.op == "gte":
+            return left >= expr.value
+        if expr.op == "lt":
+            return left < expr.value
+        return left <= expr.value
+    if isinstance(expr, In):
+        if expr.field not in payload:
+            return False
+        return any(_same(payload[expr.field], item) for item in expr.values)
+    if isinstance(expr, IsNull):
+        return expr.field not in payload or payload[expr.field] is None
+    if isinstance(expr, Not):
+        return not _matches(payload, expr.inner)
+    if isinstance(expr, And):
+        return all(_matches(payload, child) for child in expr.children)
+    if isinstance(expr, Or):
+        return any(_matches(payload, child) for child in expr.children)
+    raise ValidationError(f"unsupported filter node: {type(expr).__name__}")
+
+
+def _compile_milvus(expr: Expr) -> str:
+    if isinstance(expr, Eq):
+        return f"{expr.field} == {_literal(expr.value)}"
+    if isinstance(expr, Cmp):
+        return f"{expr.field} {_CMP_OPS[expr.op]} {_literal(expr.value)}"
+    if isinstance(expr, In):
+        rendered = ", ".join(_literal(item) for item in expr.values)
+        return f"{expr.field} in [{rendered}]"
+    if isinstance(expr, IsNull):
+        return f"{expr.field} is null"
+    if isinstance(expr, Not):
+        return f"not ({_compile_milvus(expr.inner)})"
+    if isinstance(expr, And):
+        return " and ".join(f"({_compile_milvus(child)})" for child in expr.children)
+    if isinstance(expr, Or):
+        return " or ".join(f"({_compile_milvus(child)})" for child in expr.children)
+    raise ValidationError(f"unsupported filter node: {type(expr).__name__}")
+
+
+def _equality_value(value: object) -> object:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        if not math.isfinite(value):
+            raise ValidationError("filter value must be finite")
+        return value
+    if isinstance(value, str):
+        return value
+    raise ValidationError(f"filter value must be a string, number, or bool, not {type(value).__name__}")
+
+
+def _number_value(value: object) -> int | float:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ValidationError("comparison value must be a number")
+    if isinstance(value, float) and not math.isfinite(value):
+        raise ValidationError("filter value must be finite")
+    return value
+
+
+def _value_kind(value: object) -> str:
+    if isinstance(value, bool):
+        return "bool"
+    if isinstance(value, (int, float)):
+        return "number"
+    if isinstance(value, str):
+        return "string"
+    return type(value).__name__
+
+
+def _is_number(value: object) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool)
+
+
+def _same(left: object, right: object) -> bool:
+    if left is None or right is None:
+        return False
+    if isinstance(left, bool) or isinstance(right, bool):
+        return isinstance(left, bool) and isinstance(right, bool) and left is right
+    if _is_number(left) and _is_number(right):
+        return left == right
+    return type(left) is type(right) and left == right
+
+
+def _literal(value: object) -> str:
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, int):
+        return str(value)
+    if isinstance(value, float):
+        return repr(value)
+    if isinstance(value, str):
+        escaped = value.replace("\\", "\\\\").replace('"', '\\"')
+        return f'"{escaped}"'
+    raise ValidationError(f"filter value must be a string, number, or bool, not {type(value).__name__}")

vecspan/naming.py

@@ -0,0 +1,35 @@
+"""Shared rules for collection and field names."""
+
+from __future__ import annotations
+
+import re
+
+from vecspan.errors import ValidationError
+
+_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+META_COLLECTION = "_vecspan_schemas"
+RESERVED_COLLECTIONS = frozenset({META_COLLECTION})
+_MAX_NAME_LENGTH = 255
+
+
+def check_identifier(name: object, label: str) -> str:
+    """Return `name` when it is a short ASCII identifier."""
+    if not isinstance(name, str) or not _IDENTIFIER.match(name) or len(name) > _MAX_NAME_LENGTH:
+        raise ValidationError(f"invalid {label}: {name!r}")
+    return name
+
+
+def check_field_name(name: object) -> str:
+    """Return a field name that is safe to compile into a filter."""
+    checked = check_identifier(name, "field name")
+    if checked == "id" or checked.startswith("$"):
+        raise ValidationError(f"reserved field name: {checked!r}")
+    return checked
+
+
+def check_collection_name(name: object) -> str:
+    """Return a collection name callers are allowed to create."""
+    checked = check_identifier(name, "collection name")
+    if checked in RESERVED_COLLECTIONS:
+        raise ValidationError(f"reserved collection name: {checked!r}")
+    return checked

vecspan/registry.py

@@ -0,0 +1,67 @@
+"""Resolve a connection URL to a backend."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from urllib.parse import urlsplit
+
+from vecspan.backends.memory import MemoryBackend
+from vecspan.backends.base import VectorBackend
+from vecspan.errors import ValidationError
+
+_BUILTINS = frozenset({"memory", "milvus", "qdrant", "qdrant+https"})
+_EXTRA: dict[str, Callable[..., VectorBackend]] = {}
+
+
+def register_backend(scheme: str) -> Callable[[Callable[..., VectorBackend]], Callable[..., VectorBackend]]:
+    """Register a factory for a URL scheme.
+
+    The factory is called as `factory(parsed_url, **kwargs)` and returns a
+    `VectorBackend`. Built-in schemes cannot be replaced.
+    """
+
+    def decorator(factory: Callable[..., VectorBackend]) -> Callable[..., VectorBackend]:
+        key = scheme.lower()
+        if not key or any(character in key for character in ":/"):
+            raise ValidationError(f"invalid backend scheme: {scheme!r}")
+        if key in _BUILTINS or key in _EXTRA:
+            raise ValidationError(f"backend {key!r} is already registered")
+        _EXTRA[key] = factory
+        return factory
+
+    return decorator
+
+
+def unregister_backend(scheme: str) -> None:
+    """Remove a backend registered with `register_backend`."""
+    _EXTRA.pop(scheme.lower(), None)
+
+
+def create_backend(url: str, **kwargs: object) -> VectorBackend:
+    """Build the backend addressed by `url`."""
+    if not isinstance(url, str) or "://" not in url:
+        raise ValidationError("URL must look like memory://, milvus://host:19530, or qdrant://host:6333")
+    parsed = urlsplit(url)
+    scheme = parsed.scheme.lower()
+    if scheme == "memory":
+        if kwargs:
+            raise ValidationError(f"unknown option: {sorted(kwargs)[0]}")
+        _check_memory_url(parsed)
+        return MemoryBackend()
+    if scheme == "milvus":
+        from vecspan.backends.milvus import MilvusBackend
+
+        return MilvusBackend.from_url(url, **kwargs)
+    if scheme in {"qdrant", "qdrant+https"}:
+        from vecspan.backends.qdrant import QdrantBackend
+
+        return QdrantBackend.from_url(url, **kwargs)
+    factory = _EXTRA.get(scheme)
+    if factory is None:
+        raise ValidationError(f"unknown backend scheme: {scheme!r}")
+    return factory(parsed, **kwargs)
+
+
+def _check_memory_url(parsed: object) -> None:
+    if parsed.netloc or parsed.path not in ("", "/") or parsed.query or parsed.fragment or parsed.username:
+        raise ValidationError("memory URL does not take a host, path, or options")