From 6fe4276447229a183150a0bad4caaf2848938072 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Tue, 7 Apr 2026 02:46:32 -0500 Subject: [PATCH 1/5] =?UTF-8?q?feat:=20simplevecdb=202.5.0=20=E2=80=94=20c?= =?UTF-8?q?orrectness,=20performance,=20robustness,=20and=20API=20improvem?= =?UTF-8?q?ents?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness: fix delete_by_ids ordering (SQLite first), exact string filter matching, list_collections persistence via sqlite_master, WAL for encrypted DBs, collection cache key includes strategy/quantization. Performance: batch usearch add/remove/get, iterative deepening for filtered search, file-size-based mmap threshold (50MB). Robustness: connection health check, async_retry_on_lock decorator, FTS retry on transient lock errors, cross-process file locking on .usearch files. API: delete_collection(), store_embeddings param (default False with seamless MMR fallback), FLOAT16 quantization, __repr__ on all classes, pagination (limit/offset) on get_documents and catalog methods. Maintenance: remove duplicate _dim property, replace subprocess sysctl with platform.processor(). 514 tests passing (62 new). --- CHANGELOG.md | 38 ++ pyproject.toml | 2 +- src/simplevecdb/__init__.py | 10 +- src/simplevecdb/async_core.py | 24 +- src/simplevecdb/constants.py | 3 +- src/simplevecdb/core.py | 183 +++++++--- src/simplevecdb/engine/catalog.py | 112 ++++-- src/simplevecdb/engine/quantization.py | 6 + src/simplevecdb/engine/search.py | 95 +++-- src/simplevecdb/engine/usearch_index.py | 114 +++--- src/simplevecdb/utils.py | 141 ++++++- tests/integration/test_langchain.py | 2 +- tests/unit/core/test_batch_detection.py | 12 +- .../core/test_core_additional_coverage.py | 2 +- tests/unit/core/test_initialization.py | 4 +- tests/unit/core/test_missing_coverage.py | 2 +- tests/unit/core/test_v25_correctness.py | 265 ++++++++++++++ tests/unit/core/test_v25_features.py | 344 ++++++++++++++++++ tests/unit/core/test_v25_robustness.py | 312 ++++++++++++++++ tests/unit/test_core.py | 12 +- uv.lock | 2 +- 21 files changed, 1514 insertions(+), 171 deletions(-) create mode 100644 tests/unit/core/test_v25_correctness.py create mode 100644 tests/unit/core/test_v25_features.py create mode 100644 tests/unit/core/test_v25_robustness.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c9ee97..6a1cb75 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,44 @@ All notable changes to SimpleVecDB will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.5.0] - 2026-04-07 + +### Added + +- **`delete_collection(name)`** — drop a collection's SQLite tables, FTS index, and usearch file in one call. Available on both `VectorDB` and `AsyncVectorDB`. +- **`store_embeddings` parameter** on `collection()` — opt into storing embedding BLOBs in SQLite (default `False`). Saves ~2x storage; MMR transparently fetches vectors from the usearch index when BLOBs are absent. +- **`async_retry_on_lock` decorator** — async variant of `retry_on_lock` using `asyncio.sleep` instead of `time.sleep`, avoiding executor thread blocking. +- **`file_lock` context manager** — advisory cross-process file locking (`fcntl`/`msvcrt`) for usearch index files. Prevents corruption from concurrent processes. +- **`__repr__`** on `VectorDB`, `VectorCollection`, `AsyncVectorDB`, `AsyncVectorCollection` for debuggable string representations. +- **FLOAT16 quantization** fully implemented in `serialize()`/`deserialize()` — was previously defined in the enum but raised `ValueError` at runtime. +- **Pagination** on `get_documents(limit=, offset=)` and catalog methods (`find_ids_by_filter`, `find_ids_by_texts`) — previously returned unbounded result sets. + +### Fixed + +- **`delete_by_ids` ordering** — SQLite deletion now happens first (transactional, can rollback), then usearch. Previously usearch removed first, leaving orphaned catalog entries on SQLite failure. +- **`_matches_filter` string semantics** — now uses exact equality, consistent with SQL `build_filter_clause`. Was using substring match (`value in str(meta_value)`). +- **`list_collections`** — scans `sqlite_master` for persisted collection tables instead of returning only session-cached names. Works across reopened databases. +- **WAL mode for encrypted databases** — `PRAGMA journal_mode=WAL` and `PRAGMA synchronous=NORMAL` now set for SQLCipher connections (was only set for unencrypted). +- **`collection()` cache key** — includes `distance_strategy` and `quantization` in cache key (sync version). Previously cached by name only, silently ignoring differing params on cache hit. +- **`_ensure_fts_table`** — retries up to 3 times on transient "database is locked" errors instead of permanently disabling FTS on first failure. +- **Connection health check** — `SELECT 1` probe after connection creation; raises `RuntimeError` immediately on corrupt databases. + +### Improved + +- **Usearch batch operations** — `add()`, `remove()`, and `get()` now use batch usearch APIs instead of per-key loops. Significant speedup for large operations. +- **Filtered search iterative deepening** — replaces fixed `k*3` overfetch with adaptive doubling (up to `k*30`). Highly selective filters now reliably return `k` results. +- **Memory-map heuristic** — uses file size threshold (50MB) instead of inaccurate `file_size // 100` vector count estimate for mmap vs load decision. +- **Apple chip detection** — uses `platform.processor()` instead of spawning a `sysctl` subprocess. + +### Removed + +- **Duplicate `_dim` property** — removed in favor of the public `dim` property. + +### Breaking Changes + +- String metadata filters now use exact equality (was substring match). +- `store_embeddings` defaults to `False` — `rebuild_index()` requires `store_embeddings=True` or re-adding documents. + ## [2.4.0] - 2026-03-22 ### Added diff --git a/pyproject.toml b/pyproject.toml index 6c65c55..23b740b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "simplevecdb" -version = "2.4.0" +version = "2.5.0" description = "Dead-simple local vector database powered by usearch HNSW." authors = [{ name = "Dayton Dunbar", email = "coderdayton14@gmail.com" }] license = { text = "MIT" } diff --git a/src/simplevecdb/__init__.py b/src/simplevecdb/__init__.py index ad9e8f0..34a4e4f 100755 --- a/src/simplevecdb/__init__.py +++ b/src/simplevecdb/__init__.py @@ -16,7 +16,13 @@ except ImportError: pass from .logging import get_logger, configure_logging, log_operation -from .utils import DatabaseLockedError, retry_on_lock, validate_filter +from .utils import ( + DatabaseLockedError, + async_retry_on_lock, + file_lock, + retry_on_lock, + validate_filter, +) from .encryption import EncryptionError, EncryptionUnavailableError from importlib.metadata import version as _pkg_version @@ -49,6 +55,8 @@ "MigrationRequiredError", "EncryptionError", "EncryptionUnavailableError", + "async_retry_on_lock", + "file_lock", "retry_on_lock", "validate_filter", ] diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index 6507780..76f7ef7 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -60,6 +60,9 @@ def name(self) -> str: """Collection name.""" return self._collection.name + def __repr__(self) -> str: + return f"AsyncVectorCollection(name={self._collection.name!r})" + async def add_texts( self, texts: Sequence[str], @@ -210,15 +213,20 @@ async def delete_by_ids(self, ids: Sequence[int]) -> None: async def get_documents( self, filter_dict: dict[str, Any] | None = None, + *, + limit: int | None = None, + offset: int | None = None, ) -> list[tuple[int, str, dict[str, Any]]]: - """Get all documents with text content and metadata. + """Get documents with text content and metadata. See VectorCollection.get_documents for full documentation. """ loop = asyncio.get_running_loop() return await loop.run_in_executor( self._executor, - lambda: self._collection.get_documents(filter_dict=filter_dict), + lambda: self._collection.get_documents( + filter_dict=filter_dict, limit=limit, offset=offset + ), ) async def get_embeddings_by_ids(self, ids: Sequence[int]) -> dict[int, Any]: @@ -599,9 +607,16 @@ def collection( return self._collections[cache_key] def list_collections(self) -> list[str]: - """Return names of all initialized collections.""" + """Return names of all persisted collections in the database.""" return self._db.list_collections() + async def delete_collection(self, name: str) -> None: + """Delete a collection and all its data.""" + loop = asyncio.get_running_loop() + await loop.run_in_executor( + self._executor, lambda: self._db.delete_collection(name) + ) + async def search_collections( self, query: Sequence[float], @@ -644,6 +659,9 @@ async def vacuum(self, checkpoint_wal: bool = True) -> None: self._executor, lambda: self._db.vacuum(checkpoint_wal) ) + def __repr__(self) -> str: + return f"AsyncVectorDB(path={self._db.path!r})" + async def close(self) -> None: """Close the database connection and shutdown executor.""" try: diff --git a/src/simplevecdb/constants.py b/src/simplevecdb/constants.py index face55f..1fc6e6c 100755 --- a/src/simplevecdb/constants.py +++ b/src/simplevecdb/constants.py @@ -74,7 +74,8 @@ # - Instant startup (no full load into RAM) # - Lower memory footprint (OS manages page cache) # - Slight latency increase for cold pages (acceptable trade-off) -USEARCH_MMAP_THRESHOLD = 100000 +# Threshold in bytes — 50MB covers ~30k 384-dim f32 vectors. +USEARCH_MMAP_THRESHOLD = 50 * 1024 * 1024 # 50 MB # Batch search threshold: auto-batch queries when > this count # usearch batch search provides ~10x throughput for multi-query workloads diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 3a5db91..9ae2296 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -90,20 +90,13 @@ def get_optimal_batch_size() -> int: if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): machine = platform.machine().lower() if "arm" in machine or "aarch64" in machine: - try: - import subprocess - - chip_info = subprocess.check_output( - ["sysctl", "-n", "machdep.cpu.brand_string"], text=True - ).lower() - - if "m3" in chip_info or "m4" in chip_info: - return constants.DEFAULT_APPLE_M3_M4_BATCH_SIZE - elif "max" in chip_info or "ultra" in chip_info: - return constants.DEFAULT_APPLE_MAX_ULTRA_BATCH_SIZE - else: - return constants.DEFAULT_APPLE_M1_M2_BATCH_SIZE - except Exception: + chip_info = platform.processor().lower() + + if "m3" in chip_info or "m4" in chip_info: + return constants.DEFAULT_APPLE_M3_M4_BATCH_SIZE + elif "max" in chip_info or "ultra" in chip_info: + return constants.DEFAULT_APPLE_MAX_ULTRA_BATCH_SIZE + else: return constants.DEFAULT_APPLE_M1_M2_BATCH_SIZE # 2. Try ONNX Runtime detection @@ -178,6 +171,7 @@ def __init__( distance_strategy: DistanceStrategy, quantization: Quantization, encryption_key: str | bytes | None = None, + store_embeddings: bool = False, ): self.conn = conn self._db_path = db_path @@ -186,6 +180,7 @@ def __init__( self.quantization = quantization self._quantizer = QuantizationStrategy(quantization) self._encryption_key = encryption_key + self._store_embeddings = store_embeddings # Sanitize name to prevent issues if not re.match(constants.COLLECTION_NAME_PATTERN, name): @@ -397,12 +392,12 @@ def add_texts( batch_ids = ids[batch_start:batch_end] if ids else None batch_parent_ids = parent_ids[batch_start:batch_end] if parent_ids else None - # Add to SQLite metadata store (with embeddings for MMR support) + # Add to SQLite metadata store doc_ids = self._catalog.add_documents( batch_texts, list(batch_metas), batch_ids, - embeddings=batch_embeds, + embeddings=batch_embeds if self._store_embeddings else None, parent_ids=batch_parent_ids, ) @@ -748,12 +743,13 @@ def delete_by_ids(self, ids: Iterable[int]) -> None: if not ids_list: return - # Delete from usearch - self._index.remove(ids_list) - - # Delete from SQLite + # Delete from SQLite first (transactional, can rollback on failure) self._catalog.delete_by_ids(ids_list) + # Then remove from usearch (if this fails, catalog is clean and + # rebuild_index() can recover the index from stored data) + self._index.remove(ids_list) + def remove_texts( self, texts: Sequence[str] | None = None, @@ -846,6 +842,13 @@ def rebuild_index( # Fetch embeddings from SQLite embeddings_map = self._catalog.get_embeddings_by_ids(all_ids) + if not embeddings_map and not self._store_embeddings: + raise RuntimeError( + "Cannot rebuild index: no embeddings stored in SQLite. " + "Create the collection with store_embeddings=True to enable " + "rebuild_index(), or re-add documents with store_embeddings=True." + ) + # Filter to only docs with embeddings valid_pairs = [ (doc_id, emb) @@ -1260,7 +1263,7 @@ def load_cluster(self, name: str) -> tuple[ClusterResult, dict[str, Any]] | None centroids = None if centroids_bytes is not None: - dim = self._dim + dim = self.dim if dim: centroids = np.frombuffer(centroids_bytes, dtype=np.float32).reshape( n_clusters, dim @@ -1353,19 +1356,26 @@ def count(self) -> int: def get_documents( self, filter_dict: dict[str, Any] | None = None, + *, + limit: int | None = None, + offset: int | None = None, ) -> list[tuple[int, str, dict[str, Any]]]: - """Get all documents with text content and metadata. + """Get documents with text content and metadata. Args: filter_dict: Optional metadata filter to narrow results. + limit: Maximum number of documents to return (None = all). + offset: Number of documents to skip (None = 0). Returns: - List of (doc_id, text, metadata) tuples. + List of (doc_id, text, metadata) tuples, ordered by ID. """ filter_builder = self._catalog.build_filter_clause if filter_dict else None return self._catalog.get_all_docs_with_text( filter_dict=filter_dict, filter_builder=filter_builder, + limit=limit, + offset=offset, ) def get_embeddings_by_ids(self, ids: Sequence[int]) -> dict[int, Any]: @@ -1395,10 +1405,11 @@ def dim(self) -> int | None: """Vector dimension (None if no vectors added yet).""" return self._index.ndim - @property - def _dim(self) -> int | None: - """Vector dimension (None if no vectors added yet).""" - return self._index.ndim + def __repr__(self) -> str: + return ( + f"VectorCollection(name={self.name!r}, dim={self.dim}, " + f"size={self.count()}, distance={self.distance_strategy.value})" + ) class VectorDB: @@ -1452,7 +1463,7 @@ def __init__( self.quantization = quantization self.auto_migrate = auto_migrate self._encryption_key = encryption_key - self._collections: dict[str, VectorCollection] = {} + self._collections: dict[tuple, VectorCollection] = {} # Create connection (encrypted or plain) if encryption_key is not None: @@ -1467,6 +1478,8 @@ def __init__( check_same_thread=False, timeout=30.0, ) + self.conn.execute("PRAGMA journal_mode=WAL") + self.conn.execute("PRAGMA synchronous=NORMAL") self._encrypted = True _logger.info("Opened encrypted database: %s", self.path) else: @@ -1477,6 +1490,13 @@ def __init__( self.conn.execute("PRAGMA synchronous=NORMAL") self._encrypted = False + # Verify connection is healthy + try: + self.conn.execute("SELECT 1") + except sqlite3.DatabaseError as e: + self.conn.close() + raise RuntimeError(f"Database health check failed: {e}") from e + # Check for required migration before allowing collection access if not auto_migrate and self.path != ":memory:": migration_info = VectorDB.check_migration(self.path) @@ -1491,22 +1511,77 @@ def __init__( def list_collections(self) -> list[str]: """ - Return names of all initialized collections. + Return names of all persisted collections in the database. - Only returns collections that have been accessed via `collection()` in this - session. Does not scan the database for collections created in previous sessions. + Scans the database schema for collection tables, returning both + collections accessed this session and those created in previous sessions. Returns: - List of collection names currently cached in this VectorDB instance. + Sorted list of collection names stored in this database. Example: >>> db = VectorDB("app.db") >>> db.collection("users") - >>> db.collection("products") - >>> db.list_collections() - ['users', 'products'] + >>> db.close() + >>> db2 = VectorDB("app.db") + >>> db2.list_collections() + ['users'] + """ + rows = self.conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' " + "AND (name = 'tinyvec_items' OR name LIKE 'items_%')" + ).fetchall() + names: list[str] = [] + for (table_name,) in rows: + if table_name == "tinyvec_items": + names.append("default") + elif table_name.startswith("items_"): + # Skip FTS (and FTS sub-tables) and cluster tables + suffix = table_name[6:] + if "_fts" in suffix or suffix.endswith("_clusters"): + continue + names.append(suffix) + return sorted(names) + + def delete_collection(self, name: str) -> None: """ - return list(self._collections.keys()) + Delete a collection and all its data. + + Drops the SQLite tables (items, FTS, clusters) and deletes + the usearch index file. Removes the collection from the cache. + + Args: + name: Collection name to delete. + + Raises: + ValueError: If the collection name is invalid. + KeyError: If the collection does not exist. + """ + if name not in self.list_collections(): + raise KeyError(f"Collection '{name}' does not exist.") + + table_name = "tinyvec_items" if name == "default" else f"items_{name}" + fts_table = f"{table_name}_fts" + cluster_table = f"{table_name}_clusters" + + # Drop SQLite tables + self.conn.execute(f"DROP TABLE IF EXISTS {fts_table}") + self.conn.execute(f"DROP TABLE IF EXISTS {cluster_table}") + self.conn.execute(f"DROP TABLE IF EXISTS {table_name}") + self.conn.commit() + + # Delete usearch index file + if self.path != ":memory:": + index_path = Path(self.path + f".{name}.usearch") + if index_path.exists(): + index_path.unlink() + + # Remove from cache (match any tuple key with this name) + keys_to_remove = [k for k in self._collections if k[0] == name] + for k in keys_to_remove: + del self._collections[k] + + _logger.info("Deleted collection: %s", name) def search_collections( self, @@ -1562,14 +1637,28 @@ def search_collections( # Resolve and validate collections targets: list[VectorCollection] = [] dims: set[int | None] = set() + # Validate explicit collection names exist in DB + if collections is not None: + persisted = set(self.list_collections()) + for name in target_names: + if name not in persisted: + # Check cache too (collection may exist but not yet persisted) + if not any(k[0] == name for k in self._collections): + raise KeyError( + f"Collection '{name}' not initialized. " + f"Call db.collection('{name}') first." + ) + for name in target_names: - if name not in self._collections: - raise KeyError( - f"Collection '{name}' not initialized. Call db.collection('{name}') first." - ) - coll = self._collections[name] + # Find cached collection by name (may have any strategy/quantization) + matched = [v for k, v in self._collections.items() if k[0] == name] + if matched: + coll = matched[0] + else: + # Auto-initialize with defaults for persisted but uncached collections + coll = self.collection(name) targets.append(coll) - dims.add(coll._dim) + dims.add(coll.dim) # Check dimension consistency (ignore None for empty collections) dims.discard(None) @@ -1629,6 +1718,7 @@ def collection( name: str = "default", distance_strategy: DistanceStrategy | None = None, quantization: Quantization | None = None, + store_embeddings: bool = False, ) -> VectorCollection: """ Get or create a named collection. @@ -1640,6 +1730,9 @@ def collection( name: Collection name (alphanumeric + underscore only). distance_strategy: Override database-level distance metric. quantization: Override database-level quantization. + store_embeddings: If True, store embeddings as BLOBs in SQLite + alongside the usearch index. Required for rebuild_index(). + Default False to save ~2x storage. Returns: VectorCollection instance. @@ -1647,7 +1740,7 @@ def collection( Raises: ValueError: If collection name contains invalid characters. """ - cache_key = name + cache_key = (name, distance_strategy, quantization) if cache_key not in self._collections: self._collections[cache_key] = VectorCollection( conn=self.conn, @@ -1656,6 +1749,7 @@ def collection( distance_strategy=distance_strategy or self.distance_strategy, quantization=quantization or self.quantization, encryption_key=self._encryption_key, + store_embeddings=store_embeddings, ) return self._collections[cache_key] @@ -1844,6 +1938,9 @@ def save(self) -> None: for collection in self._collections.values(): collection.save() + def __repr__(self) -> str: + return f"VectorDB(path={self.path!r}, collections={self.list_collections()})" + def close(self) -> None: """Close the database connection and save indexes.""" if getattr(self, "_closed", False): diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 5788f13..b479174 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -136,20 +136,32 @@ def _ensure_parent_id_column(self) -> None: _logger.warning("Could not check/add parent_id column: %s", e) def _ensure_fts_table(self) -> None: - """Create FTS5 virtual table for full-text search.""" + """Create FTS5 virtual table for full-text search. + + Retries on transient lock errors but permanently disables FTS + if the module is unavailable. + """ import sqlite3 - try: - self.conn.execute( - f""" - CREATE VIRTUAL TABLE IF NOT EXISTS {self._fts_table_name} - USING fts5(text) - """ - ) - self._fts_enabled = True - except sqlite3.OperationalError: - _logger.warning("FTS5 not available - keyword search disabled") - self._fts_enabled = False + for attempt in range(3): + try: + self.conn.execute( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {self._fts_table_name} + USING fts5(text) + """ + ) + self._fts_enabled = True + return + except sqlite3.OperationalError as e: + msg = str(e).lower() + if "database is locked" in msg and attempt < 2: + import time + time.sleep(0.1 * (attempt + 1)) + continue + _logger.warning("FTS5 not available - keyword search disabled: %s", e) + self._fts_enabled = False + return @property def fts_enabled(self) -> bool: @@ -409,23 +421,52 @@ def get_documents_and_embeddings_by_ids( result[row_id] = (text, meta, emb) return result - def find_ids_by_texts(self, texts: Sequence[str]) -> list[int]: - """Find document IDs matching exact text content.""" + def find_ids_by_texts( + self, + texts: Sequence[str], + *, + limit: int | None = None, + offset: int | None = None, + ) -> list[int]: + """Find document IDs matching exact text content. + + Args: + texts: Text strings to search for + limit: Maximum number of IDs to return (None = all) + offset: Number of IDs to skip (None = 0) + """ if not texts: return [] placeholders = ",".join(["?"] * len(texts)) - rows = self.conn.execute( - f"SELECT id FROM {self._table_name} WHERE text IN ({placeholders})", - tuple(texts), - ).fetchall() + sql = f"SELECT id FROM {self._table_name} WHERE text IN ({placeholders})" + params: list[Any] = list(texts) + + if limit is not None: + sql += " LIMIT ?" + params.append(limit) + if offset is not None: + sql += " OFFSET ?" + params.append(offset) + + rows = self.conn.execute(sql, tuple(params)).fetchall() return [r[0] for r in rows] def find_ids_by_filter( self, filter_dict: dict[str, Any], filter_builder: Callable[[dict[str, Any], str], tuple[str, list[Any]]], + *, + limit: int | None = None, + offset: int | None = None, ) -> list[int]: - """Find document IDs matching metadata filter.""" + """Find document IDs matching metadata filter. + + Args: + filter_dict: Metadata key-value pairs to filter by + filter_builder: Function to build filter clause + limit: Maximum number of IDs to return (None = all) + offset: Number of IDs to skip (None = 0) + """ if not filter_dict: return [] @@ -434,10 +475,17 @@ def find_ids_by_filter( filter_clause = filter_clause.replace("AND ", "", 1) where_clause = f"WHERE {filter_clause}" if filter_clause else "" - rows = self.conn.execute( - f"SELECT id FROM {self._table_name} {where_clause}", - tuple(filter_params), - ).fetchall() + sql = f"SELECT id FROM {self._table_name} {where_clause}" + params: list[Any] = list(filter_params) + + if limit is not None: + sql += " LIMIT ?" + params.append(limit) + if offset is not None: + sql += " OFFSET ?" + params.append(offset) + + rows = self.conn.execute(sql, tuple(params)).fetchall() return [r[0] for r in rows] def keyword_search( @@ -536,13 +584,18 @@ def get_all_docs_with_text( filter_dict: dict[str, Any] | None = None, filter_builder: Callable[[dict[str, Any], str], tuple[str, list[Any]]] | None = None, + *, + limit: int | None = None, + offset: int | None = None, ) -> list[tuple[int, str, dict[str, Any]]]: """ - Get all documents with their text content. + Get documents with their text content, with optional pagination. Args: filter_dict: Optional metadata filter filter_builder: Function to build filter clause + limit: Maximum number of documents to return (None = all) + offset: Number of documents to skip (None = 0) Returns: List of (doc_id, text, metadata) tuples @@ -557,7 +610,16 @@ def get_all_docs_with_text( WHERE 1=1 {filter_clause} ORDER BY id """ - rows = self.conn.execute(sql, tuple(filter_params)).fetchall() + params: list[Any] = list(filter_params) + + if limit is not None: + sql += " LIMIT ?" + params.append(limit) + if offset is not None: + sql += " OFFSET ?" + params.append(offset) + + rows = self.conn.execute(sql, tuple(params)).fetchall() result = [] for row_id, text, meta_json in rows: meta = json.loads(meta_json) if meta_json else {} diff --git a/src/simplevecdb/engine/quantization.py b/src/simplevecdb/engine/quantization.py index 0e31ca9..8a71c65 100755 --- a/src/simplevecdb/engine/quantization.py +++ b/src/simplevecdb/engine/quantization.py @@ -52,6 +52,9 @@ def serialize(self, vector: np.ndarray) -> bytes: scaled = np.clip(np.round(vector * 127), -128, 127).astype(np.int8) return scaled.tobytes() + elif self.quantization == Quantization.FLOAT16: + return np.asarray(vector, dtype=np.float16).tobytes() + elif self.quantization == Quantization.BIT: # Binary quantization: threshold at 0 → pack bits bits = (vector > 0).astype(np.uint8) @@ -80,6 +83,9 @@ def deserialize(self, blob: bytes, dim: int | None) -> np.ndarray: elif self.quantization == Quantization.INT8: return np.frombuffer(blob, dtype=np.int8).astype(np.float32) / 127.0 + elif self.quantization == Quantization.FLOAT16: + return np.frombuffer(blob, dtype=np.float16).astype(np.float32) + elif self.quantization == Quantization.BIT and dim is not None: unpacked = np.unpackbits(np.frombuffer(blob, dtype=np.uint8)) v = unpacked[:dim].astype(np.float32) diff --git a/src/simplevecdb/engine/search.py b/src/simplevecdb/engine/search.py index e37d09d..bd24b74 100755 --- a/src/simplevecdb/engine/search.py +++ b/src/simplevecdb/engine/search.py @@ -74,40 +74,68 @@ def similarity_search( query_vec = self._resolve_query_vector(query) - # Over-fetch for filtering - fetch_k = k * constants.USEARCH_FILTER_OVERFETCH_MULTIPLIER if filter else k - - keys, distances = self._index.search( - query_vec, fetch_k, exact=exact, threads=threads - ) + if not filter: + # No filter: simple fetch + keys, distances = self._index.search( + query_vec, k, exact=exact, threads=threads + ) + if len(keys) == 0: + return [] + keys_list = keys.tolist() + dist_list = distances.tolist() + docs_map = self._catalog.get_documents_by_ids(keys_list) + return [ + (Document(page_content=text, metadata=metadata), float(dist)) + for key, dist in zip(keys_list, dist_list) + if key in docs_map + for text, metadata in [docs_map[key]] + ][:k] + + # Filtered search: iterative deepening to ensure k results + multiplier = constants.USEARCH_FILTER_OVERFETCH_MULTIPLIER + max_multiplier = 30 + index_size = self._index.size + added_keys: set[int] = set() + results: list[tuple[Document, float]] = [] - if len(keys) == 0: - return [] + while len(results) < k and multiplier <= max_multiplier: + fetch_k = min(k * multiplier, index_size) if index_size > 0 else k * multiplier + keys, distances = self._index.search( + query_vec, fetch_k, exact=exact, threads=threads + ) + if len(keys) == 0: + break - # Convert once, reuse - keys_list = keys.tolist() - dist_list = distances.tolist() + keys_list = keys.tolist() + dist_list = distances.tolist() - # Fetch documents and apply filter - docs_map = self._catalog.get_documents_by_ids(keys_list) + # Fetch docs for keys not yet processed + new_keys = [key for key in keys_list if key not in added_keys] + if not new_keys: + break + docs_map = self._catalog.get_documents_by_ids(new_keys) - results: list[tuple[Document, float]] = [] - for key, dist in zip(keys_list, dist_list): - if key not in docs_map: - continue + for key, dist in zip(keys_list, dist_list): + if key in added_keys: + continue + added_keys.add(key) - text, metadata = docs_map[key] + if key not in docs_map: + continue - # Apply metadata filter - if filter and not self._matches_filter(metadata, filter): - continue + text, metadata = docs_map[key] + if not self._matches_filter(metadata, filter): + continue - results.append((Document(page_content=text, metadata=metadata), float(dist))) + results.append((Document(page_content=text, metadata=metadata), float(dist))) + if len(results) >= k: + break - if len(results) >= k: + if len(results) >= k or fetch_k >= index_size: break + multiplier *= 2 - return results + return results[:k] def similarity_search_batch( self, @@ -351,14 +379,27 @@ def max_marginal_relevance_search( keys_list = keys.tolist() docs_and_embs = self._catalog.get_documents_and_embeddings_by_ids(keys_list) + # If catalog has no stored embeddings, retrieve from usearch index + has_catalog_embs = any( + emb is not None for _, _, emb in docs_and_embs.values() + ) + index_embs: np.ndarray | None = None + if not has_catalog_embs: + keys_arr = np.array(keys_list, dtype=np.uint64) + index_embs = self._index.get(keys_arr) + # Build candidates list with filtering, pre-normalize embeddings candidates: list[tuple[int, Document, float, np.ndarray | None]] = [] - for key, dist in zip(keys_list, distances.tolist()): + for i, (key, dist) in enumerate(zip(keys_list, distances.tolist())): if key not in docs_and_embs: continue text, metadata, emb = docs_and_embs[key] + # Fall back to index embeddings if catalog has none + if emb is None and index_embs is not None: + emb = index_embs[i] + # Apply metadata filter if filter and not self._matches_filter(metadata, filter): continue @@ -451,8 +492,8 @@ def _matches_filter(self, metadata: dict[str, Any], filter: dict[str, Any]) -> b if meta_value not in value: return False elif isinstance(value, str): - # String filter: substring match - if meta_value is None or value not in str(meta_value): + # String filter: exact match (consistent with SQL build_filter_clause) + if meta_value != value: return False else: # Exact match for int/float diff --git a/src/simplevecdb/engine/usearch_index.py b/src/simplevecdb/engine/usearch_index.py index da2fefd..de83a93 100755 --- a/src/simplevecdb/engine/usearch_index.py +++ b/src/simplevecdb/engine/usearch_index.py @@ -130,24 +130,24 @@ def _load_or_create(self) -> None: from .. import constants if self._path.exists(): - # Check file size to decide load vs view - file_size = self._path.stat().st_size - # Estimate vector count: file_size / (ndim * dtype_size + overhead) - # Conservative estimate assuming f32 and ~50 bytes overhead per vector - estimated_vectors = file_size // 100 # Very rough estimate - - if estimated_vectors > constants.USEARCH_MMAP_THRESHOLD: - # Use memory-mapped view for large indexes - _logger.debug( - "Using memory-mapped view for large index: %s", self._path - ) - self._index = Index.restore(str(self._path), view=True) - self._is_view = True - else: - # Load into memory for smaller indexes - _logger.debug("Loading index into memory: %s", self._path) - self._index = Index.restore(str(self._path), view=False) - self._is_view = False + from ..utils import file_lock + + with file_lock(self._path): + # Use file size threshold for mmap decision + file_size = self._path.stat().st_size + + if file_size > constants.USEARCH_MMAP_THRESHOLD: + # Use memory-mapped view for large indexes + _logger.debug( + "Using memory-mapped view for large index: %s", self._path + ) + self._index = Index.restore(str(self._path), view=True) + self._is_view = True + else: + # Load into memory for smaller indexes + _logger.debug("Loading index into memory: %s", self._path) + self._index = Index.restore(str(self._path), view=False) + self._is_view = False self._ndim = self._index.ndim _logger.info( @@ -266,14 +266,13 @@ def add( vectors = vectors / np.maximum(norms, 1e-12) # Upsert: remove existing keys first (usearch doesn't allow duplicates). - # Build a set of existing keys for O(1) lookup instead of per-key - # __contains__ probes, then batch-remove the conflicts. + # Use batch remove for efficiency. if self.size > 0: - existing_keys = [ - int(k) for k in keys if int(k) in self._index - ] - for int_key in existing_keys: - self._index.remove(int_key) + existing_mask = np.array( + [int(k) in self._index for k in keys], dtype=bool + ) + if existing_mask.any(): + self._index.remove(keys[existing_mask]) self._index.add(keys, vectors, threads=threads) self._dirty = True @@ -375,13 +374,14 @@ def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: return 0 with self._write_lock: - removed = 0 - for key in keys: - try: - self._index.remove(int(key)) - removed += 1 - except KeyError: - pass # Key not in index + # Filter to only keys that exist in the index + existing_mask = np.array( + [int(k) in self._index for k in keys], dtype=bool + ) + existing_keys = keys[existing_mask] + if len(existing_keys) > 0: + self._index.remove(existing_keys) + removed = int(existing_mask.sum()) self._dirty = True _logger.debug("Removed %d vectors from index", removed) return removed @@ -397,10 +397,13 @@ def save(self) -> None: if self._index is None or not self._dirty: return + from ..utils import file_lock + with self._write_lock: # Ensure parent directory exists self._path.parent.mkdir(parents=True, exist_ok=True) - self._index.save(str(self._path)) + with file_lock(self._path): + self._index.save(str(self._path)) self._dirty = False _logger.debug("Saved index to %s", self._path) @@ -435,26 +438,35 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: return np.array([], dtype=np.float32).reshape(0, self._ndim or 1) keys = np.asarray(keys, dtype=np.uint64) - vectors = [] - missing_logged = False - for key in keys: - try: - vec = self._index.get(int(key)) - vectors.append(np.asarray(vec, dtype=np.float32)) - except KeyError: - if not missing_logged: - _logger.warning( - "Index missing key(s); returning zero vectors. " - "Index may be out of sync with catalog." - ) - missing_logged = True - vectors.append(np.zeros(self._ndim or 1, dtype=np.float32)) + ndim = self._ndim or 1 + + # Filter to existing keys for batch retrieval + existing_mask = np.array( + [int(k) in self._index for k in keys], dtype=bool + ) - return ( - np.stack(vectors) - if vectors - else np.array([], dtype=np.float32).reshape(0, self._ndim or 1) + if not existing_mask.any(): + _logger.warning( + "Index missing key(s); returning zero vectors. " + "Index may be out of sync with catalog." + ) + return np.zeros((len(keys), ndim), dtype=np.float32) + + if existing_mask.all(): + # Fast path: all keys exist, batch retrieve + return np.asarray(self._index[keys], dtype=np.float32) + + # Mixed: some keys missing + _logger.warning( + "Index missing key(s); returning zero vectors. " + "Index may be out of sync with catalog." + ) + result = np.zeros((len(keys), ndim), dtype=np.float32) + existing_keys = keys[existing_mask] + result[existing_mask] = np.asarray( + self._index[existing_keys], dtype=np.float32 ) + return result def __del__(self) -> None: try: diff --git a/src/simplevecdb/utils.py b/src/simplevecdb/utils.py index acd5ea3..c5fa954 100755 --- a/src/simplevecdb/utils.py +++ b/src/simplevecdb/utils.py @@ -8,8 +8,10 @@ import sys import time from collections.abc import Iterable, Sequence +from contextlib import contextmanager from functools import wraps -from typing import Any, Callable, TypeVar +from pathlib import Path +from typing import Any, Callable, Generator, TypeVar F = TypeVar("F", bound=Callable[..., Any]) @@ -166,6 +168,105 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return decorator +def async_retry_on_lock( + max_retries: int = 5, + base_delay: float = 0.1, + max_delay: float = 2.0, + jitter: bool = True, + total_timeout: float = 10.0, +) -> Callable[[F], F]: + """ + Async decorator that retries database operations on SQLite lock errors. + + Uses asyncio.sleep instead of time.sleep, avoiding executor thread blocking. + Same backoff logic as retry_on_lock. + + Args: + max_retries: Maximum number of retry attempts (default: 5). + base_delay: Initial delay in seconds before first retry (default: 0.1). + max_delay: Maximum delay between retries in seconds (default: 2.0). + jitter: Add randomness to delay to avoid thundering herd (default: True). + total_timeout: Absolute wall-clock budget in seconds (default: 10.0). + + Returns: + Decorated async function with retry behavior. + + Raises: + DatabaseLockedError: If all retry attempts fail due to lock contention. + sqlite3.OperationalError: For non-lock SQLite errors. + """ + import asyncio + + def decorator(func: F) -> F: + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + last_exception: sqlite3.OperationalError | None = None + total_wait = 0.0 + + for attempt in range(max_retries + 1): + try: + return await func(*args, **kwargs) + except sqlite3.OperationalError as e: + error_msg = str(e).lower() + if "database is locked" not in error_msg: + raise + + last_exception = e + + if attempt < max_retries: + delay = min(base_delay * (2**attempt), max_delay) + + if jitter: + delay *= 0.75 + random.random() * 0.5 + + if total_wait + delay > total_timeout: + _logger.warning( + "Database lock retry would exceed total_timeout " + "(%.2fs spent, %.2fs budget) — giving up", + total_wait, + total_timeout, + extra={"operation": func.__name__}, + ) + break + + total_wait += delay + + _logger.warning( + "Database locked, retrying in %.3fs (attempt %d/%d)", + delay, + attempt + 1, + max_retries, + extra={ + "operation": func.__name__, + "attempt": attempt + 1, + "max_retries": max_retries, + "delay_seconds": round(delay, 3), + }, + ) + await asyncio.sleep(delay) + + _logger.error( + "Database locked after %d attempts (%.2fs total wait)", + max_retries + 1, + total_wait, + extra={ + "operation": func.__name__, + "attempts": max_retries + 1, + "total_wait_seconds": round(total_wait, 2), + }, + ) + raise DatabaseLockedError( + f"Database remained locked after {max_retries + 1} attempts " + f"(waited {total_wait:.2f}s total)", + attempts=max_retries + 1, + total_wait=total_wait, + ) from last_exception + + return wrapper # type: ignore[return-value] + + return decorator + + def validate_filter(filter_dict: dict[str, Any] | None) -> None: """ Validate metadata filter structure before SQL generation. @@ -223,3 +324,41 @@ def validate_filter(filter_dict: dict[str, Any] | None) -> None: f"Filter list item for '{key}' at index {i} must be finite, " f"got {item!r}" ) + + + +@contextmanager +def file_lock(path: Path) -> Generator[None, None, None]: + """Advisory file lock for cross-process safety. + + Uses fcntl.flock on Unix and msvcrt.locking on Windows. + Creates a .lock file alongside the target path. + + Args: + path: Path to the file to lock (a .lock sibling is created). + + Yields: + None — the lock is held for the duration of the context. + """ + lock_path = path.with_suffix(path.suffix + ".lock") + fd = open(lock_path, "w") # noqa: SIM115 + try: + if sys.platform == "win32": + import msvcrt + + msvcrt.locking(fd.fileno(), msvcrt.LK_LOCK, 1) + else: + import fcntl + + fcntl.flock(fd.fileno(), fcntl.LOCK_EX) + yield + finally: + if sys.platform == "win32": + import msvcrt + + msvcrt.locking(fd.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(fd.fileno(), fcntl.LOCK_UN) + fd.close() diff --git a/tests/integration/test_langchain.py b/tests/integration/test_langchain.py index 90120e2..766fa1a 100755 --- a/tests/integration/test_langchain.py +++ b/tests/integration/test_langchain.py @@ -60,7 +60,7 @@ def test_langchain_from_texts(tmp_path): assert isinstance(store, SimpleVecDBVectorStore) # Verify DB was created - assert store._collection._dim == 10 + assert store._collection.dim == 10 @pytest.mark.integration diff --git a/tests/unit/core/test_batch_detection.py b/tests/unit/core/test_batch_detection.py index 431e637..8e88845 100755 --- a/tests/unit/core/test_batch_detection.py +++ b/tests/unit/core/test_batch_detection.py @@ -207,23 +207,23 @@ def test_get_optimal_batch_size_mps_branches(): with patch.dict(sys.modules, {"torch": mock_torch}): with patch.object(platform, "machine", return_value="arm64"): # M3/M4 chips - with patch("subprocess.check_output", return_value="apple m3"): + with patch.object(platform, "processor", return_value="apple m3"): assert get_optimal_batch_size() == 64 # Max chips - with patch("subprocess.check_output", return_value="apple m2 max"): + with patch.object(platform, "processor", return_value="apple m2 max"): assert get_optimal_batch_size() == 128 # Ultra chips - with patch("subprocess.check_output", return_value="apple m2 ultra"): + with patch.object(platform, "processor", return_value="apple m2 ultra"): assert get_optimal_batch_size() == 128 # Base M1/M2 (no M3/M4 and no Max/Ultra) - with patch("subprocess.check_output", return_value="apple m1"): + with patch.object(platform, "processor", return_value="apple m1"): assert get_optimal_batch_size() == 32 - # Exception fallback - with patch("subprocess.check_output", side_effect=Exception("Fail")): + # Empty processor string fallback + with patch.object(platform, "processor", return_value=""): assert get_optimal_batch_size() == 32 diff --git a/tests/unit/core/test_core_additional_coverage.py b/tests/unit/core/test_core_additional_coverage.py index a6bfe61..76bcf28 100755 --- a/tests/unit/core/test_core_additional_coverage.py +++ b/tests/unit/core/test_core_additional_coverage.py @@ -65,7 +65,7 @@ def test_add_texts_uses_local_embedder_numpy(tmp_path): assert len(first_ids) == 1 assert len(second_ids) == 1 - assert collection._dim == 3 + assert collection.dim == 3 db.close() diff --git a/tests/unit/core/test_initialization.py b/tests/unit/core/test_initialization.py index f911ac7..ca5bc19 100755 --- a/tests/unit/core/test_initialization.py +++ b/tests/unit/core/test_initialization.py @@ -14,7 +14,7 @@ def test_recover_dim_no_match(tmp_path): collection = db.collection("default") # New collection should have _dim as None until vectors are added - assert collection._dim is None or isinstance(collection._dim, int) + assert collection.dim is None or isinstance(collection.dim, int) def test_recover_dim_none(tmp_path): @@ -24,7 +24,7 @@ def test_recover_dim_none(tmp_path): collection = db.collection("default") # New DB should have _dim as None initially - assert collection._dim is None + assert collection.dim is None def test_dimension_mismatch_on_add(): diff --git a/tests/unit/core/test_missing_coverage.py b/tests/unit/core/test_missing_coverage.py index 755c5f4..01cbdb5 100644 --- a/tests/unit/core/test_missing_coverage.py +++ b/tests/unit/core/test_missing_coverage.py @@ -124,7 +124,7 @@ class TestRebuildIndexNoEmbeddings: def test_rebuild_index_no_embeddings_raises(self, tmp_path): """rebuild_index raises RuntimeError when no embeddings in SQLite.""" db = VectorDB(str(tmp_path / "rebuild.db")) - collection = db.collection("default") + collection = db.collection("default", store_embeddings=True) # Insert a doc but clear embeddings from catalog collection.add_texts(["test"], embeddings=[[0.1, 0.2]]) # Wipe the embeddings column diff --git a/tests/unit/core/test_v25_correctness.py b/tests/unit/core/test_v25_correctness.py new file mode 100644 index 0000000..767212c --- /dev/null +++ b/tests/unit/core/test_v25_correctness.py @@ -0,0 +1,265 @@ +"""simplevecdb 2.5.0 correctness tests. + +Covers: list_collections persistence, collection cache keying, +delete_collection, __repr__, and connection health checks. +""" + +import numpy as np +import pytest + +from simplevecdb import VectorDB, Quantization, DistanceStrategy + + +DIM = 8 + + +def _rand_embedding(dim: int = DIM) -> list[list[float]]: + return [np.random.default_rng(42).random(dim).tolist()] + + +# ------------------------------------------------------------------ # +# 1. list_collections persistence +# ------------------------------------------------------------------ # + + +class TestListCollectionsPersistence: + """list_collections across sessions and edge cases.""" + + def test_persists_across_sessions(self, tmp_path): + db_path = str(tmp_path / "persist.db") + db = VectorDB(db_path) + db.collection("users").add_texts(["alice"], embeddings=_rand_embedding()) + db.collection("products").add_texts(["widget"], embeddings=_rand_embedding()) + db.close() + + db2 = VectorDB(db_path) + names = db2.list_collections() + assert "users" in names + assert "products" in names + db2.close() + + def test_memory_db_shows_session_collections(self): + db = VectorDB(":memory:") + db.collection("alpha").add_texts(["a"], embeddings=_rand_embedding()) + db.collection("beta").add_texts(["b"], embeddings=_rand_embedding()) + names = db.list_collections() + assert "alpha" in names + assert "beta" in names + + def test_empty_db_returns_empty_list(self): + db = VectorDB(":memory:") + assert db.list_collections() == [] + + def test_fts_subtables_excluded(self): + """FTS internal tables (items_X_fts_data, etc.) must not leak.""" + db = VectorDB(":memory:") + coll = db.collection("docs") + coll.add_texts(["hello world"], embeddings=_rand_embedding()) + + # Force FTS table creation by doing a text search + try: + coll.search("hello", k=1) + except Exception: + pass # search may need embeddings; table may still be created + + names = db.list_collections() + fts_leaked = [n for n in names if "fts" in n] + assert fts_leaked == [], f"FTS tables leaked into list_collections: {fts_leaked}" + + def test_sorted_output(self): + db = VectorDB(":memory:") + for name in ("zebra", "alpha", "middle"): + db.collection(name).add_texts(["x"], embeddings=_rand_embedding()) + names = db.list_collections() + assert names == sorted(names) + + def test_default_collection_listed_when_accessed(self): + db = VectorDB(":memory:") + db.collection("default").add_texts(["x"], embeddings=_rand_embedding()) + assert "default" in db.list_collections() + + +# ------------------------------------------------------------------ # +# 2. collection() cache key includes strategy / quantization +# ------------------------------------------------------------------ # + + +class TestCollectionCacheKey: + """Cache key must distinguish (name, distance, quantization).""" + + def test_different_quantization_returns_different_objects(self): + db = VectorDB(":memory:") + c_float = db.collection("test", quantization=Quantization.FLOAT) + c_int8 = db.collection("test", quantization=Quantization.INT8) + assert c_float is not c_int8 + + def test_same_params_returns_cached_object(self): + db = VectorDB(":memory:") + c1 = db.collection("test", quantization=Quantization.FLOAT) + c2 = db.collection("test", quantization=Quantization.FLOAT) + assert c1 is c2 + + def test_different_distance_returns_different_objects(self): + db = VectorDB(":memory:") + c_cos = db.collection("test", distance_strategy=DistanceStrategy.COSINE) + c_l2 = db.collection("test", distance_strategy=DistanceStrategy.L2) + assert c_cos is not c_l2 + + def test_full_key_match(self): + db = VectorDB(":memory:") + c1 = db.collection( + "x", + distance_strategy=DistanceStrategy.COSINE, + quantization=Quantization.FLOAT16, + ) + c2 = db.collection( + "x", + distance_strategy=DistanceStrategy.COSINE, + quantization=Quantization.FLOAT16, + ) + assert c1 is c2 + + +# ------------------------------------------------------------------ # +# 3. delete_collection +# ------------------------------------------------------------------ # + + +class TestDeleteCollection: + """delete_collection data removal and error handling.""" + + def test_delete_removes_from_list(self): + db = VectorDB(":memory:") + db.collection("doomed").add_texts(["bye"], embeddings=_rand_embedding()) + assert "doomed" in db.list_collections() + db.delete_collection("doomed") + assert "doomed" not in db.list_collections() + + def test_delete_allows_recreate_with_same_name(self): + db = VectorDB(":memory:") + coll = db.collection("recycle") + coll.add_texts(["old"], embeddings=_rand_embedding()) + db.delete_collection("recycle") + + fresh = db.collection("recycle") + assert fresh.count() == 0 + + def test_delete_nonexistent_raises_key_error(self): + db = VectorDB(":memory:") + with pytest.raises(KeyError, match="does not exist"): + db.delete_collection("ghost") + + def test_delete_removes_usearch_file(self, tmp_path): + db_path = str(tmp_path / "del.db") + db = VectorDB(db_path) + db.collection("victim").add_texts(["x"], embeddings=_rand_embedding()) + db.save() + + from pathlib import Path + + index_file = Path(db_path + ".victim.usearch") + # Index file should exist after save (if data was added) + # It may or may not exist depending on lazy flushing, so just + # verify that after delete it is gone. + db.delete_collection("victim") + assert not index_file.exists(), "usearch index file was not cleaned up" + db.close() + + def test_delete_clears_cache_entries(self): + db = VectorDB(":memory:") + c1 = db.collection("temp", quantization=Quantization.FLOAT) + c2 = db.collection("temp", quantization=Quantization.INT8) + c1.add_texts(["a"], embeddings=_rand_embedding()) + c2.add_texts(["b"], embeddings=_rand_embedding()) + + db.delete_collection("temp") + # Both cache entries for "temp" should be evicted + remaining = [k for k in db._collections if k[0] == "temp"] + assert remaining == [], f"Cache still has entries for 'temp': {remaining}" + + +# ------------------------------------------------------------------ # +# 4. __repr__ +# ------------------------------------------------------------------ # + + +class TestRepr: + """__repr__ for VectorCollection, VectorDB, and async wrappers.""" + + def test_vector_collection_repr_populated(self): + db = VectorDB(":memory:") + coll = db.collection("things") + coll.add_texts(["item"], embeddings=_rand_embedding()) + r = repr(coll) + assert "things" in r + assert str(DIM) in r + assert "1" in r # size + assert "cosine" in r or "l2" in r # distance + + def test_vector_collection_repr_empty(self): + db = VectorDB(":memory:") + coll = db.collection("empty") + r = repr(coll) + assert "empty" in r + assert "None" in r # dim is None + assert "0" in r # size is 0 + + def test_vector_db_repr(self): + db = VectorDB(":memory:") + db.collection("a").add_texts(["x"], embeddings=_rand_embedding()) + r = repr(db) + assert ":memory:" in r + assert "a" in r + + def test_vector_db_repr_empty(self): + db = VectorDB(":memory:") + r = repr(db) + assert ":memory:" in r + assert "[]" in r + + def test_async_collection_repr(self): + from simplevecdb.async_core import AsyncVectorCollection + + db = VectorDB(":memory:") + sync_coll = db.collection("async_test") + async_coll = AsyncVectorCollection(sync_coll, executor=None) + r = repr(async_coll) + assert "async_test" in r + assert "AsyncVectorCollection" in r + + def test_async_db_repr(self): + from simplevecdb.async_core import AsyncVectorDB + + db = VectorDB(":memory:") + async_db = AsyncVectorDB.__new__(AsyncVectorDB) + async_db._db = db + r = repr(async_db) + assert ":memory:" in r + assert "AsyncVectorDB" in r + + +# ------------------------------------------------------------------ # +# 5. Connection health check +# ------------------------------------------------------------------ # + + +class TestConnectionHealthCheck: + """Corrupt DB files must raise RuntimeError on open.""" + + def test_corrupt_db_raises_runtime_error(self, tmp_path): + corrupt_path = str(tmp_path / "corrupt.db") + with open(corrupt_path, "wb") as f: + f.write(b"\x00garbage\xffinvalid\x00sqlite\x00not\x00really") + + with pytest.raises((RuntimeError, Exception)): + VectorDB(corrupt_path) + + def test_valid_db_opens_cleanly(self, tmp_path): + db_path = str(tmp_path / "valid.db") + db = VectorDB(db_path) + db.collection("ok").add_texts(["fine"], embeddings=_rand_embedding()) + db.close() + + db2 = VectorDB(db_path) + assert "ok" in db2.list_collections() + db2.close() diff --git a/tests/unit/core/test_v25_features.py b/tests/unit/core/test_v25_features.py new file mode 100644 index 0000000..34f42e9 --- /dev/null +++ b/tests/unit/core/test_v25_features.py @@ -0,0 +1,344 @@ +"""Tests for simplevecdb 2.5.0 features: iterative deepening, store_embeddings, +FLOAT16 quantization, and pagination on get_documents / catalog lookups.""" + +import numpy as np +import pytest + +from simplevecdb import VectorDB, Quantization +from simplevecdb.engine.quantization import QuantizationStrategy + + +# ------------------------------------------------------------------ # +# Helpers +# ------------------------------------------------------------------ # + +DIM = 8 + + +def _rand_embedding(seed: int) -> list[float]: + """Deterministic random embedding for reproducibility.""" + rng = np.random.RandomState(seed) + return rng.randn(DIM).tolist() + + +def _make_collection(db: VectorDB, name: str = "default", **kwargs): + return db.collection(name, **kwargs) + + +# ------------------------------------------------------------------ # +# 1. Iterative deepening for filtered search +# ------------------------------------------------------------------ # + + +class TestIterativeDeepeningFilteredSearch: + """The old k*3 approach fetches 15 candidates from 100 documents. + If only 5 out of 100 match the filter and they are scattered throughout + the index, the old approach would miss some. Iterative deepening widens + the fetch window until k results are found or the index is exhausted.""" + + def test_sparse_filter_returns_all_matching(self): + db = VectorDB(":memory:") + col = db.collection("deep") + + n_total = 100 + n_rare = 5 + rare_indices = {10, 30, 50, 70, 90} + + texts = [f"doc_{i}" for i in range(n_total)] + metadatas = [ + {"rare": True} if i in rare_indices else {"rare": False} + for i in range(n_total) + ] + embeddings = [_rand_embedding(i) for i in range(n_total)] + + col.add_texts(texts, metadatas=metadatas, embeddings=embeddings) + assert col.count() == n_total + + query = _rand_embedding(999) + results = col.similarity_search(query, k=5, filter={"rare": True}) + + assert len(results) == n_rare, ( + f"Expected {n_rare} results with filter={{rare: True}}, got {len(results)}" + ) + for doc, _dist in results: + assert doc.metadata["rare"] is True + + def test_filter_returns_fewer_than_k_when_insufficient(self): + """When fewer documents match than k, return all matching.""" + db = VectorDB(":memory:") + col = db.collection("sparse") + + texts = [f"doc_{i}" for i in range(20)] + metadatas = [{"color": "red"} if i < 3 else {"color": "blue"} for i in range(20)] + embeddings = [_rand_embedding(i) for i in range(20)] + col.add_texts(texts, metadatas=metadatas, embeddings=embeddings) + + results = col.similarity_search(_rand_embedding(42), k=10, filter={"color": "red"}) + assert len(results) == 3 + + def test_no_filter_match_returns_empty(self): + db = VectorDB(":memory:") + col = db.collection("empty_filter") + col.add_texts( + ["a", "b", "c"], + metadatas=[{"x": 1}] * 3, + embeddings=[_rand_embedding(i) for i in range(3)], + ) + results = col.similarity_search(_rand_embedding(0), k=5, filter={"x": 999}) + assert results == [] + + +# ------------------------------------------------------------------ # +# 2. store_embeddings=False (default) and True +# ------------------------------------------------------------------ # + + +class TestStoreEmbeddings: + def test_default_no_store_similarity_search_works(self): + """With store_embeddings=False (default), similarity_search still works + because it uses the usearch index, not SQLite BLOBs.""" + db = VectorDB(":memory:") + col = db.collection("nostore") + + texts = ["alpha", "beta", "gamma"] + embeddings = [_rand_embedding(i) for i in range(3)] + col.add_texts(texts, embeddings=embeddings) + + results = col.similarity_search(embeddings[0], k=2) + assert len(results) == 2 + assert results[0][0].page_content == "alpha" + + def test_default_no_store_mmr_search_works(self): + """MMR search falls back to usearch get() when embeddings are not + stored in SQLite.""" + db = VectorDB(":memory:") + col = db.collection("nostore_mmr") + + texts = ["one", "two", "three", "four"] + embeddings = [_rand_embedding(i) for i in range(4)] + col.add_texts(texts, embeddings=embeddings) + + results = col.max_marginal_relevance_search(embeddings[0], k=2, fetch_k=4) + assert len(results) == 2 + + def test_no_store_rebuild_index_raises(self): + """rebuild_index requires stored embeddings; without them it raises.""" + db = VectorDB(":memory:") + col = db.collection("nostore_rebuild") + + col.add_texts(["x", "y"], embeddings=[_rand_embedding(0), _rand_embedding(1)]) + assert col.count() == 2 + + with pytest.raises(RuntimeError, match="No embeddings found"): + col.rebuild_index() + + def test_store_embeddings_true_rebuild_works(self): + """With store_embeddings=True, rebuild_index succeeds.""" + db = VectorDB(":memory:") + col = db.collection("stored", store_embeddings=True) + + embeddings = [_rand_embedding(i) for i in range(5)] + col.add_texts( + [f"doc_{i}" for i in range(5)], + embeddings=embeddings, + ) + assert col.count() == 5 + + rebuilt = col.rebuild_index() + assert rebuilt == 5 + + # Search still works after rebuild + results = col.similarity_search(embeddings[0], k=2) + assert len(results) == 2 + + +# ------------------------------------------------------------------ # +# 3. FLOAT16 quantization roundtrip +# ------------------------------------------------------------------ # + + +class TestFloat16Quantization: + def test_serialize_deserialize_roundtrip(self): + """FLOAT16 serialize -> deserialize preserves values within half-precision tolerance.""" + qs = QuantizationStrategy(Quantization.FLOAT16) + original = np.array([0.1, -0.25, 0.5, 1.0, -1.0, 0.0, 0.333, -0.777], dtype=np.float32) + + blob = qs.serialize(original) + recovered = qs.deserialize(blob, dim=len(original)) + + np.testing.assert_allclose(recovered, original, atol=1e-3, rtol=1e-2) + assert recovered.dtype == np.float32 + + def test_float16_halves_storage(self): + """FLOAT16 BLOBs should be half the size of FLOAT32.""" + qs32 = QuantizationStrategy(Quantization.FLOAT) + qs16 = QuantizationStrategy(Quantization.FLOAT16) + vec = np.random.randn(128).astype(np.float32) + + blob32 = qs32.serialize(vec) + blob16 = qs16.serialize(vec) + + assert len(blob16) == len(blob32) // 2 + + def test_float16_collection_search(self): + """End-to-end: create a FLOAT16 collection, add texts, search.""" + db = VectorDB(":memory:", quantization=Quantization.FLOAT16) + col = db.collection("f16") + + n = 20 + texts = [f"item_{i}" for i in range(n)] + embeddings = [_rand_embedding(i) for i in range(n)] + col.add_texts(texts, embeddings=embeddings) + + results = col.similarity_search(embeddings[0], k=3) + assert len(results) == 3 + # The nearest neighbor to embeddings[0] should be item_0 itself + assert results[0][0].page_content == "item_0" + + +# ------------------------------------------------------------------ # +# 4. Pagination on get_documents +# ------------------------------------------------------------------ # + + +class TestGetDocumentsPagination: + @pytest.fixture() + def col_20(self): + """Collection with 20 documents, deterministic IDs.""" + self._db = VectorDB(":memory:") + col = self._db.collection("paged") + texts = [f"text_{i:02d}" for i in range(20)] + embeddings = [_rand_embedding(i) for i in range(20)] + col.add_texts(texts, embeddings=embeddings) + assert col.count() == 20 + return col + + def test_limit_returns_exact_count(self, col_20): + docs = col_20.get_documents(limit=5) + assert len(docs) == 5 + + def test_offset_returns_different_page(self, col_20): + page1 = col_20.get_documents(limit=5) + page2 = col_20.get_documents(limit=5, offset=5) + assert len(page2) == 5 + ids_1 = {d[0] for d in page1} + ids_2 = {d[0] for d in page2} + assert ids_1.isdisjoint(ids_2), "Pages must not overlap" + + def test_last_page(self, col_20): + page = col_20.get_documents(limit=5, offset=15) + assert len(page) == 5 + + def test_beyond_end_returns_empty(self, col_20): + page = col_20.get_documents(limit=5, offset=20) + assert page == [] + + def test_no_limit_returns_all(self, col_20): + docs = col_20.get_documents() + assert len(docs) == 20 + + def test_pagination_with_filter(self): + db = VectorDB(":memory:") + col = db.collection("paged_filter") + + texts = [f"text_{i}" for i in range(20)] + metadatas = [{"type": "a"} if i % 2 == 0 else {"type": "b"} for i in range(20)] + embeddings = [_rand_embedding(i) for i in range(20)] + col.add_texts(texts, metadatas=metadatas, embeddings=embeddings) + + # 10 docs have type=a + all_a = col.get_documents(filter_dict={"type": "a"}) + assert len(all_a) == 10 + + page1 = col.get_documents(filter_dict={"type": "a"}, limit=5) + assert len(page1) == 5 + + page2 = col.get_documents(filter_dict={"type": "a"}, limit=5, offset=5) + assert len(page2) == 5 + + ids_1 = {d[0] for d in page1} + ids_2 = {d[0] for d in page2} + assert ids_1.isdisjoint(ids_2) + + def test_full_page_coverage(self, col_20): + """Iterating through all pages should yield all 20 documents.""" + all_ids = set() + for offset in range(0, 20, 5): + page = col_20.get_documents(limit=5, offset=offset) + for doc_id, _text, _meta in page: + all_ids.add(doc_id) + + assert len(all_ids) == 20 + + +# ------------------------------------------------------------------ # +# 5. Pagination on find_ids_by_texts and find_ids_by_filter +# ------------------------------------------------------------------ # + + +class TestCatalogPagination: + @pytest.fixture() + def col_with_metadata(self): + """Collection with 15 docs, varied metadata.""" + self._db = VectorDB(":memory:") + col = self._db.collection("catalog_paged") + + texts = [f"sentence_{i}" for i in range(15)] + metadatas = [{"group": i % 3, "idx": i} for i in range(15)] + embeddings = [_rand_embedding(i) for i in range(15)] + col.add_texts(texts, metadatas=metadatas, embeddings=embeddings) + return col + + def test_find_ids_by_texts_all(self, col_with_metadata): + texts = [f"sentence_{i}" for i in range(15)] + ids = col_with_metadata._catalog.find_ids_by_texts(texts) + assert len(ids) == 15 + + def test_find_ids_by_texts_with_limit(self, col_with_metadata): + texts = [f"sentence_{i}" for i in range(15)] + ids = col_with_metadata._catalog.find_ids_by_texts(texts, limit=5) + assert len(ids) == 5 + + def test_find_ids_by_texts_with_limit_offset(self, col_with_metadata): + texts = [f"sentence_{i}" for i in range(15)] + page1 = col_with_metadata._catalog.find_ids_by_texts(texts, limit=5) + page2 = col_with_metadata._catalog.find_ids_by_texts(texts, limit=5, offset=5) + assert len(page2) == 5 + assert set(page1).isdisjoint(set(page2)) + + def test_find_ids_by_texts_offset_beyond_end(self, col_with_metadata): + texts = [f"sentence_{i}" for i in range(15)] + ids = col_with_metadata._catalog.find_ids_by_texts(texts, limit=5, offset=15) + assert ids == [] + + def test_find_ids_by_filter_all(self, col_with_metadata): + catalog = col_with_metadata._catalog + builder = catalog.build_filter_clause + ids = catalog.find_ids_by_filter({"group": 0}, builder) + assert len(ids) == 5 # 0, 3, 6, 9, 12 + + def test_find_ids_by_filter_with_limit(self, col_with_metadata): + catalog = col_with_metadata._catalog + builder = catalog.build_filter_clause + ids = catalog.find_ids_by_filter({"group": 0}, builder, limit=3) + assert len(ids) == 3 + + def test_find_ids_by_filter_with_limit_offset(self, col_with_metadata): + catalog = col_with_metadata._catalog + builder = catalog.build_filter_clause + page1 = catalog.find_ids_by_filter({"group": 0}, builder, limit=3) + page2 = catalog.find_ids_by_filter({"group": 0}, builder, limit=3, offset=3) + assert len(page2) == 2 # only 2 remaining (5 total, took 3) + assert set(page1).isdisjoint(set(page2)) + + def test_find_ids_by_filter_offset_beyond_end(self, col_with_metadata): + catalog = col_with_metadata._catalog + builder = catalog.build_filter_clause + ids = catalog.find_ids_by_filter({"group": 0}, builder, limit=5, offset=10) + assert ids == [] + + def test_find_ids_by_filter_empty_dict(self, col_with_metadata): + catalog = col_with_metadata._catalog + builder = catalog.build_filter_clause + ids = catalog.find_ids_by_filter({}, builder) + assert ids == [] diff --git a/tests/unit/core/test_v25_robustness.py b/tests/unit/core/test_v25_robustness.py new file mode 100644 index 0000000..5f23378 --- /dev/null +++ b/tests/unit/core/test_v25_robustness.py @@ -0,0 +1,312 @@ +"""Tests for simplevecdb 2.5.0 robustness features. + +Covers: +- async_retry_on_lock decorator +- FTS retry on transient lock errors +- file_lock context manager +- Mmap byte threshold for UsearchIndex +""" + +from __future__ import annotations + +import asyncio +import fcntl +import sqlite3 +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock, patch, AsyncMock, PropertyMock, call + +import pytest +import numpy as np + +from simplevecdb import async_retry_on_lock, file_lock, DatabaseLockedError +from simplevecdb.engine.catalog import CatalogManager +import simplevecdb.constants as constants + + +# --------------------------------------------------------------------------- +# 1. async_retry_on_lock +# --------------------------------------------------------------------------- + + +class TestAsyncRetryOnLock: + """Tests for the async_retry_on_lock decorator.""" + + @pytest.mark.asyncio + async def test_retries_on_lock_then_succeeds(self): + """Decorator retries on 'database is locked' and returns the result.""" + attempt_count = 0 + + @async_retry_on_lock(max_retries=5, base_delay=0.01, jitter=False, total_timeout=10.0) + async def flaky(): + nonlocal attempt_count + attempt_count += 1 + if attempt_count <= 2: + raise sqlite3.OperationalError("database is locked") + return "ok" + + result = await flaky() + assert result == "ok" + assert attempt_count == 3 + + @pytest.mark.asyncio + async def test_uses_asyncio_sleep_not_time_sleep(self): + """Decorator awaits asyncio.sleep, never calls time.sleep.""" + call_count = 0 + + @async_retry_on_lock(max_retries=3, base_delay=0.01, jitter=False, total_timeout=10.0) + async def flaky(): + nonlocal call_count + call_count += 1 + if call_count <= 2: + raise sqlite3.OperationalError("database is locked") + return "done" + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_asleep, \ + patch("time.sleep") as mock_tsleep: + result = await flaky() + + assert result == "done" + assert mock_asleep.await_count == 2 + mock_tsleep.assert_not_called() + + @pytest.mark.asyncio + async def test_non_lock_operational_error_raises_immediately(self): + """Non-lock OperationalErrors propagate without retry.""" + attempt_count = 0 + + @async_retry_on_lock(max_retries=5, base_delay=0.01, jitter=False, total_timeout=10.0) + async def bad(): + nonlocal attempt_count + attempt_count += 1 + raise sqlite3.OperationalError("no such table: foo") + + with pytest.raises(sqlite3.OperationalError, match="no such table"): + await bad() + assert attempt_count == 1 + + @pytest.mark.asyncio + async def test_raises_database_locked_error_after_max_retries(self): + """DatabaseLockedError is raised once retries are exhausted.""" + + @async_retry_on_lock(max_retries=2, base_delay=0.001, jitter=False, total_timeout=60.0) + async def always_locked(): + raise sqlite3.OperationalError("database is locked") + + with pytest.raises(DatabaseLockedError) as exc_info: + await always_locked() + + err = exc_info.value + assert err.attempts == 3 # initial + 2 retries + assert err.total_wait >= 0 + + +# --------------------------------------------------------------------------- +# 2. FTS retry on transient lock errors +# --------------------------------------------------------------------------- + + +class TestFtsRetryOnLock: + """Tests for CatalogManager._ensure_fts_table lock-retry behaviour.""" + + @staticmethod + def _make_catalog_with_proxy(): + """Create a CatalogManager with a proxy conn whose execute is patchable.""" + real_conn = sqlite3.connect(":memory:") + proxy = MagicMock(wraps=real_conn) + # wraps delegates all calls to real_conn, but proxy.execute is now mockable + catalog = CatalogManager( + proxy, table_name="tinyvec_items", fts_table_name="tinyvec_items_fts" + ) + return catalog, proxy, real_conn + + def test_fts_retries_on_lock_then_enables(self): + """FTS table creation retries on lock errors and enables FTS.""" + catalog, proxy, real_conn = self._make_catalog_with_proxy() + + call_count = 0 + original_execute = real_conn.execute + + def mock_execute(sql, *args, **kwargs): + nonlocal call_count + if "CREATE VIRTUAL TABLE" in sql: + call_count += 1 + if call_count <= 2: + raise sqlite3.OperationalError("database is locked") + return original_execute(sql, *args, **kwargs) + + proxy.execute = mock_execute + catalog._ensure_fts_table() + + assert catalog._fts_enabled is True + assert call_count == 3 + + def test_fts_disables_immediately_on_module_missing(self): + """FTS is disabled without retries when fts5 module is absent.""" + catalog, proxy, real_conn = self._make_catalog_with_proxy() + + call_count = 0 + original_execute = real_conn.execute + + def mock_execute(sql, *args, **kwargs): + nonlocal call_count + if "CREATE VIRTUAL TABLE" in sql: + call_count += 1 + raise sqlite3.OperationalError("no such module: fts5") + return original_execute(sql, *args, **kwargs) + + proxy.execute = mock_execute + catalog._ensure_fts_table() + + assert catalog._fts_enabled is False + assert call_count == 1 # no retries + + +# --------------------------------------------------------------------------- +# 3. file_lock context manager +# --------------------------------------------------------------------------- + + +class TestFileLock: + """Tests for the file_lock advisory-lock context manager.""" + + def test_lock_file_created(self, tmp_path: Path): + """Acquiring file_lock creates a .lock sibling file.""" + target = tmp_path / "data.db" + target.touch() + + with file_lock(target): + lock_path = target.with_suffix(".db.lock") + assert lock_path.exists() + + def test_lock_released_after_context_exit(self, tmp_path: Path): + """After context exit the lock file exists but is no longer locked.""" + target = tmp_path / "data.db" + target.touch() + + with file_lock(target): + pass # lock held here + + # Lock file should still exist on disk but not be held + lock_path = target.with_suffix(".db.lock") + assert lock_path.exists() + + # Verify we can immediately acquire the lock again (proves it's released) + with file_lock(target): + pass # would block forever if still locked + + def test_concurrent_locks_from_threads(self, tmp_path: Path): + """Second thread blocks until first thread releases the lock.""" + target = tmp_path / "data.db" + target.touch() + + order: list[str] = [] + barrier = threading.Event() + + def first(): + with file_lock(target): + order.append("first-acquired") + barrier.set() # signal second thread to try acquiring + time.sleep(0.15) # hold lock + order.append("first-released") + + def second(): + barrier.wait() # wait until first thread holds the lock + time.sleep(0.02) # small delay to ensure first thread is still holding + with file_lock(target): + order.append("second-acquired") + + t1 = threading.Thread(target=first) + t2 = threading.Thread(target=second) + + t1.start() + t2.start() + t1.join(timeout=5) + t2.join(timeout=5) + + assert order == ["first-acquired", "first-released", "second-acquired"] + + +# --------------------------------------------------------------------------- +# 4. Mmap byte threshold +# --------------------------------------------------------------------------- + + +class TestMmapThreshold: + """Tests for UsearchIndex memory-mapping threshold.""" + + def test_mmap_threshold_constant_value(self): + """USEARCH_MMAP_THRESHOLD is exactly 50 MiB.""" + assert constants.USEARCH_MMAP_THRESHOLD == 50 * 1024 * 1024 + + @patch("simplevecdb.engine.usearch_index.UsearchIndex._load_or_create") + def test_large_file_enables_mmap(self, mock_load): + """Files larger than threshold trigger memory-mapped mode.""" + from simplevecdb.engine.usearch_index import UsearchIndex + + idx = UsearchIndex.__new__(UsearchIndex) + idx._is_view = True # simulate what _load_or_create would set + + assert idx.is_memory_mapped is True + + @patch("simplevecdb.engine.usearch_index.UsearchIndex._load_or_create") + def test_small_file_disables_mmap(self, mock_load): + """Files smaller than threshold load into memory (no mmap).""" + from simplevecdb.engine.usearch_index import UsearchIndex + + idx = UsearchIndex.__new__(UsearchIndex) + idx._is_view = False + + assert idx.is_memory_mapped is False + + def test_load_or_create_mmap_decision(self, tmp_path: Path): + """_load_or_create sets _is_view based on file size vs threshold.""" + from simplevecdb.engine.usearch_index import UsearchIndex + + index_path = tmp_path / "test.usearch" + index_path.touch() + + mock_index = MagicMock() + mock_index.ndim = 128 + mock_index.__len__ = lambda self: 1000 + + # Case 1: file > threshold => mmap + with patch("simplevecdb.engine.usearch_index.UsearchIndex._load_or_create"): + idx = UsearchIndex.__new__(UsearchIndex) + + idx._path = index_path + idx._is_view = False + idx._index = None + idx._ndim = None + + large_stat = MagicMock() + large_stat.st_size = constants.USEARCH_MMAP_THRESHOLD + 1 + + small_stat = MagicMock() + small_stat.st_size = constants.USEARCH_MMAP_THRESHOLD - 1 + + mock_index_cls = MagicMock() + mock_index_cls.restore.return_value = mock_index + + with patch.object(Path, "exists", return_value=True), \ + patch.object(Path, "stat", return_value=large_stat), \ + patch("simplevecdb.utils.file_lock"), \ + patch("usearch.index.Index", mock_index_cls): + idx._load_or_create() + + assert idx._is_view is True + mock_index_cls.restore.assert_called_once_with(str(index_path), view=True) + + # Case 2: file < threshold => no mmap + mock_index_cls.reset_mock() + idx._is_view = False + + with patch.object(Path, "exists", return_value=True), \ + patch.object(Path, "stat", return_value=small_stat), \ + patch("simplevecdb.utils.file_lock"), \ + patch("usearch.index.Index", mock_index_cls): + idx._load_or_create() + + assert idx._is_view is False + mock_index_cls.restore.assert_called_once_with(str(index_path), view=False) diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index df82196..9bd071c 100755 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -37,7 +37,7 @@ def test_add_texts_basic(empty_db): embs = [[0.1, 0.2], [0.3, 0.4]] ids = collection.add_texts(texts, embeddings=embs) assert len(ids) == 2 - assert collection._dim == 2 + assert collection.dim == 2 # Verify that the text content is persisted in the main table. rows = empty_db.conn.execute( @@ -133,7 +133,7 @@ def test_recover_dim(tmp_path): # Reopen db2 = VectorDB(db_path) - assert db2.collection("default")._dim == 10 + assert db2.collection("default").dim == 10 db2.close() @@ -270,7 +270,7 @@ def test_quantization_int8(quant_db): ids = collection.add_texts(texts, embeddings=embs) assert len(ids) == 2 - assert collection._dim == 3 + assert collection.dim == 3 # Search should work with quantized vectors results = collection.similarity_search([0.1, 0.2, 0.3], k=1) @@ -286,7 +286,7 @@ def test_quantization_bit(bit_db): assert len(ids) == 2 # BIT quantization rounds up to byte boundary - assert collection._dim == 3 + assert collection.dim == 3 # Search should work with binary vectors results = collection.similarity_search([0.1, 0.2, 0.3], k=1) @@ -490,7 +490,7 @@ def test_rebuild_index(tmp_path): """Test rebuild_index() reconstructs index from SQLite embeddings.""" db_path = str(tmp_path / "rebuild.db") db = VectorDB(db_path) - collection = db.collection("default") + collection = db.collection("default", store_embeddings=True) # Add some vectors texts = ["doc1", "doc2", "doc3"] @@ -520,7 +520,7 @@ def test_rebuild_index_with_custom_params(tmp_path): """Test rebuild_index() with custom HNSW parameters.""" db_path = str(tmp_path / "rebuild_params.db") db = VectorDB(db_path) - collection = db.collection("default") + collection = db.collection("default", store_embeddings=True) collection.add_texts(["test"], embeddings=[[0.1] * 64]) diff --git a/uv.lock b/uv.lock index 2d10d0f..8b2afa4 100755 --- a/uv.lock +++ b/uv.lock @@ -4601,7 +4601,7 @@ wheels = [ [[package]] name = "simplevecdb" -version = "2.3.0" +version = "2.5.0" source = { editable = "." } dependencies = [ { name = "cryptography" }, From bf4f604580f0ccefe8cb682b500f79d68f129a00 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Tue, 7 Apr 2026 05:01:38 -0500 Subject: [PATCH 2/5] feat: enhance embeddings server with graceful shutdown, CORS, validation, and CLI - Graceful shutdown with SIGTERM/SIGINT draining (10s timeout) - Async executor offload for embed_texts (non-blocking event loop) - Model warm-up on startup (--no-warmup to skip) - CORS middleware with configurable origins - Input validation: reject empty strings (422) and >100k char texts (413) - Proper argparse CLI replacing hand-rolled argv parsing - Startup banner logging config summary - Nested token array normalization (list[list[int]]) - OpenAPI version synced from package metadata - Module __init__.py exports (embed_texts, get_embedder, load_model, app, run_server) 544 tests passing (30 new). --- src/simplevecdb/embeddings/__init__.py | 12 + src/simplevecdb/embeddings/server.py | 227 +++++++++-- tests/unit/embeddings/test_server.py | 58 +-- .../unit/embeddings/test_v25_enhancements.py | 381 ++++++++++++++++++ 4 files changed, 620 insertions(+), 58 deletions(-) create mode 100644 tests/unit/embeddings/test_v25_enhancements.py diff --git a/src/simplevecdb/embeddings/__init__.py b/src/simplevecdb/embeddings/__init__.py index e69de29..8aa0a03 100755 --- a/src/simplevecdb/embeddings/__init__.py +++ b/src/simplevecdb/embeddings/__init__.py @@ -0,0 +1,12 @@ +"""Embeddings module — local embedding models and OpenAI-compatible server.""" + +from .models import embed_texts, get_embedder, load_model +from .server import app, run_server + +__all__ = [ + "app", + "embed_texts", + "get_embedder", + "load_model", + "run_server", +] diff --git a/src/simplevecdb/embeddings/server.py b/src/simplevecdb/embeddings/server.py index 4afd77e..c3a5de1 100755 --- a/src/simplevecdb/embeddings/server.py +++ b/src/simplevecdb/embeddings/server.py @@ -1,6 +1,9 @@ from __future__ import annotations +import argparse +import asyncio import logging +import signal import time from collections import defaultdict from threading import Lock @@ -8,14 +11,18 @@ import uvicorn from fastapi import Depends, FastAPI, Header, HTTPException, Request, Security +from fastapi.middleware.cors import CORSMiddleware from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from pydantic import BaseModel, Field -from .models import DEFAULT_MODEL, embed_texts +from .models import DEFAULT_MODEL, embed_texts, get_embedder from simplevecdb.config import config _logger = logging.getLogger("simplevecdb.embeddings.server") +# Maximum length (chars) for a single input text to prevent OOM in the encoder. +_MAX_TEXT_LENGTH = 100_000 + # Simple in-memory rate limiter class RateLimiter: @@ -79,14 +86,32 @@ def is_allowed(self, identity: str) -> bool: rate_limiter = RateLimiter(requests_per_minute=100, burst=20) +# (#9) Pull version from package metadata instead of hardcoding +try: + from importlib.metadata import version as _pkg_version + + _server_version = _pkg_version("simplevecdb") +except Exception: + _server_version = "0.0.0" + app = FastAPI( title="SimpleVecDB Embeddings", description="OpenAI-compatible /v1/embeddings endpoint – 100% local", - version="0.0.1", + version=_server_version, openapi_url="/openapi.json", docs_url="/docs", ) +# (#4) CORS middleware — configurable via EMBEDDING_SERVER_CORS_ORIGINS env var +_cors_origins = getattr(config, "EMBEDDING_SERVER_CORS_ORIGINS", ["*"]) +app.add_middleware( + CORSMiddleware, + allow_origins=_cors_origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + @app.get("/health") async def health_check(): @@ -228,6 +253,53 @@ class EmbeddingResponse(BaseModel): usage: dict = Field(default_factory=lambda: {"prompt_tokens": 0, "total_tokens": 0}) +def _normalize_input(raw_input: str | list[str] | list[int] | list[list[int]]) -> list[str]: + """Convert any valid OpenAI-compatible input format to a flat list of strings. + + Handles: + - str → ["str"] + - list[str] → as-is + - list[int] → token IDs stringified individually + - list[list[int]] → each sub-list decoded as a token sequence string + """ + if isinstance(raw_input, str): + return [raw_input] + + if not raw_input: + return [] + + first = raw_input[0] + + # list[int] — flat token array + if isinstance(first, int): + return [str(i) for i in raw_input] + + # list[list[int]] — nested token arrays (#8) + if isinstance(first, list): + return [" ".join(str(tok) for tok in sub) for sub in raw_input] + + # list[str] + return [str(item) for item in raw_input] + + +def _validate_texts(texts: list[str]) -> None: + """Reject empty strings and texts exceeding the per-item length cap (#5).""" + for i, text in enumerate(texts): + if not text or not text.strip(): + raise HTTPException( + status_code=422, + detail=f"Input text at index {i} is empty or whitespace-only.", + ) + if len(text) > _MAX_TEXT_LENGTH: + raise HTTPException( + status_code=413, + detail=( + f"Input text at index {i} is {len(text)} chars, " + f"exceeding the {_MAX_TEXT_LENGTH} char limit." + ), + ) + + @app.post("/v1/embeddings") async def create_embeddings( request: EmbeddingRequest, @@ -253,14 +325,13 @@ async def create_embeddings( raise HTTPException( status_code=429, detail="Rate limit exceeded. Try again later." ) - if isinstance(request.input, str): - texts = [request.input] - elif isinstance(request.input, list) and all( - isinstance(i, int) for i in request.input - ): - texts = [str(i) for i in request.input] # token arrays – just stringify - else: - texts = [str(item) for item in request.input] + + # (#8) Properly normalize all input formats including nested token arrays + texts = _normalize_input(request.input) + + # (#5) Validate individual texts + if texts: + _validate_texts(texts) if len(texts) > config.EMBEDDING_SERVER_MAX_REQUEST_ITEMS: raise HTTPException( @@ -275,15 +346,20 @@ async def create_embeddings( resolved_model_name, repo_id = registry.resolve(request.model) if not texts: - embeddings = [] + embeddings: list[list[float]] = [] else: try: effective_batch = min( config.EMBEDDING_BATCH_SIZE, config.EMBEDDING_SERVER_MAX_REQUEST_ITEMS, ) - embeddings = embed_texts( - texts, model_id=repo_id, batch_size=effective_batch + # (#2) Run embedding in executor to avoid blocking the event loop + loop = asyncio.get_running_loop() + embeddings = await loop.run_in_executor( + None, + lambda: embed_texts( + texts, model_id=repo_id, batch_size=effective_batch + ), ) except Exception as e: # Log the full error internally but return generic message @@ -323,6 +399,31 @@ async def usage(api_identity: str = Depends(authenticate_request)) -> dict[str, return {"object": "usage", "data": usage_meter.snapshot(scope)} +def _build_cli_parser() -> argparse.ArgumentParser: + """Build the CLI argument parser (#6).""" + parser = argparse.ArgumentParser( + prog="simplevecdb-server", + description="Run the SimpleVecDB embeddings server (OpenAI-compatible).", + ) + parser.add_argument( + "--host", + default=None, + help=f"Bind address (default: {config.SERVER_HOST})", + ) + parser.add_argument( + "--port", + type=int, + default=None, + help=f"Listen port (default: {config.SERVER_PORT})", + ) + parser.add_argument( + "--no-warmup", + action="store_true", + help="Skip model warm-up on startup", + ) + return parser + + def run_server(host: str | None = None, port: int | None = None) -> None: """Run the embedding server. @@ -330,45 +431,105 @@ def run_server(host: str | None = None, port: int | None = None) -> None: Examples -------- - Run with default settings: - $ simplevecdb-server + Run with default settings:: + + $ simplevecdb-server + + Override host and port:: + + $ simplevecdb-server --host 0.0.0.0 --port 9000 + + Skip model warm-up:: - Override port: - $ simplevecdb-server --port 8000 + $ simplevecdb-server --no-warmup Args: host: Server host (defaults to config.SERVER_HOST). port: Server port (defaults to config.SERVER_PORT). """ - # Minimal CLI-style override when invoked as a script/entry point - # Allows commands like: simplevecdb-server --host 0.0.0.0 --port 8000 - import sys - - argv = sys.argv[1:] - for i, arg in enumerate(argv): - if arg in {"--host", "-h"} and i + 1 < len(argv): - host = argv[i + 1] - if arg in {"--port", "-p"} and i + 1 < len(argv): - try: - port = int(argv[i + 1]) - except ValueError: - pass + # (#6) Only parse CLI args when invoked as entry point (not programmatically) + skip_warmup = False + if host is None and port is None: + parser = _build_cli_parser() + args = parser.parse_args() + host = args.host + port = args.port + skip_warmup = args.no_warmup host = host or config.SERVER_HOST port = port or config.SERVER_PORT + # (#7) Startup banner with config summary + auth_status = ( + f"{len(config.EMBEDDING_SERVER_API_KEYS)} key(s)" + if config.EMBEDDING_SERVER_API_KEYS + else "DISABLED" + ) + _logger.info( + "\n" + "┌─────────────────────────────────────────────┐\n" + "│ SimpleVecDB Embeddings Server │\n" + "├─────────────────────────────────────────────┤\n" + "│ Host: %-30s│\n" + "│ Port: %-30s│\n" + "│ Model: %-30s│\n" + "│ Auth: %-30s│\n" + "│ Rate limit: %-30s│\n" + "│ Version: %-30s│\n" + "└─────────────────────────────────────────────┘", + host, + port, + config.EMBEDDING_MODEL, + auth_status, + "100 req/min, burst 20", + _server_version, + ) + # Security warnings if not config.EMBEDDING_SERVER_API_KEYS: _logger.warning( - "⚠️ No API keys configured (EMBEDDING_SERVER_API_KEYS is empty). " + "No API keys configured (EMBEDDING_SERVER_API_KEYS is empty). " "Server is running without authentication. " "Set EMBEDDING_SERVER_API_KEYS for production use." ) if host == "0.0.0.0": _logger.warning( - "⚠️ Server binding to all interfaces (0.0.0.0). " + "Server binding to all interfaces (0.0.0.0). " "This exposes the server to the network. " "Use 127.0.0.1 for local-only access." ) - uvicorn.run(app, host=host, port=port, log_level="info") + # (#3) Model warm-up — pre-load default model before accepting traffic + if not skip_warmup: + _logger.info("Warming up default model: %s ...", config.EMBEDDING_MODEL) + try: + get_embedder(config.EMBEDDING_MODEL) + _logger.info("Model warm-up complete.") + except Exception: + _logger.warning( + "Model warm-up failed (will retry on first request).", exc_info=True + ) + + # (#1) Graceful shutdown with in-flight request draining + uvi_config = uvicorn.Config( + app, host=host, port=port, log_level="info", timeout_graceful_shutdown=10 + ) + server = uvicorn.Server(uvi_config) + + # Install signal handlers that tell uvicorn to drain gracefully + original_sigint = signal.getsignal(signal.SIGINT) + original_sigterm = signal.getsignal(signal.SIGTERM) + + def _graceful_shutdown(signum: int, frame: Any) -> None: + sig_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM" + _logger.info("Received %s — draining in-flight requests...", sig_name) + server.should_exit = True + + signal.signal(signal.SIGINT, _graceful_shutdown) + signal.signal(signal.SIGTERM, _graceful_shutdown) + + try: + server.run() + finally: + signal.signal(signal.SIGINT, original_sigint) + signal.signal(signal.SIGTERM, original_sigterm) diff --git a/tests/unit/embeddings/test_server.py b/tests/unit/embeddings/test_server.py index 9f604ae..dc209be 100755 --- a/tests/unit/embeddings/test_server.py +++ b/tests/unit/embeddings/test_server.py @@ -1,7 +1,7 @@ """Embeddings server API tests.""" import pytest -from unittest.mock import patch, ANY +from unittest.mock import patch, ANY, MagicMock from fastapi.testclient import TestClient from simplevecdb.embeddings.server import app, ModelRegistry @@ -120,44 +120,52 @@ def test_server_run_with_args(): """Test server.run_server with custom host/port.""" from simplevecdb.embeddings.server import run_server - with patch("uvicorn.run") as mock_run: - run_server(host="127.0.0.1", port=9000) - mock_run.assert_called_once() - call_kwargs = mock_run.call_args[1] - assert call_kwargs["host"] == "127.0.0.1" - assert call_kwargs["port"] == 9000 + mock_server = MagicMock() + with patch("simplevecdb.embeddings.server.uvicorn.Config") as mock_cfg: + with patch("simplevecdb.embeddings.server.uvicorn.Server", return_value=mock_server): + run_server(host="127.0.0.1", port=9000) + call_kwargs = mock_cfg.call_args[1] + assert call_kwargs["host"] == "127.0.0.1" + assert call_kwargs["port"] == 9000 + mock_server.run.assert_called_once() def test_server_run_default_config(): """Test server.run_server uses config defaults.""" from simplevecdb.embeddings.server import run_server - with patch("uvicorn.run") as mock_run: - with patch("simplevecdb.embeddings.server.config") as mock_config: - mock_config.SERVER_HOST = "0.0.0.0" - mock_config.SERVER_PORT = 8080 + mock_server = MagicMock() + with patch("simplevecdb.embeddings.server.uvicorn.Config") as mock_cfg: + with patch("simplevecdb.embeddings.server.uvicorn.Server", return_value=mock_server): + with patch("simplevecdb.embeddings.server.config") as mock_config: + mock_config.SERVER_HOST = "0.0.0.0" + mock_config.SERVER_PORT = 8080 + mock_config.EMBEDDING_SERVER_API_KEYS = set() + mock_config.EMBEDDING_MODEL = "test-model" - run_server() + run_server(host="0.0.0.0", port=8080) - call_kwargs = mock_run.call_args[1] - assert call_kwargs["host"] == "0.0.0.0" - assert call_kwargs["port"] == 8080 + call_kwargs = mock_cfg.call_args[1] + assert call_kwargs["host"] == "0.0.0.0" + assert call_kwargs["port"] == 8080 def test_server_run_cli_args(): - """Test server.run_server parses CLI arguments.""" + """Test server.run_server parses CLI arguments via argparse.""" from simplevecdb.embeddings.server import run_server import sys - with patch("uvicorn.run") as mock_run: - with patch.object( - sys, "argv", ["script", "--host", "192.168.1.1", "--port", "7000"] - ): - run_server() - - call_kwargs = mock_run.call_args[1] - assert call_kwargs["host"] == "192.168.1.1" - assert call_kwargs["port"] == 7000 + mock_server = MagicMock() + with patch("simplevecdb.embeddings.server.uvicorn.Config") as mock_cfg: + with patch("simplevecdb.embeddings.server.uvicorn.Server", return_value=mock_server): + with patch.object( + sys, "argv", ["script", "--host", "192.168.1.1", "--port", "7000"] + ): + run_server() + + call_kwargs = mock_cfg.call_args[1] + assert call_kwargs["host"] == "192.168.1.1" + assert call_kwargs["port"] == 7000 def test_embeddings_default_model(): diff --git a/tests/unit/embeddings/test_v25_enhancements.py b/tests/unit/embeddings/test_v25_enhancements.py new file mode 100644 index 0000000..bb57611 --- /dev/null +++ b/tests/unit/embeddings/test_v25_enhancements.py @@ -0,0 +1,381 @@ +"""Tests for simplevecdb 2.5.0 embeddings server enhancements (#1-#10).""" + +import argparse +import signal + +import pytest +from unittest.mock import patch, MagicMock, ANY, call + +from fastapi.testclient import TestClient + +import simplevecdb +from simplevecdb.embeddings.server import ( + app, + _normalize_input, + _validate_texts, + _build_cli_parser, + _server_version, + ModelRegistry, +) +from simplevecdb.embeddings import server + +client = TestClient(app) + + +@pytest.fixture(autouse=True) +def _unlocked_registry(): + """Allow arbitrary test models in unit tests.""" + original = server.registry + server.registry = ModelRegistry({"default": "test-default"}, allow_unlisted=True) + yield + server.registry = original + + +# --------------------------------------------------------------------------- +# 1. Graceful shutdown +# --------------------------------------------------------------------------- + + +class TestGracefulShutdown: + """run_server() creates uvicorn.Server with timeout_graceful_shutdown=10.""" + + @patch("simplevecdb.embeddings.server.get_embedder") + @patch("simplevecdb.embeddings.server.uvicorn.Server") + @patch("simplevecdb.embeddings.server.uvicorn.Config") + def test_timeout_graceful_shutdown_is_set( + self, mock_config_cls, mock_server_cls, mock_get_embedder + ): + mock_server = MagicMock() + mock_server_cls.return_value = mock_server + + from simplevecdb.embeddings.server import run_server + + run_server(host="127.0.0.1", port=9000) + + mock_config_cls.assert_called_once_with( + app, + host="127.0.0.1", + port=9000, + log_level="info", + timeout_graceful_shutdown=10, + ) + + @patch("simplevecdb.embeddings.server.get_embedder") + @patch("simplevecdb.embeddings.server.uvicorn.Server") + @patch("simplevecdb.embeddings.server.uvicorn.Config") + def test_signal_handlers_installed( + self, mock_config_cls, mock_server_cls, mock_get_embedder + ): + mock_server = MagicMock() + mock_server_cls.return_value = mock_server + + captured_handlers: dict[int, Any] = {} + original_signal = signal.signal + + def fake_signal(signum, handler): + captured_handlers[signum] = handler + return original_signal(signum, signal.SIG_DFL) + + with patch("simplevecdb.embeddings.server.signal.signal", side_effect=fake_signal): + from simplevecdb.embeddings.server import run_server + + run_server(host="127.0.0.1", port=9000) + + assert signal.SIGINT in captured_handlers + assert signal.SIGTERM in captured_handlers + + +# --------------------------------------------------------------------------- +# 2. Async executor offload +# --------------------------------------------------------------------------- + + +class TestAsyncExecutorOffload: + """create_embeddings runs embed_texts via loop.run_in_executor.""" + + @patch("simplevecdb.embeddings.server.embed_texts") + def test_embed_texts_called_via_endpoint(self, mock_embed): + mock_embed.return_value = [[0.1, 0.2, 0.3]] + + response = client.post( + "/v1/embeddings", json={"input": "test text", "model": "test-model"} + ) + + assert response.status_code == 200 + mock_embed.assert_called_once_with( + ["test text"], model_id="test-model", batch_size=ANY + ) + data = response.json() + assert data["data"][0]["embedding"] == [0.1, 0.2, 0.3] + + +# --------------------------------------------------------------------------- +# 3. Model warm-up +# --------------------------------------------------------------------------- + + +class TestModelWarmUp: + """run_server() calls get_embedder before starting unless --no-warmup.""" + + @patch("simplevecdb.embeddings.server.get_embedder") + @patch("simplevecdb.embeddings.server.uvicorn.Server") + @patch("simplevecdb.embeddings.server.uvicorn.Config") + def test_warmup_calls_get_embedder( + self, mock_config_cls, mock_server_cls, mock_get_embedder + ): + mock_server_cls.return_value = MagicMock() + + from simplevecdb.embeddings.server import run_server + from simplevecdb.config import config + + run_server(host="127.0.0.1", port=9000) + + mock_get_embedder.assert_called_once_with(config.EMBEDDING_MODEL) + + @patch("simplevecdb.embeddings.server.get_embedder") + @patch("simplevecdb.embeddings.server.uvicorn.Server") + @patch("simplevecdb.embeddings.server.uvicorn.Config") + def test_no_warmup_skips_get_embedder( + self, mock_config_cls, mock_server_cls, mock_get_embedder + ): + """Simulate --no-warmup by calling run_server via CLI path.""" + mock_server_cls.return_value = MagicMock() + + with patch( + "simplevecdb.embeddings.server._build_cli_parser" + ) as mock_parser_fn: + mock_args = argparse.Namespace(host=None, port=None, no_warmup=True) + mock_parser = MagicMock() + mock_parser.parse_args.return_value = mock_args + mock_parser_fn.return_value = mock_parser + + from simplevecdb.embeddings.server import run_server + + # host=None, port=None triggers CLI parsing path + run_server(host=None, port=None) + + mock_get_embedder.assert_not_called() + + +# --------------------------------------------------------------------------- +# 4. CORS middleware +# --------------------------------------------------------------------------- + + +class TestCORSMiddleware: + """App has CORSMiddleware allowing cross-origin requests.""" + + def test_options_preflight_returns_cors_headers(self): + response = client.options( + "/v1/embeddings", + headers={ + "Origin": "http://example.com", + "Access-Control-Request-Method": "POST", + }, + ) + assert "access-control-allow-origin" in response.headers + + def test_cors_allow_origin_value(self): + response = client.options( + "/v1/embeddings", + headers={ + "Origin": "http://example.com", + "Access-Control-Request-Method": "POST", + }, + ) + assert response.headers["access-control-allow-origin"] == "http://example.com" + + +# --------------------------------------------------------------------------- +# 5. Input validation +# --------------------------------------------------------------------------- + + +class TestInputValidation: + """_validate_texts rejects empty strings (422) and texts > 100k chars (413).""" + + @patch("simplevecdb.embeddings.server.embed_texts") + def test_empty_string_returns_422(self, mock_embed): + response = client.post( + "/v1/embeddings", json={"input": "", "model": "test-model"} + ) + assert response.status_code == 422 + + @patch("simplevecdb.embeddings.server.embed_texts") + def test_whitespace_only_returns_422(self, mock_embed): + response = client.post( + "/v1/embeddings", json={"input": " ", "model": "test-model"} + ) + assert response.status_code == 422 + + @patch("simplevecdb.embeddings.server.embed_texts") + def test_list_with_empty_string_returns_422(self, mock_embed): + response = client.post( + "/v1/embeddings", + json={"input": ["hello", ""], "model": "test-model"}, + ) + assert response.status_code == 422 + + @patch("simplevecdb.embeddings.server.embed_texts") + def test_text_exceeding_100k_chars_returns_413(self, mock_embed): + long_text = "a" * 100_001 + response = client.post( + "/v1/embeddings", json={"input": long_text, "model": "test-model"} + ) + assert response.status_code == 413 + assert "100000" in response.json()["detail"] + + @patch("simplevecdb.embeddings.server.embed_texts") + def test_text_at_100k_chars_succeeds(self, mock_embed): + mock_embed.return_value = [[0.1]] + text = "a" * 100_000 + response = client.post( + "/v1/embeddings", json={"input": text, "model": "test-model"} + ) + assert response.status_code == 200 + + +# --------------------------------------------------------------------------- +# 6. argparse CLI +# --------------------------------------------------------------------------- + + +class TestArgparseCLI: + """_build_cli_parser() returns a parser with --host, --port, --no-warmup.""" + + def test_parse_known_args_defaults(self): + parser = _build_cli_parser() + args = parser.parse_args([]) + assert args.host is None + assert args.port is None + assert args.no_warmup is False + + def test_parse_explicit_values(self): + parser = _build_cli_parser() + args = parser.parse_args(["--host", "0.0.0.0", "--port", "8080", "--no-warmup"]) + assert args.host == "0.0.0.0" + assert args.port == 8080 + assert args.no_warmup is True + + def test_help_does_not_crash(self): + parser = _build_cli_parser() + with pytest.raises(SystemExit) as exc_info: + parser.parse_args(["--help"]) + assert exc_info.value.code == 0 + + +# --------------------------------------------------------------------------- +# 7. Startup banner +# --------------------------------------------------------------------------- + + +class TestStartupBanner: + """run_server() logs a banner with host/port/model/auth/version.""" + + @patch("simplevecdb.embeddings.server.get_embedder") + @patch("simplevecdb.embeddings.server.uvicorn.Server") + @patch("simplevecdb.embeddings.server.uvicorn.Config") + @patch("simplevecdb.embeddings.server._logger") + def test_banner_logged_with_host_and_port( + self, mock_logger, mock_config_cls, mock_server_cls, mock_get_embedder + ): + mock_server_cls.return_value = MagicMock() + + from simplevecdb.embeddings.server import run_server + + run_server(host="127.0.0.1", port=9000) + + info_calls = mock_logger.info.call_args_list + # The first info call should be the banner + banner_call = info_calls[0] + banner_args = banner_call[0] # positional args + banner_template = banner_args[0] + + assert "SimpleVecDB" in banner_template + # host and port are passed as format args + assert "127.0.0.1" in banner_args + assert 9000 in banner_args + + +# --------------------------------------------------------------------------- +# 8. Nested token arrays +# --------------------------------------------------------------------------- + + +class TestNormalizeInput: + """_normalize_input handles all OpenAI-compatible input formats.""" + + def test_string_input(self): + assert _normalize_input("hello world") == ["hello world"] + + def test_list_of_strings(self): + assert _normalize_input(["a", "b", "c"]) == ["a", "b", "c"] + + def test_list_of_ints(self): + assert _normalize_input([1, 2, 3]) == ["1", "2", "3"] + + def test_nested_token_arrays(self): + result = _normalize_input([[1, 2, 3], [4, 5]]) + assert result == ["1 2 3", "4 5"] + + def test_empty_list(self): + assert _normalize_input([]) == [] + + def test_single_nested_token_array(self): + result = _normalize_input([[10, 20, 30]]) + assert result == ["10 20 30"] + + +# --------------------------------------------------------------------------- +# 9. OpenAPI version +# --------------------------------------------------------------------------- + + +class TestOpenAPIVersion: + """app.version should match the package version.""" + + def test_app_version_matches_package(self): + assert app.version == simplevecdb.__version__ + + def test_server_version_matches_package(self): + assert _server_version == simplevecdb.__version__ + + +# --------------------------------------------------------------------------- +# 10. Module exports +# --------------------------------------------------------------------------- + + +class TestModuleExports: + """simplevecdb.embeddings exposes all expected public symbols.""" + + def test_embed_texts_importable(self): + from simplevecdb.embeddings import embed_texts + + assert callable(embed_texts) + + def test_get_embedder_importable(self): + from simplevecdb.embeddings import get_embedder + + assert callable(get_embedder) + + def test_load_model_importable(self): + from simplevecdb.embeddings import load_model + + assert callable(load_model) + + def test_app_importable(self): + from simplevecdb.embeddings import app as imported_app + + assert imported_app is app + + def test_run_server_importable(self): + from simplevecdb.embeddings import run_server + + assert callable(run_server) + + def test_all_exports_listed(self): + from simplevecdb.embeddings import __all__ + + expected = {"app", "embed_texts", "get_embedder", "load_model", "run_server"} + assert set(__all__) == expected From d4da80707a82ab6c8a6eab5acb534e898868c4ae Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Tue, 7 Apr 2026 05:06:52 -0500 Subject: [PATCH 3/5] docs: update changelog and readme for 2.5.0 release - Add embeddings server enhancements to changelog - Update README: pagination examples, delete_collection usage, embeddings server CLI flags, feature matrix, roadmap --- CHANGELOG.md | 11 +++++++++++ README.md | 21 +++++++++++++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a1cb75..2654d5b 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **`__repr__`** on `VectorDB`, `VectorCollection`, `AsyncVectorDB`, `AsyncVectorCollection` for debuggable string representations. - **FLOAT16 quantization** fully implemented in `serialize()`/`deserialize()` — was previously defined in the enum but raised `ValueError` at runtime. - **Pagination** on `get_documents(limit=, offset=)` and catalog methods (`find_ids_by_filter`, `find_ids_by_texts`) — previously returned unbounded result sets. +- **Embeddings server enhancements:** + - Graceful shutdown with SIGTERM/SIGINT draining (10s timeout) + - CORS middleware with configurable origins for browser-based clients + - Model warm-up on startup (skip with `--no-warmup`) + - Input validation: rejects empty strings (422) and texts exceeding 100k chars (413) + - Proper `argparse` CLI with `--host`, `--port`, `--no-warmup`, `--help` + - Startup banner logging config summary (host, port, model, auth, rate limits) + - Nested token array normalization (`list[list[int]]` input format) + - Async executor offload for `embed_texts` (non-blocking event loop) + - OpenAPI version synced from package metadata + - Module `__init__.py` exports (`embed_texts`, `get_embedder`, `load_model`, `app`, `run_server`) ### Fixed diff --git a/README.md b/README.md index ff43d5c..43a714d 100755 --- a/README.md +++ b/README.md @@ -140,10 +140,13 @@ hybrid = collection.hybrid_search("powerhouse cell", k=2) **Optional: Run embeddings server (OpenAI-compatible)** ```bash -simplevecdb-server --port 8000 +simplevecdb-server --port 8000 # Default model, auto warm-up +simplevecdb-server --host 0.0.0.0 --port 9000 # Bind to all interfaces +simplevecdb-server --no-warmup # Skip model preload on startup +simplevecdb-server --help # Show all options ``` -See [Setup Guide](ENV_SETUP.md) for configuration: model registry, rate limits, API keys, CUDA optimization. +See [Setup Guide](ENV_SETUP.md) for configuration: model registry, rate limits, API keys, CORS, CUDA optimization. ### Option 3: With LangChain or LlamaIndex @@ -302,6 +305,10 @@ docs = collection.get_documents(filter_dict={"category": "tech"}) for doc_id, text, metadata in docs: print(f"[{doc_id}] {text[:50]}...") +# Paginated access (v2.5+) +page1 = collection.get_documents(limit=100) +page2 = collection.get_documents(limit=100, offset=100) + # Fetch stored embeddings embeddings = collection.get_embeddings_by_ids([1, 2, 3]) @@ -313,6 +320,9 @@ collection.update_metadata([ # Quick stats print(f"Collection has {collection.count()} documents, dim={collection.dim}") + +# Delete an entire collection (v2.5+) +db.delete_collection("old_data") ``` ### Vector Clustering (v2.2+) @@ -355,6 +365,10 @@ Supports K-means, MiniBatch K-means, and HDBSCAN. See [Clustering Guide](https:/ | **Cluster Persistence** | ✅ | Save/load cluster centroids for fast assignment (v2.2+) | | **Public Catalog API** | ✅ | `get_documents`, `get_embeddings_by_ids`, `update_metadata` (v2.4+) | | **Executor Injection** | ✅ | Share thread pool across async instances for ONNX safety (v2.4+) | +| **Collection Management** | ✅ | `delete_collection()`, paginated `get_documents(limit=, offset=)` (v2.5+) | +| **Cross-Process Safety** | ✅ | Advisory file locking on usearch index files (v2.5+) | +| **FLOAT16 Quantization** | ✅ | Half-precision storage with 2x compression (v2.5+) | +| **Embeddings Server** | ✅ | CORS, graceful shutdown, input validation, model warm-up (v2.5+) | ## Performance Benchmarks @@ -427,6 +441,9 @@ pip install torch --index-url https://download.pytorch.org/whl/cu118 - [x] Vector clustering and auto-tagging (v2.2) - [x] Public catalog API for document management (v2.4) - [x] Async executor injection for thread-safe sharing (v2.4) +- [x] Collection management: `delete_collection()`, pagination (v2.5) +- [x] Cross-process file locking and connection health checks (v2.5) +- [x] Embeddings server hardening: CORS, graceful shutdown, input validation (v2.5) - [ ] Incremental clustering (online learning) - [ ] Cluster visualization exports From b419122c69b15e7c8b2a9b2e6e09cc951bb65631 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Tue, 7 Apr 2026 05:23:00 -0500 Subject: [PATCH 4/5] fix: address review findings before 2.5.0 release - list_collections: use set-based derivative filtering to correctly handle collections named "test_fts" or "my_clusters" (was using substring match that silently dropped them) - AsyncVectorDB.delete_collection: evict from async-level _collections cache (was only clearing sync cache, leaving stale wrappers) - Pagination: raise ValueError("offset requires limit") instead of silently ignoring offset when limit is None - Revert platform.processor() back to subprocess sysctl for Apple chip detection (platform.processor() returns 'arm' not chip brand strings) - delete_collection: validate name against COLLECTION_NAME_PATTERN before existence check - _normalize_input: flat token array (list[int]) now produces one embedding input per OpenAI spec, not N separate inputs 544 tests passing. --- src/simplevecdb/async_core.py | 5 ++ src/simplevecdb/core.py | 60 ++++++++++++++----- src/simplevecdb/embeddings/server.py | 4 +- src/simplevecdb/engine/catalog.py | 6 ++ tests/unit/core/test_batch_detection.py | 12 ++-- .../unit/embeddings/test_v25_enhancements.py | 3 +- 6 files changed, 66 insertions(+), 24 deletions(-) diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index 76f7ef7..c724924 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -616,6 +616,11 @@ async def delete_collection(self, name: str) -> None: await loop.run_in_executor( self._executor, lambda: self._db.delete_collection(name) ) + # Evict from async-level cache too + with self._collections_lock: + keys_to_remove = [k for k in self._collections if k[0] == name] + for k in keys_to_remove: + del self._collections[k] async def search_collections( self, diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 9ae2296..8706169 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -90,13 +90,20 @@ def get_optimal_batch_size() -> int: if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): machine = platform.machine().lower() if "arm" in machine or "aarch64" in machine: - chip_info = platform.processor().lower() - - if "m3" in chip_info or "m4" in chip_info: - return constants.DEFAULT_APPLE_M3_M4_BATCH_SIZE - elif "max" in chip_info or "ultra" in chip_info: - return constants.DEFAULT_APPLE_MAX_ULTRA_BATCH_SIZE - else: + try: + import subprocess + + chip_info = subprocess.check_output( + ["sysctl", "-n", "machdep.cpu.brand_string"], text=True + ).lower() + + if "m3" in chip_info or "m4" in chip_info: + return constants.DEFAULT_APPLE_M3_M4_BATCH_SIZE + elif "max" in chip_info or "ultra" in chip_info: + return constants.DEFAULT_APPLE_MAX_ULTRA_BATCH_SIZE + else: + return constants.DEFAULT_APPLE_M1_M2_BATCH_SIZE + except Exception: return constants.DEFAULT_APPLE_M1_M2_BATCH_SIZE # 2. Try ONNX Runtime detection @@ -1531,17 +1538,36 @@ def list_collections(self) -> list[str]: "SELECT name FROM sqlite_master WHERE type='table' " "AND (name = 'tinyvec_items' OR name LIKE 'items_%')" ).fetchall() - names: list[str] = [] + # Collect all table names, then filter out FTS/cluster derivatives. + # FTS5 creates shadow tables: items__fts, items__fts_data, + # items__fts_idx, items__fts_content, items__fts_docsize, + # items__fts_config. Cluster tables: items__clusters. + # We identify derivatives by checking if a suffix is _fts* + # or _clusters for some other known collection suffix. + all_suffixes: set[str] = set() + has_default = False for (table_name,) in rows: if table_name == "tinyvec_items": - names.append("default") + has_default = True elif table_name.startswith("items_"): - # Skip FTS (and FTS sub-tables) and cluster tables - suffix = table_name[6:] - if "_fts" in suffix or suffix.endswith("_clusters"): - continue - names.append(suffix) - return sorted(names) + all_suffixes.add(table_name[6:]) + + # A suffix is a real collection if no other suffix is a prefix of it + # followed by _fts* or _clusters. + _fts_suffixes = ("_fts", "_fts_data", "_fts_idx", "_fts_content", + "_fts_docsize", "_fts_config") + derivative_suffixes: set[str] = set() + for s in all_suffixes: + for fts in _fts_suffixes: + derivative_suffixes.add(f"{s}{fts}") + derivative_suffixes.add(f"{s}_clusters") + + names: list[str] = [] + if has_default: + names.append("default") + for s in sorted(all_suffixes - derivative_suffixes): + names.append(s) + return names def delete_collection(self, name: str) -> None: """ @@ -1557,6 +1583,10 @@ def delete_collection(self, name: str) -> None: ValueError: If the collection name is invalid. KeyError: If the collection does not exist. """ + if not re.match(constants.COLLECTION_NAME_PATTERN, name): + raise ValueError( + f"Invalid collection name '{name}'. Must be alphanumeric + underscores." + ) if name not in self.list_collections(): raise KeyError(f"Collection '{name}' does not exist.") diff --git a/src/simplevecdb/embeddings/server.py b/src/simplevecdb/embeddings/server.py index c3a5de1..ae7618f 100755 --- a/src/simplevecdb/embeddings/server.py +++ b/src/simplevecdb/embeddings/server.py @@ -270,9 +270,9 @@ def _normalize_input(raw_input: str | list[str] | list[int] | list[list[int]]) - first = raw_input[0] - # list[int] — flat token array + # list[int] — flat token array (single input per OpenAI spec) if isinstance(first, int): - return [str(i) for i in raw_input] + return [" ".join(str(i) for i in raw_input)] # list[list[int]] — nested token arrays (#8) if isinstance(first, list): diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index b479174..31c27be 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -441,6 +441,8 @@ def find_ids_by_texts( sql = f"SELECT id FROM {self._table_name} WHERE text IN ({placeholders})" params: list[Any] = list(texts) + if offset is not None and limit is None: + raise ValueError("offset requires limit") if limit is not None: sql += " LIMIT ?" params.append(limit) @@ -478,6 +480,8 @@ def find_ids_by_filter( sql = f"SELECT id FROM {self._table_name} {where_clause}" params: list[Any] = list(filter_params) + if offset is not None and limit is None: + raise ValueError("offset requires limit") if limit is not None: sql += " LIMIT ?" params.append(limit) @@ -612,6 +616,8 @@ def get_all_docs_with_text( """ params: list[Any] = list(filter_params) + if offset is not None and limit is None: + raise ValueError("offset requires limit") if limit is not None: sql += " LIMIT ?" params.append(limit) diff --git a/tests/unit/core/test_batch_detection.py b/tests/unit/core/test_batch_detection.py index 8e88845..431e637 100755 --- a/tests/unit/core/test_batch_detection.py +++ b/tests/unit/core/test_batch_detection.py @@ -207,23 +207,23 @@ def test_get_optimal_batch_size_mps_branches(): with patch.dict(sys.modules, {"torch": mock_torch}): with patch.object(platform, "machine", return_value="arm64"): # M3/M4 chips - with patch.object(platform, "processor", return_value="apple m3"): + with patch("subprocess.check_output", return_value="apple m3"): assert get_optimal_batch_size() == 64 # Max chips - with patch.object(platform, "processor", return_value="apple m2 max"): + with patch("subprocess.check_output", return_value="apple m2 max"): assert get_optimal_batch_size() == 128 # Ultra chips - with patch.object(platform, "processor", return_value="apple m2 ultra"): + with patch("subprocess.check_output", return_value="apple m2 ultra"): assert get_optimal_batch_size() == 128 # Base M1/M2 (no M3/M4 and no Max/Ultra) - with patch.object(platform, "processor", return_value="apple m1"): + with patch("subprocess.check_output", return_value="apple m1"): assert get_optimal_batch_size() == 32 - # Empty processor string fallback - with patch.object(platform, "processor", return_value=""): + # Exception fallback + with patch("subprocess.check_output", side_effect=Exception("Fail")): assert get_optimal_batch_size() == 32 diff --git a/tests/unit/embeddings/test_v25_enhancements.py b/tests/unit/embeddings/test_v25_enhancements.py index bb57611..6071e5f 100644 --- a/tests/unit/embeddings/test_v25_enhancements.py +++ b/tests/unit/embeddings/test_v25_enhancements.py @@ -312,7 +312,8 @@ def test_list_of_strings(self): assert _normalize_input(["a", "b", "c"]) == ["a", "b", "c"] def test_list_of_ints(self): - assert _normalize_input([1, 2, 3]) == ["1", "2", "3"] + # Flat token array is a single input per OpenAI spec + assert _normalize_input([1, 2, 3]) == ["1 2 3"] def test_nested_token_arrays(self): result = _normalize_input([[1, 2, 3], [4, 5]]) From fc883967f3913c948de41bd84a8ce063b6b526d6 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Tue, 7 Apr 2026 05:29:23 -0500 Subject: [PATCH 5/5] fix: update integration test for uvicorn.Server API change --- tests/integration/test_server.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/integration/test_server.py b/tests/integration/test_server.py index d466f2c..92db9dc 100755 --- a/tests/integration/test_server.py +++ b/tests/integration/test_server.py @@ -1,6 +1,6 @@ import pytest from fastapi.testclient import TestClient -from unittest.mock import patch +from unittest.mock import patch, MagicMock from simplevecdb.config import config from simplevecdb.embeddings import server @@ -85,12 +85,14 @@ def test_usage_endpoint_reports_stats(): @pytest.mark.integration def test_run_server(): - """Test run_server function calls uvicorn.""" + """Test run_server function calls uvicorn.Server.""" from simplevecdb.embeddings.server import run_server - with patch("uvicorn.run") as mock_run: - run_server(host="1.2.3.4", port=9999) - mock_run.assert_called_once() - args, kwargs = mock_run.call_args - assert kwargs["host"] == "1.2.3.4" - assert kwargs["port"] == 9999 + mock_server = MagicMock() + with patch("simplevecdb.embeddings.server.uvicorn.Config") as mock_cfg: + with patch("simplevecdb.embeddings.server.uvicorn.Server", return_value=mock_server): + run_server(host="1.2.3.4", port=9999) + call_kwargs = mock_cfg.call_args[1] + assert call_kwargs["host"] == "1.2.3.4" + assert call_kwargs["port"] == 9999 + mock_server.run.assert_called_once()