vecspan 0.1.0__py3-none-any.whl

vecspan/__init__.py

@@ -0,0 +1,50 @@
+"""One Python API for Milvus, Qdrant, and other vector databases."""
+
+from vecspan.capabilities import Capability
+from vecspan.client import VectorClient
+from vecspan.errors import (
+    BackendError,
+    CollectionAlreadyExistsError,
+    CollectionNotFoundError,
+    UnsupportedOperation,
+    ValidationError,
+    VecspanError,
+)
+from vecspan.filters import Filter
+from vecspan.registry import register_backend
+from vecspan.types import (
+    CollectionSpec,
+    Distance,
+    IdType,
+    PayloadField,
+    PayloadType,
+    Point,
+    ScrollPage,
+    SearchHit,
+    VectorField,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "BackendError",
+    "Capability",
+    "CollectionAlreadyExistsError",
+    "CollectionNotFoundError",
+    "CollectionSpec",
+    "Distance",
+    "Filter",
+    "IdType",
+    "PayloadField",
+    "PayloadType",
+    "Point",
+    "ScrollPage",
+    "SearchHit",
+    "UnsupportedOperation",
+    "ValidationError",
+    "VecspanError",
+    "VectorClient",
+    "VectorField",
+    "__version__",
+    "register_backend",
+]

vecspan/backends/__init__.py

@@ -0,0 +1,6 @@
+"""Built-in backends. Milvus and Qdrant stay lazy so importing vecspan is cheap."""
+
+from vecspan.backends.base import VectorBackend
+from vecspan.backends.memory import MemoryBackend
+
+__all__ = ["MemoryBackend", "VectorBackend"]

vecspan/backends/base.py

@@ -0,0 +1,107 @@
+"""Backend interface implemented by Milvus, Qdrant, and the in-memory store."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Sequence
+
+from vecspan.capabilities import ALL_CAPABILITIES, Capability
+from vecspan.filters import Filter
+from vecspan.types import CollectionSpec, Point, ScrollPage, SearchHit
+
+
+class VectorBackend(ABC):
+    """Storage operations behind `VectorClient`.
+
+    `VectorClient` validates arguments and checks capabilities before it calls
+    these methods. A backend still raises `VecspanError` for conditions only
+    the database can see, such as a missing collection or a missing point id.
+    """
+
+    name = "backend"
+    capabilities: frozenset[Capability] = ALL_CAPABILITIES
+
+    @abstractmethod
+    def create_collection(self, spec: CollectionSpec) -> None:
+        """Create a collection from `spec`."""
+
+    @abstractmethod
+    def drop_collection(self, name: str) -> None:
+        """Drop a collection."""
+
+    @abstractmethod
+    def has_collection(self, name: str) -> bool:
+        """Return whether the collection exists."""
+
+    @abstractmethod
+    def describe_collection(self, name: str) -> CollectionSpec:
+        """Return the schema stored for a collection."""
+
+    @abstractmethod
+    def list_collections(self) -> list[str]:
+        """Return collection names visible to callers."""
+
+    @abstractmethod
+    def upsert(self, name: str, points: Sequence[Point]) -> None:
+        """Insert points or replace them completely when the id already exists."""
+
+    @abstractmethod
+    def update_payload(self, name: str, ids: Sequence[int | str], payload: dict) -> None:
+        """Merge `payload` into existing points. Missing ids are an error."""
+
+    @abstractmethod
+    def update_vectors(
+        self,
+        name: str,
+        ids: Sequence[int | str],
+        vectors: Sequence[dict[str, list[float]]],
+    ) -> None:
+        """Replace the given vector fields on existing points and leave the rest."""
+
+    @abstractmethod
+    def delete(
+        self,
+        name: str,
+        *,
+        ids: Sequence[int | str] | None = None,
+        filter: Filter | None = None,
+    ) -> None:
+        """Delete by id or by filter. Exactly one of the two is set."""
+
+    @abstractmethod
+    def get(self, name: str, ids: Sequence[int | str], *, with_vectors: bool = False) -> list[Point]:
+        """Return the points that exist, in the same order as `ids`."""
+
+    @abstractmethod
+    def count(self, name: str, *, filter: Filter | None = None) -> int:
+        """Count points, optionally restricted by a filter."""
+
+    @abstractmethod
+    def scroll(
+        self,
+        name: str,
+        *,
+        filter: Filter | None = None,
+        limit: int = 100,
+        offset: int | str | None = None,
+        with_vectors: bool = False,
+    ) -> ScrollPage:
+        """Read a page of points. `offset` comes from a previous page."""
+
+    @abstractmethod
+    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]:
+        """Return the nearest points for one named vector field."""
+
+    def close(self) -> None:
+        """Release the underlying connection. The default does nothing."""
+        return None

vecspan/backends/memory.py

@@ -0,0 +1,289 @@
+"""In-process backend used by tests and local development."""
+
+from __future__ import annotations
+
+import copy
+import math
+import threading
+from dataclasses import dataclass
+from typing import Sequence
+
+from vecspan.backends.base import VectorBackend
+from vecspan.errors import (
+    BackendError,
+    CollectionAlreadyExistsError,
+    CollectionNotFoundError,
+    ValidationError,
+)
+from vecspan.filters import Filter, matches
+from vecspan.types import CollectionSpec, Distance, Point, ScrollPage, SearchHit
+from vecspan.validate import (
+    check_limit,
+    check_with_vectors,
+    prepare_filter,
+    prepare_ids,
+    prepare_payload,
+    prepare_points,
+    prepare_query_vector,
+    prepare_search_options,
+    prepare_vector_updates,
+)
+
+PointId = int | str
+
+
+@dataclass
+class _Record:
+    id: PointId
+    vectors: dict[str, list[float]]
+    payload: dict
+
+
+class MemoryBackend(VectorBackend):
+    """Brute-force vector store kept in the current process."""
+
+    name = "memory"
+
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self._collections: dict[str, tuple[CollectionSpec, dict[PointId, _Record]]] = {}
+        self._closed = False
+
+    def create_collection(self, spec: CollectionSpec) -> None:
+        with self._lock:
+            self._check_open()
+            if not isinstance(spec, CollectionSpec):
+                raise ValidationError("spec must be a CollectionSpec")
+            if spec.name in self._collections:
+                raise CollectionAlreadyExistsError(f"collection {spec.name!r} already exists")
+            self._collections[spec.name] = (spec, {})
+
+    def drop_collection(self, name: str) -> None:
+        with self._lock:
+            self._check_open()
+            if name not in self._collections:
+                raise CollectionNotFoundError(f"collection {name!r} does not exist")
+            del self._collections[name]
+
+    def has_collection(self, name: str) -> bool:
+        with self._lock:
+            self._check_open()
+            return name in self._collections
+
+    def describe_collection(self, name: str) -> CollectionSpec:
+        with self._lock:
+            self._check_open()
+            return self._spec(name)
+
+    def list_collections(self) -> list[str]:
+        with self._lock:
+            self._check_open()
+            return list(self._collections)
+
+    def upsert(self, name: str, points: Sequence[Point]) -> None:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            for point in prepare_points(spec, points):
+                store[point.id] = _Record(
+                    point.id,
+                    {key: list(value) for key, value in point.vectors.items()},
+                    copy.deepcopy(point.payload),
+                )
+
+    def update_payload(self, name: str, ids: Sequence[PointId], payload: dict) -> None:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            prepared_ids = prepare_ids(spec, ids)
+            patch = prepare_payload(spec, payload, allow_empty=False)
+            records = self._require(store, prepared_ids)
+            for record in records:
+                record.payload.update(copy.deepcopy(patch))
+
+    def update_vectors(
+        self,
+        name: str,
+        ids: Sequence[PointId],
+        vectors: Sequence[dict[str, list[float]]],
+    ) -> None:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            updates = prepare_vector_updates(spec, ids, vectors)
+            records = {record.id: record for record in self._require(store, [point_id for point_id, _ in updates])}
+            for point_id, vector_update in updates:
+                records[point_id].vectors.update(
+                    {key: list(value) for key, value in vector_update.items()}
+                )
+
+    def delete(
+        self,
+        name: str,
+        *,
+        ids: Sequence[PointId] | None = None,
+        filter: Filter | None = None,
+    ) -> None:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            if (ids is None) == (filter is None):
+                raise ValidationError("pass exactly one of ids or filter")
+            if ids is not None:
+                for point_id in prepare_ids(spec, ids):
+                    store.pop(point_id, None)
+                return
+            prepared = prepare_filter(spec, filter)
+            assert prepared is not None
+            doomed = [point_id for point_id, record in store.items() if matches(record.payload, prepared)]
+            for point_id in doomed:
+                del store[point_id]
+
+    def get(self, name: str, ids: Sequence[PointId], *, with_vectors: bool = False) -> list[Point]:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            include = check_with_vectors(with_vectors)
+            found: list[Point] = []
+            for point_id in prepare_ids(spec, ids):
+                record = store.get(point_id)
+                if record is not None:
+                    found.append(self._point(record, include))
+            return found
+
+    def count(self, name: str, *, filter: Filter | None = None) -> int:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            prepared = prepare_filter(spec, filter)
+            if prepared is None:
+                return len(store)
+            return sum(1 for record in store.values() if matches(record.payload, prepared))
+
+    def scroll(
+        self,
+        name: str,
+        *,
+        filter: Filter | None = None,
+        limit: int = 100,
+        offset: int | str | None = None,
+        with_vectors: bool = False,
+    ) -> ScrollPage:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            prepared = prepare_filter(spec, filter)
+            size = check_limit(limit)
+            include = check_with_vectors(with_vectors)
+            start = 0 if offset is None else offset
+            if isinstance(start, bool) or not isinstance(start, int) or start < 0:
+                raise ValidationError("scroll offset must be an integer from a previous page")
+            ordered = [
+                record
+                for _, record in sorted(store.items(), key=lambda item: _id_sort_key(item[0]))
+                if prepared is None or matches(record.payload, prepared)
+            ]
+            page = ordered[start : start + size]
+            next_offset = None if start + size >= len(ordered) else start + size
+            return ScrollPage(
+                points=[self._point(record, include) for record in page],
+                next_offset=next_offset,
+            )
+
+    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]:
+        with self._lock:
+            self._check_open()
+            spec, store = self._collection(name)
+            field, query = prepare_query_vector(spec, vector, vector_name)
+            prepared = prepare_filter(spec, filter)
+            size = check_limit(limit)
+            include = check_with_vectors(with_vectors)
+            prepare_search_options(backend_options, allow_nprobe=False)
+            ranked: list[tuple[float, PointId, _Record]] = []
+            for record in store.values():
+                if prepared is not None and not matches(record.payload, prepared):
+                    continue
+                score = _score(field.distance, query, record.vectors[field.name])
+                ranked.append((score, record.id, record))
+            ranked.sort(key=lambda item: _rank_key(field.distance, item[0], item[1]))
+            hits: list[SearchHit] = []
+            for score, _, record in ranked[:size]:
+                point = self._point(record, include)
+                hits.append(
+                    SearchHit(
+                        id=point.id,
+                        score=score,
+                        distance=field.distance,
+                        payload=point.payload,
+                        vectors=point.vectors if include else None,
+                    )
+                )
+            return hits
+
+    def close(self) -> None:
+        with self._lock:
+            self._closed = True
+            self._collections.clear()
+
+    def _check_open(self) -> None:
+        if self._closed:
+            raise BackendError("client is closed")
+
+    def _collection(self, name: str) -> tuple[CollectionSpec, dict[PointId, _Record]]:
+        try:
+            return self._collections[name]
+        except KeyError:
+            raise CollectionNotFoundError(f"collection {name!r} does not exist") from None
+
+    def _spec(self, name: str) -> CollectionSpec:
+        return self._collection(name)[0]
+
+    def _require(self, store: dict[PointId, _Record], ids: Sequence[PointId]) -> list[_Record]:
+        records: list[_Record] = []
+        for point_id in ids:
+            record = store.get(point_id)
+            if record is None:
+                raise ValidationError(f"unknown point id: {point_id!r}")
+            records.append(record)
+        return records
+
+    @staticmethod
+    def _point(record: _Record, with_vectors: bool) -> Point:
+        vectors = {key: list(value) for key, value in record.vectors.items()} if with_vectors else {}
+        return Point(id=record.id, vectors=vectors, payload=copy.deepcopy(record.payload))
+
+
+def _id_sort_key(point_id: PointId) -> tuple:
+    if isinstance(point_id, int):
+        return (0, point_id)
+    return (1, point_id)
+
+
+def _rank_key(distance: Distance, score: float, point_id: PointId) -> tuple:
+    primary = score if distance == Distance.L2 else -score
+    return (primary, _id_sort_key(point_id))
+
+
+def _score(distance: Distance, query: Sequence[float], stored: Sequence[float]) -> float:
+    if distance == Distance.DOT:
+        return _dot(query, stored)
+    if distance == Distance.COSINE:
+        denom = math.sqrt(_dot(query, query) * _dot(stored, stored))
+        if denom == 0:
+            return 0.0
+        return _dot(query, stored) / denom
+    return math.sqrt(sum((left - right) ** 2 for left, right in zip(query, stored)))
+
+
+def _dot(left: Sequence[float], right: Sequence[float]) -> float:
+    return sum(left_item * right_item for left_item, right_item in zip(left, right))