vink 0.1.0a1__py3-none-any.whl

vink/__init__.py

@@ -0,0 +1,52 @@
+"""
+Vink: Vector Incremental Nano Kit
+
+A lightweight vector database that incrementally switches from exact to
+approximate search as your data grows — without full index rebuilds.
+
+Note:
+    ANN switching is one-way — once switched to approximate search,
+    the system will never switch back to exact search.
+
+Key differentiators:
+    - **Incremental inserts**: Add vectors anytime — no rebuild per insert.
+    - **Automatic strategy switching**: No manual tuning — exact for small datasets, ANN for large.
+    - **Thread-safe**: Background ANN building doesn't block new operations.
+    - **Soft deletes + compact**: Efficient deletion with explicit storage reclamation.
+
+Features:
+    - ~100x faster using RII and Product Quantization for large datasets.
+    - Pure Python: No external services or dependencies required.
+    - Supports Euclidean (L2) and cosine (dot) product similarity.
+    - SQLite-backed persistent storage.
+
+Technical Background:
+    Vink uses Reconfigurable Inverted Index (RII) with Product Quantization (PQ)
+    for approximate nearest neighbor search. The switch from exact to approximate
+    happens when the normalized power-law complexity reaches 1.0:
+    (dim * vectors / 1M) ^ switch_exp >= 1.0. Default switch_exp is 1.0.
+
+References:
+    .. [Matsui18] Matsui et al., "Reconfigurable Inverted Index", ACM MM 2018.
+
+    .. [Jegou11] Jegou et al., "Product Quantization for Nearest Neighbor Search",
+       IEEE TPAMI 2011.
+
+    .. [Matsui15] Matsui et al., "Optimized Product Quantization for Nearest
+       Neighbor Search", CVPR 2015.
+
+See Also:
+    - RII: https://github.com/matsui528/rii
+    - nanopq: https://github.com/matsui528/nanopq
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from vink.core import VinkDB
+from vink.exceptions import *
+from vink.models import AnnConfig
+
+try:
+    __version__ = version("vink")
+except PackageNotFoundError:
+    __version__ = "0.0.0"

vink/core.py

@@ -0,0 +1,519 @@
+import shutil
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from threading import Thread
+from typing import TYPE_CHECKING, Annotated, Callable, Literal
+
+if TYPE_CHECKING:
+    from vink.latency_predictor import LatencyPredictor
+
+import numpy as np
+import regex as re
+from pydantic import Field, ValidationError
+from readerwriterlock import rwlock
+
+from vink.exceptions import InvalidInputError, VectorDimensionError
+from vink.models import AnnConfig, VectorRecords
+from vink.sql_wrapper import SQLiteWrapper
+
+# The strategies and the latency predictor are lazy imported
+from vink.strategies.base import BaseStrategy
+from vink.utils.input_validation import (
+    pretty_errors,
+    validate_arguments,
+    validate_embedding,
+    validate_id,
+)
+from vink.utils.logging import log_info, logger
+
+
+class VinkDB:
+    """
+    Vector database with hybrid exact/approximate nearest neighbor search.
+
+    VinkDB automatically switches from exact brute-force search to approximate
+    nearest neighbor (ANN) search based on dataset size, using Reconfigurable Inverted
+    Index (RII) and Product Quantization (PQ) for efficient ANN.
+
+    Note:
+        ANN switching is one-way — once switched, the system never switches back to exact search.
+
+    Features:
+
+        - Hybrid search: exact for small datasets, ANN for large datasets.
+        - Automatic strategy switching based on runtime-calibrated latency prediction.
+        - Normalized embeddings for consistent distance metrics.
+        - Supports Euclidean (L2) and cosine (dot) product similarity.
+        - Soft deletes: efficient deletion without data reorganization.
+
+    Getting ANNConfig:
+
+        To customize ANN behavior, create an ANNConfig instance:
+
+        >>> from vink import AnnConfig
+        >>> config = AnnConfig(
+        ...     num_subspaces=16,
+        ...     codebook_size=128,
+        ...     switch_latency_ms=300,
+        ...     quantizer="pq"
+        ... )
+        >>> db = VinkDB(dir_path="./data", dim=384, ann_config=config)
+
+        For help with AnnConfig parameters, call AnnConfig.help()
+    """
+
+    @validate_arguments
+    def __init__(
+        self,
+        dir_path: str | Path,
+        dim:  Annotated[int, Field(ge=16)],
+        metric: Literal["euclidean", "cosine"] = "euclidean",
+        force_exact: bool = False,
+        ann_config: AnnConfig | None = None,
+        embedding_callback: Callable | None = None,
+        overwrite: bool = False,
+        verbose: bool = False,
+    ) -> None:
+        """
+        Initialize a VinkDB instance.
+
+        Note:
+            The only editable attributes after initialization are:
+                - ann_config
+                - embedding_callback
+                - verbose
+            Everything else is read-only properties.
+
+        Args:
+            dir_path (str | Path): Directory path to store vector data. Contains the pickled index
+                and SQLite database for vector records.
+                Use ":memory:" for volatile in-memory storage.
+            dim (int): Dimension of the vectors. Must be higher than 16.
+            metric (Literal["euclidean", "cosine"], optional): Distance metric to use.
+                Defaults to "euclidean".
+            force_exact (bool, optional): If True, only exact calculation is used.
+                If False, switches between exact and ANN based on runtime calibration.
+                Defaults to False.
+            ann_config (AnnConfig, optional): Configuration for approximate nearest neighbor search.
+                Used during switching and compacting. Defaults to ANNConfig with standard settings.
+                Only applicable when force_exact is False.
+            embedding_callback (Callable, optional): Callback function to generate embeddings
+                from content. If provided, 'embedding' key is optional in
+                vector records as it will be generated via this callback. Defaults to None.
+            overwrite (bool, optional): Overwrite existing index if exists. Defaults to False.
+            verbose (bool, optional): Enable verbose output. Defaults to False.
+        """
+        # Determine if in-memory mode
+        self._in_memory = isinstance(dir_path, str) and dir_path.strip() == ":memory:"
+
+        self._dir_path = Path(dir_path)
+        self._dim = dim
+        self._metric = metric
+        self._ann_config = ann_config
+        self._force_exact = force_exact
+        self.embedding_callback = embedding_callback
+        self.verbose = verbose
+
+        self._strategy: BaseStrategy | None = None
+        self._latency_predictor: "LatencyPredictor" | None = None
+
+        # Default the config with standard settings if force_exact is not true
+        if not (self._force_exact or self._ann_config):
+            self._ann_config = AnnConfig()
+
+        self._validate_config()
+
+        if not self._in_memory:
+            if overwrite and self.dir_path.exists():
+                shutil.rmtree(self.dir_path)
+
+            self.dir_path.mkdir(parents=True, exist_ok=True)
+            self._records_db_path = str(self.dir_path / "records.sqlite")
+        else:
+            self._records_db_path = ":memory:"
+
+        self._records_db = SQLiteWrapper(
+            self._records_db_path,
+            index_config={
+              "dimension": str(self.dim),
+              "metric": self.metric,
+              "strategy": "exact",
+            }
+        )
+
+        # Threading components for ANN auto-switch
+        self._ann_building = False
+        self._rwlock = rwlock.RWLockFair()
+
+        self.load()
+
+    @property
+    def dir_path(self) -> Path:
+        return self._dir_path
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    @property
+    def metric(self) -> str:
+        return self._metric
+
+    @property
+    def force_exact(self) -> bool:
+        return self._force_exact
+
+    @property
+    def in_memory(self) -> bool:
+        return self._in_memory
+
+    @property
+    def strategy(self) -> str:
+        """The internal indexing strategy currently active, formatted in snake_case."""
+        if self._strategy is None:
+            return "exact_search"
+
+        strategy_name = self._strategy.__class__.__name__
+        parts = re.split(r"(?<!^)(?=\p{LU})", strategy_name)
+        return "_".join([p.lower() for p in parts])
+
+    def count(self, status: Literal["active", "deleted"] | None = None) -> int:
+        """Count vectors in the database.
+
+        Args:
+            status (Literal["active", "deleted"], optional): Which vectors to count.
+                Count all if not provided.
+
+        Returns:
+            int: Count of vectors.
+        """
+        return self._records_db.count(status)
+
+    def stats(self) -> dict:  # pragma: no cover
+        """Return database statistics and metadata.
+
+        Returns:
+            dict: Database metadata including version, dimension, metric, strategy,
+                last_saved_at, last_deleted_at, active_count, deleted_count,
+                and other stored metadata.
+        """
+        return {
+            "version": self._records_db["version"],
+            "dimension": self._records_db["dimension"],
+            "metric": self._records_db["metric"],
+            "strategy": self._records_db["strategy"],
+            "last_saved_at": self._records_db["last_saved_at"],
+            "last_deleted_at": self._records_db["last_deleted_at"],
+            "active_count": self.count("active"),
+            "deleted_count": self.count("deleted"),
+        }
+
+    def _validate_config(self) -> None:
+        """
+        Internal handshake to verify embedding dimensions and PQ constraints.
+        """
+        # Callback Handshake validation
+        if self.embedding_callback:
+            try:
+                raw_vec = self.embedding_callback("vink_warmup_test")
+
+                # This handles casting, shape normalization (1, d), and L2 projection
+                validated_vec = validate_embedding(raw_vec, dim=self.dim, metric=self.metric)
+
+                if validated_vec.shape[-1] != self._dim:
+                    raise VectorDimensionError(
+                        f"Embedding callback output dimension ({validated_vec.shape[-1]}) "
+                        f"does not match VinkDB dimension ({self._dim})."
+                    )
+            except (VectorDimensionError, InvalidInputError):
+                # Let these specific errors bubble up for the test/user
+                raise
+            except Exception as e:
+                raise InvalidInputError("Embedding callback crashed during handshake") from e
+
+        if not self._force_exact:
+            self._ann_config.validate_vector_dim(self._dim)
+
+    @validate_arguments
+    def add(self, vector_records: list[dict]) -> list[str]:
+        """Add vectors to the index.
+
+        Args:
+            vector_records (list[dict]): List of dicts with 'content', 'metadata',
+                and 'embedding' keys. 'id' is optional
+                If not provided, a UUIDv7 will be auto-generated.
+
+        Note:
+            The first batch (when database is empty) is limited to 10,000 vectors to avoid
+            expensive initial index operations. This constraint only applies to the first add()
+            call. Subsequent batches can be any size.
+
+        Returns:
+            list[str]: List of assigned UUIDv7 IDs.
+
+        Raises:
+            InvalidInputError: If validation fails or if the first batch exceeds 10,000 vectors.
+        """
+        if not vector_records:
+            log_info(self.verbose, "Input is empty, returning empty list.")
+            return []
+
+        try:
+            records = VectorRecords(
+                dim=self.dim,
+                metric=self._metric,
+                records=vector_records,
+                embedding_callback=self.embedding_callback,
+            )
+        except ValidationError as e:
+            raise InvalidInputError(f"Invalid vector records: {pretty_errors(e)}") from None
+
+        log_info(self.verbose, "Adding {} vector records to index.", len(vector_records))
+
+        if self._ann_building:
+            assigned_ids = [r.id for r in records.records]
+            self._records_db.insert(records, is_buffer=True)
+            log_info(
+                self.verbose, "Successfully added {} records to buffer.", len(assigned_ids),
+            )
+            return assigned_ids
+
+        assigned_ids = self._strategy.add(records)
+
+        if self.strategy == "exact_search":
+            # Check if switch should be triggered based on new count
+            if self._should_switch():
+                self._ann_building = True
+                Thread(target=self._prepare_approx_strategy, daemon=True).start()
+
+        log_info(
+            self.verbose, "Successfully added {} records to index.", len(assigned_ids)
+        )
+        return assigned_ids
+
+    @validate_arguments
+    def soft_delete(self, ids: list[str]) -> None:
+        """Soft-delete vectors from the index by their IDs (marks as deleted).
+
+        Args:
+            ids (list[str]): List of UUIDv7 IDs to soft-delete.
+        """
+        log_info(self.verbose, "Soft-deleting {} vectors from index.", len(ids))
+
+        id_bytes = [validate_id(id_str) for id_str in ids]
+
+        # If ANN is building, write to buffer for replay after switch
+        if self.strategy != "approximate_search" and self._ann_building:
+            self._records_db.soft_delete(id_bytes)
+            self._records_db["last_deleted_at"] = datetime.now(timezone.utc).isoformat()
+            log_info(
+                self.verbose, "Marked {} vectors for soft-deletion in buffer.", len(ids)
+            )
+            return
+
+        self._strategy.soft_delete(id_bytes)
+        self._records_db["last_deleted_at"] = datetime.now(timezone.utc).isoformat()
+
+    def compact(self) -> None:
+        """Hard-delete soft-deleted records and rebuild the index.
+
+        Note:
+            For ApproximateSearch, the ANN index is rebuilt from scratch which can take
+            20-200+ seconds depending on data size. This operation should be called
+            during maintenance windows or off-peak hours.
+            If not enough vectors remain to retrain the codec, rebuild is skipped.
+        """
+        log_info(self.verbose, "Compacting database...")
+        self._strategy.compact()
+        log_info(self.verbose, "Compaction complete.")
+
+    def save(self) -> None:
+        """Save the index to disk."""
+        log_info(self.verbose, "Saving index to {}.", self._dir_path)
+        self._strategy.save()
+        self._records_db["last_saved_at"] = datetime.now(timezone.utc).isoformat()
+        log_info(self.verbose, "Index saved successfully.")
+
+    def load(self, overwrite: bool = False) -> None:
+        """Load the index from disk.
+
+        Args:
+            overwrite (bool): If True, replace in-memory state with loaded data.
+                Defaults to False.
+        """
+        log_info(self.verbose, "Loading index from {}.", self._dir_path)
+
+        if self._strategy is None:
+            params = {
+                "db": self._records_db,
+                "dir_path": self._dir_path if not self._in_memory else None,
+                "dim": self._dim,
+                "in_memory": self._in_memory,
+                "metric": self.metric,
+                "verbose": self.verbose,
+            }
+            if self.strategy == "exact_search":
+                from vink.strategies.exact_search import ExactSearch
+                strategy_class = ExactSearch
+            else:
+                from vink.strategies.approximate_search import ApproximateSearch
+                strategy_class = ApproximateSearch
+                params["ann_config"] = self._ann_config
+
+            self._strategy = strategy_class(**params)
+
+        self._strategy.load(overwrite=overwrite)
+
+        # Lazy init predictor only if strategy is exact
+        if self.strategy == "exact_search":
+            from vink.latency_predictor import LatencyPredictor
+            self._latency_predictor = LatencyPredictor(dim=self._dim)
+
+        log_info(self.verbose, "Index loaded successfully.")
+
+    @validate_arguments
+    def search(
+        self,
+        query_vec: list[float] | np.ndarray,
+        top_k: int = 10,
+        include_vectors: bool = False,
+        filters: list[str] | None = None,
+    ) -> list[dict]:
+        """Search for k nearest neighbors using the configured metric.
+
+        Args:
+            query_vec (list[float] | np.ndarray): The query vector as a list of floats,
+                1D numpy array (d,), or 2D numpy array (1, d).
+            top_k (int, optional): Number of nearest neighbors to return. Defaults to 10.
+            include_vectors (bool, optional): If True, include 'embedding' key in results.
+                Defaults to False.
+            filters (list[str] | None, optional): Filter expressions to apply before scoring.
+                E.g., ["category == 'science'", "price >= 10"].
+
+        Returns:
+            list[dict]: List of dicts with 'id', 'content', 'metadata', 'distance',
+                and optionally 'embedding' (if include_vectors is True).
+        """
+        log_info(
+            self.verbose,
+            "Searching for {} nearest neighbors using {}.",
+            top_k,
+            self.strategy,
+        )
+
+        start = time.perf_counter()
+
+        validated_query = validate_embedding(query_vec, dim=self.dim, metric=self.metric,)
+        results = self._strategy.search(
+            validated_query, top_k=top_k, include_vectors=include_vectors, filters=filters
+        )
+
+        elapsed_ms = (time.perf_counter() - start) * 1000
+
+        # Tune predictor with actual latency
+        if self._latency_predictor is not None:
+            self._latency_predictor.tune(self.count(), elapsed_ms)
+
+        log_info(
+            self.verbose,
+            "Found {} results for query in {} ms.",
+            len(results),
+            round(elapsed_ms, 2)
+        )
+        return results
+
+    def _should_switch(self) -> bool:
+        """
+        Check if ANN switch should be triggered based on dual conditions:
+
+        1. Sufficiency: num_vectors >= min_required (num_subspaces * codebook_size)
+        2. latency: predicted latency exceeds this threshold
+        """
+        if self._force_exact:
+            return False
+
+        n_vecs = self.count()
+        cfg = self._ann_config
+        min_required = cfg.num_subspaces * cfg.codebook_size
+
+        if n_vecs < min_required:
+            return False
+
+        # Use predictor for proactive switching
+        predicted_latency = self._latency_predictor.predict(n_vecs)
+        return predicted_latency >= cfg.switch_latency_ms
+
+    def _prepare_approx_strategy(self) -> None:
+        """
+        Build ANN strategy in a background daemon thread.
+
+        Runs in a daemon thread so add()/search() remain unblocked.
+        Replays buffered records after the strategy switch completes.
+        """
+        self._strategy._ensure_cache()
+        vectors = self._strategy.active_vectors_arr
+        ids = self._strategy.active_ids_arr
+
+        from vink.strategies.approximate_search import ApproximateSearch
+
+        approx_strategy = ApproximateSearch(
+            db=self._records_db,
+            dir_path=self._dir_path if not self._in_memory else None,
+            dim=self._dim,
+            in_memory=self._in_memory,
+            metric=self.metric,
+            verbose=self.verbose,
+            ann_config=self._ann_config,
+        )
+        approx_strategy.fit(vectors, ids)
+
+        log_info(self.verbose, "ANN index fit complete, switching strategy.")
+
+        # Automatically switch after successful build
+        self._switch_to_approx_strategy(approx_strategy)
+
+    def _switch_to_approx_strategy(self, strategy) -> None:
+        """Switch to approximate search and auto dumps buffer."""
+        with self._rwlock.gen_wlock():
+            self._strategy = strategy
+
+        self._ann_building = False
+
+        cursor = self._records_db.conn.cursor()
+        buffer_rows = cursor.execute("""
+            SELECT id, embedding FROM vec_records
+            WHERE buffered = TRUE AND deleted = FALSE
+        """).fetchall()
+
+        if not buffer_rows:
+            return
+
+        records = [
+            {
+                "id": row[0],
+                "embedding": np.frombuffer(row[1], dtype=np.float32),
+
+                # Not used, kept for validation
+                "content": "",
+                "metadata": {},
+            }
+            for row in buffer_rows
+        ]
+
+        strategy.add(VectorRecords(dim=self.dim, metric=self.metric, records=records), is_buffer=True)
+        self._records_db.clear_buffer()
+        log_info(
+            self.verbose,
+            "Buffer dump: added {} vectors to ANN index.",
+            len(records),
+        )
+
+    def __enter__(self) -> "VinkDB":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if exc_type:
+            logger.error(f"Transaction failed: {exc_val}")
+            return False # Tell python to reraise it
+        self.save()

vink/exceptions.py

@@ -0,0 +1,40 @@
+class VinkDBError(Exception):
+    """Base exception for all Vink errors."""
+
+    pass
+
+
+class InvalidInputError(VinkDBError):
+    """Raised when one or multiple invalid input(s) are encountered."""
+
+    pass
+
+
+class VectorDimensionError(InvalidInputError):
+    """Raised when vector dimensions don't match the index configuration."""
+
+    pass
+
+
+class InvalidIdError(InvalidInputError):
+    """Raised when an invalid UUIDv7 is provided."""
+
+    pass
+
+
+class IndexNotFittedError(Exception):
+    """Raised when an operation requiring learned quantization is called on an unitialized index."""
+
+    pass
+
+
+class FilterError(InvalidInputError):
+    """Raised when a filter expression fails to parse."""
+
+    pass
+
+
+class DatabaseCorruptedError(VinkDBError):
+    """Raised when database files are corrupted."""
+
+    pass