fromscratch-pyvecdb 0.3.0__py3-none-any.whl

pyvecdb/persistence.py

@@ -0,0 +1,149 @@
+"""Save and load vector stores to/from JSON files.
+
+Think of persistence like saving your game -- you can close the program,
+come back later, and pick up right where you left off. Without persistence,
+all your vectors would vanish when the program stops!
+
+The JSON format looks like this:
+
+.. code-block:: json
+
+    {
+        "dimensions": 3,
+        "records": [
+            {
+                "record_id": "song-1",
+                "vector": [0.8, 0.2, 0.5],
+                "metadata": {"genre": "rock"}
+            }
+        ]
+    }
+
+Example:
+    >>> from pyvecdb.store import VectorStore
+    >>> from pyvecdb.persistence import save_store, load_store
+    >>> store = VectorStore(dimensions=2)
+    >>> store.insert("a", [1.0, 2.0])
+    >>> save_store(store, "/tmp/my_vectors.json")
+    >>> loaded = load_store("/tmp/my_vectors.json")
+    >>> loaded.count
+    1
+
+"""
+
+import json
+from pathlib import Path
+from typing import cast
+
+from pyvecdb.atomic import atomic_write_text
+from pyvecdb.errors import PersistenceError, VecDBError
+from pyvecdb.record import VectorRecord
+from pyvecdb.store import VectorStore
+
+# Every real save file has these top-level keys; anything else isn't one of ours.
+_REQUIRED_KEYS = frozenset({"dimensions", "records"})
+
+
+def store_to_dict(
+    store: VectorStore,
+) -> dict[str, int | list[dict[str, str | list[float] | dict[str, str | int | float | bool]]]]:
+    """Convert a VectorStore to a plain dictionary.
+
+    Args:
+        store: The vector store to serialize.
+
+    Returns:
+        A dictionary with 'dimensions' and 'records' keys.
+
+    """
+    return {
+        "dimensions": store.dimensions,
+        "records": [r.to_dict() for r in store.list_records()],
+    }
+
+
+def dict_to_store(data: dict[str, object]) -> VectorStore:
+    """Create a VectorStore from a dictionary.
+
+    Args:
+        data: A dictionary with 'dimensions' and 'records' keys.
+
+    Returns:
+        A new VectorStore populated with the records.
+
+    """
+    dimensions = int(data["dimensions"])  # type: ignore[arg-type]
+    store = VectorStore(dimensions=dimensions)
+    raw_records = data.get("records", [])
+    if not isinstance(raw_records, list):
+        msg = "Records must be a list"
+        raise TypeError(msg)
+    for raw in raw_records:  # pyright: ignore[reportUnknownVariableType]
+        if not isinstance(raw, dict):
+            msg = "Each record must be a dict"
+            raise TypeError(msg)
+        rec_data: dict[str, object] = {str(k): v for k, v in raw.items()}  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
+        record = VectorRecord.from_dict(rec_data)
+        store.insert(record.record_id, record.vector, record.metadata)
+    return store
+
+
+def save_store(store: VectorStore, path: str | Path) -> None:
+    """Save a vector store to a JSON file.
+
+    Args:
+        store: The vector store to save.
+        path: The file path to write to.
+
+    """
+    file_path = Path(path)
+    data = store_to_dict(store)
+    atomic_write_text(file_path, json.dumps(data, indent=2))
+
+
+def load_store(path: str | Path) -> VectorStore:
+    """Load a vector store from a JSON file.
+
+    Args:
+        path: The file path to read from.
+
+    Returns:
+        A new VectorStore populated with the saved records.
+
+    Raises:
+        PersistenceError: If the file is missing, unreadable, not valid JSON,
+            or not a valid save file.
+
+    """
+    file_path = Path(path)
+    try:
+        raw = file_path.read_text(encoding="utf-8")
+    # A file that isn't UTF-8 text (an image, a zip, a corrupted save) fails
+    # while decoding, not while reading, so it needs catching too.
+    except (OSError, UnicodeDecodeError) as exc:
+        msg = f"Could not read {file_path}: {exc}"
+        raise PersistenceError(msg) from exc
+    try:
+        data: object = json.loads(raw)
+    # JSONDecodeError is a kind of ValueError; a plain ValueError also means a
+    # number too long for Python to read, and RecursionError means lists or
+    # objects nested too deep. Each one is a file we can't load.
+    except (ValueError, RecursionError) as exc:
+        msg = f"Invalid JSON in {file_path}: {exc}"
+        raise PersistenceError(msg) from exc
+    if not isinstance(data, dict):
+        msg = f"{file_path} is not a valid save file: the top level must be an object"
+        raise PersistenceError(msg)
+    saved: dict[str, object] = {str(k): v for k, v in cast("dict[object, object]", data).items()}
+    missing = _REQUIRED_KEYS - saved.keys()
+    if missing:
+        msg = f"{file_path} is not a valid save file: missing {sorted(missing)}"
+        raise PersistenceError(msg)
+    try:
+        return dict_to_store(saved)
+    # OverflowError: a number too big to be a float, or Infinity dimensions.
+    # VecDBError: a record the store itself refuses -- an empty vector, the
+    # wrong number of dimensions, or the same ID twice.
+    except (KeyError, TypeError, ValueError, AttributeError, OverflowError, VecDBError) as exc:
+        msg = f"{file_path} is not a valid save file: {exc}"
+        raise PersistenceError(msg) from exc

pyvecdb/record.py

@@ -0,0 +1,94 @@
+"""Vector record -- a vector with an ID and extra info.
+
+Think of a ``VectorRecord`` like a trading card:
+
+- The **ID** is the card number (unique identifier)
+- The **vector** is the card's stats (a list of numbers)
+- The **metadata** is the extra info printed on the card (name, team, year)
+
+When you search a vector database, you get back records -- so you know
+*which* item matched, not just *where* it is in space.
+
+Example:
+    >>> from pyvecdb.record import VectorRecord
+    >>> card = VectorRecord(
+    ...     record_id="card-42",
+    ...     vector=[0.8, 0.2, 0.5],
+    ...     metadata={"name": "Pikachu", "type": "electric"},
+    ... )
+    >>> card.record_id
+    'card-42'
+    >>> card.dimensions
+    3
+
+"""
+
+from dataclasses import dataclass, field
+
+from pyvecdb.errors import EmptyVectorError
+
+
+@dataclass
+class VectorRecord:
+    """A vector paired with an ID and optional metadata.
+
+    Attributes:
+        record_id: A unique string identifier for this record.
+        vector: The list of floats representing this record's position in space.
+        metadata: Optional key-value pairs with extra info about this record.
+
+    """
+
+    record_id: str
+    vector: list[float]
+    metadata: dict[str, str | int | float | bool] = field(default_factory=lambda: {})
+
+    def __post_init__(self) -> None:
+        """Validate the record after creation."""
+        if not self.vector:
+            msg = "Vector must have at least one component"
+            raise EmptyVectorError(msg)
+
+    @property
+    def dimensions(self) -> int:
+        """Return the number of dimensions in this record's vector."""
+        return len(self.vector)
+
+    def to_dict(self) -> dict[str, str | list[float] | dict[str, str | int | float | bool]]:
+        """Convert this record to a plain dictionary for serialization.
+
+        Returns:
+            A dictionary with 'record_id', 'vector', and 'metadata' keys.
+
+        """
+        return {
+            "record_id": self.record_id,
+            "vector": list(self.vector),
+            "metadata": dict(self.metadata),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, object]) -> VectorRecord:
+        """Create a VectorRecord from a dictionary.
+
+        Args:
+            data: A dictionary with 'record_id', 'vector', and optionally 'metadata'.
+
+        Returns:
+            A new VectorRecord instance.
+
+        """
+        record_id = str(data["record_id"])
+        raw_vector = data["vector"]
+        if not isinstance(raw_vector, list):
+            msg = "Vector must be a list"
+            raise TypeError(msg)
+        vector = [float(x) for x in raw_vector]  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
+        raw_metadata = data.get("metadata")
+        if not isinstance(raw_metadata, dict):
+            raw_metadata = {}
+        metadata: dict[str, str | int | float | bool] = {}
+        for k, v in raw_metadata.items():  # pyright: ignore[reportUnknownVariableType]
+            if isinstance(v, (str, int, float, bool)):
+                metadata[str(k)] = v  # pyright: ignore[reportUnknownArgumentType]
+        return cls(record_id=record_id, vector=vector, metadata=metadata)

pyvecdb/search.py

@@ -0,0 +1,165 @@
+"""K-Nearest Neighbors (KNN) search for finding similar vectors.
+
+Imagine you're standing in a huge room full of people, and you want to find
+the 5 people standing closest to you. That's KNN search! Instead of people,
+we have vectors (lists of numbers), and "close" means "similar."
+
+This module provides the core search function that:
+
+1. Computes the distance from your query vector to every vector in the store
+2. Sorts by distance (closest first)
+3. Returns the top K results
+
+Example:
+    >>> from pyvecdb.store import VectorStore
+    >>> from pyvecdb.search import knn_search
+    >>> store = VectorStore(dimensions=2)
+    >>> store.insert("a", [1.0, 0.0])
+    >>> store.insert("b", [0.0, 1.0])
+    >>> store.insert("c", [0.9, 0.1])
+    >>> results = knn_search(store, query=[1.0, 0.0], k=2)
+    >>> [r.record_id for r, _ in results]
+    ['a', 'c']
+
+"""
+
+from pyvecdb.distance import (
+    cosine_similarity,
+    dot_product,
+    euclidean_distance,
+    manhattan_distance,
+)
+from pyvecdb.errors import DimensionMismatchError, EmptyStoreError
+from pyvecdb.filtering import MetadataFilter, filter_records
+from pyvecdb.record import VectorRecord
+from pyvecdb.store import VectorStore
+
+# Supported metric names
+METRIC_COSINE = "cosine"
+METRIC_EUCLIDEAN = "euclidean"
+METRIC_DOT = "dot"
+METRIC_MANHATTAN = "manhattan"
+
+VALID_METRICS = {METRIC_COSINE, METRIC_EUCLIDEAN, METRIC_DOT, METRIC_MANHATTAN}
+
+
+def _compute_score(query: list[float], candidate: list[float], metric: str) -> float:
+    """Compute a similarity/distance score between two vectors.
+
+    For cosine and dot product, higher is better (more similar).
+    For euclidean and manhattan, lower is better (closer together).
+
+    We negate distance metrics so that sorting by descending score
+    always puts the best matches first.
+
+    Args:
+        query: The query vector.
+        candidate: The candidate vector to compare against.
+        metric: One of 'cosine', 'euclidean', 'dot', or 'manhattan'.
+
+    Returns:
+        A score where higher means more similar.
+
+    """
+    if metric == METRIC_COSINE:
+        return cosine_similarity(query, candidate)
+    if metric == METRIC_EUCLIDEAN:
+        return -euclidean_distance(query, candidate)
+    if metric == METRIC_MANHATTAN:
+        return -manhattan_distance(query, candidate)
+    # dot product -- delegate so the same dimension check applies as the other metrics
+    return dot_product(query, candidate)
+
+
+def knn_search(
+    store: VectorStore,
+    query: list[float],
+    k: int = 5,
+    metric: str = METRIC_COSINE,
+    metadata_filter: MetadataFilter | None = None,
+) -> list[tuple[VectorRecord, float]]:
+    """Find the K most similar vectors to a query vector.
+
+    This is the heart of a vector database! It checks every vector in the
+    store, computes how similar it is to your query, and returns the top K.
+
+    Args:
+        store: The vector store to search.
+        query: The query vector to find neighbors for.
+        k: How many results to return (default 5).
+        metric: Distance metric -- 'cosine', 'euclidean', 'dot', or 'manhattan'.
+        metadata_filter: Optional filter to narrow results by metadata.
+
+    Returns:
+        A list of (record, score) tuples, sorted by similarity (best first).
+
+    Raises:
+        EmptyStoreError: If the store has no records.
+        DimensionMismatchError: If the query has wrong dimensions.
+        ValueError: If the metric is not recognized or k < 1.
+
+    """
+    if store.count == 0:
+        msg = "Cannot search an empty store"
+        raise EmptyStoreError(msg)
+
+    if len(query) != store.dimensions:
+        msg = f"Query has {len(query)} dimensions, store expects {store.dimensions}"
+        raise DimensionMismatchError(msg)
+
+    if metric not in VALID_METRICS:
+        msg = f"Unknown metric '{metric}'. Choose from: {', '.join(sorted(VALID_METRICS))}"
+        raise ValueError(msg)
+
+    min_k = 1
+    if k < min_k:
+        msg = f"k must be at least {min_k}, got {k}"
+        raise ValueError(msg)
+
+    # Get all records, optionally filtered
+    candidates = store.list_records()
+    if metadata_filter is not None:
+        candidates = filter_records(candidates, metadata_filter)
+
+    if not candidates:
+        return []
+
+    # Score each candidate
+    scored: list[tuple[VectorRecord, float]] = []
+    for record in candidates:
+        score = _compute_score(query, record.vector, metric)
+        scored.append((record, score))
+
+    # Sort by score descending (higher = better match)
+    scored.sort(key=lambda pair: pair[1], reverse=True)
+
+    return scored[:k]
+
+
+def batch_search(
+    store: VectorStore,
+    queries: list[list[float]],
+    k: int = 5,
+    metric: str = METRIC_COSINE,
+    metadata_filter: MetadataFilter | None = None,
+) -> list[list[tuple[VectorRecord, float]]]:
+    """Search for multiple query vectors at once.
+
+    Like asking "find the 5 closest pizza shops" for each of your friends
+    at the same time!
+
+    Args:
+        store: The vector store to search.
+        queries: A list of query vectors.
+        k: How many results per query (default 5).
+        metric: Distance metric to use.
+        metadata_filter: Optional metadata filter applied to all queries.
+
+    Returns:
+        A list of result lists, one per query.
+
+    """
+    return [
+        knn_search(store, query=q, k=k, metric=metric, metadata_filter=metadata_filter)
+        for q in queries
+    ]

pyvecdb/store.py

@@ -0,0 +1,250 @@
+"""Vector store -- a container for vector records.
+
+Think of a ``VectorStore`` like a filing cabinet where every folder
+(record) contains a list of numbers (vector) and a label (ID).
+The cabinet has one rule: every folder must contain the same number
+of pages (dimensions).
+
+The store provides CRUD operations:
+
+- **Create**: Insert a new vector record
+- **Read**: Get a record by its ID
+- **Update**: Replace a record's vector or metadata
+- **Delete**: Remove a record by its ID
+- **List**: Get all records in the store
+
+Example:
+    >>> from pyvecdb.store import VectorStore
+    >>> store = VectorStore(dimensions=3)
+    >>> store.insert("song-1", [0.8, 0.2, 0.5], metadata={"genre": "rock"})
+    >>> store.get("song-1").metadata["genre"]
+    'rock'
+    >>> store.count
+    1
+
+"""
+
+from pyvecdb.errors import (
+    DimensionMismatchError,
+    DuplicateIDError,
+    RecordNotFoundError,
+)
+from pyvecdb.record import VectorRecord
+
+
+def _copy(record: VectorRecord) -> VectorRecord:
+    """Return a copy of *record* with its own vector and metadata.
+
+    The store only ever hands out copies, like a library that lends you a
+    photocopy: scribble on it all you like, the original stays clean.
+    """
+    return VectorRecord(
+        record_id=record.record_id,
+        vector=list(record.vector),
+        metadata=dict(record.metadata),
+    )
+
+
+class VectorStore:
+    """A container that stores vector records with the same dimensionality.
+
+    All vectors in a store must have the same number of dimensions.
+    Think of it like a spreadsheet where every row has the same number of columns.
+
+    Attributes:
+        dimensions: The number of dimensions each vector must have.
+
+    """
+
+    def __init__(self, dimensions: int) -> None:
+        """Create a new empty vector store.
+
+        Args:
+            dimensions: The number of dimensions for vectors in this store.
+
+        Raises:
+            ValueError: If dimensions is less than 1.
+
+        """
+        min_dims = 1
+        if dimensions < min_dims:
+            msg = f"Dimensions must be at least {min_dims}, got {dimensions}"
+            raise ValueError(msg)
+        self._dimensions = dimensions
+        self._records: dict[str, VectorRecord] = {}
+
+    @property
+    def dimensions(self) -> int:
+        """Return the dimensionality of this store."""
+        return self._dimensions
+
+    @property
+    def count(self) -> int:
+        """Return the number of records in this store."""
+        return len(self._records)
+
+    def _check_dimensions(self, vector: list[float]) -> None:
+        """Raise DimensionMismatchError if the vector has the wrong length."""
+        if len(vector) != self._dimensions:
+            msg = f"Expected {self._dimensions} dimensions, got {len(vector)}"
+            raise DimensionMismatchError(msg)
+
+    def insert(
+        self,
+        record_id: str,
+        vector: list[float],
+        metadata: dict[str, str | int | float | bool] | None = None,
+    ) -> VectorRecord:
+        """Insert a new vector record into the store.
+
+        Args:
+            record_id: A unique identifier for this record.
+            vector: The vector (list of floats) to store.
+            metadata: Optional key-value pairs of extra information.
+
+        Returns:
+            A copy of the newly created VectorRecord.
+
+        Raises:
+            DuplicateIDError: If a record with this ID already exists.
+            DimensionMismatchError: If the vector has the wrong number of dimensions.
+
+        """
+        if record_id in self._records:
+            msg = f"Record with ID '{record_id}' already exists"
+            raise DuplicateIDError(msg)
+        self._check_dimensions(vector)
+        record = VectorRecord(
+            record_id=record_id,
+            vector=list(vector),
+            metadata=dict(metadata) if metadata is not None else {},
+        )
+        self._records[record_id] = record
+        return _copy(record)
+
+    def get(self, record_id: str) -> VectorRecord:
+        """Get a record by its ID.
+
+        Args:
+            record_id: The ID of the record to retrieve.
+
+        Returns:
+            A copy of the matching VectorRecord. Changing the copy doesn't
+            change the store; use ``update`` for that.
+
+        Raises:
+            RecordNotFoundError: If no record with this ID exists.
+
+        """
+        if record_id not in self._records:
+            msg = f"No record found with ID '{record_id}'"
+            raise RecordNotFoundError(msg)
+        return _copy(self._records[record_id])
+
+    def update(
+        self,
+        record_id: str,
+        vector: list[float] | None = None,
+        metadata: dict[str, str | int | float | bool] | None = None,
+    ) -> VectorRecord:
+        """Update an existing record's vector and/or metadata.
+
+        Args:
+            record_id: The ID of the record to update.
+            vector: New vector (optional -- keeps old if not provided).
+            metadata: New metadata (optional -- keeps old if not provided).
+
+        Returns:
+            A copy of the updated VectorRecord.
+
+        Raises:
+            RecordNotFoundError: If no record with this ID exists.
+            DimensionMismatchError: If the new vector has wrong dimensions.
+
+        """
+        if record_id not in self._records:
+            msg = f"No record found with ID '{record_id}'"
+            raise RecordNotFoundError(msg)
+        record = self._records[record_id]
+        if vector is not None:
+            self._check_dimensions(vector)
+            record.vector = list(vector)
+        if metadata is not None:
+            record.metadata = dict(metadata)
+        return _copy(record)
+
+    def delete(self, record_id: str) -> None:
+        """Remove a record from the store.
+
+        Args:
+            record_id: The ID of the record to delete.
+
+        Raises:
+            RecordNotFoundError: If no record with this ID exists.
+
+        """
+        if record_id not in self._records:
+            msg = f"No record found with ID '{record_id}'"
+            raise RecordNotFoundError(msg)
+        del self._records[record_id]
+
+    def list_records(self) -> list[VectorRecord]:
+        """Return all records in the store as a list.
+
+        Returns:
+            Copies of all the VectorRecord objects, in insertion order.
+
+        """
+        return [_copy(record) for record in self._records.values()]
+
+    def contains(self, record_id: str) -> bool:
+        """Check whether a record with the given ID exists.
+
+        Args:
+            record_id: The ID to check for.
+
+        Returns:
+            True if the record exists, False otherwise.
+
+        """
+        return record_id in self._records
+
+    def clear(self) -> None:
+        """Remove all records from the store."""
+        self._records.clear()
+
+    def batch_insert(
+        self,
+        records: list[tuple[str, list[float], dict[str, str | int | float | bool] | None]],
+    ) -> list[VectorRecord]:
+        """Insert multiple records at once.
+
+        Args:
+            records: A list of (record_id, vector, metadata) tuples.
+
+        Returns:
+            Copies of the newly created VectorRecord objects.
+
+        This is all-or-nothing: if *any* record is invalid, nothing is inserted
+        and the store is left exactly as it was -- like a shopping cart that
+        refuses to check out unless every item scans correctly.
+
+        Raises:
+            DuplicateIDError: If any ID already exists or is duplicated in the batch.
+            DimensionMismatchError: If any vector has wrong dimensions.
+
+        """
+        # Validate the whole batch first so a bad record never leaves a half-done insert.
+        seen_ids: set[str] = set()
+        for record_id, vector, _metadata in records:
+            if record_id in self._records or record_id in seen_ids:
+                msg = f"Record with ID '{record_id}' already exists"
+                raise DuplicateIDError(msg)
+            seen_ids.add(record_id)
+            self._check_dimensions(vector)
+
+        # Everything checks out -- now it is safe to insert.
+        results: list[VectorRecord] = []
+        for record_id, vector, metadata in records:
+            results.append(self.insert(record_id, vector, metadata))
+        return results