diff --git a/.gitignore b/.gitignore index 89ca3aa..9ea06ef 100755 --- a/.gitignore +++ b/.gitignore @@ -68,3 +68,7 @@ simplevecdb_plan.md AGENTS.md NEXT_UPDATES.md pro_pack/ + + +# Local notes +IMPORTANT.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b0e503..f5d49e9 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,112 @@ 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.6.2] - 2026-06-06 + +### Correctness and contract fixes + +Hardening of the index-rebuild, search, clustering, and integration layers +surfaced by a code review. Two intentional behavior changes are noted under +“Changed”. + +#### Fixed + +- **`rebuild_index` no longer bricks a collection on failure** — if building or + swapping the new HNSW index raises after the live index is closed, the + collection re-opens the intact on-disk index instead of holding a closed one. +- **Catalog write lock released on connection error** — a raising + `connection.__enter__` no longer leaks the catalog lock (which could deadlock + the database). +- **Max-Marginal-Relevance respects the distance metric** — MMR on `l2` + collections used a cosine-specific relevance formula that swamped the + diversity term; it now uses a bounded, metric-appropriate relevance. +- **`similarity_search_batch` fills `k` under filters and accepts text queries** + — large filtered batches no longer silently under-deliver, and a text query in + a large batch behaves the same as in a small one. +- **Clustering handles impossible `n_clusters`** — `ClusterEngine.cluster_vectors` + raises a clear error when `n_clusters` exceeds the number of vectors; + `Collection.cluster()` caps `n_clusters` to the number of vectors actually + clustered (the sample when `sample_size` is set, fixing a latent error when + `n_clusters > sample_size`). +- **Metadata filter keys match literally** — a filter key containing a dot + (e.g. `{"a.b": x}`) now matches the literal top-level key `a.b` instead of the + nested JSON path `a → b`, consistent with the Python filter path. Keys + containing a double-quote are rejected. +- **BIT-quantized vector retrieval unpacks correctly** — `UsearchIndex.get()` + (used by the MMR fallback) returned packed bytes for BIT indexes instead of + the unpacked ±1 float vectors; it now unpacks them. +- **`rebuild_index` no longer blocks the database during the HNSW build** — the + expensive build runs without the shared lock (held only to snapshot and swap); + writes that land during the build are folded into the new index before the + swap. +- **Embedding server caps request body size** — an ASGI middleware rejects + request bodies larger than the server's own accept limits before they are + buffered/parsed, closing an unauthenticated memory-exhaustion vector (only + relevant with the `[server]` extra exposed on a network). A missing + encryption salt sidecar now logs a warning instead of silently falling back + to the shared legacy salt. +- **Robustness pass** — malformed FTS5 keyword queries raise `ValueError` instead + of a raw SQLite error; the cluster-state table is created eagerly so a + rolled-back first `save_cluster` cannot desync it; a non-integer + `EMBEDDING_BATCH_SIZE`/`EMBEDDING_SERVER_MAX_REQUEST_ITEMS` env value warns and + falls back instead of crashing import; `vacuum()` holds the DB lock; a failed + index add after the catalog commit is logged (divergence visibility); hybrid + search applies the Python metadata filter on the keyword side too (SQL/Python + parity); `logging.configure_logging` swaps handlers atomically. +- **LangChain `asimilarity_search_with_score`** offloads to a thread instead of + blocking the event loop. + +#### Changed + +- **`AsyncVectorCollection.increment_metadata` now returns `int`** (1 if the row + existed and was updated, 0 otherwise), matching the synchronous API; it + previously discarded the value and returned `None`. +- **LlamaIndex metadata filters fail loudly on unsupported shapes** — the + `SimpleVecDBLlamaStore` adapter now maps comparison operators + (`$gt/$gte/$lt/$lte/$ne/$in/$nin`) instead of silently treating them as + equality, and raises `NotImplementedError` for `OR`/`NOT` conditions and + unsupported operators rather than returning wrong results. +- **LangChain relevance scoring now works** — `SimpleVecDBVectorStore` implements + `_select_relevance_score_fn`, so `similarity_search_with_relevance_scores` and + `as_retriever(search_type="similarity_score_threshold")` return metric-aware + `[0, 1]` relevance (higher = better). `similarity_search_with_score` still + returns the raw distance (FAISS/Chroma convention), now documented as such. + +### Clustering and hierarchy fixes + +Internal correctness and performance work on the clustering and hierarchy +layers. No public API changes; existing databases are unaffected. + +#### Fixed + +- **`load_cluster` survives empty k-means clusters** — when k-means leaves a + requested cluster empty (common with duplicate vectors or `n_clusters` near + the number of distinct points), the stored `n_clusters` is smaller than the + number of centroid rows. The centroid reshape now derives its row count from + the stored buffer rather than `n_clusters`, which previously raised + `ValueError` on load. +- **`assign_to_cluster` matches metadata keys literally** — a `metadata_key` + containing `.` or `[` is now matched as a literal top-level key (via + `json_each`) instead of being misread as a nested JSON path, which had caused + every already-assigned document to be re-assigned on each call. +- **`cluster(algorithm="hdbscan", sample_size=…)` raises instead of silently + dropping documents** — HDBSCAN produces no centroids, so out-of-sample + documents cannot be assigned. The combination now raises a clear `ValueError` + rather than clustering only the sample. + +#### Performance + +- **BLAS-backed out-of-sample centroid assignment** — nearest-centroid + assignment uses the `‖c‖² − 2·x·c` expansion (a single matmul) instead of + materialising the dense `(n_vectors, n_centroids, dim)` broadcast temporary + that could exhaust memory on large collections. +- **Unassigned-id lookup pushed into SQLite** — `assign_to_cluster(doc_ids=None)` + finds documents lacking the cluster key with one `json_each` query instead of + loading and JSON-parsing every row's text and metadata. +- **Bounded ancestor-walk for cycle detection** — `set_parent` detects + parent/child cycles by walking the ancestor chain with a depth-bounded + recursive CTE instead of materialising the entire descendant subtree. + ## [2.6.1] - 2026-05-10 ### Storage, mutation, and eventing improvements diff --git a/README.md b/README.md index 3ae1874..b7e85fc 100755 --- a/README.md +++ b/README.md @@ -7,17 +7,17 @@ Buy Me a Coffee at ko-fi.com -**The dead-simple, local-first vector database.** +**A local-first, embedded vector database backed by SQLite and usearch.** -SimpleVecDB brings **Chroma-like simplicity** to a single **SQLite file**. Built on `usearch` HNSW indexing, it offers high-performance vector search, quantization, and zero infrastructure headaches. Perfect for local RAG, offline agents, and indie hackers who need production-grade vector search without the operational overhead. +SimpleVecDB pairs **Chroma-like ergonomics** with a **file-based** store — a SQLite database for metadata and text alongside a `usearch` HNSW index per collection. It provides high-performance vector search, quantization, and hybrid retrieval with no separate services to run. It fits local RAG pipelines, offline agents, and any application that needs production-grade vector search without the operational overhead of a hosted database. ## Why SimpleVecDB? -- **Zero Infrastructure** — Just a `.db` file. No Docker, no Redis, no cloud bills. -- **Blazing Fast** — 10-100x faster search via usearch HNSW. Adaptive: brute-force for <10k vectors (perfect recall), HNSW for larger collections. -- **Truly Portable** — Runs anywhere SQLite runs: Linux, macOS, Windows, even WASM. -- **Async Ready** — Full async/await support with optional executor injection for thread-safe ONNX/usearch sharing. -- **Batteries Included** — Optional FastAPI embeddings server + LangChain/LlamaIndex integrations via `[integrations]` extra. +- **Zero Infrastructure** — Local files on disk: a SQLite database plus a `usearch` index. No Docker, no Redis, no external services. +- **High Performance** — usearch HNSW indexing with adaptive search: brute-force under 10k vectors (perfect recall), HNSW above that. +- **Portable** — Runs anywhere SQLite runs: Linux, macOS, Windows, and WASM. +- **Async Support** — A complete async/await surface with optional executor injection for thread-safe ONNX/usearch sharing. +- **Integrations Included** — Optional FastAPI embeddings server and LangChain/LlamaIndex adapters via the `[integrations]` extra. - **Production Ready** — Hybrid search (BM25 + vector), metadata filtering, multi-collection support, and automatic hardware acceleration. ### When to Choose SimpleVecDB @@ -67,7 +67,7 @@ pip install "simplevecdb[server]" **Verify Installation:** ```bash -python -c "from simplevecdb import VectorDB; print('SimpleVecDB installed successfully!')" +python -c "import simplevecdb; print(simplevecdb.__version__)" ``` ## Quickstart @@ -157,8 +157,8 @@ in the [Setup Guide](ENV_SETUP.md). ### Option 3: LangChain or LlamaIndex -Already wired into one of the big RAG frameworks? Drop SimpleVecDB in -as the vector store: +Already using one of the major RAG frameworks? Use SimpleVecDB as the +vector store: ```bash pip install "simplevecdb[integrations]" @@ -208,7 +208,7 @@ A few of the things SimpleVecDB does well — see - **Quantization** — `FLOAT32`, `FLOAT16`, `INT8`, `BIT` for 1×–32× compression. - **Multi-collection + cross-collection search** — isolated namespaces in - one `.db` file, with merged ranked search across them. + one database, with merged ranked search across them. - **Mongo-style filters** — `$eq $ne $gt $gte $lt $lte $in $nin $exists $between` on metadata, edges, and events. - **Memory primitives (v2.6.1)** — pending-vector buffer with atomic @@ -300,7 +300,7 @@ ideas in [GitHub Discussions](https://github.com/coderdayton/simplevecdb/discuss ## Contributing -Contributions are welcome! Whether you're fixing bugs, improving documentation, or proposing new features: +Contributions are welcome — bug fixes, documentation improvements, and new feature proposals alike: 1. Read [CONTRIBUTING.md](CONTRIBUTING.md) for development setup 2. Check existing [Issues](https://github.com/coderdayton/simplevecdb/issues) and [Discussions](https://github.com/coderdayton/simplevecdb/discussions) @@ -318,12 +318,12 @@ Contributions are welcome! Whether you're fixing bugs, improving documentation, - [GitHub Releases](https://github.com/coderdayton/simplevecdb/releases) — Changelog and updates - [Examples Gallery](https://coderdayton.github.io/SimpleVecDB/examples/) — Community-contributed notebooks -## Other Ways to Support +## Support the Project -- ☕ **[Buy me a coffee](https://ko-fi.com/xbbvii)** - One-time donation -- ⭐ **Star the repo** - Helps with visibility -- 🐛 **Report bugs** - Improve the project for everyone -- 📝 **Contribute** - See [CONTRIBUTING.md](CONTRIBUTING.md) +- **[Sponsor on Ko-fi](https://ko-fi.com/xbbvii)** — one-time donations +- **Star the repository** — helps with visibility +- **[Report issues](https://github.com/coderdayton/simplevecdb/issues)** — bug reports and feedback +- **[Contribute](CONTRIBUTING.md)** — development setup and guidelines ## License diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 9b0e503..f5d49e9 100755 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -5,6 +5,112 @@ 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.6.2] - 2026-06-06 + +### Correctness and contract fixes + +Hardening of the index-rebuild, search, clustering, and integration layers +surfaced by a code review. Two intentional behavior changes are noted under +“Changed”. + +#### Fixed + +- **`rebuild_index` no longer bricks a collection on failure** — if building or + swapping the new HNSW index raises after the live index is closed, the + collection re-opens the intact on-disk index instead of holding a closed one. +- **Catalog write lock released on connection error** — a raising + `connection.__enter__` no longer leaks the catalog lock (which could deadlock + the database). +- **Max-Marginal-Relevance respects the distance metric** — MMR on `l2` + collections used a cosine-specific relevance formula that swamped the + diversity term; it now uses a bounded, metric-appropriate relevance. +- **`similarity_search_batch` fills `k` under filters and accepts text queries** + — large filtered batches no longer silently under-deliver, and a text query in + a large batch behaves the same as in a small one. +- **Clustering handles impossible `n_clusters`** — `ClusterEngine.cluster_vectors` + raises a clear error when `n_clusters` exceeds the number of vectors; + `Collection.cluster()` caps `n_clusters` to the number of vectors actually + clustered (the sample when `sample_size` is set, fixing a latent error when + `n_clusters > sample_size`). +- **Metadata filter keys match literally** — a filter key containing a dot + (e.g. `{"a.b": x}`) now matches the literal top-level key `a.b` instead of the + nested JSON path `a → b`, consistent with the Python filter path. Keys + containing a double-quote are rejected. +- **BIT-quantized vector retrieval unpacks correctly** — `UsearchIndex.get()` + (used by the MMR fallback) returned packed bytes for BIT indexes instead of + the unpacked ±1 float vectors; it now unpacks them. +- **`rebuild_index` no longer blocks the database during the HNSW build** — the + expensive build runs without the shared lock (held only to snapshot and swap); + writes that land during the build are folded into the new index before the + swap. +- **Embedding server caps request body size** — an ASGI middleware rejects + request bodies larger than the server's own accept limits before they are + buffered/parsed, closing an unauthenticated memory-exhaustion vector (only + relevant with the `[server]` extra exposed on a network). A missing + encryption salt sidecar now logs a warning instead of silently falling back + to the shared legacy salt. +- **Robustness pass** — malformed FTS5 keyword queries raise `ValueError` instead + of a raw SQLite error; the cluster-state table is created eagerly so a + rolled-back first `save_cluster` cannot desync it; a non-integer + `EMBEDDING_BATCH_SIZE`/`EMBEDDING_SERVER_MAX_REQUEST_ITEMS` env value warns and + falls back instead of crashing import; `vacuum()` holds the DB lock; a failed + index add after the catalog commit is logged (divergence visibility); hybrid + search applies the Python metadata filter on the keyword side too (SQL/Python + parity); `logging.configure_logging` swaps handlers atomically. +- **LangChain `asimilarity_search_with_score`** offloads to a thread instead of + blocking the event loop. + +#### Changed + +- **`AsyncVectorCollection.increment_metadata` now returns `int`** (1 if the row + existed and was updated, 0 otherwise), matching the synchronous API; it + previously discarded the value and returned `None`. +- **LlamaIndex metadata filters fail loudly on unsupported shapes** — the + `SimpleVecDBLlamaStore` adapter now maps comparison operators + (`$gt/$gte/$lt/$lte/$ne/$in/$nin`) instead of silently treating them as + equality, and raises `NotImplementedError` for `OR`/`NOT` conditions and + unsupported operators rather than returning wrong results. +- **LangChain relevance scoring now works** — `SimpleVecDBVectorStore` implements + `_select_relevance_score_fn`, so `similarity_search_with_relevance_scores` and + `as_retriever(search_type="similarity_score_threshold")` return metric-aware + `[0, 1]` relevance (higher = better). `similarity_search_with_score` still + returns the raw distance (FAISS/Chroma convention), now documented as such. + +### Clustering and hierarchy fixes + +Internal correctness and performance work on the clustering and hierarchy +layers. No public API changes; existing databases are unaffected. + +#### Fixed + +- **`load_cluster` survives empty k-means clusters** — when k-means leaves a + requested cluster empty (common with duplicate vectors or `n_clusters` near + the number of distinct points), the stored `n_clusters` is smaller than the + number of centroid rows. The centroid reshape now derives its row count from + the stored buffer rather than `n_clusters`, which previously raised + `ValueError` on load. +- **`assign_to_cluster` matches metadata keys literally** — a `metadata_key` + containing `.` or `[` is now matched as a literal top-level key (via + `json_each`) instead of being misread as a nested JSON path, which had caused + every already-assigned document to be re-assigned on each call. +- **`cluster(algorithm="hdbscan", sample_size=…)` raises instead of silently + dropping documents** — HDBSCAN produces no centroids, so out-of-sample + documents cannot be assigned. The combination now raises a clear `ValueError` + rather than clustering only the sample. + +#### Performance + +- **BLAS-backed out-of-sample centroid assignment** — nearest-centroid + assignment uses the `‖c‖² − 2·x·c` expansion (a single matmul) instead of + materialising the dense `(n_vectors, n_centroids, dim)` broadcast temporary + that could exhaust memory on large collections. +- **Unassigned-id lookup pushed into SQLite** — `assign_to_cluster(doc_ids=None)` + finds documents lacking the cluster key with one `json_each` query instead of + loading and JSON-parsing every row's text and metadata. +- **Bounded ancestor-walk for cycle detection** — `set_parent` detects + parent/child cycles by walking the ancestor chain with a depth-bounded + recursive CTE instead of materialising the entire descendant subtree. + ## [2.6.1] - 2026-05-10 ### Storage, mutation, and eventing improvements diff --git a/docs/ENV_SETUP.md b/docs/ENV_SETUP.md index 81a5f2d..7f279fa 100755 --- a/docs/ENV_SETUP.md +++ b/docs/ENV_SETUP.md @@ -53,6 +53,7 @@ Configuration for `simplevecdb-server`. | `SERVER_HOST` | Host to bind the server to. | `0.0.0.0` | | `SERVER_PORT` | Port to bind the server to. | `53287` (Code default) / `8000` (Example) | | `EMBEDDING_SERVER_MAX_REQUEST_ITEMS` | Max number of prompts allowed per `/v1/embeddings` request (protects latency). | `max(32, EMBEDDING_BATCH_SIZE)` | +| `EMBEDDING_SERVER_MAX_BODY_BYTES` | Max raw request body size in bytes; larger bodies are rejected (413) before being buffered/parsed, preventing memory exhaustion. | _Derived from the request-item and text-length limits (min 1 MiB)_ | | `EMBEDDING_SERVER_API_KEYS` | Comma-separated API keys to require `Authorization: Bearer`/`X-API-Key`. | _Disabled (unauthenticated)_ | When `EMBEDDING_SERVER_API_KEYS` is set, SimpleVecDB also tracks request counts and token usage per key. Call `GET /v1/usage` with the same key to retrieve your stats. diff --git a/docs/Features.md b/docs/Features.md index a846c16..1062bfe 100644 --- a/docs/Features.md +++ b/docs/Features.md @@ -6,8 +6,9 @@ release-by-release detail, see the [Changelog](CHANGELOG.md). ## Storage & schema -- **Single-file SQLite** — one `.db` file (or `:memory:`) holds everything: - documents, vectors, FTS5 index, edges, events, TTL, clusters. +- **File-based storage** — a `.db` file (or `:memory:`) holds documents, the + FTS5 index, edges, events, TTL, and clusters; vectors live in a per-collection + `.usearch` HNSW index file alongside it. - **Multi-collection** — isolated namespaces per database via `db.collection("name")`. Each collection has its own quantization, distance metric, and (optional) embedding storage. diff --git a/docs/index.md b/docs/index.md index bdef15c..e92692e 100755 --- a/docs/index.md +++ b/docs/index.md @@ -5,17 +5,17 @@ [![License: MIT](https://img.shields.io/github/license/coderdayton/simplevecdb)](LICENSE) [![GitHub Stars](https://img.shields.io/github/stars/coderdayton/simplevecdb?style=social)](https://github.com/coderdayton/simplevecdb) -**The dead-simple, local-first vector database.** +**A local-first, embedded vector database backed by SQLite and usearch.** -SimpleVecDB brings **Chroma-like simplicity** to a single **SQLite file**. Built on **usearch HNSW** (v2.0+), it offers 10-100x faster vector search, quantization, and zero infrastructure headaches. Perfect for local RAG, offline agents, and indie hackers who need production-grade vector search without the operational overhead. +SimpleVecDB pairs **Chroma-like ergonomics** with a **file-based** store — a SQLite database for metadata and text alongside a `usearch` HNSW index per collection. It provides high-performance vector search, quantization, and hybrid retrieval with no separate services to run. It fits local RAG pipelines, offline agents, and any application that needs production-grade vector search without the operational overhead of a hosted database. ## Why SimpleVecDB? -- **Zero Infrastructure** — Just a `.db` file. No Docker, no Redis, no cloud bills. -- **Blazing Fast** — 10-100x faster with HNSW indexing, sub-millisecond queries on 100k+ vectors. -- **Truly Portable** — Runs anywhere Python runs: Linux, macOS, Windows. -- **Async Ready** — Full async/await support for web servers and concurrent workloads. -- **Batteries Included** — Optional FastAPI embeddings server + LangChain/LlamaIndex integrations. +- **Zero Infrastructure** — Local files on disk: a SQLite database plus a `usearch` index. No Docker, no Redis, no external services. +- **High Performance** — usearch HNSW indexing with adaptive search: brute-force under 10k vectors (perfect recall), HNSW above that. +- **Portable** — Runs anywhere SQLite runs: Linux, macOS, Windows, and WASM. +- **Async Support** — A complete async/await surface with optional executor injection for thread-safe ONNX/usearch sharing. +- **Integrations Included** — Optional FastAPI embeddings server and LangChain/LlamaIndex adapters via the `[integrations]` extra. - **Production Ready** — Hybrid search (BM25 + vector), metadata filtering, multi-collection support, and automatic hardware acceleration. ### When to Choose SimpleVecDB @@ -179,7 +179,7 @@ See **[Examples](examples.md)** for complete RAG workflows with Ollama. ### Multi-Collection Support -Organize vectors by domain within a single database file: +Organize vectors by domain within a single database: ```python from simplevecdb import VectorDB, Quantization @@ -272,7 +272,7 @@ See [Clustering Guide](guides/clustering.md) for algorithms, metrics, and use ca | Feature | Status | Description | | :------------------------ | :----- | :--------------------------------------------------------- | -| **Single-File Storage** | ✅ | SQLite `.db` file + `.usearch` index files | +| **File-Based Storage** | ✅ | SQLite `.db` file + `.usearch` index files | | **Multi-Collection** | ✅ | Isolated namespaces per database | | **HNSW Indexing** | ✅ | 10-100x faster approximate nearest neighbor (usearch) | | **Vector Search** | ✅ | Cosine, Euclidean, Inner Product metrics | diff --git a/pyproject.toml b/pyproject.toml index 8e8e5e8..58c436d 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "simplevecdb" -version = "2.6.1" +version = "2.6.2" 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/async_core.py b/src/simplevecdb/async_core.py index 161a203..e274bb2 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -301,9 +301,12 @@ async def increment_metadata( self, doc_id: int, deltas: dict[str, int | float], - ) -> None: - """Atomically apply numeric deltas to JSON metadata counters.""" - await self._run(self._collection.increment_metadata, doc_id, deltas) + ) -> int: + """Atomically apply numeric deltas to JSON metadata counters. + + Returns 1 if the row existed and was updated, 0 otherwise. + """ + return await self._run(self._collection.increment_metadata, doc_id, deltas) async def add_edge( self, diff --git a/src/simplevecdb/config.py b/src/simplevecdb/config.py index f430b1b..fb6d1e5 100755 --- a/src/simplevecdb/config.py +++ b/src/simplevecdb/config.py @@ -76,19 +76,41 @@ class Config: EMBEDDING_MODEL_REGISTRY_LOCKED: bool = _parse_bool_env( os.getenv("EMBEDDING_MODEL_REGISTRY_LOCKED"), True ) - # Auto-detect optimal batch size if not explicitly set + # Auto-detect optimal batch size if not explicitly set. Tolerate a malformed + # env value (e.g. "auto") instead of crashing every import of the package. _batch_size_env = os.getenv("EMBEDDING_BATCH_SIZE") - EMBEDDING_BATCH_SIZE: int = ( - int(_batch_size_env) - if _batch_size_env is not None - else get_optimal_batch_size() - ) + try: + EMBEDDING_BATCH_SIZE: int = ( + int(_batch_size_env) + if _batch_size_env is not None + else get_optimal_batch_size() + ) + except ValueError: + import warnings as _warnings + + _warnings.warn( + f"Invalid EMBEDDING_BATCH_SIZE={_batch_size_env!r}; using auto-detected size.", + stacklevel=2, + ) + EMBEDDING_BATCH_SIZE = get_optimal_batch_size() _request_limit_env = os.getenv("EMBEDDING_SERVER_MAX_REQUEST_ITEMS") or os.getenv( "EMBEDDING_SERVER_MAX_BATCH" ) - EMBEDDING_SERVER_MAX_REQUEST_ITEMS: int = ( - int(_request_limit_env) if _request_limit_env else max(32, EMBEDDING_BATCH_SIZE) - ) + try: + EMBEDDING_SERVER_MAX_REQUEST_ITEMS: int = ( + int(_request_limit_env) + if _request_limit_env + else max(32, EMBEDDING_BATCH_SIZE) + ) + except ValueError: + import warnings as _warnings + + _warnings.warn( + f"Invalid EMBEDDING_SERVER_MAX_REQUEST_ITEMS={_request_limit_env!r}; " + "using a default.", + stacklevel=2, + ) + EMBEDDING_SERVER_MAX_REQUEST_ITEMS = max(32, EMBEDDING_BATCH_SIZE) EMBEDDING_SERVER_API_KEYS: set[str] = _parse_api_keys( os.getenv("EMBEDDING_SERVER_API_KEYS") ) diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index a48de35..2198179 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -198,6 +198,9 @@ def __init__( # collections sharing the same sqlite3.Connection serialize their # transactional access from Python. self._lock: threading.RLock = lock if lock is not None else threading.RLock() + # Serializes rebuild_index() calls against each other WITHOUT holding the + # shared DB lock during the (slow) HNSW build. Distinct from self._lock. + self._rebuild_lock = threading.Lock() # Shared transaction depth; defaults to a per-collection state when # the parent VectorDB didn't pass one (e.g. legacy direct ctor use). self._tx_state: _TxState = tx_state if tx_state is not None else _TxState() @@ -395,10 +398,22 @@ def add_texts( parent_ids=batch_parent_ids, ) - # Add to usearch index - self._index.add( - np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads - ) + # Add to usearch index. The catalog rows above are already + # committed, so if this fails the two stores diverge (rows present, + # vectors missing). Log it so the divergence is visible instead of + # silent; recovery is rebuild_index() (needs store_embeddings=True). + try: + self._index.add( + np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads + ) + except Exception: + _logger.error( + "Index add failed for %d docs after catalog commit; catalog " + "and index have diverged. Run rebuild_index() to resync " + "(requires store_embeddings=True).", + len(doc_ids), + ) + raise all_ids.extend(doc_ids) @@ -844,37 +859,61 @@ def rebuild_index( """ _logger.info("Rebuilding usearch index for collection '%s'...", self.name) - # Serialize the entire fetch + build + swap on the connection-level - # lock so concurrent add/delete operations cannot mutate the catalog - # mid-rebuild and produce a stale or inconsistent snapshot. The lock - # is reentrant; CatalogManager read methods reacquire it but that - # is harmless under RLock. - with self._lock: - return self._rebuild_index_locked( - connectivity, expansion_add, expansion_search + # Serialize rebuilds against each other with a dedicated lock so two + # concurrent rebuild_index() calls cannot race on the .rebuild file — + # WITHOUT holding the shared DB lock during the slow HNSW build below. + with self._rebuild_lock: + # Phase 1 (DB lock): snapshot ids + embeddings from SQLite. + with self._lock: + snapshot = self._rebuild_snapshot() + if snapshot is None: + return 0 + keys, vectors, ndim, old_path = snapshot + + # Phase 2 (no DB lock): build the new index from the snapshot. The + # HNSW build is the expensive step and does not need the DB lock, so + # searches and writes on this and other collections are not blocked. + rebuild_path = old_path.with_suffix(old_path.suffix + ".rebuild") + if rebuild_path.exists(): + rebuild_path.unlink() + new_index = UsearchIndex( + index_path=str(rebuild_path), + ndim=ndim, + distance_strategy=self.distance_strategy, + quantization=self.quantization, + connectivity=connectivity + if connectivity is not None + else constants.USEARCH_DEFAULT_CONNECTIVITY, + expansion_add=expansion_add + if expansion_add is not None + else constants.USEARCH_DEFAULT_EXPANSION_ADD, + expansion_search=expansion_search + if expansion_search is not None + else constants.USEARCH_DEFAULT_EXPANSION_SEARCH, ) + new_index.add(keys, vectors) - def _rebuild_index_locked( + # Phase 3 (DB lock): fold in writes that landed during the build, + # then atomically swap the rebuilt index into place. + with self._lock: + return self._rebuild_commit( + new_index, keys, ndim, old_path, rebuild_path + ) + + def _rebuild_snapshot( self, - connectivity: int | None, - expansion_add: int | None, - expansion_search: int | None, - ) -> int: - # Precondition: caller must already hold ``self._lock``. ``rebuild_index`` - # is the only public entry point and acquires it before delegating; this - # private helper relies on RLock re-entrancy so the catalog reads below - # are serialized against concurrent add/delete on the shared connection. - # Routing the read through CatalogManager keeps the lock invariant - # explicit instead of bare ``self.conn.execute(...)``. - all_ids = self._catalog.list_all_ids() + ) -> tuple[np.ndarray, np.ndarray, int, Path] | None: + """Snapshot (keys, vectors, ndim, old_path) from SQLite under the DB lock. + Returns None when the collection has no documents; raises if no usable + embeddings are stored. The caller must hold ``self._lock``. + """ + all_ids = self._catalog.list_all_ids() if not all_ids: _logger.warning("No documents found in collection") - return 0 + return None - # 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. " @@ -882,13 +921,11 @@ def _rebuild_index_locked( "rebuild_index(), or re-add documents with store_embeddings=True." ) - # Filter to only docs with embeddings valid_pairs = [ (doc_id, emb) for doc_id in all_ids if (emb := embeddings_map.get(doc_id)) is not None ] - if not valid_pairs: raise RuntimeError( "No embeddings found in SQLite. Cannot rebuild index. " @@ -897,61 +934,79 @@ def _rebuild_index_locked( keys = np.array([doc_id for doc_id, _ in valid_pairs], dtype=np.uint64) vectors = np.array([emb for _, emb in valid_pairs], dtype=np.float32) + return keys, vectors, vectors.shape[1], self._index._path - # Determine dimension - ndim = vectors.shape[1] - - # Atomic rebuild: build the new index at a sibling path, save it - # durably, then os.replace() it onto the live path. The old index - # remains intact and recoverable until the final rename succeeds. - old_path = self._index._path - self._index.close() - - rebuild_path = old_path.with_suffix(old_path.suffix + ".rebuild") - if rebuild_path.exists(): - # Clean up remnant from a prior failed rebuild - rebuild_path.unlink() + def _rebuild_commit( + self, + new_index: UsearchIndex, + keys: np.ndarray, + ndim: int, + old_path: Path, + rebuild_path: Path, + ) -> int: + """Fold in concurrent catalog writes, then atomically swap. Holds DB lock.""" + # Catch up mutations that landed during the unlocked build so the new + # index reflects the current catalog, not just the snapshot. + snap_ids = {int(k) for k in keys} + current_ids = set(self._catalog.list_all_ids()) + added = current_ids - snap_ids + removed = snap_ids - current_ids + + add_pairs: list[tuple[int, Any]] = [] + if added: + emap = self._catalog.get_embeddings_by_ids(list(added)) + add_pairs = [(i, emb) for i in added if (emb := emap.get(i)) is not None] + if len(add_pairs) < len(added): + _logger.warning( + "rebuild_index: %d of %d docs added during the build have no " + "stored embeddings and were not indexed; catalog and index " + "will diverge for them (store_embeddings disabled?).", + len(added) - len(add_pairs), + len(added), + ) + if add_pairs: + new_index.add( + np.array([i for i, _ in add_pairs], dtype=np.uint64), + np.array([emb for _, emb in add_pairs], dtype=np.float32), + ) + if removed: + new_index.remove(np.array(sorted(removed), dtype=np.uint64)) - new_index = UsearchIndex( - index_path=str(rebuild_path), - ndim=ndim, - distance_strategy=self.distance_strategy, - quantization=self.quantization, - connectivity=connectivity - if connectivity is not None - else constants.USEARCH_DEFAULT_CONNECTIVITY, - expansion_add=expansion_add - if expansion_add is not None - else constants.USEARCH_DEFAULT_EXPANSION_ADD, - expansion_search=expansion_search - if expansion_search is not None - else constants.USEARCH_DEFAULT_EXPANSION_SEARCH, - ) - new_index.add(keys, vectors) - new_index.save() + # Count what was actually indexed, not what merely appeared in the + # catalog (added docs without stored embeddings are skipped above). + total = len(snap_ids) + len(add_pairs) - len(removed) - # Atomically swap the rebuilt index into place. Until this rename, - # the old index file at old_path is still the canonical copy. - os.replace(str(rebuild_path), str(old_path)) + # Atomic swap: the old index file stays canonical until os.replace(). + self._index.close() try: - dir_fd = os.open(str(old_path.parent), os.O_RDONLY) + new_index.save() + os.replace(str(rebuild_path), str(old_path)) try: - os.fsync(dir_fd) - finally: - os.close(dir_fd) - except OSError: - pass - - # Repoint the rebuilt index at the canonical path so future saves - # land at old_path rather than the now-vanished rebuild_path. - new_index._path = old_path - self._index = new_index + dir_fd = os.open(str(old_path.parent), os.O_RDONLY) + try: + os.fsync(dir_fd) + finally: + os.close(dir_fd) + except OSError: + pass + new_index._path = old_path + self._index = new_index + except BaseException: + # The live index was already closed and the swap failed; old_path + # still holds a valid index (original or rebuilt). Re-open it so the + # collection stays usable instead of holding a closed (bricked) index. + self._index = UsearchIndex( + index_path=str(old_path), + ndim=ndim, + distance_strategy=self.distance_strategy, + quantization=self.quantization, + ) + self._search._index = self._index + raise - # Update search engine reference self._search._index = self._index - - _logger.info("Rebuilt index with %d vectors", len(keys)) - return len(keys) + _logger.info("Rebuilt index with %d vectors", total) + return total # ------------------------------------------------------------------ # # Hierarchical Relationships @@ -1107,11 +1162,28 @@ def cluster( vectors = self._index.get(np.array(doc_ids, dtype=np.uint64)) + # Cap n_clusters to the number of vectors actually clustered (the + # sample when sampling, else the full set). This is a friendly + # public-API convenience; the low-level ClusterEngine.cluster_vectors() + # raises instead, for direct callers that want strictness. + _clustered_count = len(doc_ids) + if ( + sample_size is not None + and sample_size < len(doc_ids) + and algorithm != "hdbscan" + ): + _clustered_count = sample_size effective_n_clusters = n_clusters if n_clusters is not None and algorithm in ("kmeans", "minibatch_kmeans"): - effective_n_clusters = min(n_clusters, len(doc_ids)) + effective_n_clusters = min(n_clusters, _clustered_count) if sample_size and sample_size < len(doc_ids): + if algorithm == "hdbscan": + raise ValueError( + "sample_size is not supported with the 'hdbscan' algorithm: " + "HDBSCAN produces no centroids, so out-of-sample documents " + "cannot be assigned to clusters. Cluster the full set instead." + ) rng = np.random.default_rng(random_state) sample_indices = rng.choice(len(doc_ids), sample_size, replace=False) sample_ids = [doc_ids[i] for i in sample_indices] @@ -1320,8 +1392,13 @@ def load_cluster(self, name: str) -> tuple[ClusterResult, dict[str, Any]] | None if centroids_bytes is not None: dim = self.dim if dim: + # Derive the row count from the buffer (-1) rather than the + # stored ``n_clusters``: k-means can leave a requested cluster + # empty (common with duplicate vectors or n_clusters near the + # number of distinct points), so n_clusters_found < centroid + # rows and a reshape(n_clusters, dim) would raise ValueError. centroids = np.frombuffer(centroids_bytes, dtype=np.float32).reshape( - n_clusters, dim + -1, dim ) result = ClusterResult( @@ -1382,13 +1459,13 @@ def assign_to_cluster( ) if doc_ids is None: - all_ids = list(self._index.keys()) - # Get all documents to check for metadata key existence - all_docs = self._catalog.get_all_docs_with_text() - assigned_ids = { - doc_id for doc_id, _, meta in all_docs if metadata_key in meta - } - doc_ids = [d for d in all_ids if d not in assigned_ids] + # Push the "already assigned?" test into SQLite so we don't load + # and JSON-parse every row's text + metadata just to find the + # unassigned ids. Intersect with the index keys so we only try to + # assign documents that actually have a vector. + index_keys = set(self._index.keys()) + unassigned = self._catalog.find_ids_without_metadata_key(metadata_key) + doc_ids = [d for d in unassigned if d in index_keys] if not doc_ids: return 0 @@ -2316,7 +2393,7 @@ def get( class VectorDB: """ - Dead-simple local vector database powered by usearch HNSW. + Local-first, embedded vector database powered by usearch HNSW. SQLite stores metadata and text; usearch stores vectors in separate .usearch files per collection. Provides Chroma-like API with built-in @@ -2792,10 +2869,13 @@ def vacuum(self, checkpoint_wal: bool = True) -> None: Args: checkpoint_wal: If True (default), also truncate the WAL file. """ - if checkpoint_wal: - self.conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") - self.conn.execute("VACUUM") - self.conn.execute("PRAGMA optimize") + # Hold the DB lock: wal_checkpoint(TRUNCATE) and VACUUM require exclusive + # access, and other threads share this sqlite3 connection. + with self._lock: + if checkpoint_wal: + self.conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + self.conn.execute("VACUUM") + self.conn.execute("PRAGMA optimize") def save(self) -> None: """Save all collection indexes to disk.""" diff --git a/src/simplevecdb/embeddings/server.py b/src/simplevecdb/embeddings/server.py index 33714b3..b445997 100755 --- a/src/simplevecdb/embeddings/server.py +++ b/src/simplevecdb/embeddings/server.py @@ -4,6 +4,7 @@ import asyncio import hmac import logging +import os import signal import time from collections import defaultdict @@ -25,6 +26,89 @@ _MAX_TEXT_LENGTH = 100_000 +def _default_max_body_bytes() -> int: + """Body cap derived from the server's own accept limits (items × text len).""" + try: + items = int(config.EMBEDDING_SERVER_MAX_REQUEST_ITEMS) + except Exception: + items = 2048 + return max(items * (_MAX_TEXT_LENGTH + 16) + 65536, 1 << 20) + + +# Cap the raw request body BEFORE Pydantic buffers it: the per-item count is +# only enforced after the whole body is parsed, so without this an +# unauthenticated client could exhaust memory with one giant body. Override via +# EMBEDDING_SERVER_MAX_BODY_BYTES. +try: + _MAX_BODY_BYTES = int( + os.getenv("EMBEDDING_SERVER_MAX_BODY_BYTES") or _default_max_body_bytes() + ) +except ValueError: + _logger.warning( + "Invalid EMBEDDING_SERVER_MAX_BODY_BYTES=%r; using the derived default.", + os.getenv("EMBEDDING_SERVER_MAX_BODY_BYTES"), + ) + _MAX_BODY_BYTES = _default_max_body_bytes() +# Never let an override drop the cap to ~0 (which would reject every request). +_MAX_BODY_BYTES = max(_MAX_BODY_BYTES, 1 << 20) + + +class _MaxBodySizeMiddleware: + """ASGI middleware that rejects request bodies larger than ``max_bytes``.""" + + def __init__(self, app: Any, max_bytes: int) -> None: + self._app = app + self._max_bytes = max_bytes + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + if scope["type"] != "http": + await self._app(scope, receive, send) + return + + # Fast path: reject on a declared Content-Length over the cap. + for name, value in scope.get("headers", []): + if name == b"content-length": + try: + declared = int(value) + except ValueError: + break + if declared > self._max_bytes: + await self._reject(send) + return + break + + # Slow path: count bytes for chunked/unknown-length bodies. + received = 0 + + async def limited_receive() -> Any: + nonlocal received + message = await receive() + if message["type"] == "http.request": + received += len(message.get("body", b"")) + if received > self._max_bytes: + return {"type": "http.disconnect"} + return message + + await self._app(scope, limited_receive, send) + + async def _reject(self, send: Any) -> None: + body = ( + b'{"error":{"message":"Request body too large",' + b'"type":"payload_too_large","code":413}}' + ) + await send( + { + "type": "http.response.start", + "status": 413, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode()), + ], + } + ) + await send({"type": "http.response.body", "body": body}) + + def _validate_request_item_cap() -> None: """Reject EMBEDDING_SERVER_MAX_REQUEST_ITEMS > _MAX_ENCODE_BATCH at startup. @@ -129,6 +213,9 @@ def is_allowed(self, identity: str) -> bool: docs_url="/docs", ) +# Reject oversized request bodies before they are buffered/parsed. +app.add_middleware(_MaxBodySizeMiddleware, max_bytes=_MAX_BODY_BYTES) + # (#4) CORS middleware — configurable via EMBEDDING_SERVER_CORS_ORIGINS env var. # Default is no CORS (no allow_origins, no credentials) so the server is safe # to deploy without explicit CORS configuration. Operators that want CORS set diff --git a/src/simplevecdb/encryption.py b/src/simplevecdb/encryption.py index d9715d4..d3db50f 100755 --- a/src/simplevecdb/encryption.py +++ b/src/simplevecdb/encryption.py @@ -120,8 +120,11 @@ def _normalize_key(key: str | bytes, salt: bytes | None = None) -> bytes: salt_to_use = salt if salt is not None else _NORMALIZE_KEY_SALT - # Cache key includes both the raw passphrase bytes and the salt so the - # same passphrase yields different cache entries for different DBs. + # Cache key = (raw passphrase bytes, salt) so the same passphrase yields + # distinct entries per salt/DB. We deliberately do NOT hash the passphrase + # here: a fast hash (sha256) trips weak-password-hash scanners, and a slow + # KDF would defeat this cache, whose sole purpose is to AVOID re-running the + # 600k-iter PBKDF2. The cache is process-local and LRU-bounded. key_bytes = key.encode("utf-8") if isinstance(key, str) else bytes(key) cache_key = (key_bytes, salt_to_use) @@ -186,7 +189,15 @@ def _resolve_salt( return _NORMALIZE_KEY_SALT if not create_if_missing: - # Legacy resource — created before per-DB salts existed. + # Legacy resource — created before per-DB salts existed, or a sidecar + # that was removed. Either way per-DB salt protection is not active; + # surface it (instead of silently using the shared fixed salt) so + # operators can migrate, or notice a deleted sidecar. + _logger.warning( + "No salt sidecar for %s; using the legacy shared salt. Per-DB salt " + "protection is not active for this resource.", + resource_path, + ) return _NORMALIZE_KEY_SALT salt = secrets.token_bytes(SALT_SIZE) @@ -322,6 +333,11 @@ def create_encrypted_connection( # database (sidecar present). Normalize every key shape to a # 32-byte derived value and feed it as ``x'hex'`` so we never # interpolate raw passphrase characters into SQL. + # + # At-rest key strength comes from the application-layer PBKDF2 + # (PBKDF2_ITERATIONS = 600k) used to derive this 32-byte key. The + # ``x'hex'`` form is a raw key, so SQLCipher runs no internal KDF + # and its build-dependent ``kdf_iter`` default does not apply here. salt = _resolve_salt(db_path_obj, create_if_missing=is_new_db) normalized_key = _normalize_key(key, salt=salt) conn.execute(f"PRAGMA key = \"x'{normalized_key.hex()}'\"") diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 835e3a0..8dcc9e3 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -119,9 +119,15 @@ def __init__( def __enter__(self): self._lock.acquire() - if self._tx.depth == 0: - self._conn.__enter__() - self._owns_conn = True + try: + if self._tx.depth == 0: + self._conn.__enter__() + self._owns_conn = True + except BaseException: + # __exit__ is not called if __enter__ raises; release the lock + # ourselves so a connection-level error cannot leak it. + self._lock.release() + raise return self def __exit__(self, exc_type, exc, tb): @@ -318,6 +324,10 @@ def create_tables(self) -> None: self._ensure_embedding_column() self._ensure_parent_id_column() self._ensure_fts_table() + # Create the cluster-state table eagerly so a first save_cluster inside a + # rolled-back transaction cannot leave _cluster_table_ready set without + # the table actually existing. + self._ensure_cluster_table() # 2.6.1 auxiliary tables (pending vectors, edges, events, TTL). # Each is idempotent (CREATE TABLE IF NOT EXISTS), so existing 2.6.0 # databases gain them transparently on first open. @@ -939,9 +949,25 @@ def keyword_search( ORDER BY score ASC LIMIT ? """ + import sqlite3 # noqa: PLC0415 + params = (query,) + tuple(filter_params) + (k,) - with self._lock: - rows = self.conn.execute(sql, params).fetchall() + try: + with self._lock: + rows = self.conn.execute(sql, params).fetchall() + except sqlite3.OperationalError as exc: + # FTS5 raises OperationalError on a malformed MATCH query (unbalanced + # quotes, a bare operator, ...). Surface a clear caller-facing error + # instead of the raw SQLite message; re-raise unrelated op errors. + msg = str(exc).lower() + # Match only FTS5-query failure shapes; do NOT match a bare + # "syntax error", which can come from an unrelated SQL bug and must + # not be mislabeled as the user's query. + if any(s in msg for s in ("fts5", "unterminated", "malformed")): + raise ValueError( + f"Invalid full-text search query {query!r}: {exc}" + ) from exc + raise return [(int(row[0]), float(row[1])) for row in rows] def build_filter_clause( @@ -981,7 +1007,10 @@ def build_filter_clause( clauses: list[str] = [] params: list[Any] = [] for key, value in normalized.items(): - json_path = f"$.{key}" + # Quote the path label so a literal key like "a.b" matches the + # top-level member, not the nested path a -> b (matches the Python + # _matches_filter semantics and find_ids_without_metadata_key). + json_path = f'$."{key}"' text_extract = f"json_extract({metadata_column}, ?)" num_extract = f"CAST({text_extract} AS REAL)" type_extract = f"json_type({metadata_column}, ?)" @@ -1143,6 +1172,29 @@ def get_all_docs_with_text( result.append((int(row_id), text, meta)) return result + def find_ids_without_metadata_key(self, key: str) -> list[int]: + """Return ids whose metadata object lacks the top-level ``key``. + + Pushes the key-existence test into SQLite so callers that only need the + unassigned ids avoid loading and JSON-parsing every row's text + + metadata. ``json_each`` enumerates the literal top-level members, so the + bound ``key`` is matched exactly as stored. Unlike a ``$.`` JSON + path this is correct for keys containing ``.`` or ``[`` (which a path + would misread as nested access) and treats a present-but-null value as + assigned — matching the Python ``key in meta`` test it replaces. NULL or + ``{}`` metadata yields no members, so such rows count as unassigned. + ``key`` is bound as a parameter and cannot inject SQL. + """ + with self._lock: + rows = self.conn.execute( + f"SELECT id FROM {self._table_name} " + f"WHERE NOT EXISTS (" + f"SELECT 1 FROM json_each({self._table_name}.metadata) " + f"WHERE key = ?)", + (key,), + ).fetchall() + return [int(r[0]) for r in rows] + def update_metadata_batch(self, updates: list[tuple[int, dict[str, Any]]]) -> int: """ Update metadata for multiple documents in a single transaction. @@ -2123,13 +2175,36 @@ def set_parent(self, doc_id: int, parent_id: int | None) -> bool: # writer cannot create a cycle-forming edge between the check and the # UPDATE. The lock serializes; `with self.conn:` wraps the UPDATE in # an implicit transaction that commits on success. + from .. import constants + with self._writable(): if parent_id is not None: if parent_id == doc_id: raise ValueError("A document cannot be its own parent") - descendants = self.get_descendants(doc_id) - descendant_ids = {d[0] for d in descendants} - if parent_id in descendant_ids: + # A cycle forms iff doc_id is already an ancestor of parent_id + # (then doc_id -> parent_id -> ... -> doc_id). Walk *up* the + # ancestor chain of parent_id — bounded by MAX_HIERARCHY_DEPTH, + # ids only, with an early LIMIT 1 — instead of materialising + # doc_id's entire descendant subtree (text + metadata) just to + # test one membership. + cycle = self.conn.execute( + f""" + WITH RECURSIVE ancestors(id, depth) AS ( + SELECT parent_id, 1 FROM {self._table_name} + WHERE id = ? AND parent_id IS NOT NULL + + UNION ALL + + SELECT t.parent_id, a.depth + 1 + FROM {self._table_name} t + JOIN ancestors a ON t.id = a.id + WHERE t.parent_id IS NOT NULL AND a.depth < ? + ) + SELECT 1 FROM ancestors WHERE id = ? LIMIT 1 + """, + (parent_id, constants.MAX_HIERARCHY_DEPTH, doc_id), + ).fetchone() + if cycle is not None: raise ValueError( f"Cannot set parent: document {parent_id} is a descendant of {doc_id}" ) diff --git a/src/simplevecdb/engine/clustering.py b/src/simplevecdb/engine/clustering.py index 0dca1bd..af4107f 100755 --- a/src/simplevecdb/engine/clustering.py +++ b/src/simplevecdb/engine/clustering.py @@ -56,6 +56,16 @@ def cluster_vectors( algorithm=algorithm, ) + if ( + algorithm in ("kmeans", "minibatch_kmeans") + and n_clusters is not None + and n_clusters > len(vectors) + ): + raise ValueError( + f"n_clusters ({n_clusters}) cannot exceed the number of vectors " + f"({len(vectors)})" + ) + if algorithm == "hdbscan": labels, centroids, inertia = self._hdbscan(vectors, min_cluster_size) elif algorithm == "minibatch_kmeans": @@ -255,5 +265,12 @@ def assign_to_nearest_centroid( centroids: np.ndarray, ) -> np.ndarray: """Assign vectors to nearest centroid (for out-of-sample assignment).""" - distances = np.linalg.norm(vectors[:, np.newaxis] - centroids, axis=2) - return np.argmin(distances, axis=1).astype(np.int32) + # ||x - c||^2 = ||x||^2 - 2 x·c + ||c||^2. The ||x||^2 term is constant + # across centroids for a given row, so it does not change the argmin and + # is dropped. This avoids the (n_vectors, n_centroids, dim) broadcast + # temporary that ``vectors[:, None] - centroids`` would materialise + # (tens of GB — and an OOM — on large collections) and replaces the + # Python-level broadcast with a single BLAS-backed matmul. + centroid_sq = np.einsum("ij,ij->i", centroids, centroids) + distances_sq = centroid_sq[np.newaxis, :] - 2.0 * (vectors @ centroids.T) + return np.argmin(distances_sq, axis=1).astype(np.int32) diff --git a/src/simplevecdb/engine/search.py b/src/simplevecdb/engine/search.py index 813d4ba..3df709e 100755 --- a/src/simplevecdb/engine/search.py +++ b/src/simplevecdb/engine/search.py @@ -171,8 +171,15 @@ def similarity_search_batch( validate_filter(filter) - # For small query counts, sequential search avoids batch overhead - if len(queries) <= constants.USEARCH_BATCH_THRESHOLD: + # The native batch path requires pre-embedded vector queries and does a + # single fixed over-fetch, so it can neither auto-embed text queries nor + # re-fetch to fill k under a selective filter. Route through the per-query + # path (which handles both) for small batches, any filter, or text queries. + if ( + len(queries) <= constants.USEARCH_BATCH_THRESHOLD + or filter is not None + or any(isinstance(q, str) for q in queries) + ): return [ self.similarity_search(q, k, filter, exact=exact, threads=threads) for q in queries @@ -349,6 +356,11 @@ def hybrid_search( if cid not in docs_map: continue text, metadata = docs_map[cid] + # Defensive parity with the vector side: the SQL filter already + # excluded non-matches, but apply the Python check too so a + # SQL/Python grammar divergence can't admit a wrong candidate. + if filter and not self._matches_filter(metadata, filter): + continue rrf_scores[cid] = rrf_scores.get(cid, 0.0) + 1.0 / (rrf_k + kw_rank + 1) if cid not in doc_lookup: doc_lookup[cid] = Document(page_content=text, metadata=metadata) @@ -453,6 +465,7 @@ def max_marginal_relevance_search( sel_matrix: np.ndarray | None = ( emb[np.newaxis, :].copy() if emb is not None else None ) + is_l2 = self._distance_strategy == DistanceStrategy.L2 while len(selected) < k and unselected: best_score = -float("inf") @@ -461,9 +474,15 @@ def max_marginal_relevance_search( for pos, idx in enumerate(unselected): _, _, dist, emb = candidates[idx] - # Relevance: convert distance to similarity (lower distance = higher similarity) - # For cosine distance in [0, 2], similarity = 1 - distance/2 - relevance = 1.0 - dist / 2.0 + # Relevance: convert distance to similarity (lower distance = higher). + if is_l2: + # usearch returns squared L2 in [0, inf); map to a bounded, + # monotonically-decreasing similarity so large distances don't + # swamp the diversity (redundancy) term. + relevance = 1.0 / (1.0 + dist**0.5) + else: + # Cosine distance in [0, 2]: similarity = 1 - distance/2 + relevance = 1.0 - dist / 2.0 # Redundancy: max similarity to any already-selected doc redundancy = 0.0 diff --git a/src/simplevecdb/engine/usearch_index.py b/src/simplevecdb/engine/usearch_index.py index 38b39f5..9a1df7f 100755 --- a/src/simplevecdb/engine/usearch_index.py +++ b/src/simplevecdb/engine/usearch_index.py @@ -359,7 +359,8 @@ def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: Note: usearch HNSW doesn't support true deletion efficiently. Keys are marked as deleted but space isn't reclaimed until rebuild. - For heavy delete workloads, consider periodic rebuild(). + For heavy delete workloads, periodically call + ``VectorCollection.rebuild_index()`` to reclaim space and recall. Args: keys: Keys to remove @@ -470,6 +471,18 @@ def keys(self) -> list[int]: return [] return [int(k) for k in self._index.keys] + def _vectors_from_index(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: + """Fetch stored vectors for keys, unpacking BIT-quantized bytes to ±1 floats. + + For BIT quantization usearch stores packed bytes (ndim/8 per vector); a + plain float cast would yield the wrong shape and meaningless values, so + the bits are unpacked back to the float dimension. + """ + raw = self._index[keys] + if self._quantization == Quantization.BIT: + return _unpack_bits(np.asarray(raw, dtype=np.uint8), self._ndim or 1) + return np.asarray(raw, dtype=np.float32) + def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: """ Retrieve vectors by their keys. @@ -498,7 +511,7 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: if existing_mask.all(): # Fast path: all keys exist, batch retrieve - return np.asarray(self._index[keys], dtype=np.float32) + return self._vectors_from_index(keys) # Mixed: some keys missing _logger.warning( @@ -507,7 +520,7 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: ) 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) + result[existing_mask] = self._vectors_from_index(existing_keys) return result def __del__(self) -> None: diff --git a/src/simplevecdb/integrations/langchain.py b/src/simplevecdb/integrations/langchain.py index f18dc48..2ce18e6 100755 --- a/src/simplevecdb/integrations/langchain.py +++ b/src/simplevecdb/integrations/langchain.py @@ -1,4 +1,4 @@ -from collections.abc import Iterable +from collections.abc import Callable, Iterable from typing import Any try: @@ -141,7 +141,12 @@ def similarity_search_with_score( **kwargs: Any, ) -> list[tuple[LangChainDocument, float]]: """ - Return with scores (distances). + Return docs with their raw distance (lower = more similar). + + This matches the LangChain FAISS/Chroma convention. For a [0, 1] + relevance score (higher = better) — used by + ``as_retriever(search_type="similarity_score_threshold")`` — use + ``similarity_search_with_relevance_scores`` instead. Args: query: Text query string. @@ -149,7 +154,7 @@ def similarity_search_with_score( **kwargs: Additional arguments (e.g., filter). Returns: - List of (Document, score) tuples. + List of (Document, distance) tuples. """ if self.embedding: query_vec = self.embedding.embed_query(query) @@ -168,6 +173,22 @@ def similarity_search_with_score( for doc, score in results ] + def _select_relevance_score_fn(self) -> Callable[[float], float]: + """Map this collection's distance to a [0, 1] relevance (higher = better). + + Powers ``similarity_search_with_relevance_scores`` and + ``as_retriever(search_type="similarity_score_threshold")``; + ``similarity_search_with_score`` returns the raw distance, this inverts it. + """ + from simplevecdb.types import DistanceStrategy # noqa: PLC0415 + + if self._collection.distance_strategy == DistanceStrategy.L2: + # usearch returns squared L2 in [0, inf): bounded and decreasing. + return lambda distance: 1.0 / (1.0 + max(distance, 0.0) ** 0.5) + # Cosine distance in [0, 2] -> relevance in [0, 1]; clamp against any + # floating overshoot so LangChain doesn't flag out-of-range scores. + return lambda distance: max(0.0, min(1.0, 1.0 - distance / 2.0)) + def delete(self, ids: list[str] | None = None, **kwargs: Any) -> None: """ Delete documents by ID. @@ -273,6 +294,14 @@ async def asimilarity_search(self, *args, **kwargs): return await asyncio.to_thread(self.similarity_search, *args, **kwargs) + async def asimilarity_search_with_score(self, *args, **kwargs): + import asyncio + + # Base class would call the sync version on the event loop; offload it. + return await asyncio.to_thread( + self.similarity_search_with_score, *args, **kwargs + ) + async def amax_marginal_relevance_search( self, *args, diff --git a/src/simplevecdb/integrations/llamaindex.py b/src/simplevecdb/integrations/llamaindex.py index d285dee..dba05e6 100755 --- a/src/simplevecdb/integrations/llamaindex.py +++ b/src/simplevecdb/integrations/llamaindex.py @@ -224,6 +224,12 @@ def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: if internal_id is not None: self._collection.delete_by_ids([internal_id]) self._id_map.pop(internal_id, None) + else: + _logger.warning( + "delete(ref_doc_id=%r): no matching document found; nothing " + "was deleted.", + ref_doc_id, + ) def delete_nodes( self, @@ -257,13 +263,47 @@ def _filters_to_dict( ) -> dict[str, Any] | None: if filters is None: return None + + from llama_index.core.vector_stores.types import ( # noqa: PLC0415 + FilterCondition, + FilterOperator, + ) + + # The underlying engine ANDs all conditions; any non-AND condition + # (OR, NOT) is not representable, so fail loudly instead of silently + # returning AND semantics. + if getattr(filters, "condition", None) not in (None, FilterCondition.AND): + raise NotImplementedError( + "Only AND metadata filter conditions are supported; got " + f"{getattr(filters, 'condition', None)}." + ) + + # Map LlamaIndex operators onto the engine's Mongo-style grammar. + op_map = { + FilterOperator.GT: "$gt", + FilterOperator.GTE: "$gte", + FilterOperator.LT: "$lt", + FilterOperator.LTE: "$lte", + FilterOperator.NE: "$ne", + FilterOperator.IN: "$in", + FilterOperator.NIN: "$nin", + } + result: dict[str, Any] = {} - if hasattr(filters, "filters"): - for filter_item in filters.filters: # type: ignore[attr-defined] - if hasattr(filter_item, "key") and hasattr(filter_item, "value"): - key = getattr(filter_item, "key") - value = getattr(filter_item, "value") - result[key] = value + for filter_item in getattr(filters, "filters", None) or []: + if not (hasattr(filter_item, "key") and hasattr(filter_item, "value")): + continue + key = filter_item.key + value = filter_item.value + operator = getattr(filter_item, "operator", None) + if operator is None or operator == FilterOperator.EQ: + result[key] = value + elif operator in op_map: + result[key] = {op_map[operator]: value} + else: + raise NotImplementedError( + f"Unsupported metadata filter operator: {operator}" + ) return result or None def _build_query_result( diff --git a/src/simplevecdb/logging.py b/src/simplevecdb/logging.py index 2fb63b3..93e5daa 100755 --- a/src/simplevecdb/logging.py +++ b/src/simplevecdb/logging.py @@ -104,9 +104,6 @@ def configure_logging( logger.setLevel(level) - # Remove existing handlers to avoid duplicates - logger.handlers.clear() - # Create handler if handler is None: handler = logging.StreamHandler() @@ -116,7 +113,9 @@ def configure_logging( handler.setFormatter(formatter) handler.setLevel(level) - logger.addHandler(handler) + # Replace handlers in a single assignment so a concurrent logging call never + # observes a window with no handlers (clear()+addHandler() leaves one). + logger.handlers[:] = [handler] @contextmanager diff --git a/src/simplevecdb/utils.py b/src/simplevecdb/utils.py index 890c936..38d709a 100755 --- a/src/simplevecdb/utils.py +++ b/src/simplevecdb/utils.py @@ -400,6 +400,11 @@ def validate_filter(filter_dict: dict[str, Any] | None) -> None: raise ValueError( f"Filter keys must be strings, got {type(key).__name__}: {key!r}" ) + if '"' in key: + raise ValueError( + f"Filter keys must not contain a double-quote character: {key!r} " + "(such keys cannot be represented as a JSON path)." + ) # Normalize tuple shorthand for validation; the actual SQL builder # also normalizes, so this is just for the error path here. value = _normalize_filter_value(key, raw_value) diff --git a/tests/unit/integrations/test_llamaindex_filters_tier1.py b/tests/unit/integrations/test_llamaindex_filters_tier1.py new file mode 100644 index 0000000..9dbd84d --- /dev/null +++ b/tests/unit/integrations/test_llamaindex_filters_tier1.py @@ -0,0 +1,70 @@ +"""Tier-1 fix: LlamaIndex metadata filter operators/conditions are honored.""" + +from __future__ import annotations + +import pytest + +try: + import llama_index # noqa: F401 +except ImportError: + pytest.skip("llama-index not installed", allow_module_level=True) + +from llama_index.core.vector_stores.types import ( + FilterCondition, + FilterOperator, + MetadataFilter, + MetadataFilters, +) + +from simplevecdb.integrations.llamaindex import SimpleVecDBLlamaStore + + +def _store() -> SimpleVecDBLlamaStore: + return SimpleVecDBLlamaStore(db_path=":memory:") + + +def test_gt_operator_maps_to_dollar_gt() -> None: + f = MetadataFilters( + filters=[MetadataFilter(key="score", value=0.5, operator=FilterOperator.GT)] + ) + assert _store()._filters_to_dict(f) == {"score": {"$gt": 0.5}} + + +def test_eq_operator_stays_bare_value() -> None: + f = MetadataFilters( + filters=[MetadataFilter(key="tag", value="x", operator=FilterOperator.EQ)] + ) + assert _store()._filters_to_dict(f) == {"tag": "x"} + + +def test_or_condition_raises_not_implemented() -> None: + f = MetadataFilters( + filters=[ + MetadataFilter(key="a", value=1, operator=FilterOperator.EQ), + MetadataFilter(key="b", value=2, operator=FilterOperator.EQ), + ], + condition=FilterCondition.OR, + ) + with pytest.raises(NotImplementedError): + _store()._filters_to_dict(f) + + +def test_in_operator_maps_to_dollar_in() -> None: + f = MetadataFilters( + filters=[ + MetadataFilter(key="tag", value=["a", "b"], operator=FilterOperator.IN) + ] + ) + assert _store()._filters_to_dict(f) == {"tag": {"$in": ["a", "b"]}} + + +def test_not_condition_raises_not_implemented() -> None: + f = MetadataFilters( + filters=[ + MetadataFilter(key="a", value=1, operator=FilterOperator.EQ), + MetadataFilter(key="b", value=2, operator=FilterOperator.EQ), + ], + condition=FilterCondition.NOT, + ) + with pytest.raises(NotImplementedError): + _store()._filters_to_dict(f) diff --git a/tests/unit/test_clustering.py b/tests/unit/test_clustering.py index b643c94..4a11b1d 100755 --- a/tests/unit/test_clustering.py +++ b/tests/unit/test_clustering.py @@ -121,6 +121,19 @@ def test_cluster_with_sample_size(self, db_path: Path, dim: int): db.close() + def test_sample_size_with_hdbscan_raises(self, db_path: Path, dim: int): + """sample_size + hdbscan raises instead of silently dropping docs.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(10, 3, dim) + collection.add_texts(texts, embeddings=embeddings.tolist()) + + with pytest.raises(ValueError, match="sample_size is not supported"): + collection.cluster(algorithm="hdbscan", sample_size=9) + + db.close() + def test_cluster_result_summary(self, db_path: Path, dim: int): """ClusterResult.summary() returns cluster counts.""" db = VectorDB(db_path) @@ -574,6 +587,104 @@ def test_assign_to_cluster(self, db_path: Path, dim: int): db.close() + def test_save_load_cluster_with_empty_clusters(self, db_path: Path, dim: int): + """load_cluster survives k-means runs that leave a requested cluster empty. + + With only 2 distinct vectors but n_clusters=4, k-means yields 4 centroid + rows but fewer non-empty labels, so the stored n_clusters is < centroid + rows. The reshape must derive the row count from the buffer, not the + stored n_clusters (which previously raised ValueError on load). + """ + db = VectorDB(db_path) + collection = db.collection("test") + + a = np.zeros(dim, dtype=np.float32) + b = np.zeros(dim, dtype=np.float32) + b[0] = 10.0 + embeddings = np.array([a, a, a, b, b, b], dtype=np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(6)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=4, algorithm="kmeans", random_state=42) + assert result.centroids is not None + assert result.centroids.shape[0] == 4 + # Fewer distinct clusters than requested -> stored n_clusters < rows. + assert result.n_clusters < result.centroids.shape[0] + + collection.save_cluster("degenerate", result) + loaded = collection.load_cluster("degenerate") + + assert loaded is not None + loaded_result, _ = loaded + assert loaded_result.centroids is not None + assert loaded_result.centroids.shape == result.centroids.shape + + db.close() + + def test_assign_to_cluster_only_assigns_unassigned(self, db_path: Path, dim: int): + """assign_to_cluster(doc_ids=None) targets only docs lacking the key.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, dim).astype(np.float32) + embeddings[:5, 0] += 10.0 + embeddings[5:, 1] += 10.0 + collection.add_texts( + [f"doc_{i}" for i in range(10)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + collection.assign_cluster_metadata(result) # all 10 now have "cluster" + collection.save_cluster("saved", result) + + new_embs = np.random.randn(3, dim).astype(np.float32) + new_embs[:, 0] += 10.0 + collection.add_texts(["new_a", "new_b", "new_c"], embeddings=new_embs.tolist()) + + # Only the 3 new (unassigned) docs should be touched. + assigned = collection.assign_to_cluster("saved") + assert assigned == 3 + + db.close() + + def test_assign_to_cluster_handles_dotted_metadata_key( + self, db_path: Path, dim: int + ): + """A metadata_key with '.' is matched as a literal key, not a JSON path. + + The unassigned-id lookup must test literal top-level key existence. A + ``$.cluster.v2`` JSON path would read the dotted key as nested access, + find nothing, and re-assign every already-assigned doc. + """ + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(0) + embeddings = np.random.randn(8, dim).astype(np.float32) + embeddings[:4, 0] += 10.0 + embeddings[4:, 1] += 10.0 + collection.add_texts( + [f"doc_{i}" for i in range(8)], embeddings=embeddings.tolist() + ) + + key = "cluster.v2" + result = collection.cluster(n_clusters=2, random_state=0) + collection.assign_cluster_metadata(result, metadata_key=key) + collection.save_cluster("saved", result) + + new_embs = np.random.randn(2, dim).astype(np.float32) + new_embs[:, 0] += 10.0 + collection.add_texts(["new_a", "new_b"], embeddings=new_embs.tolist()) + + # The 8 existing docs already carry the dotted key -> only the 2 new + # docs are unassigned. + assigned = collection.assign_to_cluster("saved", metadata_key=key) + assert assigned == 2 + + db.close() + def test_assign_to_cluster_raises_for_unknown(self, db_path: Path): """assign_to_cluster raises ValueError for unknown cluster.""" db = VectorDB(db_path) @@ -704,3 +815,19 @@ def test_assign_to_nearest_centroid(self): vectors = np.array([[1, 1], [9, 9], [0.5, 0.5]], dtype=np.float32) labels = engine.assign_to_nearest_centroid(vectors, centroids) assert list(labels) == [0, 1, 0] + + def test_assign_to_nearest_centroid_matches_bruteforce(self): + """Optimised assignment equals the explicit pairwise-distance argmin.""" + from simplevecdb.engine.clustering import ClusterEngine + + rng = np.random.default_rng(7) + vectors = rng.standard_normal((200, 16)).astype(np.float32) + centroids = rng.standard_normal((12, 16)).astype(np.float32) + + engine = ClusterEngine() + got = engine.assign_to_nearest_centroid(vectors, centroids) + + expected = np.argmin( + np.linalg.norm(vectors[:, np.newaxis] - centroids, axis=2), axis=1 + ).astype(np.int32) + assert np.array_equal(got, expected) diff --git a/tests/unit/test_hierarchy.py b/tests/unit/test_hierarchy.py index 35ad94a..7c0f26a 100755 --- a/tests/unit/test_hierarchy.py +++ b/tests/unit/test_hierarchy.py @@ -391,6 +391,34 @@ def test_set_parent_cycle_detection(self, db_path: Path, dim: int): db.close() + def test_set_parent_cycle_detection_deep_chain(self, db_path: Path, dim: int): + """Cycle check walks multiple ancestor levels without false positives.""" + db = VectorDB(db_path) + collection = db.collection("test") + + # Chain: a -> b -> c -> d + a = collection.add_texts(["a"], embeddings=[self.make_embedding(dim)])[0] + b = collection.add_texts( + ["b"], embeddings=[self.make_embedding(dim)], parent_ids=[a] + )[0] + c = collection.add_texts( + ["c"], embeddings=[self.make_embedding(dim)], parent_ids=[b] + )[0] + d = collection.add_texts( + ["d"], embeddings=[self.make_embedding(dim)], parent_ids=[c] + )[0] + + # Re-parenting d directly under a (its existing ancestor) is NOT a cycle. + assert collection.set_parent(d, a) is True + + # Rebuild the deep chain and reject a multi-level cycle: a -> d would + # make a a child of its own descendant. + collection.set_parent(d, c) + with pytest.raises(ValueError, match="descendant"): + collection.set_parent(a, d) + + db.close() + class TestHierarchyMigration: """Test that existing databases get parent_id column added.""" diff --git a/tests/unit/test_tier1_fixes.py b/tests/unit/test_tier1_fixes.py new file mode 100644 index 0000000..0c2e397 --- /dev/null +++ b/tests/unit/test_tier1_fixes.py @@ -0,0 +1,241 @@ +"""Tier-1 review fixes — TDD regression tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from simplevecdb import VectorDB + + +def _emb(n: int, dim: int, seed: int = 0) -> list[list[float]]: + rng = np.random.default_rng(seed) + return rng.standard_normal((n, dim)).astype(np.float32).tolist() + + +class TestClusteringGuard: + def test_cluster_vectors_n_clusters_exceeds_raises_clear_error(self) -> None: + """cluster_vectors with n_clusters > n_vectors gives a descriptive error.""" + pytest.importorskip("sklearn") + from simplevecdb.engine.clustering import ClusterEngine + + engine = ClusterEngine() + vectors = np.asarray(_emb(5, 16), dtype=np.float32) + with pytest.raises(ValueError, match="cannot exceed"): + engine.cluster_vectors( + vectors, doc_ids=[0, 1, 2, 3, 4], algorithm="kmeans", n_clusters=10 + ) + + +class TestAsyncIncrementReturnsInt: + @pytest.mark.asyncio + async def test_returns_counter_update_result(self) -> None: + """async increment_metadata returns the int result (1 updated / 0 missing).""" + from simplevecdb import AsyncVectorDB + + async with AsyncVectorDB(":memory:") as db: + col = db.collection("t") + ids = await col.add_texts(["doc"], embeddings=_emb(1, 16)) + updated = await col.increment_metadata(ids[0], {"hits": 1}) + assert updated == 1 + missing = await col.increment_metadata(999999, {"hits": 1}) + assert missing == 0 + + +class TestCatalogWritableLockRelease: + def test_lock_released_when_conn_enter_fails(self) -> None: + """If conn.__enter__ raises, the writable lock must be released, not leaked.""" + import sqlite3 + import threading + import types + + from simplevecdb.engine.catalog import _CatalogWritable + + lock = threading.RLock() + + class BadConn: + def __enter__(self): + raise sqlite3.OperationalError("simulated busy") + + def __exit__(self, *exc): + return False + + tx = types.SimpleNamespace(depth=0) + writable = _CatalogWritable(lock, BadConn(), tx) # type: ignore[arg-type] + + with pytest.raises(sqlite3.OperationalError): + writable.__enter__() + + # Probe from a *different* thread (RLock is reentrant; same thread would + # falsely re-acquire and mask a leak). + probe: dict[str, bool] = {} + + def _probe() -> None: + got = lock.acquire(blocking=False) + probe["got"] = got + if got: + lock.release() + + t = threading.Thread(target=_probe) + t.start() + t.join() + assert probe["got"], "lock leaked (still held) after __enter__ failure" + + +class TestMMRMetricAware: + def test_l2_mmr_respects_diversity(self) -> None: + """On an L2 collection MMR must still apply diversity (metric-aware relevance).""" + from simplevecdb.types import DistanceStrategy + + db = VectorDB(":memory:", distance_strategy=DistanceStrategy.L2) + col = db.collection("t") + + def v(x: float, y: float) -> list[float]: + a = np.zeros(4, dtype=np.float32) + a[0], a[1] = x, y + return a.tolist() + + # C1a/C1b are a redundant near-duplicate pair; D is the diverse option, + # slightly farther from the query. + col.add_texts( + ["C1a", "C1b", "D"], embeddings=[v(1.0, 0.0), v(1.2, 0.0), v(0.0, 0.5)] + ) + res = col.max_marginal_relevance_search( + v(10.0, 0.0), k=2, fetch_k=10, lambda_mult=0.3 + ) + got = [d.page_content for d in res] + assert "D" in got, f"MMR ignored diversity on L2 (picked redundant pair): {got}" + + +class TestBatchSearchFilterFallback: + def _populate(self, col, dim: int = 8) -> None: + texts, embs, metas = [], [], [] + # 18 docs aligned with the query (near), 2 orthogonal "keep" docs (far). + for i in range(18): + e = np.zeros(dim, dtype=np.float32) + e[0] = 1.0 + e[1] = 0.001 * i + texts.append(f"near{i}") + embs.append(e.tolist()) + metas.append({"keep": False}) + for j in range(2): + e = np.zeros(dim, dtype=np.float32) + e[1] = 1.0 + e[2] = 0.001 * j + texts.append(f"keep{j}") + embs.append(e.tolist()) + metas.append({"keep": True}) + col.add_texts(texts, embeddings=embs, metadatas=metas) + + def test_large_filtered_batch_returns_full_k(self) -> None: + """A >threshold batch with a selective filter must still return k per query.""" + db = VectorDB(":memory:") + col = db.collection("t") + self._populate(col) + q = np.zeros(8, dtype=np.float32) + q[0] = 1.0 + queries = [q.tolist()] * 11 # > USEARCH_BATCH_THRESHOLD (10) -> native path + results = col.similarity_search_batch(queries, k=2, filter={"keep": True}) + assert len(results) == 11 + for r in results: + assert len(r) == 2, f"filtered batch under-delivered: {len(r)}" + assert all(doc.metadata.get("keep") is True for doc, _ in r) + + def test_text_query_outcome_consistent_across_batch_size(self) -> None: + """A text query must behave the same in a small (routed) and a large batch, + not crash only in the large (native) path.""" + db = VectorDB(":memory:") + col = db.collection("t") + self._populate(col) + + def outcome(n: int) -> str: + try: + col.similarity_search_batch(["some text query"] * n, k=2) # type: ignore[list-item] + return "ok" + except Exception as exc: # noqa: BLE001 - classifying the failure mode + m = str(exc).lower() + if any( + s in m + for s in ( + "inhomogeneous", + "could not convert", + "setting an array element", + ) + ): + return "numpy-crash" + return "routed-error" + + small = outcome(3) # <= USEARCH_BATCH_THRESHOLD -> always per-query routed + large = outcome(11) # > threshold + assert small == large, ( + f"text query differs by batch size: small={small} large={large}" + ) + + +class TestRebuildIndexRecovery: + def test_failed_rebuild_keeps_collection_usable( + self, tmp_path: Path, monkeypatch + ) -> None: + """If the rebuild fails after closing the live index, the collection must + re-open the intact on-disk index and stay usable, not brick.""" + from simplevecdb.engine.usearch_index import UsearchIndex + + db = VectorDB(tmp_path / "r.db") + col = db.collection("t", store_embeddings=True) + embs = _emb(20, 16) + col.add_texts([f"d{i}" for i in range(20)], embeddings=embs) + assert len(col.similarity_search(embs[0], k=3)) >= 1 + + orig_add = UsearchIndex.add + + def add_hook(self, *a, **k): # noqa: ANN001, ANN002, ANN003 + # Fail only while building the new (.rebuild) index, after the live + # index has already been closed. + if ".rebuild" in str(self._path): + raise RuntimeError("simulated rebuild add failure") + return orig_add(self, *a, **k) + + monkeypatch.setattr(UsearchIndex, "add", add_hook) + with pytest.raises(RuntimeError): + col.rebuild_index() + + after = col.similarity_search(embs[0], k=3) + assert len(after) >= 1, "collection bricked after failed rebuild" + + +class TestFilterLiteralKey: + def test_sql_filter_matches_literal_top_level_key(self) -> None: + """build_filter_clause must match a literal key 'a.b', not the nested path a->b.""" + db = VectorDB(":memory:") + col = db.collection("t") + col.add_texts( + ["lit", "nested"], + embeddings=_emb(2, 8), + metadatas=[{"a.b": "X"}, {"a": {"b": "X"}}], + ) + cat = col._catalog + ids = cat.find_ids_by_filter({"a.b": "X"}, cat.build_filter_clause) + id_to_text = {doc_id: text for doc_id, text, _ in col.get_documents(limit=100)} + matched = {id_to_text[i] for i in ids} + assert matched == {"lit"}, f"expected literal-key match only, got {matched}" + + def test_quote_in_key_rejected(self) -> None: + db = VectorDB(":memory:") + col = db.collection("t") + col.add_texts(["d"], embeddings=_emb(1, 8)) + cat = col._catalog + with pytest.raises(ValueError): + cat.find_ids_by_filter({'bad"key': 1}, cat.build_filter_clause) + + +class TestClusterSampleSizeCap: + def test_n_clusters_capped_to_sample_size(self, tmp_path: Path) -> None: + """With sampling, n_clusters caps to the sampled count instead of erroring.""" + pytest.importorskip("sklearn") + db = VectorDB(tmp_path / "cs.db") + col = db.collection("t") + col.add_texts([f"d{i}" for i in range(30)], embeddings=_emb(30, 16)) + result = col.cluster(n_clusters=20, sample_size=10, random_state=0) + assert result.n_clusters <= 10 diff --git a/tests/unit/test_tier2_fixes.py b/tests/unit/test_tier2_fixes.py new file mode 100644 index 0000000..8538c81 --- /dev/null +++ b/tests/unit/test_tier2_fixes.py @@ -0,0 +1,117 @@ +"""Tier-2 review fixes — TDD regression tests.""" + +from __future__ import annotations + +import numpy as np + +from simplevecdb import Quantization, VectorDB + + +def _emb(n: int, dim: int, seed: int = 0) -> list[list[float]]: + rng = np.random.default_rng(seed) + return rng.standard_normal((n, dim)).astype(np.float32).tolist() + + +class TestBitGetUnpacks: + def test_bit_index_get_returns_unpacked_float_vectors(self) -> None: + """UsearchIndex.get() on a BIT index must unpack bytes to ±1 floats.""" + db = VectorDB(":memory:") + col = db.collection("t", quantization=Quantization.BIT) + dim = 16 + ids = col.add_texts([f"d{i}" for i in range(4)], embeddings=_emb(4, dim)) + got = col._index.get(np.array(ids, dtype=np.uint64)) + assert got.shape == (4, dim), f"expected (4, {dim}), got {got.shape}" + assert set(np.unique(got)).issubset({-1.0, 1.0}), ( + f"non ±1 values: {np.unique(got)}" + ) + + +class TestRebuildConcurrentCatchup: + def test_rebuild_catches_up_writes_during_build( + self, tmp_path, monkeypatch + ) -> None: + """Writes that land during the (unlocked) rebuild build are folded into + the new index before the swap, so none are lost.""" + from simplevecdb.engine.usearch_index import UsearchIndex + + db = VectorDB(tmp_path / "r.db") + col = db.collection("t", store_embeddings=True) + ids = col.add_texts([f"d{i}" for i in range(10)], embeddings=_emb(10, 8)) + + orig_add = UsearchIndex.add + fired = {"v": False} + new_id: dict[str, int] = {} + + def add_hook(self, keys, vectors, **kwargs): # noqa: ANN001, ANN002, ANN003 + result = orig_add(self, keys, vectors, **kwargs) + # Simulate a writer landing during the build window (once). + if ".rebuild" in str(self._path) and not fired["v"]: + fired["v"] = True + new_id["v"] = col.add_texts( + ["concurrent"], embeddings=_emb(1, 8, seed=99) + )[0] + col.delete_by_ids([ids[0]]) + return result + + monkeypatch.setattr(UsearchIndex, "add", add_hook) + col.rebuild_index() + + keys = set(col._index.keys()) + assert new_id["v"] in keys, "concurrent add lost during rebuild" + assert ids[0] not in keys, "concurrent delete not reflected after rebuild" + + +class TestLangChainRelevanceScoreFn: + def test_cosine_relevance_higher_for_closer(self) -> None: + import pytest + + pytest.importorskip("langchain_core") + from simplevecdb.integrations.langchain import SimpleVecDBVectorStore + + store = SimpleVecDBVectorStore(db_path=":memory:") # default cosine + fn = store._select_relevance_score_fn() + assert fn(0.0) == 1.0 # closest -> max relevance + assert fn(2.0) == 0.0 # farthest cosine distance -> min relevance + assert fn(0.5) > fn(1.5) # monotonically decreasing + + def test_l2_relevance_bounded_and_decreasing(self) -> None: + import pytest + + pytest.importorskip("langchain_core") + from simplevecdb.integrations.langchain import SimpleVecDBVectorStore + from simplevecdb.types import DistanceStrategy + + store = SimpleVecDBVectorStore( + db_path=":memory:", distance_strategy=DistanceStrategy.L2 + ) + fn = store._select_relevance_score_fn() + assert fn(0.0) == 1.0 + assert 0.0 < fn(100.0) < fn(1.0) <= 1.0 + + +class TestLangChainRelevanceEndToEnd: + def test_relevance_scores_in_range_from_real_search(self) -> None: + """End-to-end: similarity_search_with_relevance_scores returns [0,1] + scores (higher=closer) using real backend distances, not just the lambda.""" + import pytest + + pytest.importorskip("langchain_core") + from langchain_core.embeddings import Embeddings + + from simplevecdb.integrations.langchain import SimpleVecDBVectorStore + + class _FakeEmb(Embeddings): + _M = {"near": [1.0, 0.0, 0.0, 0.0], "far": [0.0, 1.0, 0.0, 0.0]} + + def embed_documents(self, texts): # noqa: ANN001, ANN201 + return [self._M[t] for t in texts] + + def embed_query(self, text): # noqa: ANN001, ANN201 + return [1.0, 0.0, 0.0, 0.0] + + store = SimpleVecDBVectorStore(embedding=_FakeEmb()) + store.add_texts(["near", "far"]) + scored = store.similarity_search_with_relevance_scores("q", k=2) + by_text = {doc.page_content: score for doc, score in scored} + assert all(0.0 <= s <= 1.0 for s in by_text.values()), by_text + assert by_text["near"] > by_text["far"], by_text diff --git a/tests/unit/test_tier3_fixes.py b/tests/unit/test_tier3_fixes.py new file mode 100644 index 0000000..c75d948 --- /dev/null +++ b/tests/unit/test_tier3_fixes.py @@ -0,0 +1,65 @@ +"""Tier-3 security fixes — TDD regression tests.""" + +from __future__ import annotations + +import asyncio + +import pytest + + +class TestServerBodyCap: + def test_oversize_content_length_rejected(self) -> None: + server = pytest.importorskip("simplevecdb.embeddings.server") + captured: list[dict] = [] + + async def downstream(scope, receive, send): # noqa: ANN001 + await send({"type": "http.response.start", "status": 200, "headers": []}) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(msg): # noqa: ANN001 + captured.append(msg) + + mw = server._MaxBodySizeMiddleware(downstream, max_bytes=100) + scope = {"type": "http", "headers": [(b"content-length", b"100000")]} + asyncio.run(mw(scope, receive, send)) + assert captured and captured[0]["status"] == 413 + + def test_small_body_passes_through(self) -> None: + server = pytest.importorskip("simplevecdb.embeddings.server") + captured: list[dict] = [] + + async def downstream(scope, receive, send): # noqa: ANN001 + await send({"type": "http.response.start", "status": 200, "headers": []}) + + async def receive(): + return {"type": "http.request", "body": b"hi", "more_body": False} + + async def send(msg): # noqa: ANN001 + captured.append(msg) + + mw = server._MaxBodySizeMiddleware(downstream, max_bytes=1000) + scope = {"type": "http", "headers": [(b"content-length", b"2")]} + asyncio.run(mw(scope, receive, send)) + assert captured[0]["status"] == 200 + + +class TestBodyCapEnvGuard: + def test_bad_or_zero_env_does_not_crash_and_is_floored(self, monkeypatch) -> None: + """A malformed or zero EMBEDDING_SERVER_MAX_BODY_BYTES must not crash the + server import and must never drop the cap below the 1 MiB floor.""" + import importlib + + server = pytest.importorskip("simplevecdb.embeddings.server") + try: + monkeypatch.setenv("EMBEDDING_SERVER_MAX_BODY_BYTES", "not-an-int") + importlib.reload(server) + assert server._MAX_BODY_BYTES >= (1 << 20) + + monkeypatch.setenv("EMBEDDING_SERVER_MAX_BODY_BYTES", "0") + importlib.reload(server) + assert server._MAX_BODY_BYTES >= (1 << 20) + finally: + monkeypatch.delenv("EMBEDDING_SERVER_MAX_BODY_BYTES", raising=False) + importlib.reload(server) diff --git a/tests/unit/test_tier4_fixes.py b/tests/unit/test_tier4_fixes.py new file mode 100644 index 0000000..ca06442 --- /dev/null +++ b/tests/unit/test_tier4_fixes.py @@ -0,0 +1,51 @@ +"""Tier-4 polish/robustness fixes — regression tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from simplevecdb import VectorDB + + +def _emb(n: int, dim: int, seed: int = 0) -> list[list[float]]: + rng = np.random.default_rng(seed) + return rng.standard_normal((n, dim)).astype(np.float32).tolist() + + +class TestFTS5SyntaxError: + def test_malformed_query_raises_valueerror(self) -> None: + """A malformed FTS5 MATCH query surfaces as ValueError, not OperationalError.""" + db = VectorDB(":memory:") + col = db.collection("t") + col.add_texts(["hello world", "foo bar"], embeddings=_emb(2, 8)) + if not col._catalog.fts_enabled: + pytest.skip("FTS5 not available") + with pytest.raises(ValueError): + col.keyword_search('"unbalanced', k=5) + + +class TestClusterTableRollback: + def test_save_cluster_works_after_rolled_back_transaction( + self, tmp_path: Path + ) -> None: + """The cluster table is created eagerly, so a rolled-back first use cannot + leave the ready-flag set without the table existing.""" + pytest.importorskip("sklearn") + db = VectorDB(tmp_path / "c.db") + col = db.collection("t") + col.add_texts([f"d{i}" for i in range(6)], embeddings=_emb(6, 8)) + result = col.cluster(n_clusters=2, random_state=0) + + try: + with db.transaction(): + col.save_cluster("inside_tx", result) + raise RuntimeError("force rollback") + except RuntimeError: + pass + + # Table still exists -> this save succeeds and round-trips. + col.save_cluster("after", result) + assert col.load_cluster("after") is not None diff --git a/uv.lock b/uv.lock index bd1eaac..68d7b76 100755 --- a/uv.lock +++ b/uv.lock @@ -4705,7 +4705,7 @@ wheels = [ [[package]] name = "simplevecdb" -version = "2.6.1" +version = "2.6.2" source = { editable = "." } dependencies = [ { name = "cryptography" },