diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml index b9e2e90..b960d76 100755 --- a/.github/FUNDING.yml +++ b/.github/FUNDING.yml @@ -1,6 +1,5 @@ github: [coderdayton] custom: [ - "https://simplevecdb.lemonsqueezy.com/", - "https://buymeacoffee.com/coderdayton", + "https://ko-fi.com/xbbvii", ] diff --git a/.gitignore b/.gitignore index 9c0b1eb..89ca3aa 100755 --- a/.gitignore +++ b/.gitignore @@ -1,25 +1,61 @@ # Python / uv .venv/ __pycache__/ -*.pyc -.env -dist/ -build/ +*.py[cod] *.egg-info/ +.eggs/ +build/ +dist/ +.tox/ +.nox/ + +# Tooling caches .mypy_cache/ .pytest_cache/ .ruff_cache/ +.cache/ +.hypothesis/ +cython_debug/ -# IDE -.vscode/ -.idea/ +# Coverage +.coverage +.coverage.* +coverage.xml +*.cover +htmlcov/ # Jupyter -.ipynb_checkpoints +.ipynb_checkpoints/ + +# Environment / secrets +.env +.envrc +.direnv/ -# Databases +# Databases (SimpleVecDB writes WAL/SHM sidecars) *.db +*.db-journal +*.db-shm +*.db-wal *.sqlite +*.sqlite3 + +# Editor / IDE +.vscode/ +.idea/ +.history/ +*.iml +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db +desktop.ini + +# Docs build +site/ # Agentic CLI tools (per-developer state) .opencode/ @@ -27,12 +63,8 @@ opencode.json .claude/ .codex -# Project specific +# Project-specific scratch simplevecdb_plan.md AGENTS.md -htmlcov/ -site/ -htmlcov/ -.coverage NEXT_UPDATES.md -pro_pack/ \ No newline at end of file +pro_pack/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 9d9f81b..9b0e503 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,109 @@ 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.1] - 2026-05-10 + +### Storage, mutation, and eventing improvements + +This release closes ten long-standing gaps in the catalog layer with a coherent +set of additive primitives. No public API breaks; existing 2.6.0 databases +upgrade transparently (the new tables are created on first open). + +#### New features + +- **Native vector update via pending buffer** — `collection.update_embedding(id, vector)` + writes a row to a per-collection `_pending_vectors` overlay inside one SQL + transaction; the new vector becomes visible to reads immediately and is + promoted to the HNSW index on `collection.pending.flush()`. Removes the + HNSW remove+re-add churn previously required for in-place updates. +- **Bulk vector math** — `collection.pending.update_many([(id, vec), …])` and + `collection.pending.blend_toward(ids, centroid, alpha)` for batched edits. +- **Atomic transaction boundary** — `with db.transaction() as tx: …` and + `with collection.tx(): …` wrap a SAVEPOINT around catalog writes + (metadata, counters, edges, events, TTL, and `update_embedding`'s + pending overlay) so a raised exception rolls all SQL writes back. + Coarse vector mutations (`add_texts`, `delete`) are NOT rolled back — + use `update_embedding` + `pending.flush()` for vector changes that + must be commit-gated. Nested contexts share a single savepoint stack + via the new `_TxState` helper. +- **Weighted directed edges** — new `collection.edges` namespace with + `add_edge / get_edges / update_edge / delete_edge / prune` over a + per-collection `_edges` table. Numeric columns (`weight`, `bonus`, `hits`, + `last_touch`) are addressable by the new range-filter grammar; deltas + (`dweight=+0.02, dhits=+1`) compile to a single atomic SQL UPDATE. +- **Atomic counter increments** — `collection.increment_metadata(id, {"hits": 1, "drift": 0.02})` + applies a dict of numeric deltas to JSON metadata in one statement using + chained `json_set(... json_extract + ?)` calls. WAL-atomic; safe under + concurrent writers. +- **Mongo-style range filters** — `filter={"score": {"$gt": 0.5, "$lte": 0.9}}` + on `similarity_search`, `keyword_search`, `hybrid_search`, `edges.get_edges`, + and `events.read`. Supported operators: `$eq $ne $gt $gte $lt $lte $in $nin + $exists $between`. Tuple shorthand (`("range", lo, hi)`, `(">", x)`) is + normalised into the operator-dict form. +- **Append-only change feed** — every mutating method now appends one row to + a per-collection `_events` table (kind, doc_id, payload, monotonic seq). + `collection.events.read(since=, kind=, limit=)`, + `collection.events.subscribe(since=, poll_interval=)`, and + `collection.events.prune(before_seq=)` expose the feed; cross-process + visibility comes from the existing WAL mode. +- **TTL / expiry hooks** — `collection.ttl.set(id, seconds=…, on_expire="delete"|"callback")`, + `collection.ttl.clear(id)`, and `collection.ttl.sweep()` over a + `_ttl` table; `start_background(interval=…)` runs the sweep in a daemon + thread (off by default). +- **Incremental rebuild scheduler** — `collection.maintenance.rebuild_if_needed(max_pending=, max_deleted=)` + triggers a full `rebuild_index()` only when the configured pending / + tombstone / wall-time thresholds are crossed. +- **Multi-process write safety** — added `PRAGMA busy_timeout=5000` and + `PRAGMA foreign_keys=ON` at every connection-open site (encrypted and + unencrypted). The native 5 s wait window reduces `DatabaseLockedError` + pressure under contention; foreign keys cascade-delete pending / + edges / TTL rows when a doc is deleted. The events table is + intentionally FK-less so the audit trail survives deletions. +- **Async wrappers** — `AsyncVectorCollection` gains async equivalents of the + new methods (`update_embedding`, `flush_pending`, `increment_metadata`, + `add_edge`, `update_edge`, `delete_edge`, `get_edges`, `set_ttl`, + `clear_ttl`, `sweep_ttl`, `read_events`, `last_event_seq`, + `rebuild_if_needed`). + +#### New types & constants + +- `simplevecdb.types`: `Edge`, `Event`, `TTLEntry` frozen dataclasses. +- `simplevecdb.constants`: `PENDING_FLUSH_DEFAULT_BATCH=1000`, + `EVENTS_POLL_INTERVAL_S=0.1`, `EVENTS_RETENTION_LIMIT=100_000`, + `TTL_SWEEP_DEFAULT_INTERVAL_S=60.0`, `REBUILD_PENDING_THRESHOLD=5_000`, + `REBUILD_TOMBSTONE_THRESHOLD=5_000`, `REBUILD_MIN_INTERVAL_S=3600.0`, + `SQLITE_BUSY_TIMEOUT_MS=5000`. + +#### Test coverage + +- `tests/unit/test_v26_1_features.py` — 25 tests covering the five must-have + primitives end-to-end: `update_embedding` + pending buffer + flush; edges + CRUD with atomic deltas, range filtering, and prune; `increment_metadata` + under 800-thread contention (exact total preserved); transaction rollback + and commit semantics; Mongo-style and tuple-shorthand range filters in + `similarity_search`; events append on every mutation; TTL sweep with + `delete` and `callback` paths; threshold-driven rebuild scheduler. + +#### Removed + +- **`sqlite-vec` dependency** dropped from `pyproject.toml`. The package was + never imported and the v1.x → v2.0 auto-migration code path could not have + worked without explicitly loading the extension at connection time. +- **`MigrationRequiredError`**, **`VectorDB.check_migration()`**, and the + **`auto_migrate=`** constructor flag have been removed. Databases written + by `simplevecdb < 2.0.0` (sqlite-vec backend) are no longer auto-migrated + on open. To upgrade a v1.x database, dump the rows with a v1.x install and + re-ingest them through the v2 API; or stay on the last release that shipped + the migration path (anything ≤ v2.6.1's predecessor). +- The catalog helpers `check_legacy_sqlite_vec`, `get_legacy_vectors`, and + `drop_legacy_vec_table` are gone alongside the migration entry point. + +#### Out of scope + +- No external pub/sub for events — polling only. +- No multi-master writer support; single-writer + many readers remains the + recommended topology. + ## [2.6.0] - 2026-05-06 ### Review pass 3 — final correctness/security pass before tag diff --git a/README.md b/README.md index 43a714d..3ae1874 100755 --- a/README.md +++ b/README.md @@ -5,6 +5,8 @@ [![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) +Buy Me a Coffee at ko-fi.com + **The dead-simple, local-first vector database.** 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. @@ -70,48 +72,52 @@ python -c "from simplevecdb import VectorDB; print('SimpleVecDB installed succes ## Quickstart -SimpleVecDB is **just a vector storage layer**—it doesn't include an LLM or generate embeddings. This design keeps it lightweight and flexible. Choose your integration path: - -### Option 1: With OpenAI (Simplest) +SimpleVecDB is just a storage and search layer — it doesn't ship an LLM +and won't generate embeddings for you. Bring whichever embedding source +you already use; three common ones below. -Best for: Quick prototypes, production apps with OpenAI subscriptions. +### Option 1: OpenAI embeddings ```python from simplevecdb import VectorDB from openai import OpenAI -db = VectorDB("knowledge.db") -collection = db.collection("docs") client = OpenAI() - -texts = ["Paris is the capital of France.", "Mitochondria powers cells."] -embeddings = [ - client.embeddings.create(model="text-embedding-3-small", input=t).data[0].embedding - for t in texts +db = VectorDB("notes.db") +notes = db.collection("personal") + +def embed(text: str) -> list[float]: + return ( + client.embeddings + .create(model="text-embedding-3-small", input=text) + .data[0].embedding + ) + +entries = [ + ("Cherry MX silent reds bottom out around 45g — quieter than browns", "keyboards"), + ("Sourdough hydration sweet spot is ~75% with this flour", "baking"), + ("EXPLAIN ANALYZE showed seq scan; ANALYZE on the table fixed it", "work"), + ("Passport renewal took 3 weeks, not the advertised 6–8", "admin"), ] -collection.add_texts( - texts=texts, - embeddings=embeddings, - metadatas=[{"category": "geography"}, {"category": "biology"}] +notes.add_texts( + texts=[t for t, _ in entries], + embeddings=[embed(t) for t, _ in entries], + metadatas=[{"tag": tag} for _, tag in entries], ) -# Search -query_emb = client.embeddings.create( - model="text-embedding-3-small", - input="capital of France" -).data[0].embedding +hits = notes.similarity_search(embed("how loud are silent reds"), k=2) +for doc, score in hits: + print(f"{score:.3f} {doc.page_content}") -results = collection.similarity_search(query_emb, k=1) -print(results[0][0].page_content) # "Paris is the capital of France." - -# Filter by metadata -filtered = collection.similarity_search(query_emb, k=10, filter={"category": "geography"}) +work = notes.similarity_search( + embed("query plan slow"), + k=5, + filter={"tag": "work"}, +) ``` -### Option 2: Fully Local (Privacy-First) - -Best for: Offline apps, sensitive data, zero API costs. +### Option 2: Fully local (no network, no API key) ```bash pip install "simplevecdb[server]" @@ -121,36 +127,38 @@ pip install "simplevecdb[server]" from simplevecdb import VectorDB from simplevecdb.embeddings.models import embed_texts -db = VectorDB("local.db") -collection = db.collection("docs") - -texts = ["Paris is the capital of France.", "Mitochondria powers cells."] -embeddings = embed_texts(texts) # Local HuggingFace models - -collection.add_texts(texts=texts, embeddings=embeddings) +db = VectorDB("notes.db") +notes = db.collection("personal") -# Search -query_emb = embed_texts(["capital of France"])[0] -results = collection.similarity_search(query_emb, k=1) +texts = [ + "Cherry MX silent reds bottom out around 45g", + "Sourdough hydration sweet spot is ~75% with this flour", + "EXPLAIN ANALYZE showed seq scan; ANALYZE on the table fixed it", +] +notes.add_texts(texts=texts, embeddings=embed_texts(texts)) -# Hybrid search (BM25 + vector) -hybrid = collection.hybrid_search("powerhouse cell", k=2) +vec = notes.similarity_search(embed_texts(["quieter switches"])[0], k=2) +mixed = notes.hybrid_search("postgres slow query", k=3) ``` -**Optional: Run embeddings server (OpenAI-compatible)** +If you'd rather hit an HTTP endpoint than import the embedding models +directly, the bundled server speaks the same shape as OpenAI's +embeddings API: ```bash -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 +simplevecdb-server --port 8000 # default model, auto warm-up +simplevecdb-server --host 0.0.0.0 --port 9000 +simplevecdb-server --no-warmup # skip the model preload +simplevecdb-server --help ``` -See [Setup Guide](ENV_SETUP.md) for configuration: model registry, rate limits, API keys, CORS, CUDA optimization. +Server tuning (model registry, rate limits, API keys, CORS, CUDA) lives +in the [Setup Guide](ENV_SETUP.md). -### Option 3: With LangChain or LlamaIndex +### Option 3: LangChain or LlamaIndex -Best for: Existing RAG pipelines, framework-based workflows. +Already wired into one of the big RAG frameworks? Drop SimpleVecDB in +as the vector store: ```bash pip install "simplevecdb[integrations]" @@ -161,214 +169,64 @@ from simplevecdb.integrations.langchain import SimpleVecDBVectorStore from langchain_openai import OpenAIEmbeddings store = SimpleVecDBVectorStore( - db_path="langchain.db", - embedding=OpenAIEmbeddings(model="text-embedding-3-small") + db_path="notes.db", + embedding=OpenAIEmbeddings(model="text-embedding-3-small"), ) -store.add_texts(["Paris is the capital of France."]) -results = store.similarity_search("capital of France", k=1) -hybrid = store.hybrid_search("France capital", k=3) # BM25 + vector +store.add_texts([ + "Cherry MX silent reds bottom out around 45g", + "EXPLAIN ANALYZE showed seq scan; ANALYZE on the table fixed it", +]) +store.similarity_search("quieter switches", k=1) +store.hybrid_search("postgres performance", k=3) ``` -**LlamaIndex:** +LlamaIndex is the same shape: ```python from simplevecdb.integrations.llamaindex import SimpleVecDBLlamaStore from llama_index.embeddings.openai import OpenAIEmbedding store = SimpleVecDBLlamaStore( - db_path="llama.db", - embedding=OpenAIEmbedding(model="text-embedding-3-small") -) -``` - -See **[Examples](https://coderdayton.github.io/SimpleVecDB/examples/)** for complete RAG workflows with Ollama. - -## Core Features - -### Multi-Collection Support - -Organize vectors by domain within a single database file: - -```python -from simplevecdb import VectorDB, Quantization - -db = VectorDB("app.db") -users = db.collection("users", quantization=Quantization.FLOAT16) # 2x memory savings -products = db.collection("products", quantization=Quantization.BIT) # 32x compression - -# Isolated namespaces -users.add_texts(["Alice likes hiking"], embeddings=[[0.1]*384]) -products.add_texts(["Hiking boots"], embeddings=[[0.9]*384]) -``` - -### Search Capabilities - -```python -# Vector similarity (cosine/L2) - adaptive search by default -results = collection.similarity_search(query_vector, k=10) - -# Force exact search for perfect recall (brute-force) -results = collection.similarity_search(query_vector, k=10, exact=True) - -# Force HNSW approximate search (faster, may miss some results) -results = collection.similarity_search(query_vector, k=10, exact=False) - -# Parallel search with explicit thread count -results = collection.similarity_search(query_vector, k=10, threads=8) - -# Batch search - 10x throughput for multiple queries -queries = [query1, query2, query3] # List of embedding vectors -batch_results = collection.similarity_search_batch(queries, k=10) - -# Keyword search (BM25) -results = collection.keyword_search("exact phrase", k=10) - -# Hybrid (BM25 + vector fusion) -results = collection.hybrid_search("machine learning", k=10) -results = collection.hybrid_search("ML concepts", query_vector=my_vector, k=10) - -# Metadata filtering -results = collection.similarity_search( - query_vector, - k=10, - filter={"category": "technical", "verified": True} -) -``` - -> **Tip:** LangChain and LlamaIndex integrations support all search methods. - -### Encryption (v2.1+) - -Protect sensitive data with AES-256 at-rest encryption: - -```bash -pip install "simplevecdb[encryption]" -``` - -```python -from simplevecdb import VectorDB - -# Create encrypted database -db = VectorDB("secure.db", encryption_key="your-secret-key") -collection = db.collection("confidential") - -collection.add_texts(["sensitive data"], embeddings=[[0.1]*384]) -db.close() - -# Reopen requires same key -db = VectorDB("secure.db", encryption_key="your-secret-key") -``` - -### Streaming Insert (v2.1+) - -Memory-efficient ingestion for large datasets: - -```python -def load_documents(): - for line in open("large_file.jsonl"): - doc = json.loads(line) - yield (doc["text"], doc.get("metadata"), doc.get("embedding")) - -for progress in collection.add_texts_streaming(load_documents(), batch_size=1000): - print(f"Processed {progress['docs_processed']} documents") -``` - -### Document Hierarchies (v2.1+) - -Organize documents in parent-child relationships: - -```python -# Add parent document -parent_ids = collection.add_texts(["Main document"], embeddings=[[0.1]*384]) - -# Add children -child_ids = collection.add_texts( - ["Chunk 1", "Chunk 2"], - embeddings=[[0.11]*384, [0.12]*384], - parent_ids=[parent_ids[0], parent_ids[0]] + db_path="notes.db", + embedding=OpenAIEmbedding(model="text-embedding-3-small"), ) - -# Navigate hierarchy -children = collection.get_children(parent_ids[0]) -parent = collection.get_parent(child_ids[0]) -descendants = collection.get_descendants(parent_ids[0]) ``` -### Document Management (v2.4+) +End-to-end notebooks (including a fully local Ollama RAG) live in the +[examples gallery](https://coderdayton.github.io/SimpleVecDB/examples/). + +## Feature Highlights + +A few of the things SimpleVecDB does well — see +[`docs/Features.md`](docs/Features.md) for the comprehensive list. + +- **Vector + keyword + hybrid search** — cosine / L2 similarity, BM25 + via SQLite FTS5, and Reciprocal Rank Fusion in one collection. +- **Adaptive HNSW** — brute-force for <10k vectors (perfect recall), + `usearch` HNSW above that. Override per query with `exact=True/False`. +- **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. +- **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 + flush, weighted directed edges, append-only event feed, TTL with + delete/callback sweep, and a threshold-driven rebuild scheduler. +- **Atomic counters & transactions (v2.6.1)** — `increment_metadata` for + JSON deltas in one statement; SAVEPOINT-backed `db.transaction()` / + `collection.tx()` rolling all catalog writes back on error. +- **Async, encryption, clustering, hierarchies** — full async surface + (with executor injection), SQLCipher AES-256, K-means / MiniBatch + K-means / HDBSCAN, parent/child relationships. +- **Framework integrations** — drop-in `LangChain` and `LlamaIndex` + adapters via the `[integrations]` extra; optional FastAPI embeddings + server via `[server]`. + +For full method-level coverage, see [the Features doc](docs/Features.md) +or the [API reference](https://coderdayton.github.io/SimpleVecDB/api/core). -Query and update documents without touching private internals: - -```python -# Get all documents (with optional metadata filter) -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]) - -# Batch update metadata (shallow merge) -collection.update_metadata([ - (1, {"reviewed": True}), - (2, {"reviewed": True, "score": 0.95}), -]) - -# 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+) - -Discover natural groupings in your embeddings: - -```python -# Cluster documents and auto-generate tags -result = collection.cluster(n_clusters=5) -tags = collection.auto_tag(result, method="tfidf") -collection.assign_cluster_metadata(result, tags) - -# Save for fast assignment of new documents -collection.save_cluster("categories", result) -collection.assign_to_cluster("categories", new_doc_ids) -``` - -Supports K-means, MiniBatch K-means, and HDBSCAN. See [Clustering Guide](https://coderdayton.github.io/SimpleVecDB/guides/clustering) for details. - -## Feature Matrix - -| Feature | Status | Description | -| :------------------------ | :----- | :----------------------------------------------------------- | -| **Single-File Storage** | ✅ | SQLite `.db` file or in-memory mode | -| **Multi-Collection** | ✅ | Isolated namespaces per database | -| **HNSW Indexing** | ✅ | usearch HNSW for 10-100x faster search | -| **Adaptive Search** | ✅ | Auto brute-force for <10k vectors, HNSW for larger | -| **Vector Search** | ✅ | Cosine, Euclidean metrics (L1 removed in v2.0) | -| **Hybrid Search** | ✅ | BM25 + vector fusion (Reciprocal Rank Fusion) | -| **Quantization** | ✅ | FLOAT32, FLOAT16, INT8, BIT for 2-32x compression | -| **Parallel Operations** | ✅ | `threads` parameter for add/search | -| **Metadata Filtering** | ✅ | SQL `WHERE` clause support | -| **Framework Integration** | ✅ | LangChain \& LlamaIndex adapters via `[integrations]` extra | -| **Hardware Acceleration** | ✅ | Auto-detects CUDA/MPS/CPU + SIMD via usearch | -| **Local Embeddings** | ✅ | HuggingFace models via `[server]` extras | -| **Built-in Encryption** | ✅ | SQLCipher AES-256 at-rest encryption via `[encryption]` extras | -| **Streaming Insert** | ✅ | Memory-efficient large-scale ingestion with progress callbacks | -| **Document Hierarchies** | ✅ | Parent/child relationships for chunked docs | -| **Vector Clustering** | ✅ | K-means, MiniBatch K-means, HDBSCAN with auto-tagging (v2.2+) | -| **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 @@ -388,6 +246,7 @@ Supports K-means, MiniBatch K-means, and HDBSCAN. See [Clustering Guide](https:/ ## Documentation +- **[Features](docs/Features.md)** — Comprehensive list of every capability, grouped by area - **[Setup Guide](https://coderdayton.github.io/SimpleVecDB/ENV_SETUP)** — Environment variables, server configuration, authentication - **[API Reference](https://coderdayton.github.io/SimpleVecDB/api/core)** — Complete class/method documentation with type signatures - **[Benchmarks](https://coderdayton.github.io/SimpleVecDB/benchmarks)** — Quantization strategies, batch sizes, hardware optimization @@ -430,24 +289,14 @@ pip install torch --index-url https://download.pytorch.org/whl/cu118 ## Roadmap -- [x] Hybrid Search (BM25 + Vector) -- [x] Multi-collection support -- [x] HNSW indexing (usearch backend) -- [x] Adaptive search (brute-force/HNSW) -- [x] SQLCipher encryption (at-rest data protection) -- [x] Streaming insert API for large-scale ingestion -- [x] Hierarchical document relationships (parent/child) -- [x] Cross-collection search -- [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) +What's on the near-term radar: + - [ ] Incremental clustering (online learning) - [ ] Cluster visualization exports -Vote on features or propose new ones in [GitHub Discussions](https://github.com/coderdayton/simplevecdb/discussions). +For shipped capabilities, see [`docs/Features.md`](docs/Features.md) and the +release-by-release [Changelog](CHANGELOG.md). Vote on these or propose new +ideas in [GitHub Discussions](https://github.com/coderdayton/simplevecdb/discussions). ## Contributing @@ -469,30 +318,13 @@ 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 -## Sponsors - -SimpleVecDB is independently developed and maintained. If you or your company use it in production, please consider sponsoring to ensure its continued development and support. - -**Company Sponsors** - -_Become the first company sponsor!_ [Support on GitHub →](https://github.com/sponsors/coderdayton) - -**Individual Supporters** +## Other Ways to Support -_Join the list of supporters!_ [Support on GitHub →](https://github.com/sponsors/coderdayton) - - - -### Other Ways to Support - -- 🍵 **[Buy me a coffee](https://www.buymeacoffee.com/coderdayton)** - One-time donation -- 💎 **[Get the Pro Pack](https://simplevecdb.lemonsqueezy.com/)** - Production deployment templates & recipes +- ☕ **[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) -**Why sponsor?** Your support ensures SimpleVecDB stays maintained, secure, and compatible with the latest Python/SQLite versions. - ## License [MIT License](LICENSE) — Free for personal and commercial use. diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 9d9f81b..9b0e503 100755 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -5,6 +5,109 @@ 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.1] - 2026-05-10 + +### Storage, mutation, and eventing improvements + +This release closes ten long-standing gaps in the catalog layer with a coherent +set of additive primitives. No public API breaks; existing 2.6.0 databases +upgrade transparently (the new tables are created on first open). + +#### New features + +- **Native vector update via pending buffer** — `collection.update_embedding(id, vector)` + writes a row to a per-collection `_pending_vectors` overlay inside one SQL + transaction; the new vector becomes visible to reads immediately and is + promoted to the HNSW index on `collection.pending.flush()`. Removes the + HNSW remove+re-add churn previously required for in-place updates. +- **Bulk vector math** — `collection.pending.update_many([(id, vec), …])` and + `collection.pending.blend_toward(ids, centroid, alpha)` for batched edits. +- **Atomic transaction boundary** — `with db.transaction() as tx: …` and + `with collection.tx(): …` wrap a SAVEPOINT around catalog writes + (metadata, counters, edges, events, TTL, and `update_embedding`'s + pending overlay) so a raised exception rolls all SQL writes back. + Coarse vector mutations (`add_texts`, `delete`) are NOT rolled back — + use `update_embedding` + `pending.flush()` for vector changes that + must be commit-gated. Nested contexts share a single savepoint stack + via the new `_TxState` helper. +- **Weighted directed edges** — new `collection.edges` namespace with + `add_edge / get_edges / update_edge / delete_edge / prune` over a + per-collection `_edges` table. Numeric columns (`weight`, `bonus`, `hits`, + `last_touch`) are addressable by the new range-filter grammar; deltas + (`dweight=+0.02, dhits=+1`) compile to a single atomic SQL UPDATE. +- **Atomic counter increments** — `collection.increment_metadata(id, {"hits": 1, "drift": 0.02})` + applies a dict of numeric deltas to JSON metadata in one statement using + chained `json_set(... json_extract + ?)` calls. WAL-atomic; safe under + concurrent writers. +- **Mongo-style range filters** — `filter={"score": {"$gt": 0.5, "$lte": 0.9}}` + on `similarity_search`, `keyword_search`, `hybrid_search`, `edges.get_edges`, + and `events.read`. Supported operators: `$eq $ne $gt $gte $lt $lte $in $nin + $exists $between`. Tuple shorthand (`("range", lo, hi)`, `(">", x)`) is + normalised into the operator-dict form. +- **Append-only change feed** — every mutating method now appends one row to + a per-collection `_events` table (kind, doc_id, payload, monotonic seq). + `collection.events.read(since=, kind=, limit=)`, + `collection.events.subscribe(since=, poll_interval=)`, and + `collection.events.prune(before_seq=)` expose the feed; cross-process + visibility comes from the existing WAL mode. +- **TTL / expiry hooks** — `collection.ttl.set(id, seconds=…, on_expire="delete"|"callback")`, + `collection.ttl.clear(id)`, and `collection.ttl.sweep()` over a + `_ttl` table; `start_background(interval=…)` runs the sweep in a daemon + thread (off by default). +- **Incremental rebuild scheduler** — `collection.maintenance.rebuild_if_needed(max_pending=, max_deleted=)` + triggers a full `rebuild_index()` only when the configured pending / + tombstone / wall-time thresholds are crossed. +- **Multi-process write safety** — added `PRAGMA busy_timeout=5000` and + `PRAGMA foreign_keys=ON` at every connection-open site (encrypted and + unencrypted). The native 5 s wait window reduces `DatabaseLockedError` + pressure under contention; foreign keys cascade-delete pending / + edges / TTL rows when a doc is deleted. The events table is + intentionally FK-less so the audit trail survives deletions. +- **Async wrappers** — `AsyncVectorCollection` gains async equivalents of the + new methods (`update_embedding`, `flush_pending`, `increment_metadata`, + `add_edge`, `update_edge`, `delete_edge`, `get_edges`, `set_ttl`, + `clear_ttl`, `sweep_ttl`, `read_events`, `last_event_seq`, + `rebuild_if_needed`). + +#### New types & constants + +- `simplevecdb.types`: `Edge`, `Event`, `TTLEntry` frozen dataclasses. +- `simplevecdb.constants`: `PENDING_FLUSH_DEFAULT_BATCH=1000`, + `EVENTS_POLL_INTERVAL_S=0.1`, `EVENTS_RETENTION_LIMIT=100_000`, + `TTL_SWEEP_DEFAULT_INTERVAL_S=60.0`, `REBUILD_PENDING_THRESHOLD=5_000`, + `REBUILD_TOMBSTONE_THRESHOLD=5_000`, `REBUILD_MIN_INTERVAL_S=3600.0`, + `SQLITE_BUSY_TIMEOUT_MS=5000`. + +#### Test coverage + +- `tests/unit/test_v26_1_features.py` — 25 tests covering the five must-have + primitives end-to-end: `update_embedding` + pending buffer + flush; edges + CRUD with atomic deltas, range filtering, and prune; `increment_metadata` + under 800-thread contention (exact total preserved); transaction rollback + and commit semantics; Mongo-style and tuple-shorthand range filters in + `similarity_search`; events append on every mutation; TTL sweep with + `delete` and `callback` paths; threshold-driven rebuild scheduler. + +#### Removed + +- **`sqlite-vec` dependency** dropped from `pyproject.toml`. The package was + never imported and the v1.x → v2.0 auto-migration code path could not have + worked without explicitly loading the extension at connection time. +- **`MigrationRequiredError`**, **`VectorDB.check_migration()`**, and the + **`auto_migrate=`** constructor flag have been removed. Databases written + by `simplevecdb < 2.0.0` (sqlite-vec backend) are no longer auto-migrated + on open. To upgrade a v1.x database, dump the rows with a v1.x install and + re-ingest them through the v2 API; or stay on the last release that shipped + the migration path (anything ≤ v2.6.1's predecessor). +- The catalog helpers `check_legacy_sqlite_vec`, `get_legacy_vectors`, and + `drop_legacy_vec_table` are gone alongside the migration entry point. + +#### Out of scope + +- No external pub/sub for events — polling only. +- No multi-master writer support; single-writer + many readers remains the + recommended topology. + ## [2.6.0] - 2026-05-06 ### Review pass 3 — final correctness/security pass before tag diff --git a/docs/Features.md b/docs/Features.md new file mode 100644 index 0000000..a846c16 --- /dev/null +++ b/docs/Features.md @@ -0,0 +1,206 @@ +# Features + +A complete tour of what SimpleVecDB offers, grouped by capability. For +quick install + first query, start with the [README](../README.md). For +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. +- **Multi-collection** — isolated namespaces per database via + `db.collection("name")`. Each collection has its own quantization, + distance metric, and (optional) embedding storage. +- **WAL mode + 5 s busy timeout** — concurrent readers don't block writers, + and `PRAGMA busy_timeout=5000` cuts `DatabaseLockedError` pressure under + contention. +- **Foreign keys cascade** — deleting a doc cascade-cleans its pending + vectors, edges, and TTL rows. The events feed is intentionally FK-less so + the audit trail survives deletions. +- **Encryption (SQLCipher AES-256)** — `VectorDB(path, encryption_key=…)`, + via `[encryption]` extra. Salt and key derivation hardened in v2.6.0. +- **Cross-process safety** — advisory file locking on `usearch` index files + prevents two processes from corrupting the same index. +- **Vacuum** — `db.vacuum()` reclaims disk space; truncates WAL by default. + +## Indexing & vectors + +- **HNSW via usearch** — 10–100× faster than brute force on collections >10k. +- **Adaptive search** — brute-force for <10k vectors (perfect recall), HNSW + above that. Override per query with `exact=True` / `exact=False`. +- **Quantization** — `FLOAT32`, `FLOAT16` (2× compression), `INT8` (4×), + `BIT` (32×). +- **Distance metrics** — `COSINE`, `L2`. (L1 removed in v2.0.) +- **Hardware acceleration** — auto-detects CUDA / MPS / CPU + SIMD via + `usearch`. +- **Index rebuild** — `collection.rebuild_index(connectivity=, expansion_add=, + expansion_search=)` for tuning; `collection.maintenance.rebuild_if_needed( + max_pending=, max_deleted=)` triggers only when thresholds are crossed. + +## Search + +- **Vector similarity** — `collection.similarity_search(vec | text, k=, filter=, + exact=, threads=)`. +- **Batch vector search** — `similarity_search_batch(queries, k=)` for + ~10× throughput. +- **Keyword (BM25)** — `collection.keyword_search(query, k=, filter=)` + backed by SQLite FTS5. +- **Hybrid (BM25 + vector)** — `collection.hybrid_search(query, k=, + query_vector=, vector_k=, keyword_k=, rrf_k=)` using Reciprocal Rank + Fusion. +- **Max-Marginal-Relevance** — `collection.max_marginal_relevance_search( + query, k=, fetch_k=, lambda_mult=)` for diversified results. +- **Cross-collection search** — `db.search_collections(query, collections=, + k=, filter=, normalize_scores=, parallel=)` merges and re-ranks across + collections. +- **Range / set filters (v2.6.1)** — Mongo-style operators in `filter=`: + `$eq $ne $gt $gte $lt $lte $in $nin $exists $between`. Tuple shorthand + (`(">", 0.5)`, `("range", lo, hi)`) is normalised into the operator-dict + form. Works on `similarity_search`, `keyword_search`, `hybrid_search`, + `edges.get_edges`, and `events.read`. + +## Mutation & updates + +- **`add_texts`** — batch insert with optional metadata, embeddings, + parent IDs, and explicit thread count. +- **`add_texts_streaming`** — generator-driven ingestion for large + datasets, with progress callbacks. +- **`delete_by_ids` / `remove_texts(filter=)`** — point and bulk deletes. +- **`update_metadata([(id, patch), …])`** — shallow-merge metadata batch. +- **`update_embedding(id, vector)` (v2.6.1)** — buffers a vector update in + a per-collection `_pending_vectors` overlay. New vector becomes visible + to reads immediately; promoted to HNSW on `pending.flush()`. Removes + the HNSW remove+re-add churn previously required for in-place edits. +- **Bulk vector math (v2.6.1)** — `collection.pending.update_many([(id, + vec), …])` and `collection.pending.blend_toward(ids, centroid, alpha)`. +- **Atomic counters (v2.6.1)** — `collection.increment_metadata(id, + {"hits": 1, "drift": 0.02})` applies dict-of-deltas to JSON metadata in + one statement; WAL-atomic and safe under concurrent writers. +- **Transactions (v2.6.1)** — `with db.transaction() as tx: …` and + `with collection.tx(): …` wrap a SAVEPOINT around catalog writes + (metadata, counters, edges, events, TTL, and `update_embedding`'s + pending overlay). A raised exception rolls all SQL writes back. Coarse + vector mutations (`add_texts`, `delete_by_ids`) are NOT rolled back — + use `update_embedding` + `pending.flush()` for vector changes that + must be commit-gated. + +## Relationships + +- **Document hierarchies (v2.1+)** — `add_texts(..., parent_ids=…)`, + `set_parent`, `get_children`, `get_parent`, `get_descendants`, + `get_ancestors`. Useful for chunked-doc retrieval where children are + the search target but the parent provides context. +- **Weighted directed edges (v2.6.1)** — `collection.edges` namespace: + - `add_edge(src, dst, kind=, weight=, bonus=, hits=, metadata=)` + - `update_edge(src, dst, kind=, dweight=, dbonus=, dhits=)` — deltas + compile to a single atomic SQL UPDATE. + - `get_edges(src=, dst=, kind=, filter=, limit=)` — supports + range/set filters on numeric columns. + - `delete_edge`, `prune_edges`. Edges have their own `last_touch` + timestamp. + +## Lifecycle & memory primitives (v2.6.1) + +These primitives turn the database into a substrate for retrieval-with- +memory systems (frecency, decay, change feeds, expiry). + +- **TTL / expiry** — + - `collection.ttl.set(doc_id, seconds=… | expires_at=…, on_expire= + "delete" | "callback")` + - `collection.ttl.clear(doc_id)` + - `collection.ttl.sweep(now=, limit=)` returns `(deleted_ids, + callback_ids)`. + - `collection.ttl.start_background(interval=…)` runs the sweep in a + daemon thread (off by default). +- **Append-only event feed** — every mutating method appends one row to a + per-collection `_events` table (kind, doc_id, payload, monotonic seq). + - `collection.events.read(since=, kind=, limit=)` + - `collection.events.subscribe(since=, poll_interval=)` + - `collection.events.prune(before_seq=)` + - `collection.events.last_seq()` +- **Incremental rebuild scheduler** — `collection.maintenance.rebuild_if_needed( + max_pending=, max_deleted=)` triggers a full `rebuild_index()` only when + the configured pending / tombstone / wall-time thresholds are crossed. + +## Clustering (v2.2+) + +- **Algorithms** — K-means, MiniBatch K-means, HDBSCAN. +- **Workflow** — `cluster() → auto_tag() → assign_cluster_metadata()`. +- **Auto-tag methods** — `keywords`, `tfidf`, or a `custom_callback`. +- **Persistence** — `save_cluster`, `load_cluster`, `list_clusters`, + `delete_cluster`, `assign_to_cluster` for fast assignment of new + documents. +- **Discovery** — `get_cluster_members(cluster_id)`. +- See the [Clustering Guide](guides/clustering.md) for tuning advice. + +## Document management + +- `collection.get_documents(filter_dict=, limit=, offset=)` — + paginated catalog access. +- `collection.get_embeddings_by_ids([…])` — fetch stored embeddings (when + `store_embeddings=True`). +- `collection.count()`, `collection.dim`. +- `db.list_collections()`, `db.delete_collection(name)`. + +## Async API + +- `AsyncVectorDB` and `AsyncVectorCollection` mirror the entire sync + surface — every method listed above has an async equivalent. +- **Executor injection (v2.4+)** — pass `executor=ThreadPoolExecutor(...)` + to share a pool across async instances (important for ONNX / usearch + thread-safety). +- **Lifecycle** — `async with AsyncVectorDB(...)` drains the executor + with `wait=True` before closing the SQLite connection, so pool threads + finish before the underlying connection goes away. +- For a manual smoke runner that walks the entire async surface, see + `scripts/exercise_async_collection.py` in the repository. + +## Integrations + +- **LangChain** — `db.as_langchain(embeddings, collection_name=…)` returns + a `VectorStore`-compatible adapter. Supports all search methods. +- **LlamaIndex** — `db.as_llama_index(collection_name=…)` returns a + `BasePydanticVectorStore`-compatible adapter. +- **FastAPI embeddings server** — `[server]` extra adds a local HTTP + server with HuggingFace models, CORS, graceful shutdown, input + validation, and model warm-up (v2.5+). + +## Types & constants + +- `simplevecdb.types`: `Document`, `DistanceStrategy`, `Quantization`, + `Edge`, `Event`, `TTLEntry` (frozen dataclasses where applicable). +- `simplevecdb.constants` (v2.6.1) — tunables exposed as named + constants: + - `PENDING_FLUSH_DEFAULT_BATCH = 1000` + - `EVENTS_POLL_INTERVAL_S = 0.1` + - `EVENTS_RETENTION_LIMIT = 100_000` + - `TTL_SWEEP_DEFAULT_INTERVAL_S = 60.0` + - `REBUILD_PENDING_THRESHOLD = 5_000` + - `REBUILD_TOMBSTONE_THRESHOLD = 5_000` + - `REBUILD_MIN_INTERVAL_S = 3600.0` + - `SQLITE_BUSY_TIMEOUT_MS = 5000` + +## Performance snapshot + +10,000 vectors, 384 dimensions, k=10 search: + +| Quantization | Storage | Query | Compression | +| :----------- | :------ | :----- | :---------- | +| FLOAT32 | 36.0 MB | 0.20 ms | 1× | +| FLOAT16 | 28.7 MB | 0.20 ms | 2× | +| INT8 | 25.0 MB | 0.16 ms | 4× | +| BIT | 21.8 MB | 0.08 ms | 32× | + +Full benchmarks and tuning guidance: +[Benchmarks](benchmarks.md). + +## Roadmap + +Implemented features track the [Changelog](CHANGELOG.md). Currently on the +near-term radar: + +- Incremental clustering (online learning) +- Cluster visualization exports + +Vote on these or propose new ones in +[GitHub Discussions](https://github.com/coderdayton/simplevecdb/discussions). diff --git a/docs/LICENSE b/docs/LICENSE deleted file mode 100755 index e69de29..0000000 diff --git a/docs/api/core.md b/docs/api/core.md index cc7a83c..8031891 100755 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -12,7 +12,6 @@ The main database class for managing vector collections. - search_collections - vacuum - close - - check_migration ## VectorCollection diff --git a/docs/examples.md b/docs/examples.md index 285e497..5ab44ef 100755 --- a/docs/examples.md +++ b/docs/examples.md @@ -1,112 +1,97 @@ # Examples -## RAG with LangChain +## RAG notebooks -[View Notebook](https://github.com/coderdayton/simplevecdb/blob/main/examples/rag/langchain_rag.ipynb) +End-to-end RAG pipelines against a real LLM: -## RAG with LlamaIndex +- [LangChain](https://github.com/coderdayton/simplevecdb/blob/main/examples/rag/langchain_rag.ipynb) +- [LlamaIndex](https://github.com/coderdayton/simplevecdb/blob/main/examples/rag/llama_rag.ipynb) +- [Ollama (fully local)](https://github.com/coderdayton/simplevecdb/blob/main/examples/rag/ollama_rag.ipynb) -[View Notebook](https://github.com/coderdayton/simplevecdb/blob/main/examples/rag/llama_rag.ipynb) +## Storage & search -## RAG with Ollama LLM - -[View Notebook](https://github.com/coderdayton/simplevecdb/blob/main/examples/rag/ollama_rag.ipynb) - -## Quick Start Examples - -### Basic Usage +### Basic usage ```python from simplevecdb import VectorDB, Quantization -db = VectorDB("vectors.db") -collection = db.collection("docs", quantization=Quantization.FLOAT16) +db = VectorDB("notes.db") +notes = db.collection("personal", quantization=Quantization.FLOAT16) -# Add documents with embeddings -texts = ["Paris is the capital of France", "Berlin is in Germany"] -embeddings = [[0.1] * 384, [0.2] * 384] # Your embedding model output -collection.add_texts(texts, embeddings=embeddings) +notes.add_texts( + [ + "Cherry MX silent reds bottom out around 45g", + "Sourdough hydration sweet spot is ~75% with this flour", + ], + embeddings=[your_embedder(t) for t in texts], +) -# Search -query_embedding = [0.1] * 384 -results = collection.similarity_search(query_embedding, k=5) -for doc, score in results: - print(f"{doc.page_content} (score: {score:.4f})") +for doc, score in notes.similarity_search(query_vec, k=5): + print(f"{score:.3f} {doc.page_content}") ``` -### Keyword & Hybrid Search +### Keyword + hybrid search ```python -from simplevecdb import VectorDB - -db = VectorDB("local.db") -collection = db.collection("default") -collection.add_texts( - ["banana is yellow", "grapes are purple"], - embeddings=[[0.1, 0.2] * 192, [0.3, 0.4] * 192] +notes.keyword_search("postgres", k=3) # BM25 over FTS5 +notes.hybrid_search("slow query", k=3) # BM25 + vector RRF +notes.hybrid_search( + "slow query", + k=3, + query_vector=v, + vector_k=20, + keyword_k=20, + rrf_k=60, ) - -# BM25 keyword search -bm25 = collection.keyword_search("banana", k=1) - -# Hybrid search (BM25 + vectors with RRF) -hybrid = collection.hybrid_search("yellow fruit", k=2) ``` -### Batch Search (v2.0+) +### Batch search ```python -from simplevecdb import VectorDB - -db = VectorDB("vectors.db") -collection = db.collection("docs") - -# Add some documents... -collection.add_texts(texts, embeddings=embeddings) - -# Search multiple queries at once (~10x faster than sequential) -queries = [embedding1, embedding2, embedding3] -results = collection.similarity_search_batch(queries, k=10) - -for i, query_results in enumerate(results): - print(f"Query {i}: {len(query_results)} results") +results = notes.similarity_search_batch([v1, v2, v3], k=10) +for q, hits in zip(("a", "b", "c"), results): + print(q, len(hits)) ``` -### Force Exact Search +### Adaptive vs exact ```python -# Adaptive (default): brute-force for <10k, HNSW for larger -results = collection.similarity_search(query, k=10) - -# Force brute-force for perfect recall -results = collection.similarity_search(query, k=10, exact=True) - -# Force HNSW for speed on small collections -results = collection.similarity_search(query, k=10, exact=False) +notes.similarity_search(q, k=10) # adaptive — brute < 10k, HNSW above +notes.similarity_search(q, k=10, exact=True) # always brute (perfect recall) +notes.similarity_search(q, k=10, exact=False) # always HNSW ``` -### Metadata Filtering +### Metadata + range filters (v2.6.1) + +Equality and `$in` work everywhere; the v2.6.1 operators add range and +existence checks. Filters apply to `similarity_search`, +`keyword_search`, `hybrid_search`, `edges.get_edges`, and +`events.read`. ```python -collection.add_texts( - ["doc1", "doc2", "doc3"], +notes.add_texts( + [...], embeddings=[...], metadatas=[ - {"category": "tech", "year": 2024}, - {"category": "science", "year": 2023}, - {"category": "tech", "year": 2023}, - ] + {"tag": "work", "year": 2026, "score": 0.91}, + {"tag": "baking", "year": 2024, "score": 0.40}, + ], ) -# Filter by metadata -results = collection.similarity_search( - query, - k=10, - filter={"category": "tech"} -) +notes.similarity_search(q, k=10, filter={"tag": "work"}) + +notes.similarity_search(q, k=10, filter={ + "score": {"$gt": 0.5, "$lte": 0.95}, + "tag": {"$in": ["work", "research"]}, + "archived": {"$exists": False}, +}) + +# Tuple shorthand normalises to the operator-dict form +notes.similarity_search(q, k=10, filter={"score": (">", 0.5)}) +notes.similarity_search(q, k=10, filter={"year": ("range", 2024, 2026)}) ``` -### Encrypted Database (v2.1+) +### Encrypted database ```bash pip install "simplevecdb[encryption]" @@ -115,242 +100,187 @@ pip install "simplevecdb[encryption]" ```python from simplevecdb import VectorDB -# Create encrypted database with SQLCipher db = VectorDB("secure.db", encryption_key="your-secret-key") -collection = db.collection("confidential") - -collection.add_texts( - ["sensitive financial data", "private health records"], - embeddings=[[0.1]*384, [0.2]*384], - metadatas=[{"type": "finance"}, {"type": "health"}] +db.collection("confidential").add_texts( + ["sensitive note"], + embeddings=[[0.1] * 384], ) db.close() -# Reopen requires the same encryption key -db = VectorDB("secure.db", encryption_key="your-secret-key") -results = db.collection("confidential").similarity_search([0.1]*384, k=1) - -# Wrong key will fail -try: - bad_db = VectorDB("secure.db", encryption_key="wrong-key") -except Exception as e: - print(f"Access denied: {e}") +VectorDB("secure.db", encryption_key="your-secret-key") # ok +VectorDB("secure.db", encryption_key="wrong-key") # raises ``` -### Streaming Insert (v2.1+) - -Memory-efficient ingestion for large datasets: +### Streaming insert ```python import json from simplevecdb import VectorDB -db = VectorDB("large_dataset.db") -collection = db.collection("docs") +db = VectorDB("dump.db") +notes = db.collection("dump") -# Generator function for memory efficiency -def load_documents(): - with open("large_file.jsonl") as f: +def feed(): + with open("dump.jsonl") as f: for line in f: - doc = json.loads(line) - yield ( - doc["text"], - doc.get("metadata", {}), - doc["embedding"] - ) - -# Stream with progress tracking -for progress in collection.add_texts_streaming(load_documents(), batch_size=1000): - print(f"Batch {progress['batch_num']}: {progress['docs_processed']} total docs") - -# With callback instead of iteration -def log_progress(p): - if p['docs_processed'] % 10000 == 0: - print(f"Milestone: {p['docs_processed']} docs") - -list(collection.add_texts_streaming( - load_documents(), - batch_size=500, - on_progress=log_progress -)) -``` + row = json.loads(line) + yield row["text"], row.get("metadata", {}), row["embedding"] -### Document Hierarchies (v2.1+) +for progress in notes.add_texts_streaming(feed(), batch_size=1000): + print(f"batch {progress['batch_num']}: {progress['docs_processed']} docs") +``` -Organize documents in parent-child trees for chunked documents, threads, or nested content: +### Document hierarchies ```python -from simplevecdb import VectorDB - -db = VectorDB("hierarchical.db") -collection = db.collection("docs") - -# Add parent documents (e.g., full articles) -parent_ids = collection.add_texts( - ["Chapter 1: Introduction to ML", "Chapter 2: Neural Networks"], - embeddings=[[0.1]*384, [0.2]*384], - metadatas=[{"type": "chapter", "num": 1}, {"type": "chapter", "num": 2}] +parents = collection.add_texts( + ["Chapter 1", "Chapter 2"], + embeddings=[[0.1] * 384, [0.2] * 384], ) - -# Add children with parent references (e.g., chunks/sections) -child_ids = collection.add_texts( - ["Section 1.1: What is ML?", "Section 1.2: History of ML", "Section 2.1: Perceptrons"], - embeddings=[[0.11]*384, [0.12]*384, [0.21]*384], - metadatas=[ - {"type": "section", "chapter": 1}, - {"type": "section", "chapter": 1}, - {"type": "section", "chapter": 2} - ], - parent_ids=[parent_ids[0], parent_ids[0], parent_ids[1]] +children = collection.add_texts( + ["1.1 intro", "1.2 history", "2.1 perceptrons"], + embeddings=[[0.11] * 384, [0.12] * 384, [0.21] * 384], + parent_ids=[parents[0], parents[0], parents[1]], ) -# Navigate the hierarchy -children = collection.get_children(parent_ids[0]) -print(f"Chapter 1 has {len(children)} sections") - -parent = collection.get_parent(child_ids[0]) -print(f"Section belongs to: {parent.page_content}") +collection.get_children(parents[0]) +collection.get_parent(children[0]) +collection.get_descendants(parents[0]) +collection.get_ancestors(children[0]) +collection.set_parent(children[2], parents[0]) # reparent +``` -# Get all descendants (recursive) -all_descendants = collection.get_descendants(parent_ids[0]) +### Async usage -# Get ancestors (path to root) -ancestors = collection.get_ancestors(child_ids[0]) +```python +import asyncio +from simplevecdb import AsyncVectorDB -# Reparent a document -collection.set_parent(child_ids[2], parent_ids[0]) # Move section to Chapter 1 +async def main(): + async with AsyncVectorDB("notes.db") as db: + notes = db.collection("personal") + await notes.add_texts(texts, embeddings=embeddings) + return await notes.similarity_search_batch(queries, k=10) -# Search within a subtree -results = collection.similarity_search( - [0.1]*384, - k=10, - filter={"type": "section", "chapter": 1} -) +asyncio.run(main()) ``` -### Async Usage +### Cross-collection search ```python -import asyncio -from simplevecdb import AsyncVectorDB +from simplevecdb import VectorDB -async def main(): - async with AsyncVectorDB("vectors.db") as db: - collection = db.collection("docs") - - # Add documents - await collection.add_texts(texts, embeddings=embeddings) - - # Batch search - results = await collection.similarity_search_batch(queries, k=10) - - return results - -results = asyncio.run(main()) +db = VectorDB("app.db") +db.collection("users").add_texts([...], embeddings=[...]) +db.collection("products").add_texts([...], embeddings=[...]) +db.collection("tickets").add_texts([...], embeddings=[...]) + +for doc, score, name in db.search_collections(query, k=5): + print(f"[{name}] {doc.page_content} {score:.3f}") + +db.search_collections(query, collections=["users", "products"], k=3) +db.search_collections(query, k=10, filter={"category": "software"}) +db.list_collections() ``` -### Cross-Collection Search (v2.2+) +## Memory primitives (v2.6.1) + +The 2.6.1 release adds primitives for retrieval-with-memory systems — +in-place vector updates, atomic counters, edges, expiry, and an +append-only change feed. Full reference: +[Features](Features.md). -Search across multiple collections with unified ranking: +### Pending vector buffer + +`update_embedding` writes to a per-collection overlay; the new vector +becomes visible to reads immediately and is promoted into HNSW on +`pending.flush()`. Removes the HNSW remove+re-add churn previously +required for in-place updates. ```python -from simplevecdb import VectorDB +collection.update_embedding(doc_id, new_vector, source="recompute") +collection.update_embedding(other_id, vec2) +collection.pending.flush(max_batch=512) # promote both into HNSW -db = VectorDB("multi_tenant.db") +collection.pending.update_many([(id1, v1), (id2, v2)]) +collection.pending.blend_toward([id1, id2, id3], centroid=c, alpha=0.1) +``` -# Create domain-specific collections -users = db.collection("users") -products = db.collection("products") -support = db.collection("support_tickets") +### Atomic counters -# Populate collections -users.add_texts( - ["Alice - software engineer", "Bob - data scientist"], - embeddings=[[0.1]*384, [0.2]*384], - metadatas=[{"role": "eng"}, {"role": "data"}] -) -products.add_texts( - ["Python IDE Pro", "Data Analysis Suite"], - embeddings=[[0.15]*384, [0.25]*384], - metadatas=[{"category": "software"}, {"category": "software"}] -) -support.add_texts( - ["How to install Python IDE?", "Data export not working"], - embeddings=[[0.12]*384, [0.22]*384] -) +`increment_metadata` applies a dict-of-deltas to JSON metadata in one +SQL statement — WAL-atomic and safe under concurrent writers. -# Search across ALL collections at once -query = [0.1]*384 -results = db.search_collections(query, k=5) - -for doc, score, collection_name in results: - print(f"[{collection_name}] {doc.page_content} ({score:.3f})") -# Output: -# [users] Alice - software engineer (0.999) -# [support] How to install Python IDE? (0.893) -# [products] Python IDE Pro (0.871) -# ... - -# Search only specific collections -results = db.search_collections( - query, - collections=["users", "products"], - k=3 -) +```python +collection.increment_metadata(doc_id, {"hits": 1, "drift": 0.02}) +``` -# Apply metadata filter across all searched collections -results = db.search_collections( - query, - k=10, - filter={"category": "software"} +### Weighted directed edges + +```python +collection.edges.add_edge(src, dst, kind="cites", weight=0.8, hits=1) +collection.edges.update_edge(src, dst, kind="cites", dweight=+0.05, dhits=+1) + +collection.edges.get_edges( + src=src, + filter={"weight": {"$gt": 0.5}, "hits": ("range", 1, 10)}, ) +collection.edges.delete_edge(src, dst, kind="cites") -# List available collections -print(db.list_collections()) # ['users', 'products', 'support_tickets'] +# Bulk threshold prune +collection.edges.prune(kind="cites", max_weight=0.1, idle_before=cutoff_ts) ``` -**Async cross-collection search:** +### TTL / expiry ```python -import asyncio -from simplevecdb.async_core import AsyncVectorDB - -async def search_all(): - db = AsyncVectorDB("app.db") - - # Initialize collections - db.collection("users") - db.collection("products") - - # Search across collections - results = await db.search_collections([0.1]*384, k=10) - return results - -results = asyncio.run(search_all()) +import time + +collection.ttl.set(doc_id, seconds=3600, on_expire="delete") +collection.ttl.set(other, expires_at=time.time()+5, on_expire="callback") + +deleted, callbacks = collection.ttl.sweep() # one-shot +collection.ttl.start_background(interval=60.0) # daemon thread +collection.ttl.clear(doc_id) ``` -## Benchmark Scripts +### Append-only event feed -### Backend Benchmark +Every mutating call appends one row (kind, doc_id, payload, monotonic +seq). Useful for change feeds, replication, and audit trails: -Compare HNSW vs brute-force performance: +```python +seq = collection.events.last_seq() +# … do some writes … +for ev in collection.events.read(since=seq, kind="edge_add", limit=200): + print(ev.seq, ev.kind, ev.doc_id, ev.payload) -```bash -python examples/backend_benchmark.py +for ev in collection.events.subscribe(since=seq, poll_interval=0.5): + handle(ev) # blocking generator + +collection.events.prune(before_seq=seq - 100_000) ``` -### Quantization Benchmark +### Transactions -Test different quantization levels: +`db.transaction()` and `collection.tx()` wrap a SAVEPOINT around +catalog writes (metadata, counters, edges, events, TTL, and the +pending overlay). A raised exception rolls all SQL writes back. Coarse +vector mutations (`add_texts`, `delete_by_ids`) are *not* rolled back — +use `update_embedding` + `pending.flush()` for vector changes that +must be commit-gated. -```bash -python examples/quant_benchmark.py +```python +with db.transaction() as tx: + tx["personal"].increment_metadata(1, {"hits": 1}) + tx["personal"].edges.add_edge(1, 2, kind="cites", weight=0.6) + # any exception below rolls both writes back ``` -### Embedding Performance - -Benchmark local embedding generation: +## Benchmark scripts ```bash -python examples/embeddings/perf_benchmark.py +python examples/backend_benchmark.py # HNSW vs brute-force +python examples/quant_benchmark.py # quantization tradeoffs +python examples/embeddings/perf_benchmark.py # local embedding throughput ``` diff --git a/examples/rag/langchain_rag.ipynb b/examples/rag/langchain_rag.ipynb index b559cc8..67d4b57 100755 --- a/examples/rag/langchain_rag.ipynb +++ b/examples/rag/langchain_rag.ipynb @@ -64,7 +64,7 @@ "texts = [\n", " \"LangChain is awesome for building LLM applications\",\n", " \"SimpleVecDB rocks as a lightweight vector store\",\n", - " \"RAG combines retrieval with generation for better answers\"\n", + " \"RAG combines retrieval with generation for better answers\",\n", "]\n", "\n", "collection.add_texts(texts)\n", @@ -90,7 +90,7 @@ "# Initialize OpenAI embeddings\n", "embeddings = OpenAIEmbeddings(\n", " base_url=f\"http://{config.SERVER_HOST}:{config.SERVER_PORT}/v1\",\n", - " api_key=convert_to_secret_str(\"your-embedding-api-key\")\n", + " api_key=convert_to_secret_str(\"your-embedding-api-key\"),\n", ")\n", "\n", "# Create LangChain vector store from SimpleVecDB (specify collection name)\n", @@ -142,11 +142,14 @@ "\n", "# Build RAG chain: retrieve docs -> format context -> question -> LLM\n", "chain = (\n", - "\t\t{\"context\": retriever | (lambda docs: \"\\n\\n\".join(doc.page_content for doc in docs)),\n", - "\t\t \"question\": RunnablePassthrough()}\n", - "\t\t| prompt\n", - "\t\t| llm\n", - "\t\t| StrOutputParser()\n", + " {\n", + " \"context\": retriever\n", + " | (lambda docs: \"\\n\\n\".join(doc.page_content for doc in docs)),\n", + " \"question\": RunnablePassthrough(),\n", + " }\n", + " | prompt\n", + " | llm\n", + " | StrOutputParser()\n", ")\n", "\n", "print(\"✓ RAG chain components ready\")" diff --git a/examples/rag/llama_rag.ipynb b/examples/rag/llama_rag.ipynb index cd80bb8..92c3061 100755 --- a/examples/rag/llama_rag.ipynb +++ b/examples/rag/llama_rag.ipynb @@ -74,7 +74,7 @@ "\n", "Settings.embed_model = OpenAIEmbedding(\n", " api_base=f\"http://{config.SERVER_HOST}:{config.SERVER_PORT}/v1\",\n", - " api_key=\"your-api-key\"\n", + " api_key=\"your-api-key\",\n", ")" ] }, @@ -105,13 +105,19 @@ "\n", "# Create sample text files\n", "with open(os.path.join(data_dir, \"doc1.txt\"), \"w\") as f:\n", - " f.write(\"SimpleVecDB is a lightweight vector database powered by usearch HNSW indexing. It's designed to be simple and easy to use.\")\n", + " f.write(\n", + " \"SimpleVecDB is a lightweight vector database powered by usearch HNSW indexing. It's designed to be simple and easy to use.\"\n", + " )\n", "\n", "with open(os.path.join(data_dir, \"doc2.txt\"), \"w\") as f:\n", - " f.write(\"Vector databases store embeddings and allow for efficient similarity search. They are essential for RAG applications.\")\n", + " f.write(\n", + " \"Vector databases store embeddings and allow for efficient similarity search. They are essential for RAG applications.\"\n", + " )\n", "\n", "with open(os.path.join(data_dir, \"doc3.txt\"), \"w\") as f:\n", - " f.write(\"LlamaIndex is a data framework for building LLM applications. It provides tools for data ingestion, indexing, and querying.\")\n", + " f.write(\n", + " \"LlamaIndex is a data framework for building LLM applications. It provides tools for data ingestion, indexing, and querying.\"\n", + " )\n", "\n", "print(f\"✓ Sample documents created in {data_dir}\")\n", "print(f\"✓ Files in directory: {os.listdir(data_dir)}\")" diff --git a/examples/rag/ollama_rag.ipynb b/examples/rag/ollama_rag.ipynb index 49ef1cc..794ff93 100755 --- a/examples/rag/ollama_rag.ipynb +++ b/examples/rag/ollama_rag.ipynb @@ -117,7 +117,7 @@ " \"SimpleVecDB uses usearch HNSW indexing, enabling fast similarity searches on embedded documents stored in a single SQLite file.\",\n", " \"Quantization reduces the memory footprint of vectors by using fewer bits per dimension. BIT quantization uses only 1 bit per dimension, offering 32x compression compared to float32.\",\n", " \"Local-first AI means running models and storing data entirely on your own hardware, without cloud dependencies. This ensures privacy, reduces costs, and works offline.\",\n", - " \"Python is a high-level programming language known for its simplicity and readability. It's widely used in data science, machine learning, and web development.\"\n", + " \"Python is a high-level programming language known for its simplicity and readability. It's widely used in data science, machine learning, and web development.\",\n", "]\n", "\n", "# Add documents with metadata to the collection\n", @@ -182,31 +182,33 @@ "metadata": {}, "outputs": [], "source": [ - "def rag_query(question: str, k: int = 3, model: str = \"llama3.2:3b\", verbose: bool = True) -> str:\n", + "def rag_query(\n", + " question: str, k: int = 3, model: str = \"llama3.2:3b\", verbose: bool = True\n", + ") -> str:\n", " \"\"\"\n", " Answer a question using RAG.\n", - " \n", + "\n", " Args:\n", " question: The question to answer\n", " k: Number of documents to retrieve\n", " model: Ollama model to use\n", " verbose: Print retrieval details\n", - " \n", + "\n", " Returns:\n", " Generated answer\n", " \"\"\"\n", " # Step 1: Retrieve relevant documents from the collection\n", " results = collection.similarity_search(query=question, k=k)\n", - " \n", + "\n", " if verbose:\n", " print(f\"Retrieved {len(results)} documents:\")\n", " for i, (doc, score) in enumerate(results, 1):\n", " print(f\" {i}. Score: {score:.4f} | {doc.page_content[:80]}...\")\n", " print()\n", - " \n", + "\n", " # Step 2: Build context from retrieved documents\n", " context = \"\\n\\n\".join([doc.page_content for doc, _ in results])\n", - " \n", + "\n", " # Step 3: Create prompt with context\n", " prompt = f\"\"\"You are a helpful assistant. Answer the question based ONLY on the provided context. If the context doesn't contain enough information, say so.\n", "\n", @@ -216,21 +218,22 @@ "Question: {question}\n", "\n", "Answer:\"\"\"\n", - " \n", + "\n", " # Step 4: Generate answer with Ollama\n", " if verbose:\n", " print(f\"Generating answer with {model}...\\n\")\n", - " \n", + "\n", " response = ollama.generate(\n", " model=model,\n", " prompt=prompt,\n", " options={\n", " \"temperature\": 0.1, # Low temperature for more factual answers\n", - " \"num_predict\": 256 # Limit response length\n", - " }\n", + " \"num_predict\": 256, # Limit response length\n", + " },\n", " )\n", - " \n", - " return response['response']\n", + "\n", + " return response[\"response\"]\n", + "\n", "\n", "print(\"✓ RAG function ready\")" ] @@ -365,7 +368,9 @@ " \"Python's GIL (Global Interpreter Lock) means only one thread executes Python bytecode at a time.\",\n", "]\n", "\n", - "python_collection.add_texts(python_docs, metadatas=[{\"topic\": \"python\"} for _ in python_docs])\n", + "python_collection.add_texts(\n", + " python_docs, metadatas=[{\"topic\": \"python\"} for _ in python_docs]\n", + ")\n", "\n", "print(f\"✓ Created second collection: {python_collection.name}\")\n", "print(f\" Documents in '{collection.name}': (check via SQL)\")\n", @@ -411,7 +416,7 @@ "print(f\" Embedding dimension: {collection._dim}\")\n", "print(f\" Quantization: {collection.quantization}\")\n", "print(f\" Database size: {db_size_kb:.2f} KB\")\n", - "print(f\" Average size per doc: {db_size_kb/doc_count:.2f} KB\")\n", + "print(f\" Average size per doc: {db_size_kb / doc_count:.2f} KB\")\n", "\n", "# For comparison, show what float32 would use\n", "if collection._dim:\n", @@ -419,7 +424,7 @@ " print(\"\\nComparison:\")\n", " print(f\" BIT quantization: {db_size_kb:.2f} KB\")\n", " print(f\" FLOAT32 (uncompressed): ~{float32_size:.2f} KB\")\n", - " print(f\" Compression ratio: {float32_size/db_size_kb:.1f}x\")\n", + " print(f\" Compression ratio: {float32_size / db_size_kb:.1f}x\")\n", "\n", "conn.close()" ] diff --git a/lefthook.yml b/lefthook.yml index 1f74d78..4047c17 100644 --- a/lefthook.yml +++ b/lefthook.yml @@ -1,12 +1,12 @@ # Lefthook config — replaces .pre-commit-config.yaml. # -# Split: -# pre-commit → fast feedback only (version sync + ruff --fix). -# pre-push → heavyweight gates (mypy, bandit, full pytest+coverage). +# Both stages run the same gate set; pre-commit autofixes where it can +# (ruff format / ruff check --fix re-stage via ``stage_fixed: true``) +# while pre-push runs the same checks in non-mutating mode so a push +# can never carry a regression that wasn't visible at commit time. # -# Hooks within a stage run sequentially because ruff --fix can rewrite -# files mid-run and ``stage_fixed: true`` re-stages them, so the user -# does not have to re-add manually. +# Hooks within a stage run sequentially because the autofixers can +# rewrite files mid-run; later jobs see the post-fix tree. pre-commit: jobs: @@ -16,16 +16,41 @@ pre-commit: - "src/simplevecdb/__init__.py" run: uv run python3 scripts/check_version_sync.py + - name: ruff-format + glob: "*.py" + stage_fixed: true + run: uv run ruff format . + - name: ruff glob: "*.py" stage_fixed: true run: uv run ruff check . --fix + - name: mypy + glob: "*.py" + run: uv run mypy . + + - name: bandit + glob: "*.py" + run: uv run bandit -r src/ -ll -c .bandit + + - name: pytest + glob: "*.py" + run: uv run pytest tests/ -q + pre-push: jobs: - name: version-sync run: uv run python3 scripts/check_version_sync.py + - name: ruff-format + glob: "*.py" + run: uv run ruff format --check . + + - name: ruff + glob: "*.py" + run: uv run ruff check . + - name: mypy glob: "*.py" run: uv run mypy . diff --git a/mkdocs.yml b/mkdocs.yml index 79e2ccc..201d982 100755 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,6 +34,7 @@ plugins: nav: - Home: index.md + - Features: Features.md - Examples: examples.md - Benchmarks: benchmarks.md - Changelog: CHANGELOG.md diff --git a/pyproject.toml b/pyproject.toml index 3165327..8e8e5e8 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "simplevecdb" -version = "2.6.0" +version = "2.6.1" description = "Dead-simple local vector database powered by usearch HNSW." authors = [{ name = "Dayton Dunbar", email = "coderdayton14@gmail.com" }] license = { text = "MIT" } @@ -37,7 +37,6 @@ classifiers = [ dependencies = [ "numpy>=1.24", "usearch>=2.16.3", - "sqlite-vec>=0.1.6", "scikit-learn>=1.3.0", # Clustering (K-means, MiniBatch K-means, HDBSCAN) "hdbscan>=0.8.33", # Density-based clustering "sqlcipher3-binary>=0.5.0", # Encryption support diff --git a/scripts/check_version_sync.py b/scripts/check_version_sync.py index fb11c38..2f38416 100755 --- a/scripts/check_version_sync.py +++ b/scripts/check_version_sync.py @@ -16,12 +16,8 @@ from pathlib import Path -_PYPROJECT_VERSION_RE = re.compile( - r'^version\s*=\s*["\']([^"\']+)["\']', re.MULTILINE -) -_CHANGELOG_HEADING_RE = re.compile( - r"^##\s+\[(?P[^\]]+)\]", re.MULTILINE -) +_PYPROJECT_VERSION_RE = re.compile(r'^version\s*=\s*["\']([^"\']+)["\']', re.MULTILINE) +_CHANGELOG_HEADING_RE = re.compile(r"^##\s+\[(?P[^\]]+)\]", re.MULTILINE) _SEMVER_RE = re.compile(r"^\d+\.\d+\.\d+([.-][a-zA-Z0-9]+)*$") diff --git a/scripts/exercise_async_collection.py b/scripts/exercise_async_collection.py new file mode 100644 index 0000000..1a6f7e6 --- /dev/null +++ b/scripts/exercise_async_collection.py @@ -0,0 +1,358 @@ +"""Exercise every async collection function in simplevecdb. + +Walks AsyncVectorCollection's full surface — CRUD, search, hierarchy, +edges, counters, pending vectors, TTL, events, maintenance, clustering — +plus the AsyncVectorDB-level multi-collection helpers. Designed to +double as a smoke test for the 2.6.1 "advanced memory" APIs. + +Run: + uv run python /tmp/exercise_async_collection.py +""" + +from __future__ import annotations + +import asyncio +import os +import tempfile +import time +import traceback +from typing import Any, Awaitable, Callable + +import numpy as np + +from simplevecdb.async_core import AsyncVectorCollection, AsyncVectorDB + +DIM = 32 +SEED = 0 + + +def _vec(rng: np.random.Generator) -> list[float]: + """Return a unit-ish random vector of length DIM.""" + v = rng.standard_normal(DIM).astype(np.float32) + v /= np.linalg.norm(v) + 1e-9 + return v.tolist() + + +def _seed_corpus( + n: int = 12, +) -> tuple[list[str], list[dict[str, Any]], list[list[float]]]: + rng = np.random.default_rng(SEED) + # Two semantic clusters separated by metadata "topic". + docs = [ + ("python async tasks scheduler", "code"), + ("event loops and coroutines", "code"), + ("threadpool executor patterns", "code"), + ("vector databases for retrieval", "code"), + ("hnsw graph approximate nearest neighbor", "code"), + ("embedding models sentence transformers", "code"), + ("baking sourdough bread overnight", "food"), + ("croissant lamination butter technique", "food"), + ("ramen broth tonkotsu collagen", "food"), + ("pizza neapolitan flour hydration", "food"), + ("knife sharpening whetstone angles", "food"), + ("fermentation kimchi gochugaru salt", "food"), + ][:n] + texts = [t for t, _ in docs] + metas = [{"topic": tag, "hits": 0, "drift": 0.0} for _, tag in docs] + vectors = [_vec(rng) for _ in range(n)] + return texts, metas, vectors + + +class Reporter: + """Tiny pass/fail accumulator with consistent formatting.""" + + def __init__(self) -> None: + self.results: list[tuple[str, bool, str]] = [] + + async def run( + self, + label: str, + coro: Callable[[], Awaitable[Any]], + ) -> Any: + try: + value = await coro() + except Exception as exc: # noqa: BLE001 — we report and keep going + tb = traceback.format_exc(limit=2).strip().splitlines()[-1] + self.results.append((label, False, f"{type(exc).__name__}: {exc} ({tb})")) + print(f" ✗ {label}: {type(exc).__name__}: {exc}") + return None + summary = self._summarize(value) + self.results.append((label, True, summary)) + print(f" ✓ {label} → {summary}") + return value + + @staticmethod + def _summarize(value: Any) -> str: + if value is None: + return "ok" + if isinstance(value, (list, tuple)): + return f"{type(value).__name__}[{len(value)}]" + if isinstance(value, dict): + return f"dict[{len(value)}]" + if isinstance(value, (int, float, bool, str)): + return repr(value) + return type(value).__name__ + + def report(self) -> int: + passed = sum(1 for _, ok, _ in self.results if ok) + total = len(self.results) + failed = total - passed + print() + print("=" * 64) + print(f" RESULT: {passed}/{total} passed, {failed} failed") + print("=" * 64) + if failed: + for label, ok, detail in self.results: + if not ok: + print(f" FAIL {label}: {detail}") + return failed + + +async def exercise_crud_and_search( + r: Reporter, + col: AsyncVectorCollection, + texts: list[str], + metas: list[dict[str, Any]], + vectors: list[list[float]], +) -> list[int]: + print("\n[1] CRUD + search") + ids = await r.run( + "add_texts", + lambda: col.add_texts(texts, metadatas=metas, embeddings=vectors), + ) + assert ids is not None and len(ids) == len(texts) + await r.run("count", col.count) + await r.run("dim (property)", lambda: asyncio.sleep(0, result=col.dim)) + await r.run( + "get_documents (filter)", + lambda: col.get_documents(filter_dict={"topic": "code"}, limit=5), + ) + await r.run( + "get_embeddings_by_ids", + lambda: col.get_embeddings_by_ids(ids[:3]), + ) + await r.run( + "update_metadata", + lambda: col.update_metadata([(ids[0], {"starred": True})]), + ) + await r.run( + "similarity_search (vector)", + lambda: col.similarity_search(vectors[0], k=3), + ) + await r.run( + "similarity_search_batch", + lambda: col.similarity_search_batch(vectors[:2], k=3), + ) + await r.run( + "keyword_search", + lambda: col.keyword_search("python async", k=3), + ) + await r.run( + "hybrid_search", + lambda: col.hybrid_search("python async", k=3, query_vector=vectors[0]), + ) + await r.run( + "max_marginal_relevance_search", + lambda: col.max_marginal_relevance_search(vectors[0], k=3, fetch_k=8), + ) + return ids + + +async def exercise_hierarchy( + r: Reporter, col: AsyncVectorCollection, ids: list[int] +) -> None: + print("\n[2] Hierarchy") + root, mid, leaf = ids[0], ids[1], ids[2] + await r.run("set_parent (mid→root)", lambda: col.set_parent(mid, root)) + await r.run("set_parent (leaf→mid)", lambda: col.set_parent(leaf, mid)) + await r.run("get_parent", lambda: col.get_parent(leaf)) + await r.run("get_children", lambda: col.get_children(root)) + await r.run("get_descendants", lambda: col.get_descendants(root)) + await r.run("get_ancestors", lambda: col.get_ancestors(leaf)) + await r.run("set_parent (clear)", lambda: col.set_parent(leaf, None)) + + +async def exercise_edges( + r: Reporter, col: AsyncVectorCollection, ids: list[int] +) -> None: + print("\n[3] Edges") + a, b, c = ids[0], ids[1], ids[2] + await r.run( + "add_edge (related)", + lambda: col.add_edge(a, b, kind="related", weight=0.8, hits=1), + ) + await r.run( + "add_edge (cites)", + lambda: col.add_edge(a, c, kind="cites", weight=0.4), + ) + await r.run( + "update_edge (delta)", + lambda: col.update_edge(a, b, kind="related", dweight=0.1, dhits=2), + ) + await r.run( + "get_edges (by src)", + lambda: col.get_edges(src=a), + ) + await r.run( + "delete_edge", + lambda: col.delete_edge(a, c, kind="cites"), + ) + + +async def exercise_counters_pending_ttl_events( + r: Reporter, col: AsyncVectorCollection, ids: list[int] +) -> None: + print("\n[4] Counters / pending / TTL / events") + target = ids[0] + + await r.run( + "increment_metadata", + lambda: col.increment_metadata(target, {"hits": 3, "drift": 0.05}), + ) + + rng = np.random.default_rng(SEED + 1) + new_vec = _vec(rng) + await r.run( + "update_embedding (buffer)", + lambda: col.update_embedding(target, new_vec, source="exerciser"), + ) + await r.run("flush_pending", lambda: col.flush_pending(max_batch=128)) + + seq_before = await r.run("last_event_seq (before)", col.last_event_seq) + + # Already-expired TTL → sweep should harvest it. + await r.run( + "set_ttl (expires_at past)", + lambda: col.set_ttl(ids[-1], expires_at=time.time() - 1, on_expire="callback"), + ) + await r.run( + "set_ttl (seconds future)", + lambda: col.set_ttl(ids[-2], seconds=3600, on_expire="delete"), + ) + await r.run("sweep_ttl", lambda: col.sweep_ttl(limit=100)) + await r.run("clear_ttl", lambda: col.clear_ttl(ids[-2])) + + await r.run( + "read_events (since)", + lambda: col.read_events(since=int(seq_before or 0), limit=50), + ) + + +async def exercise_maintenance(r: Reporter, col: AsyncVectorCollection) -> None: + print("\n[5] Maintenance") + await r.run( + "rebuild_if_needed", + lambda: col.rebuild_if_needed(max_pending=10_000, max_deleted=10_000), + ) + await r.run("rebuild_index", col.rebuild_index) + await r.run("save", col.save) + + +async def exercise_clustering(r: Reporter, col: AsyncVectorCollection) -> None: + print("\n[6] Clustering") + cluster_result = await r.run( + "cluster (kmeans, n=2)", + lambda: col.cluster(n_clusters=2, algorithm="kmeans", min_cluster_size=2), + ) + if cluster_result is None: + return + tags = await r.run( + "auto_tag", + lambda: col.auto_tag(cluster_result, method="keywords", n_keywords=3), + ) + await r.run( + "assign_cluster_metadata", + lambda: col.assign_cluster_metadata(cluster_result, tags), + ) + await r.run("get_cluster_members(0)", lambda: col.get_cluster_members(0)) + await r.run( + "save_cluster", + lambda: col.save_cluster("snapshot", cluster_result, metadata={"by": "test"}), + ) + await r.run("list_clusters", col.list_clusters) + await r.run("load_cluster", lambda: col.load_cluster("snapshot")) + await r.run( + "assign_to_cluster", + lambda: col.assign_to_cluster("snapshot", [1, 2]), + ) + await r.run("delete_cluster", lambda: col.delete_cluster("snapshot")) + + # Validation path also lives in async wrapper. + async def _bad(): + try: + await col.cluster(algorithm="bogus") + except ValueError as exc: + msg = str(exc)[:40] + return f"ValueError({msg}…)" + return "no error" + + await r.run("cluster (invalid algo → ValueError)", _bad) + + +async def exercise_db_level( + r: Reporter, db: AsyncVectorDB, query_vec: list[float] +) -> None: + print("\n[7] DB-level helpers") + # Add a second collection so search_collections has somewhere to fan out. + other = db.collection("other", store_embeddings=True) + rng = np.random.default_rng(SEED + 2) + await other.add_texts( + ["another world", "second collection"], + embeddings=[_vec(rng), _vec(rng)], + ) + await r.run( + "list_collections", + lambda: asyncio.sleep(0, result=db.list_collections()), + ) + await r.run( + "search_collections", + lambda: db.search_collections(query_vec, k=3), + ) + await r.run("vacuum", db.vacuum) + await r.run( + "delete_collection (other)", + lambda: db.delete_collection("other"), + ) + + +async def exercise_deletion( + r: Reporter, col: AsyncVectorCollection, ids: list[int] +) -> None: + print("\n[8] Deletion paths") + await r.run( + "delete_by_ids", + lambda: col.delete_by_ids([ids[-1]]), + ) + await r.run( + "remove_texts (by filter)", + lambda: col.remove_texts(filter={"topic": "food"}), + ) + await r.run("count (post-delete)", col.count) + + +async def main() -> int: + tmp = tempfile.mkdtemp(prefix="async_exerciser_") + db_path = os.path.join(tmp, "advanced_memory.db") + print(f"DB: {db_path}") + + r = Reporter() + texts, metas, vectors = _seed_corpus() + + async with AsyncVectorDB(db_path, max_workers=4) as db: + col = db.collection("memory", store_embeddings=True) + ids = await exercise_crud_and_search(r, col, texts, metas, vectors) + if ids: + await exercise_hierarchy(r, col, ids) + await exercise_edges(r, col, ids) + await exercise_counters_pending_ttl_events(r, col, ids) + await exercise_maintenance(r, col) + await exercise_clustering(r, col) + await exercise_db_level(r, db, vectors[0]) + if ids: + await exercise_deletion(r, col, ids) + + return r.report() + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/scripts/track_metrics.py b/scripts/track_metrics.py index 9a03999..9f0ddf3 100755 --- a/scripts/track_metrics.py +++ b/scripts/track_metrics.py @@ -9,18 +9,19 @@ REPO_NAME = "tinyvecdb" GITHUB_API_URL = "https://api.github.com" + def get_github_metrics(token=None): headers = {"Accept": "application/vnd.github.v3+json"} if token: headers["Authorization"] = f"token {token}" - + url = f"{GITHUB_API_URL}/repos/{REPO_OWNER}/{REPO_NAME}" - + try: req = urllib.request.Request(url, headers=headers) with urllib.request.urlopen(req) as response: data = json.loads(response.read().decode()) - + return { "stars": data.get("stargazers_count", 0), "forks": data.get("forks_count", 0), @@ -31,12 +32,14 @@ def get_github_metrics(token=None): print(f"Error fetching repo metrics: {e}") return None + def get_sponsor_count(token=None): # Note: Accurate sponsor count requires GraphQL API and a token with read:user/read:org scope if not token: return "N/A (Requires GITHUB_TOKEN)" - - query = """ + + query = ( + """ { user(login: "%s") { sponsorshipsAsMaintainer { @@ -44,39 +47,45 @@ def get_sponsor_count(token=None): } } } - """ % REPO_OWNER - + """ + % REPO_OWNER + ) + url = "https://api.github.com/graphql" - headers = { - "Authorization": f"bearer {token}", - "Content-Type": "application/json" - } + headers = {"Authorization": f"bearer {token}", "Content-Type": "application/json"} data = json.dumps({"query": query}).encode("utf-8") - + try: req = urllib.request.Request(url, data=data, headers=headers, method="POST") with urllib.request.urlopen(req) as response: result = json.loads(response.read().decode()) - return result.get("data", {}).get("user", {}).get("sponsorshipsAsMaintainer", {}).get("totalCount", 0) + return ( + result.get("data", {}) + .get("user", {}) + .get("sponsorshipsAsMaintainer", {}) + .get("totalCount", 0) + ) except Exception as e: print(f"Error fetching sponsor count: {e}") return "Error" + def main(): token = os.environ.get("GITHUB_TOKEN") - + print(f"--- Metrics for {REPO_OWNER}/{REPO_NAME} ---") print(f"Date: {datetime.now().isoformat()}") - + metrics = get_github_metrics(token) if metrics: print(f"⭐ Stars: {metrics['stars']}") print(f"🍴 Forks: {metrics['forks']}") print(f"👀 Watchers: {metrics['watchers']}") print(f"🐞 Issues: {metrics['open_issues']}") - + sponsors = get_sponsor_count(token) print(f"💖 Sponsors: {sponsors}") + if __name__ == "__main__": main() diff --git a/src/simplevecdb/__init__.py b/src/simplevecdb/__init__.py index 8999e8a..9d66825 100755 --- a/src/simplevecdb/__init__.py +++ b/src/simplevecdb/__init__.py @@ -6,13 +6,13 @@ Document, DistanceStrategy, Quantization, - MigrationRequiredError, StreamingProgress, ProgressCallback, ) from .core import VectorDB, VectorCollection, get_optimal_batch_size from .async_core import AsyncVectorDB, AsyncVectorCollection from .config import config + try: from .integrations import SimpleVecDBVectorStore, SimpleVecDBLlamaStore except ImportError: @@ -56,7 +56,6 @@ "log_operation", # Error handling "DatabaseLockedError", - "MigrationRequiredError", "EncryptionError", "EncryptionUnavailableError", "async_retry_on_lock", diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index e5a7d65..161a203 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -25,6 +25,7 @@ from __future__ import annotations import asyncio +import functools from concurrent.futures import ThreadPoolExecutor from collections.abc import Sequence from threading import Lock @@ -62,6 +63,12 @@ def name(self) -> str: def __repr__(self) -> str: return f"AsyncVectorCollection(name={self._collection.name!r})" + async def _run(self, fn, /, *args, **kwargs): + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._executor, functools.partial(fn, *args, **kwargs) + ) + async def add_texts( self, texts: Sequence[str], @@ -76,13 +83,14 @@ async def add_texts( See VectorCollection.add_texts for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.add_texts( - texts, metadatas, embeddings, ids, - parent_ids=parent_ids, threads=threads, - ), + return await self._run( + self._collection.add_texts, + texts, + metadatas, + embeddings, + ids, + parent_ids=parent_ids, + threads=threads, ) async def similarity_search( @@ -99,12 +107,13 @@ async def similarity_search( See VectorCollection.similarity_search for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.similarity_search( - query, k, filter, exact=exact, threads=threads - ), + return await self._run( + self._collection.similarity_search, + query, + k, + filter, + exact=exact, + threads=threads, ) async def similarity_search_batch( @@ -121,12 +130,13 @@ async def similarity_search_batch( See VectorCollection.similarity_search_batch for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.similarity_search_batch( - queries, k, filter, exact=exact, threads=threads - ), + return await self._run( + self._collection.similarity_search_batch, + queries, + k, + filter, + exact=exact, + threads=threads, ) async def keyword_search( @@ -140,10 +150,11 @@ async def keyword_search( See VectorCollection.keyword_search for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.keyword_search(query, k, filter), + return await self._run( + self._collection.keyword_search, + query, + k, + filter, ) async def hybrid_search( @@ -162,18 +173,15 @@ async def hybrid_search( See VectorCollection.hybrid_search for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.hybrid_search( - query, - k, - filter, - query_vector=query_vector, - vector_k=vector_k, - keyword_k=keyword_k, - rrf_k=rrf_k, - ), + return await self._run( + self._collection.hybrid_search, + query, + k, + filter, + query_vector=query_vector, + vector_k=vector_k, + keyword_k=keyword_k, + rrf_k=rrf_k, ) async def max_marginal_relevance_search( @@ -189,12 +197,13 @@ async def max_marginal_relevance_search( See VectorCollection.max_marginal_relevance_search for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.max_marginal_relevance_search( - query, k, fetch_k, lambda_mult, filter - ), + return await self._run( + self._collection.max_marginal_relevance_search, + query, + k, + fetch_k, + lambda_mult, + filter, ) async def delete_by_ids(self, ids: Sequence[int]) -> None: @@ -203,11 +212,7 @@ async def delete_by_ids(self, ids: Sequence[int]) -> None: See VectorCollection.delete_by_ids for full documentation. """ - loop = asyncio.get_running_loop() - await loop.run_in_executor( - self._executor, - lambda: self._collection.delete_by_ids(ids), - ) + await self._run(self._collection.delete_by_ids, ids) async def get_documents( self, @@ -220,12 +225,11 @@ async def get_documents( 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, limit=limit, offset=offset - ), + return await self._run( + 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]: @@ -233,40 +237,25 @@ async def get_embeddings_by_ids(self, ids: Sequence[int]) -> dict[int, Any]: See VectorCollection.get_embeddings_by_ids for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.get_embeddings_by_ids(list(ids)), + return await self._run( + self._collection.get_embeddings_by_ids, + list(ids), ) - async def update_metadata( - self, updates: list[tuple[int, dict[str, Any]]] - ) -> int: + async def update_metadata(self, updates: list[tuple[int, dict[str, Any]]]) -> int: """Update metadata for multiple documents (shallow merge). See VectorCollection.update_metadata for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.update_metadata(updates), - ) + return await self._run(self._collection.update_metadata, updates) async def count(self) -> int: """Count documents in collection.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - self._collection.count, - ) + return await self._run(self._collection.count) async def save(self) -> None: """Save collection to disk.""" - loop = asyncio.get_running_loop() - await loop.run_in_executor( - self._executor, - self._collection.save, - ) + await self._run(self._collection.save) @property def dim(self) -> int | None: @@ -283,10 +272,183 @@ async def remove_texts( See VectorCollection.remove_texts for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.remove_texts(texts, filter), + return await self._run(self._collection.remove_texts, texts, filter) + + # ───────────────────────────────────────────────────────────────────────── + # 2.6.1 — pending vectors, counters, edges, events, TTL (Async) + # ───────────────────────────────────────────────────────────────────────── + + async def update_embedding( + self, + doc_id: int, + vector: Any, + *, + source: str | None = None, + ) -> None: + """Buffer a vector update; promoted to HNSW on flush_pending().""" + await self._run( + self._collection.update_embedding, + doc_id, + vector, + source=source, + ) + + async def flush_pending(self, *, max_batch: int | None = None) -> int: + """Flush buffered vector updates into the HNSW index.""" + return await self._run(self._collection.pending.flush, max_batch=max_batch) + + 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) + + async def add_edge( + self, + src: int, + dst: int, + *, + kind: str = "", + weight: float = 0.0, + bonus: float = 0.0, + hits: int = 0, + metadata: dict | None = None, + ) -> int: + return await self._run( + self._collection.edges.add_edge, + src, + dst, + kind=kind, + weight=weight, + bonus=bonus, + hits=hits, + metadata=metadata, + ) + + async def update_edge( + self, + src: int, + dst: int, + *, + kind: str = "", + weight: float | None = None, + bonus: float | None = None, + hits: int | None = None, + metadata: dict | None = None, + dweight: float = 0.0, + dbonus: float = 0.0, + dhits: int = 0, + ) -> int: + return await self._run( + self._collection.edges.update_edge, + src, + dst, + kind=kind, + weight=weight, + bonus=bonus, + hits=hits, + metadata=metadata, + dweight=dweight, + dbonus=dbonus, + dhits=dhits, + ) + + async def delete_edge( + self, + src: int, + dst: int, + *, + kind: str = "", + ) -> int: + return await self._run( + self._collection.edges.delete_edge, + src, + dst, + kind=kind, + ) + + async def get_edges( + self, + src: int | None = None, + dst: int | None = None, + *, + kind: str | None = None, + filter: dict[str, Any] | None = None, + limit: int | None = None, + ) -> list: + return await self._run( + self._collection.edges.get_edges, + src=src, + dst=dst, + kind=kind, + filter=filter, + limit=limit, + ) + + async def set_ttl( + self, + doc_id: int, + *, + seconds: float | None = None, + expires_at: float | None = None, + on_expire: str = "delete", + ) -> None: + await self._run( + self._collection.ttl.set, + doc_id, + seconds=seconds, + expires_at=expires_at, + on_expire=on_expire, + ) + + async def clear_ttl(self, doc_id: int) -> None: + await self._run(self._collection.ttl.clear, doc_id) + + async def sweep_ttl( + self, + *, + now: float | None = None, + limit: int = 1000, + ) -> tuple[list[int], list[int]]: + return await self._run( + self._collection.ttl.sweep, + now=now, + limit=limit, + ) + + async def read_events( + self, + *, + since: int = 0, + kind: str | None = None, + limit: int = 500, + ) -> list: + return await self._run( + self._collection.events.read, + since=since, + kind=kind, + limit=limit, + ) + + async def last_event_seq(self) -> int: + return await self._run(self._collection.events.last_seq) + + async def rebuild_if_needed( + self, + *, + max_pending: int | None = None, + max_deleted: int | None = None, + ) -> bool: + kwargs: dict[str, Any] = {} + if max_pending is not None: + kwargs["max_pending"] = max_pending + if max_deleted is not None: + kwargs["max_deleted"] = max_deleted + return await self._run( + self._collection.maintenance.rebuild_if_needed, + **kwargs, ) # ───────────────────────────────────────────────────────────────────────── @@ -301,58 +463,43 @@ async def rebuild_index( expansion_search: int | None = None, ) -> int: """Rebuild the HNSW index. See VectorCollection.rebuild_index.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.rebuild_index( - connectivity=connectivity, - expansion_add=expansion_add, - expansion_search=expansion_search, - ), + return await self._run( + self._collection.rebuild_index, + connectivity=connectivity, + expansion_add=expansion_add, + expansion_search=expansion_search, ) async def get_children(self, doc_id: int) -> list: """Get direct children of a document.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.get_children(doc_id), - ) + return await self._run(self._collection.get_children, doc_id) async def get_parent(self, doc_id: int): """Get parent document, or None.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.get_parent(doc_id), - ) + return await self._run(self._collection.get_parent, doc_id) - async def get_descendants( - self, doc_id: int, max_depth: int | None = None - ) -> list: + async def get_descendants(self, doc_id: int, max_depth: int | None = None) -> list: """Get all descendants recursively.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.get_descendants(doc_id, max_depth), + return await self._run( + self._collection.get_descendants, + doc_id, + max_depth, ) - async def get_ancestors( - self, doc_id: int, max_depth: int | None = None - ) -> list: + async def get_ancestors(self, doc_id: int, max_depth: int | None = None) -> list: """Get all ancestors to root.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.get_ancestors(doc_id, max_depth), + return await self._run( + self._collection.get_ancestors, + doc_id, + max_depth, ) async def set_parent(self, doc_id: int, parent_id: int | None) -> bool: """Set or remove parent relationship.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.set_parent(doc_id, parent_id), + return await self._run( + self._collection.set_parent, + doc_id, + parent_id, ) # ───────────────────────────────────────────────────────────────────────── @@ -387,17 +534,14 @@ async def cluster( narrowed = cast(Literal["kmeans", "minibatch_kmeans", "hdbscan"], algorithm) - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.cluster( - n_clusters, - narrowed, - filter=filter, - sample_size=sample_size, - min_cluster_size=min_cluster_size, - random_state=random_state, - ), + return await self._run( + self._collection.cluster, + n_clusters, + narrowed, + filter=filter, + sample_size=sample_size, + min_cluster_size=min_cluster_size, + random_state=random_state, ) async def auto_tag( @@ -413,15 +557,12 @@ async def auto_tag( See VectorCollection.auto_tag for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.auto_tag( - cluster_result, - method=method, - n_keywords=n_keywords, - custom_callback=custom_callback, - ), + return await self._run( + self._collection.auto_tag, + cluster_result, + method=method, + n_keywords=n_keywords, + custom_callback=custom_callback, ) async def assign_cluster_metadata( @@ -437,15 +578,12 @@ async def assign_cluster_metadata( See VectorCollection.assign_cluster_metadata for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.assign_cluster_metadata( - cluster_result, - tags, - metadata_key=metadata_key, - tag_key=tag_key, - ), + return await self._run( + self._collection.assign_cluster_metadata, + cluster_result, + tags, + metadata_key=metadata_key, + tag_key=tag_key, ) async def get_cluster_members( @@ -459,12 +597,10 @@ async def get_cluster_members( See VectorCollection.get_cluster_members for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.get_cluster_members( - cluster_id, metadata_key=metadata_key - ), + return await self._run( + self._collection.get_cluster_members, + cluster_id, + metadata_key=metadata_key, ) async def save_cluster( @@ -479,12 +615,11 @@ async def save_cluster( See VectorCollection.save_cluster for full documentation. """ - loop = asyncio.get_running_loop() - await loop.run_in_executor( - self._executor, - lambda: self._collection.save_cluster( - name, cluster_result, metadata=metadata - ), + await self._run( + self._collection.save_cluster, + name, + cluster_result, + metadata=metadata, ) async def load_cluster( @@ -496,11 +631,7 @@ async def load_cluster( See VectorCollection.load_cluster for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.load_cluster(name), - ) + return await self._run(self._collection.load_cluster, name) async def list_clusters(self) -> list[dict[str, Any]]: """ @@ -508,11 +639,7 @@ async def list_clusters(self) -> list[dict[str, Any]]: See VectorCollection.list_clusters for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - self._collection.list_clusters, - ) + return await self._run(self._collection.list_clusters) async def delete_cluster(self, name: str) -> bool: """ @@ -520,11 +647,7 @@ async def delete_cluster(self, name: str) -> bool: See VectorCollection.delete_cluster for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.delete_cluster(name), - ) + return await self._run(self._collection.delete_cluster, name) async def assign_to_cluster( self, @@ -538,16 +661,14 @@ async def assign_to_cluster( See VectorCollection.assign_to_cluster for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._collection.assign_to_cluster( - name, doc_ids, metadata_key=metadata_key - ), + return await self._run( + self._collection.assign_to_cluster, + name, + doc_ids, + metadata_key=metadata_key, ) - class AsyncVectorDB: """ Async wrapper for VectorDB. @@ -581,9 +702,18 @@ def __init__( executor: ThreadPoolExecutor | None = None, **kwargs: Any, ): - self._db = VectorDB(path=path, distance_strategy=distance_strategy, quantization=quantization, **kwargs) + self._db = VectorDB( + path=path, + distance_strategy=distance_strategy, + quantization=quantization, + **kwargs, + ) self._owns_executor = executor is None - self._executor = executor if executor is not None else ThreadPoolExecutor(max_workers=max_workers) + self._executor = ( + executor + if executor is not None + else ThreadPoolExecutor(max_workers=max_workers) + ) self._collections: dict[tuple, AsyncVectorCollection] = {} self._collections_lock = Lock() # Thread-safe collection caching @@ -623,16 +753,19 @@ def collection( ) return self._collections[cache_key] + async def _run(self, fn, /, *args, **kwargs): + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._executor, functools.partial(fn, *args, **kwargs) + ) + def list_collections(self) -> list[str]: """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) - ) + await self._run(self._db.delete_collection, name) # Evict from async-level cache too — match any tuple whose first # element is this name (the cache key now includes store_embeddings). with self._collections_lock: @@ -655,17 +788,14 @@ async def search_collections( See VectorDB.search_collections for full documentation. """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self._executor, - lambda: self._db.search_collections( - query, - collections, - k, - filter, - normalize_scores=normalize_scores, - parallel=parallel, - ), + return await self._run( + self._db.search_collections, + query, + collections, + k, + filter, + normalize_scores=normalize_scores, + parallel=parallel, ) async def vacuum(self, checkpoint_wal: bool = True) -> None: @@ -677,10 +807,7 @@ async def vacuum(self, checkpoint_wal: bool = True) -> None: Args: checkpoint_wal: If True (default), also truncate the WAL file. """ - loop = asyncio.get_running_loop() - await loop.run_in_executor( - self._executor, lambda: self._db.vacuum(checkpoint_wal) - ) + await self._run(self._db.vacuum, checkpoint_wal) def __repr__(self) -> str: return f"AsyncVectorDB(path={self._db.path!r})" diff --git a/src/simplevecdb/constants.py b/src/simplevecdb/constants.py index 10a5e13..ae2c75a 100755 --- a/src/simplevecdb/constants.py +++ b/src/simplevecdb/constants.py @@ -104,3 +104,31 @@ # Maximum sample size for sklearn.metrics.silhouette_score. The metric is # O(n²) in time and memory, so we sample on large datasets. SILHOUETTE_MAX_SAMPLE = 10_000 + +# ============================================================================ +# 2.6.1 catalog extensions +# ============================================================================ + +# Default flush batch when promoting pending vectors to the HNSW index. +PENDING_FLUSH_DEFAULT_BATCH = 1000 + +# Polling interval used by collection.events.subscribe() when no new +# events are available. WAL mode means cross-process commits show up +# on the next iteration. +EVENTS_POLL_INTERVAL_S = 0.1 + +# Soft cap on stored events. Set to None to disable startup pruning. +EVENTS_RETENTION_LIMIT = 100_000 + +# Default cadence for opt-in TTL background sweeper threads. +TTL_SWEEP_DEFAULT_INTERVAL_S = 60.0 + +# Heuristic thresholds that trigger maintenance.rebuild_if_needed. +REBUILD_PENDING_THRESHOLD = 5_000 +REBUILD_TOMBSTONE_THRESHOLD = 5_000 +REBUILD_MIN_INTERVAL_S = 3600.0 + +# SQLite native lock-wait window (ms). Lower retry pressure on multi-writer +# workloads — busy_timeout means SQLite blocks the caller in C rather +# than surfacing "database is locked" to Python. +SQLITE_BUSY_TIMEOUT_MS = 5000 diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index e58f003..a48de35 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -8,15 +8,17 @@ from __future__ import annotations import logging +import math import os import re import sqlite3 import tempfile import threading +import time import numpy as np import uuid from collections import defaultdict -from collections.abc import Generator, Iterable, Sequence +from collections.abc import Generator, Iterable, Iterator, Sequence from typing import Any, TYPE_CHECKING from pathlib import Path import platform @@ -25,7 +27,8 @@ Document, DistanceStrategy, Quantization, - MigrationRequiredError, + Edge, + Event, StreamingProgress, ProgressCallback, ClusterResult, @@ -34,7 +37,7 @@ from .utils import _import_optional from .engine.quantization import QuantizationStrategy from .engine.search import SearchEngine -from .engine.catalog import CatalogManager +from .engine.catalog import CatalogManager, _TxState from .engine.usearch_index import UsearchIndex from .engine.clustering import ClusterEngine, ClusterAlgorithm from . import constants @@ -181,6 +184,7 @@ def __init__( encryption_key: str | bytes | None = None, store_embeddings: bool = False, lock: threading.RLock | None = None, + tx_state: _TxState | None = None, ): self.conn = conn self._db_path = db_path @@ -194,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() + # 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() # Sanitize name to prevent issues if not re.match(constants.COLLECTION_NAME_PATTERN, name): @@ -204,10 +211,8 @@ def __init__( # Table names if name == "default": self._table_name = "tinyvec_items" - self._legacy_vec_table = "vec_index" # For migration else: self._table_name = f"items_{name}" - self._legacy_vec_table = f"vectors_{name}" # For migration self._fts_table_name = f"{self._table_name}_fts" @@ -219,11 +224,14 @@ def __init__( # Initialize components — share the connection lock with the catalog # so add_documents / delete_by_ids / etc. all serialize properly. + # The optional _tx_state is also shared with the parent VectorDB + # so db.transaction() can suspend per-call commits everywhere. self._catalog = CatalogManager( conn=self.conn, table_name=self._table_name, fts_table_name=self._fts_table_name, lock=self._lock, + tx_state=getattr(self, "_tx_state", None), ) self._catalog.create_tables() @@ -258,9 +266,6 @@ def __init__( distance_strategy=self.distance_strategy, ) - # Check for and perform migration from sqlite-vec - self._migrate_from_sqlite_vec_if_needed() - def _resolve_index_path(self) -> str | None: """ Resolve the actual index path, handling encryption. @@ -288,52 +293,6 @@ def _resolve_index_path(self) -> str | None: return self._index_path - def _migrate_from_sqlite_vec_if_needed(self) -> None: - """Auto-migrate from sqlite-vec to usearch on first connection.""" - if not self._catalog.check_legacy_sqlite_vec(self._legacy_vec_table): - return - - _logger.info( - "Detected legacy sqlite-vec data in collection '%s'. Migrating to usearch...", - self.name, - ) - - try: - # Get legacy vectors - legacy_data = self._catalog.get_legacy_vectors(self._legacy_vec_table) - if not legacy_data: - _logger.warning("No vectors found in legacy table") - self._catalog.drop_legacy_vec_table(self._legacy_vec_table) - return - - # Deserialize and add to usearch - keys = [] - vectors = [] - for rowid, blob in legacy_data: - vec = np.frombuffer(blob, dtype=np.float32) - keys.append(rowid) - vectors.append(vec) - - keys_arr = np.array(keys, dtype=np.uint64) - vectors_arr = np.array(vectors, dtype=np.float32) - - self._index.add(keys_arr, vectors_arr) - self._index.save() - - # Drop legacy table - self._catalog.drop_legacy_vec_table(self._legacy_vec_table) - - _logger.info( - "Migration complete: %d vectors migrated to usearch", len(keys) - ) - - except Exception as e: - _logger.error("Migration failed: %s", e) - raise RuntimeError( - f"Failed to migrate from sqlite-vec: {e}. " - "You may need to manually migrate or restore from backup." - ) from e - def add_texts( self, texts: Sequence[str], @@ -437,7 +396,9 @@ def add_texts( ) # Add to usearch index - self._index.add(np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads) + self._index.add( + np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads + ) all_ids.extend(doc_ids) @@ -605,9 +566,7 @@ def _process_streaming_batch( "Internal error: streaming batch reached persistence with " "unresolved auto-embedding placeholders." ) - embeds_resolved: list[Sequence[float]] = [ - e for e in embeds if e is not None - ] + embeds_resolved: list[Sequence[float]] = [e for e in embeds if e is not None] # Validate vectors before any persistence — see add_texts for rationale. emb_np = np.asarray(embeds_resolved, dtype=np.float32) @@ -891,7 +850,9 @@ def rebuild_index( # 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) + return self._rebuild_index_locked( + connectivity, expansion_add, expansion_search + ) def _rebuild_index_locked( self, @@ -956,9 +917,15 @@ def _rebuild_index_locked( 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, + 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() @@ -1488,6 +1455,135 @@ def update_metadata(self, updates: list[tuple[int, dict[str, Any]]]) -> int: """ return self._catalog.update_metadata_batch(updates) + def update_embedding( + self, + doc_id: int, + vector: Sequence[float] | "np.ndarray", + *, + source: str | None = None, + ) -> None: + """Buffer a native vector update without HNSW remove+re-add (gap 1). + + Writes the new vector to the per-collection pending buffer in a + single SQL transaction. The HNSW index is **not** modified until + `pending.flush()` runs (manually or via threshold-driven flush + during normal writes), so this method is cheap regardless of how + many edits a single doc receives between flushes. + + Until flush, similarity searches still rank the doc using its + previous vector — by design. Use `pending.flush()` to promote + buffered vectors when their effect on retrieval matters. + + Args: + doc_id: Existing document id. + vector: Float vector with the same dim as the index. + source: Optional free-form tag stored alongside the buffered + vector (useful for tracing which subsystem queued it). + + Raises: + ValueError: If vector dimension differs from the index dim. + """ + arr = np.asarray(vector, dtype=np.float32) + if arr.ndim != 1: + raise ValueError( + f"update_embedding expects a 1-D vector, got shape {arr.shape}" + ) + if self._index.ndim is not None and arr.shape[0] != self._index.ndim: + raise ValueError( + f"Vector dim {arr.shape[0]} != index dim {self._index.ndim}" + ) + # NaN/inf would survive the buffer and corrupt distance math on flush. + if not np.all(np.isfinite(arr)): + raise ValueError("update_embedding vector must be finite (no NaN/inf)") + self._catalog.upsert_pending_vector(doc_id, arr.tobytes(), source) + + @property + def pending(self) -> "_PendingNamespace": + """Sub-namespace for buffered vector updates (gaps 1 & 6).""" + ns = self.__dict__.get("_pending_ns") + if ns is None: + ns = _PendingNamespace(self) + self.__dict__["_pending_ns"] = ns + return ns + + def increment_metadata(self, doc_id: int, deltas: dict[str, float | int]) -> int: + """Atomically add numeric deltas to metadata counters (gap 4). + + One UPDATE statement applies all deltas via chained json_set, so + concurrent writers cannot lose increments — every call is a + SQLite-atomic read-modify-write. Missing keys are treated as 0. + + Example: + >>> collection.increment_metadata( + ... doc_id, {"retrieval_count": 1, "drift_total": 0.02} + ... ) + 1 + + Args: + doc_id: Target document id. + deltas: Map of counter name -> numeric delta (int or float). + Keys must be safe identifiers (alphanumeric + underscore, + not starting with a digit). + + Returns: + 1 if the row exists and was updated, 0 otherwise. + """ + return self._catalog.increment_metadata(doc_id, deltas) + + @property + def counters(self) -> "_CountersNamespace": + """Sub-namespace for atomic counter operations (gap 4).""" + ns = self.__dict__.get("_counters_ns") + if ns is None: + ns = _CountersNamespace(self) + self.__dict__["_counters_ns"] = ns + return ns + + @property + def edges(self) -> "_EdgesNamespace": + """Sub-namespace for weighted directed edges (gap 3).""" + ns = self.__dict__.get("_edges_ns") + if ns is None: + ns = _EdgesNamespace(self) + self.__dict__["_edges_ns"] = ns + return ns + + @property + def events(self) -> "_EventsNamespace": + """Sub-namespace for the change feed (gap 7).""" + ns = self.__dict__.get("_events_ns") + if ns is None: + ns = _EventsNamespace(self) + self.__dict__["_events_ns"] = ns + return ns + + @property + def ttl(self) -> "_TTLNamespace": + """Sub-namespace for TTL/expiry (gap 8).""" + ns = self.__dict__.get("_ttl_ns") + if ns is None: + ns = _TTLNamespace(self) + self.__dict__["_ttl_ns"] = ns + return ns + + def tx(self) -> "_CollectionTransaction": + """Single-collection transaction convenience (gap 2). + + Equivalent to opening a `db.transaction()` at the parent + VectorDB, but the context manager yields this collection + directly instead of a dict-of-collections proxy. + """ + return _CollectionTransaction(self) + + @property + def maintenance(self) -> "_MaintenanceNamespace": + """Sub-namespace for periodic rebuild scheduling (gap 9).""" + ns = self.__dict__.get("_maint_ns") + if ns is None: + ns = _MaintenanceNamespace(self) + self.__dict__["_maint_ns"] = ns + return ns + @property def dim(self) -> int | None: """Vector dimension (None if no vectors added yet).""" @@ -1503,6 +1599,721 @@ def __repr__(self) -> str: ) +class _PendingNamespace: + """Buffered vector updates exposed as `collection.pending` (gaps 1 & 6). + + The buffer stores new vectors in a SQL table; HNSW only sees them + after `flush()` promotes the batch in a single locked operation. + Multiple updates to the same doc_id between flushes coalesce + (last-write-wins). + """ + + __slots__ = ("_collection",) + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + + def update( + self, + doc_id: int, + vector: Sequence[float] | "np.ndarray", + *, + source: str | None = None, + ) -> None: + """Alias for `VectorCollection.update_embedding`.""" + self._collection.update_embedding(doc_id, vector, source=source) + + def update_many( + self, + updates: Sequence[tuple[int, Sequence[float]]], + *, + source: str | None = None, + ) -> None: + """Buffer many vector updates in one SQL transaction.""" + if not updates: + return + ndim = self._collection._index.ndim + rows: list[tuple[int, bytes, str | None]] = [] + for doc_id, vec in updates: + arr = np.asarray(vec, dtype=np.float32) + if arr.ndim != 1: + raise ValueError( + f"update_many: vector for {doc_id} must be 1-D, " + f"got shape {arr.shape}" + ) + if ndim is not None and arr.shape[0] != ndim: + raise ValueError( + f"update_many: vector dim {arr.shape[0]} for {doc_id} " + f"!= index dim {ndim}" + ) + rows.append((int(doc_id), arr.tobytes(), source)) + self._collection._catalog.upsert_pending_vectors_many(rows) + + def blend_toward( + self, + doc_ids: Sequence[int], + centroid: Sequence[float], + alpha: float, + *, + source: str | None = "blend_toward", + ) -> int: + """Drift each doc's vector toward `centroid` by alpha. + + v' = (1 - alpha) * v + alpha * centroid + + Reads the current vector from the pending buffer first (so + repeated blends compose) and falls back to the HNSW index. The + blended vector is written back to the buffer; HNSW is untouched + until flush. + + Args: + doc_ids: Documents to drift. + centroid: Target vector (same dim as the index). + alpha: Mixing weight in [0, 1]. 0 keeps the original; 1 + replaces with the centroid entirely. + source: Tag stored on the buffered rows. + + Returns: + Number of vectors blended and re-buffered. + """ + if not doc_ids: + return 0 + if not 0.0 <= alpha <= 1.0: + raise ValueError(f"alpha must be in [0, 1], got {alpha!r}") + c = np.asarray(centroid, dtype=np.float32) + if c.ndim != 1: + raise ValueError(f"centroid must be 1-D, got shape {c.shape}") + + ids_arr = np.asarray(list(doc_ids), dtype=np.uint64) + # Pull current vectors: pending wins, fallback to HNSW. + cat = self._collection._catalog + idx = self._collection._index + rows: list[tuple[int, bytes, str | None]] = [] + # Bulk fetch from HNSW once; fall through to per-id pending overlay. + idx_vecs = idx.get(ids_arr) + for i, doc_id in enumerate(doc_ids): + pending = cat.get_pending_vector(int(doc_id)) + if pending is not None: + v = np.frombuffer(pending, dtype=np.float32) + else: + v = np.asarray(idx_vecs[i], dtype=np.float32) + if v.shape[0] != c.shape[0]: + raise ValueError( + f"blend_toward: vector dim {v.shape[0]} for {doc_id} " + f"!= centroid dim {c.shape[0]}" + ) + blended = ((1.0 - alpha) * v + alpha * c).astype(np.float32) + rows.append((int(doc_id), blended.tobytes(), source)) + cat.upsert_pending_vectors_many(rows) + return len(rows) + + def size(self) -> int: + """Number of buffered vector updates not yet flushed.""" + return self._collection._catalog.count_pending_vectors() + + def flush(self, *, max_batch: int | None = None) -> int: + """Promote buffered vectors to the HNSW index. + + Reads up to `max_batch` rows in enqueue order, calls + `UsearchIndex.add()` (which already does remove+add atomically + under a single lock), then deletes the flushed pending rows in + the same SQL transaction. Returns the count flushed. + """ + cat = self._collection._catalog + idx = self._collection._index + rows = cat.list_pending_vectors(limit=max_batch) + if not rows: + return 0 + ids = np.asarray([r[0] for r in rows], dtype=np.uint64) + # Reconstruct float32 matrix from BLOBs. + ndim = idx.ndim + if ndim is None: + # Empty index — infer from first row. + ndim = len(np.frombuffer(rows[0][1], dtype=np.float32)) + mat = ( + np.frombuffer(b"".join(r[1] for r in rows), dtype=np.float32) + .reshape(len(rows), ndim) + .copy() + ) + # add() takes the write lock and does remove+add per existing key. + idx.add(ids, mat) + cat.delete_pending_vectors([int(i) for i in ids]) + # Track flush count for the rebuild scheduler. + try: + self._collection.maintenance.record_flush(len(rows)) + except Exception: + _logger.debug("flush counter record failed", exc_info=True) + return len(rows) + + +class _TTLNamespace: + """Sub-namespace exposed as `collection.ttl` (gap 8). + + Stores per-doc expiry timestamps. `sweep()` is the explicit cleanup + pass; `start_background()` launches an opt-in daemon thread that + calls `sweep()` on a fixed cadence. The thread is off by default — + nothing wakes up unless the caller asks. + """ + + __slots__ = ("_collection", "_thread", "_stop_event") + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + self._thread: threading.Thread | None = None + self._stop_event: threading.Event | None = None + + def set( + self, + doc_id: int, + expires_at: float | None = None, + *, + seconds: float | None = None, + on_expire: str = "delete", + ) -> int: + """Set a TTL by absolute timestamp or relative seconds. + + Either pass `expires_at` (unix seconds) or `seconds` + (relative to now). Both is an error. + """ + if expires_at is None and seconds is None: + raise ValueError("Must pass expires_at or seconds") + if expires_at is not None and seconds is not None: + raise ValueError("Pass exactly one of expires_at / seconds") + if on_expire not in ("delete", "callback"): + raise ValueError( + f"on_expire must be 'delete' or 'callback', got {on_expire!r}" + ) + if seconds is not None: + secs = float(seconds) + if not math.isfinite(secs): + raise ValueError(f"ttl.set seconds must be finite, got {seconds!r}") + resolved = time.time() + secs + else: + assert expires_at is not None # narrowed by guards above + resolved = float(expires_at) + if not math.isfinite(resolved): + raise ValueError(f"ttl.set expires_at must be finite, got {resolved!r}") + return self._collection._catalog.set_ttl( + doc_id, + resolved, + on_expire=on_expire, + ) + + def clear(self, doc_id: int) -> int: + """Remove a TTL entry. Returns 1 if removed, 0 otherwise.""" + return self._collection._catalog.clear_ttl(doc_id) + + def sweep( + self, + *, + now: float | None = None, + limit: int = 1000, + ) -> tuple[list[int], list[int]]: + """Apply expired entries; returns (deleted_ids, callback_ids). + + Also removes the doc from the in-memory HNSW index when an + entry is deleted. Callers handling on_expire='callback' rows + are responsible for any further action on those ids. + """ + deleted, callback_ids = self._collection._catalog.sweep_ttl( + now=now, limit=limit + ) + if deleted: + try: + self._collection._index.remove(deleted) + except Exception: + # Catalog rows are already gone; index drift means + # subsequent searches may surface phantom hits until a + # rebuild_index() runs. Log loud enough to be noticed. + _logger.warning( + "ttl.sweep: catalog deleted %d ids but HNSW remove " + "failed; index is now divergent — call " + "rebuild_index() to resync", + len(deleted), + exc_info=True, + ) + return deleted, callback_ids + + def start_background( + self, + *, + interval: float = constants.TTL_SWEEP_DEFAULT_INTERVAL_S, + ) -> None: + """Launch a daemon thread that calls sweep() every `interval`s. + + Idempotent: a second call is a no-op. The thread runs until + `stop_background()` is called or the process exits. + """ + interval_f = float(interval) + if not math.isfinite(interval_f) or interval_f <= 0: + raise ValueError( + f"ttl.start_background interval must be a positive finite " + f"number, got {interval!r}" + ) + if self._thread is not None and self._thread.is_alive(): + return + stop_event = threading.Event() + coll = self._collection + + def _loop() -> None: + while not stop_event.is_set(): + try: + self.sweep() + except Exception: + _logger.warning( + "ttl background sweep failed for collection %s", + coll.name, + exc_info=True, + ) + stop_event.wait(interval_f) + + thread = threading.Thread( + target=_loop, + name=f"simplevecdb-ttl-{coll.name}", + daemon=True, + ) + self._stop_event = stop_event + self._thread = thread + thread.start() + + def stop_background(self) -> None: + """Stop the background sweeper. Idempotent. + + Waits up to 5s for the loop to observe the stop event. If the + thread is still alive after the timeout (e.g. blocked in a + long-running sweep), `_thread` is left in place so a subsequent + start_background() does not spawn a duplicate sweeper. + """ + if self._stop_event is not None: + self._stop_event.set() + if self._thread is not None and self._thread.is_alive(): + self._thread.join(timeout=5.0) + if self._thread.is_alive(): + _logger.warning( + "ttl.stop_background: sweeper for %s did not exit " + "within 5s; leaving handle in place to prevent " + "duplicate threads on next start_background()", + self._collection.name, + ) + return + self._thread = None + self._stop_event = None + + +class _EventsNamespace: + """Append-only change feed exposed as `collection.events` (gap 7). + + Backed by the per-collection `_events` SQL table. SQLite WAL mode + (already enabled at connection open) lets cross-process readers see + committed rows immediately, so multi-process subscribers work + without an external bus. + """ + + __slots__ = ("_collection",) + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + + def append( + self, + kind: str, + *, + doc_id: int | None = None, + payload: dict | None = None, + ) -> int: + """Append an event (caller-driven). Returns the assigned seq.""" + return self._collection._catalog.append_event( + kind, doc_id=doc_id, payload=payload + ) + + def last_seq(self) -> int: + """Highest seq currently in the feed (0 if empty).""" + return self._collection._catalog.last_event_seq() + + def read( + self, + *, + since: int = 0, + kind: str | None = None, + limit: int | None = None, + ) -> list["Event"]: + """Read events with seq > since.""" + rows = self._collection._catalog.read_events( + since=since, kind=kind, limit=limit + ) + + return [ + Event(seq=r[0], ts=r[1], kind=r[2], doc_id=r[3], payload=r[4]) for r in rows + ] + + def subscribe( + self, + *, + since: int = 0, + kind: str | None = None, + poll_interval: float = constants.EVENTS_POLL_INTERVAL_S, + batch: int = 500, + ) -> "Iterator[Event]": + """Generator yielding events as they appear. Caller controls exit. + + Uses simple polling on a background-friendly cadence; WAL mode + means cross-process commits become visible to this reader on + the next iteration without explicit synchronization. + """ + last = int(since) + while True: + events = self.read(since=last, kind=kind, limit=batch) + if events: + for e in events: + yield e + last = events[-1].seq + if len(events) == batch: + # Drained a full batch; loop again immediately to keep up. + continue + time.sleep(poll_interval) + + def prune(self, *, before_seq: int) -> int: + """Drop events with seq < before_seq.""" + return self._collection._catalog.prune_events(before_seq=before_seq) + + +class _MaintenanceNamespace: + """Threshold-driven rebuild scheduler (gap 9). + + Wraps `VectorCollection.rebuild_index()` with a heuristic gate so + callers can opportunistically rebuild without rolling their own + bookkeeping. The triggers compose with OR — any one being true + fires a rebuild. + + Triggers (overridable per call): + * max_pending — total pending flushes since last rebuild + * max_deleted — tombstones in usearch (size mismatch with catalog) + * max_age_s — wall-clock seconds since last rebuild + """ + + __slots__ = ("_collection", "_pending_flushes", "_last_rebuild_ts") + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + self._pending_flushes = 0 + self._last_rebuild_ts = time.time() + + def record_flush(self, count: int = 1) -> None: + """Bump the pending-flush counter (called by pending.flush()).""" + self._pending_flushes += int(count) + + def suggest_rebuild( + self, + *, + max_pending: int = constants.REBUILD_PENDING_THRESHOLD, + max_deleted: int = constants.REBUILD_TOMBSTONE_THRESHOLD, + max_age_s: float = constants.REBUILD_MIN_INTERVAL_S, + ) -> tuple[bool, str | None]: + """Return (should_rebuild, reason).""" + if self._pending_flushes >= max_pending: + return True, f"pending_flushes={self._pending_flushes}" + # Tombstones: usearch reports deletions as size shrinkage but + # the index file may still be larger; approximate as + # catalog_count vs index.size disagreement scaled by deletions. + try: + cat_count = self._collection._catalog.count() + idx_size = self._collection._index.size + tombstones = max(0, idx_size - cat_count) + if tombstones >= max_deleted: + return True, f"tombstones={tombstones}" + except Exception: + pass + age = time.time() - self._last_rebuild_ts + if age >= max_age_s and self._pending_flushes > 0: + return True, f"age_s={age:.0f}" + return False, None + + def rebuild_if_needed( + self, + *, + max_pending: int = constants.REBUILD_PENDING_THRESHOLD, + max_deleted: int = constants.REBUILD_TOMBSTONE_THRESHOLD, + max_age_s: float = constants.REBUILD_MIN_INTERVAL_S, + ) -> bool: + """Run rebuild_index() iff a trigger fired. Returns True if it ran.""" + should, reason = self.suggest_rebuild( + max_pending=max_pending, + max_deleted=max_deleted, + max_age_s=max_age_s, + ) + if not should: + return False + _logger.info( + "Rebuilding %s index (reason=%s)", + self._collection.name, + reason, + ) + self._collection.rebuild_index() + self._pending_flushes = 0 + self._last_rebuild_ts = time.time() + try: + self._collection.events.append("rebuild", payload={"reason": reason}) + except Exception: + _logger.debug("rebuild event append failed", exc_info=True) + return True + + +class _DBTransaction: + """Context manager backing `VectorDB.transaction()` (gap 2). + + Acquires the DB-level RLock, opens a SAVEPOINT, and increments the + shared transaction-depth counter so every catalog method this DB + owns skips its per-call commit. On success the SAVEPOINT is + released; on exception it's rolled back and the depth is reset. + + Yields a mapping-like object so callers can do + `tx["collection_name"]` to operate on individual collections. + """ + + __slots__ = ("_db", "_savepoint_name", "_entered") + + def __init__(self, db: "VectorDB") -> None: + self._db = db + self._savepoint_name: str | None = None + self._entered = False + + def __enter__(self) -> "_DBTransaction": + self._db._lock.acquire() + try: + depth = self._db._tx_state.depth + name = f"simplevecdb_tx_{depth + 1}" + self._db.conn.execute(f"SAVEPOINT {name}") + self._db._tx_state.depth = depth + 1 + self._savepoint_name = name + self._entered = True + except Exception: + self._db._lock.release() + raise + return self + + def __exit__(self, exc_type, exc, tb) -> None: + try: + name = self._savepoint_name + assert name is not None + try: + if exc_type is None: + self._db.conn.execute(f"RELEASE SAVEPOINT {name}") + else: + self._db.conn.execute(f"ROLLBACK TO SAVEPOINT {name}") + self._db.conn.execute(f"RELEASE SAVEPOINT {name}") + finally: + self._db._tx_state.depth = max(0, self._db._tx_state.depth - 1) + # Outermost commit: if depth fell to 0, finalize the + # implicit Python sqlite3 transaction so changes flush. + if self._db._tx_state.depth == 0 and exc_type is None: + try: + self._db.conn.commit() + except Exception: + _logger.error("outer transaction commit failed", exc_info=True) + raise + finally: + self._db._lock.release() + + def __getitem__(self, name: str) -> "VectorCollection": + return self._db.collection(name) + + def collection(self, name: str) -> "VectorCollection": + return self._db.collection(name) + + +class _CollectionTransaction(_DBTransaction): + """Single-collection wrapper that yields the collection directly.""" + + __slots__ = ("_collection",) + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + # Reuse the shared tx_state and lock from the collection. + # _DBTransaction expects ._db; we create a shim. + super().__init__(_CollectionTxShim(collection)) # type: ignore[arg-type] + + def __enter__(self) -> "VectorCollection": # type: ignore[override] + super().__enter__() + return self._collection + + +class _CollectionTxShim: + """Minimal proxy emulating the VectorDB attributes _DBTransaction reads.""" + + __slots__ = ("_lock", "_tx_state", "conn") + + def __init__(self, collection: "VectorCollection") -> None: + self._lock = collection._lock + self._tx_state = collection._tx_state + self.conn = collection.conn + + +class _EdgesNamespace: + """Sub-namespace exposed as `collection.edges` (gap 3). + + Provides the four canonical primitives — `add_edge`, `get_edges`, + `update_edge`, `delete_edge` — plus `upsert` and `prune` convenience. + Numeric attributes (weight, bonus, hits, last_touch) are real + columns and are addressable through the range-filter grammar + (`filter={"weight": {"$lt": 0.1}}`). + """ + + __slots__ = ("_collection",) + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + + def add_edge( + self, + src_id: int, + dst_id: int, + *, + kind: str = "", + weight: float = 0.0, + bonus: float = 0.0, + hits: int = 0, + metadata: dict | None = None, + ) -> int: + """Insert a new edge. Use upsert() if collisions are expected.""" + return self._collection._catalog.add_edge( + src_id, + dst_id, + kind=kind, + weight=weight, + bonus=bonus, + hits=hits, + metadata=metadata, + ) + + def upsert( + self, + src_id: int, + dst_id: int, + *, + kind: str = "", + weight: float | None = None, + bonus: float | None = None, + hits: int | None = None, + metadata: dict | None = None, + ) -> int: + """Create-or-update; preserves existing fields where args are None.""" + return self._collection._catalog.upsert_edge( + src_id, + dst_id, + kind=kind, + weight=weight, + bonus=bonus, + hits=hits, + metadata=metadata, + ) + + def update_edge( + self, + src_id: int, + dst_id: int, + *, + kind: str = "", + weight: float | None = None, + bonus: float | None = None, + hits: int | None = None, + metadata: dict | None = None, + dweight: float = 0.0, + dbonus: float = 0.0, + dhits: int = 0, + ) -> int: + """Set absolutes and/or apply atomic deltas. See catalog.update_edge.""" + return self._collection._catalog.update_edge( + src_id, + dst_id, + kind=kind, + weight=weight, + bonus=bonus, + hits=hits, + metadata=metadata, + dweight=dweight, + dbonus=dbonus, + dhits=dhits, + ) + + def delete_edge(self, src_id: int, dst_id: int, *, kind: str = "") -> int: + """Drop a single edge by (src, dst, kind).""" + return self._collection._catalog.delete_edge(src_id, dst_id, kind=kind) + + def get_edges( + self, + src: int | None = None, + dst: int | None = None, + *, + kind: str | None = None, + filter: dict | None = None, + limit: int | None = None, + ) -> list["Edge"]: + """Read edges. `src`/`dst` constrain to outgoing/incoming/specific. + + `filter` accepts the same grammar as similarity_search; numeric + keys (weight/bonus/hits/last_touch) map to direct column + comparisons, anything else queries the JSON metadata column. + """ + rows = self._collection._catalog.get_edges( + src_id=src, + dst_id=dst, + kind=kind, + filter=filter, + limit=limit, + ) + + return [ + Edge( + src_id=r[0], + dst_id=r[1], + kind=r[2], + weight=r[3], + hits=r[4], + bonus=r[5], + last_touch=r[6], + metadata=r[7], + ) + for r in rows + ] + + def prune( + self, + *, + kind: str | None = None, + max_weight: float | None = None, + idle_before: float | None = None, + ) -> int: + """Bulk-delete edges by weight ceiling and/or age cutoff.""" + return self._collection._catalog.prune_edges( + kind=kind, + max_weight=max_weight, + idle_before=idle_before, + ) + + +class _CountersNamespace: + """Sub-namespace exposed as `collection.counters` (gap 4).""" + + __slots__ = ("_collection",) + + def __init__(self, collection: "VectorCollection") -> None: + self._collection = collection + + def increment(self, doc_id: int, deltas: dict[str, float | int]) -> int: + """Alias for `VectorCollection.increment_metadata`.""" + return self._collection._catalog.increment_metadata(doc_id, deltas) + + def increment_many(self, updates: list[tuple[int, dict[str, float | int]]]) -> int: + """Apply many counter increments in one transaction.""" + return self._collection._catalog.increment_metadata_many(updates) + + def get( + self, doc_id: int, key: str, default: float | int = 0 + ) -> float | int | None: + """Read a single numeric counter value (None if row missing).""" + return self._collection._catalog.get_metadata_counter(doc_id, key, default) + + class VectorDB: """ Dead-simple local vector database powered by usearch HNSW. @@ -1523,11 +2334,12 @@ class VectorDB: def __init__( self, path: str | Path = ":memory:", - distance_strategy: DistanceStrategy = DistanceStrategy(constants.DEFAULT_DISTANCE_STRATEGY), + distance_strategy: DistanceStrategy = DistanceStrategy( + constants.DEFAULT_DISTANCE_STRATEGY + ), quantization: Quantization = Quantization(constants.DEFAULT_QUANTIZATION), *, encryption_key: str | bytes | None = None, - auto_migrate: bool = False, ): """Initialize the vector database. @@ -1537,13 +2349,8 @@ def __init__( quantization: Default vector compression strategy. encryption_key: Optional passphrase or 32-byte key for at-rest encryption. Encrypts both SQLite (via SQLCipher) and usearch index files (via AES-256-GCM). - auto_migrate: If True, automatically migrate v1.x sqlite-vec data - to usearch. If False (default), raise MigrationRequiredError - when legacy data is detected. Use check_migration() to preview. Raises: - MigrationRequiredError: If auto_migrate=False and legacy sqlite-vec - data is detected. Contains details about what needs migration. EncryptionUnavailableError: If encryption_key provided but encryption dependencies are missing. EncryptionError: If encrypted database cannot be opened (wrong key). @@ -1552,7 +2359,6 @@ def __init__( self.path = str(path) self.distance_strategy = distance_strategy self.quantization = quantization - self.auto_migrate = auto_migrate self._encryption_key = encryption_key self._collections: dict[tuple, VectorCollection] = {} # Single RLock serializing both the _collections cache (avoid @@ -1560,6 +2366,9 @@ def __init__( # Python-level transaction context. Shared with every VectorCollection # and CatalogManager constructed by this VectorDB. self._lock = threading.RLock() + # Shared transaction-depth counter. Bumped by VectorDB.transaction() + # so all catalogs in this DB suspend per-call commits. + self._tx_state: _TxState = _TxState() # Create connection (encrypted or plain) if encryption_key is not None: @@ -1576,6 +2385,11 @@ def __init__( ) self.conn.execute("PRAGMA journal_mode=WAL") self.conn.execute("PRAGMA synchronous=NORMAL") + # Native lock-wait window so SQLite blocks the caller in C + # rather than surfacing 'database is locked' immediately + # under multi-writer load (gap 10). + self.conn.execute(f"PRAGMA busy_timeout={constants.SQLITE_BUSY_TIMEOUT_MS}") + self.conn.execute("PRAGMA foreign_keys=ON") self._encrypted = True _logger.info("Opened encrypted database: %s", self.path) else: @@ -1584,6 +2398,8 @@ def __init__( ) self.conn.execute("PRAGMA journal_mode=WAL") self.conn.execute("PRAGMA synchronous=NORMAL") + self.conn.execute(f"PRAGMA busy_timeout={constants.SQLITE_BUSY_TIMEOUT_MS}") + self.conn.execute("PRAGMA foreign_keys=ON") self._encrypted = False # Verify connection is healthy @@ -1593,17 +2409,29 @@ def __init__( 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) - if migration_info["needs_migration"]: - self.conn.close() - raise MigrationRequiredError( - path=self.path, - collections=migration_info["collections"], - total_vectors=migration_info["total_vectors"], - migration_info=migration_info, - ) + def transaction(self) -> "_DBTransaction": + """Atomic write context spanning all collections (gap 2). + + Wraps the work in a single SQLite SAVEPOINT and bumps the shared + transaction-depth counter so every catalog method skips its + per-call commit. SQL-side mutations (metadata, counters, edges, + events, TTL, and `update_embedding`'s pending-vector overlay) are + atomic: a raised exception triggers ROLLBACK TO SAVEPOINT and all + SQL writes are reverted. + + The HNSW index is not rolled back. Coarse vector mutations + (`add_texts`, `delete`, etc.) call into usearch directly and + their effects persist even if the surrounding SAVEPOINT rolls + back. Use `update_embedding` + `pending.flush()` for vector + changes whose visibility you want gated on commit. + + Example: + >>> with db.transaction() as tx: + ... tx["docs"].update_embedding(id, vec) + ... tx["docs"].increment_metadata(id, {"hits": 1}) + ... tx["docs"].edges.add_edge(src, dst, weight=0.7) + """ + return _DBTransaction(self) def list_collections(self) -> list[str]: """ @@ -1643,14 +2471,23 @@ def list_collections(self) -> list[str]: 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") + # followed by an auxiliary suffix. 2.6.1 added _pending_vectors, + # _edges, _events, _ttl alongside the existing FTS / cluster ones. + _fts_suffixes = ( + "_fts", + "_fts_data", + "_fts_idx", + "_fts_content", + "_fts_docsize", + "_fts_config", + ) + _aux_suffixes = ("_clusters", "_pending_vectors", "_edges", "_events", "_ttl") 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") + for aux in _aux_suffixes: + derivative_suffixes.add(f"{s}{aux}") names: list[str] = [] if has_default: @@ -1681,6 +2518,10 @@ def delete_collection(self, name: str) -> None: table_name = "tinyvec_items" if name == "default" else f"items_{name}" fts_table = f"{table_name}_fts" cluster_table = f"{table_name}_clusters" + pending_table = f"{table_name}_pending_vectors" + edges_table = f"{table_name}_edges" + events_table = f"{table_name}_events" + ttl_table = f"{table_name}_ttl" # Hold the lock for the full delete: drop tables, remove files, and # evict cached collections atomically. The existence check runs @@ -1707,10 +2548,17 @@ def delete_collection(self, name: str) -> None: exc_info=True, ) - # Drop SQLite tables + # Drop SQLite tables. Auxiliary tables (gap 2.6.1) reference + # the main table via FK ON DELETE CASCADE, so dropping the + # main table cleans the children too — but DROP TABLE doesn't + # cascade in SQLite, so drop them explicitly first. with self.conn: 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 {pending_table}") + self.conn.execute(f"DROP TABLE IF EXISTS {edges_table}") + self.conn.execute(f"DROP TABLE IF EXISTS {events_table}") + self.conn.execute(f"DROP TABLE IF EXISTS {ttl_table}") self.conn.execute(f"DROP TABLE IF EXISTS {table_name}") # Delete usearch index file (and encrypted variant if present), @@ -1727,7 +2575,8 @@ def delete_collection(self, name: str) -> None: pass except OSError: _logger.debug( - "Could not unlink %s during delete_collection", p, + "Could not unlink %s during delete_collection", + p, exc_info=True, ) @@ -1831,7 +2680,10 @@ def _search_one(coll: VectorCollection) -> list[tuple[Document, float, str]]: # Execute searches all_results: list[tuple[Document, float, str]] = [] if parallel and len(targets) > 1: - from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError + from concurrent.futures import ( + ThreadPoolExecutor, + TimeoutError as FuturesTimeoutError, + ) with ThreadPoolExecutor(max_workers=min(len(targets), 8)) as executor: futures = [executor.submit(_search_one, coll) for coll in targets] @@ -1907,6 +2759,7 @@ def collection( encryption_key=self._encryption_key, store_embeddings=store_embeddings, lock=self._lock, + tx_state=self._tx_state, ) return self._collections[cache_key] @@ -1929,152 +2782,6 @@ def as_llama_index(self, collection_name: str = "default") -> SimpleVecDBLlamaSt return SimpleVecDBLlamaStore(db_path=self.path, collection_name=collection_name) - # ------------------------------------------------------------------ # - # Convenience - # ------------------------------------------------------------------ # - @staticmethod - def check_migration(path: str | Path) -> dict[str, Any]: - """ - Check if a database needs migration from sqlite-vec (dry-run). - - Use this before opening a v1.x database to understand what will - be migrated. Does not modify the database. - - Args: - path: Path to the SQLite database file - - Returns: - Dict with migration info: - - needs_migration: bool - - collections: list of collection names with legacy data - - total_vectors: estimated total vector count - - estimated_size_mb: approximate usearch index size - - rollback_notes: instructions for reverting if needed - - Example: - >>> info = VectorDB.check_migration("mydb.db") - >>> if info["needs_migration"]: - ... print(f"Will migrate {info['total_vectors']} vectors") - ... print(info["rollback_notes"]) - """ - path = str(path) - if path == ":memory:" or not Path(path).exists(): - return { - "needs_migration": False, - "collections": [], - "total_vectors": 0, - "estimated_size_mb": 0.0, - "rollback_notes": "", - } - - try: - conn = sqlite3.connect(path, check_same_thread=False) - except sqlite3.DatabaseError: - # Database may be encrypted or corrupted - cannot check migration - return { - "needs_migration": False, - "collections": [], - "total_vectors": 0, - "estimated_size_mb": 0.0, - "rollback_notes": "", - } - - try: - # Check for legacy sqlite-vec tables - tables = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table'" - ).fetchall() - except sqlite3.DatabaseError: - # Database is encrypted or corrupted - cannot check migration - conn.close() - return { - "needs_migration": False, - "collections": [], - "total_vectors": 0, - "estimated_size_mb": 0.0, - "rollback_notes": "", - } - - try: - table_names = {t[0] for t in tables} - - legacy_collections = [] - total_vectors = 0 - total_bytes = 0 - - # Check default collection - if "vec_index" in table_names: - try: - count = conn.execute("SELECT COUNT(*) FROM vec_index").fetchone()[0] - if count > 0: - legacy_collections.append("default") - total_vectors += count - # Estimate: rowid(8) + embedding blob - row = conn.execute( - "SELECT embedding FROM vec_index LIMIT 1" - ).fetchone() - if row and row[0]: - dim = len(row[0]) // 4 - total_bytes += count * dim * 4 # float32 - except Exception: - pass - - # Check named collections (vectors_{name}) - for table in table_names: - if table.startswith("vectors_") and table != "vec_index": - # Validate table name from sqlite_master (defense-in-depth) - if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", table): - continue - collection_name = table[8:] # Remove "vectors_" prefix - try: - count = conn.execute( - f"SELECT COUNT(*) FROM {table}" - ).fetchone()[0] - if count > 0: - legacy_collections.append(collection_name) - total_vectors += count - row = conn.execute( - f"SELECT embedding FROM {table} LIMIT 1" - ).fetchone() - if row and row[0]: - dim = len(row[0]) // 4 - total_bytes += count * dim * 4 - except Exception: - pass - - estimated_mb = total_bytes / (1024 * 1024) - - rollback_notes = "" - if legacy_collections: - rollback_notes = f""" -MIGRATION ROLLBACK INSTRUCTIONS: -================================ -1. BEFORE upgrading, backup your database: - cp {path} {path}.backup - -2. If migration fails or you need to revert: - - Delete the new .usearch files: {path}.*.usearch - - Restore from backup: cp {path}.backup {path} - - Downgrade to simplevecdb<2.0.0 - -3. After successful migration, the legacy sqlite-vec tables are dropped. - Keep your backup until you've verified the migration worked correctly. - -4. New storage layout after migration: - - {path} (SQLite: metadata, text, FTS, embeddings) - - {path}..usearch (usearch HNSW index per collection) -""" - - return { - "needs_migration": len(legacy_collections) > 0, - "collections": legacy_collections, - "total_vectors": total_vectors, - "estimated_size_mb": round(estimated_mb, 2), - "rollback_notes": rollback_notes.strip(), - } - finally: - conn.close() - def vacuum(self, checkpoint_wal: bool = True) -> None: """ Reclaim disk space by rebuilding the SQLite database file. @@ -2115,8 +2822,11 @@ def close(self) -> None: for col in self._collections.values(): ephemeral = getattr(col, "_ephemeral_index_path", None) if ephemeral: - for p in (Path(ephemeral), Path(ephemeral + ".tmp"), - Path(ephemeral + ".lock")): + for p in ( + Path(ephemeral), + Path(ephemeral + ".tmp"), + Path(ephemeral + ".lock"), + ): try: p.unlink() except FileNotFoundError: diff --git a/src/simplevecdb/embeddings/models.py b/src/simplevecdb/embeddings/models.py index 3f8d96c..ecb14b4 100755 --- a/src/simplevecdb/embeddings/models.py +++ b/src/simplevecdb/embeddings/models.py @@ -25,6 +25,7 @@ def _validate_repo_id(repo_id: str) -> None: "Expected 'namespace/name' with [A-Za-z0-9_.-] characters only." ) + _logger = logging.getLogger("simplevecdb.embeddings.models") # HEAD timeout for HuggingFace snapshot_download. The download itself relies diff --git a/src/simplevecdb/embeddings/server.py b/src/simplevecdb/embeddings/server.py index 89d400d..33714b3 100755 --- a/src/simplevecdb/embeddings/server.py +++ b/src/simplevecdb/embeddings/server.py @@ -337,7 +337,9 @@ 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]: +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: diff --git a/src/simplevecdb/encryption.py b/src/simplevecdb/encryption.py index bb85baf..d9715d4 100755 --- a/src/simplevecdb/encryption.py +++ b/src/simplevecdb/encryption.py @@ -294,9 +294,7 @@ def create_encrypted_connection( db_path_obj = Path(path_str) salt_path = db_path_obj.with_name(db_path_obj.name + _SALT_SIDECAR_SUFFIX) is_new_db = not db_path_obj.exists() or db_path_obj.stat().st_size == 0 - has_sidecar = ( - salt_path.exists() and salt_path.stat().st_size == SALT_SIZE - ) + has_sidecar = salt_path.exists() and salt_path.stat().st_size == SALT_SIZE use_legacy_passphrase_path = (not is_new_db) and (not has_sidecar) try: @@ -353,6 +351,13 @@ def create_encrypted_connection( # Set performance optimizations (same as non-encrypted) conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") + # Native lock-wait window (gap 10): SQLite blocks the caller in + # C for up to busy_timeout ms before surfacing "database is + # locked", trimming retry pressure under multi-writer load. + from . import constants as _c + + conn.execute(f"PRAGMA busy_timeout={_c.SQLITE_BUSY_TIMEOUT_MS}") + conn.execute("PRAGMA foreign_keys=ON") return conn # type: ignore[return-value] @@ -491,9 +496,7 @@ def encrypt_file( # file, fsync, restrict permissions, then os.replace() onto the # target path. A crash mid-write leaves only the temp file; the # target is never torn. - _atomic_write_bytes( - output_path, header + nonce + ciphertext, mode=0o600 - ) + _atomic_write_bytes(output_path, header + nonce + ciphertext, mode=0o600) _logger.debug( "Encrypted %d bytes -> %d bytes", diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 1487ab2..835e3a0 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -9,6 +9,7 @@ import json import logging +import math import re import threading from typing import Any, TYPE_CHECKING, Callable @@ -16,7 +17,7 @@ from ..utils import _batched -from ..utils import validate_filter, retry_on_lock +from ..utils import validate_filter, retry_on_lock, normalize_filter if TYPE_CHECKING: import sqlite3 @@ -36,6 +37,190 @@ def _validate_table_name(name: str) -> None: ) +def _coerce_scalar(arg: Any) -> Any: + """Coerce a Python scalar to the value SQLite stores via json_extract. + + json_extract returns 0/1 for JSON booleans, so an `$eq True` filter + must compare against 1 (not Python True, which sqlite3 would bind as + integer 1 anyway, but we make it explicit). + """ + if isinstance(arg, bool): + return 1 if arg else 0 + return arg + + +# Identifier safety for keys interpolated into SQL (atomic counters, etc.). +_SAFE_IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _validate_identifier(name: str, what: str = "identifier") -> None: + """Reject names that would be unsafe to inline into SQL string literals.""" + if not isinstance(name, str) or not _SAFE_IDENT_RE.match(name): + raise ValueError( + f"Invalid {what} '{name}'. Must match {_SAFE_IDENT_RE.pattern}." + ) + + +def _check_finite_edge_field(value: Any, field: str) -> None: + """Reject NaN/inf for numeric edge columns (weight/bonus/dweight/dbonus). + + A NaN value would persist into the column and silently break every + subsequent range filter (NaN comparisons return false in SQL), so + fail at the boundary instead. + """ + if value is None: + return + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError( + f"edge {field} must be int or float, got {type(value).__name__}" + ) + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"edge {field} must be finite, got {value!r}") + + +class _TxState: + """Shared per-VectorDB transaction depth counter (gap 2). + + Used by VectorDB.transaction() to mark all collections/catalogs as + operating inside an outer SAVEPOINT. Catalog write helpers consult + this to decide whether to commit on exit. + """ + + __slots__ = ("depth",) + + def __init__(self) -> None: + self.depth: int = 0 + + +class _CatalogWritable: + """Context manager replacing `with self._lock, self.conn:`. + + Behaviour: + * Always takes the catalog write lock. + * Outside a transaction, also enters the connection's implicit-tx + context; on exit it commits (or rolls back on error). + * Inside a transaction (`tx_state.depth > 0`) it skips the conn + context — the outer SAVEPOINT handles atomicity, so we must NOT + commit at every catalog call. + """ + + __slots__ = ("_lock", "_conn", "_tx", "_owns_conn") + + def __init__( + self, + lock: threading.RLock, + conn: "sqlite3.Connection", + tx_state: _TxState, + ) -> None: + self._lock = lock + self._conn = conn + self._tx = tx_state + self._owns_conn = False + + def __enter__(self): + self._lock.acquire() + if self._tx.depth == 0: + self._conn.__enter__() + self._owns_conn = True + return self + + def __exit__(self, exc_type, exc, tb): + try: + if self._owns_conn: + # Defer to the connection's own commit/rollback semantics. + self._conn.__exit__(exc_type, exc, tb) + finally: + self._lock.release() + return False + + +# Edge column names that participate in column-direct filtering. Anything +# else falls through to JSON metadata filtering. +_EDGE_NUMERIC_COLUMNS: frozenset[str] = frozenset( + {"weight", "bonus", "hits", "last_touch"} +) + + +def _split_edge_filter( + filter_dict: dict[str, Any] | None, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Partition a filter into (edge-column filter, metadata filter). + + Filter keys that name a numeric edge column compile to direct + column comparisons; everything else routes through json_extract on + the edge's metadata column. + """ + if not filter_dict: + return {}, {} + edge_part: dict[str, Any] = {} + meta_part: dict[str, Any] = {} + for k, v in filter_dict.items(): + if k in _EDGE_NUMERIC_COLUMNS: + edge_part[k] = v + else: + meta_part[k] = v + return edge_part, meta_part + + +def _compile_edge_column_filter(edge_filter: dict[str, Any], params: list[Any]) -> str: + """Compile filters against literal edge columns (no json_extract). + + Mirrors the operator grammar from utils.normalize_filter / validate_filter + but emits SQL referencing the column directly. + """ + from ..utils import normalize_filter, validate_filter + + validate_filter(edge_filter) + normalized = normalize_filter(edge_filter) or {} + pieces: list[str] = [] + for col, value in normalized.items(): + # col is one of _EDGE_NUMERIC_COLUMNS — safe to inline. + if isinstance(value, dict): + for op, arg in value.items(): + if op == "$eq": + pieces.append(f"{col} = ?") + params.append(arg) + elif op == "$ne": + pieces.append(f"{col} != ?") + params.append(arg) + elif op == "$gt": + pieces.append(f"{col} > ?") + params.append(arg) + elif op == "$gte": + pieces.append(f"{col} >= ?") + params.append(arg) + elif op == "$lt": + pieces.append(f"{col} < ?") + params.append(arg) + elif op == "$lte": + pieces.append(f"{col} <= ?") + params.append(arg) + elif op == "$in": + placeholders = ",".join("?" for _ in arg) + pieces.append(f"{col} IN ({placeholders})") + params.extend(arg) + elif op == "$nin": + placeholders = ",".join("?" for _ in arg) + pieces.append(f"{col} NOT IN ({placeholders})") + params.extend(arg) + elif op == "$between": + lo, hi = arg + pieces.append(f"{col} BETWEEN ? AND ?") + params.extend([lo, hi]) + elif op == "$exists": + pieces.append(f"{col} IS NOT NULL" if arg else f"{col} IS NULL") + else: + raise ValueError(f"Unsupported edge operator '{op}'") + elif isinstance(value, list): + placeholders = ",".join("?" for _ in value) + pieces.append(f"{col} IN ({placeholders})") + params.extend(value) + else: + pieces.append(f"{col} = ?") + params.append(value) + return " AND ".join(pieces) + + class CatalogManager: """ Handles SQLite metadata and FTS operations. @@ -60,8 +245,13 @@ def __init__( table_name: str, fts_table_name: str, lock: threading.RLock | None = None, + tx_state: "_TxState | None" = None, ): - # Defense-in-depth: validate table names + # Defense-in-depth: validate table names. After this point every + # f-string SQL site that interpolates self._table_name (or its + # _pending_vectors / _edges / _events / _ttl / _clusters + # suffixes) is safe by construction; do not reintroduce arbitrary + # name interpolation downstream. _validate_table_name(table_name) _validate_table_name(fts_table_name) @@ -77,6 +267,22 @@ def __init__( # — two threads entering it simultaneously interleave their writes # under one implicit transaction. The lock prevents that. self._lock: threading.RLock = lock if lock is not None else threading.RLock() + # Optional shared cross-collection transaction state. When the + # state's depth > 0, _writable() suppresses inner conn commits so + # the outer SAVEPOINT controls atomicity. + self._tx_state: _TxState = tx_state if tx_state is not None else _TxState() + + def _writable(self): + """Acquire the write lock and (when no outer txn) the conn context. + + Outside a transaction, this is equivalent to `with self._lock, + self.conn:` — every catalog write commits when the block exits. + + Inside a transaction (caller has called VectorDB.transaction()), + the connection's commit-on-exit is suppressed so the SAVEPOINT + opened by the transaction owns atomicity. + """ + return _CatalogWritable(self._lock, self.conn, self._tx_state) def create_tables(self) -> None: """Create metadata and FTS tables if they don't exist.""" @@ -112,6 +318,104 @@ def create_tables(self) -> None: self._ensure_embedding_column() self._ensure_parent_id_column() self._ensure_fts_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. + self._ensure_pending_vectors_table() + self._ensure_edges_table() + self._ensure_events_table() + self._ensure_ttl_table() + + def _ensure_pending_vectors_table(self) -> None: + """Buffer of vector updates flushed to usearch in batches (gap 1).""" + self.conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {self._table_name}_pending_vectors ( + doc_id INTEGER PRIMARY KEY + REFERENCES {self._table_name}(id) ON DELETE CASCADE, + embedding BLOB NOT NULL, + source TEXT, + enqueued_at REAL NOT NULL + ) + """ + ) + + def _ensure_edges_table(self) -> None: + """Weighted directed edges between documents (gap 3).""" + self.conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {self._table_name}_edges ( + src_id INTEGER NOT NULL + REFERENCES {self._table_name}(id) ON DELETE CASCADE, + dst_id INTEGER NOT NULL + REFERENCES {self._table_name}(id) ON DELETE CASCADE, + kind TEXT NOT NULL DEFAULT '', + weight REAL NOT NULL DEFAULT 0.0, + hits INTEGER NOT NULL DEFAULT 0, + bonus REAL NOT NULL DEFAULT 0.0, + last_touch REAL NOT NULL, + metadata TEXT, + PRIMARY KEY (src_id, dst_id, kind) + ) + """ + ) + self.conn.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_{self._table_name}_edges_dst + ON {self._table_name}_edges(dst_id, kind) + """ + ) + self.conn.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_{self._table_name}_edges_weight + ON {self._table_name}_edges(kind, weight) + """ + ) + self.conn.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_{self._table_name}_edges_last_touch + ON {self._table_name}_edges(kind, last_touch) + """ + ) + + def _ensure_events_table(self) -> None: + """Append-only change feed (gap 7). Subscribers poll WHERE seq > ?.""" + self.conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {self._table_name}_events ( + seq INTEGER PRIMARY KEY AUTOINCREMENT, + ts REAL NOT NULL, + kind TEXT NOT NULL, + doc_id INTEGER, + payload TEXT + ) + """ + ) + self.conn.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_{self._table_name}_events_kind_seq + ON {self._table_name}_events(kind, seq) + """ + ) + + def _ensure_ttl_table(self) -> None: + """TTL/expiry entries (gap 8). Sweep deletes rows whose expires_at <= now().""" + self.conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {self._table_name}_ttl ( + doc_id INTEGER PRIMARY KEY + REFERENCES {self._table_name}(id) ON DELETE CASCADE, + expires_at REAL NOT NULL, + on_expire TEXT NOT NULL DEFAULT 'delete' + ) + """ + ) + self.conn.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_{self._table_name}_ttl_expires + ON {self._table_name}_ttl(expires_at) + """ + ) def _ensure_embedding_column(self) -> None: """Add embedding column if missing (migration for v2.0.0).""" @@ -174,6 +478,7 @@ def _ensure_fts_table(self) -> None: 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) @@ -189,7 +494,7 @@ def _upsert_fts_rows(self, ids: Sequence[int], texts: Sequence[str]) -> None: """Update FTS index for given document IDs. Internal helper. Must be called inside an active transaction - (``with self._lock, self.conn:``) so the FTS shadow table stays in + (``with self._writable()``) so the FTS shadow table stays in sync with the main table on crash. Args: @@ -302,7 +607,7 @@ def add_documents( real_ids: list[int] = [-1] * len(ids_list) - with self._lock, self.conn: + with self._writable(): if explicit_rows: self.conn.executemany( f""" @@ -380,7 +685,7 @@ def delete_by_ids(self, ids: Iterable[int]) -> list[int]: placeholders = ",".join("?" for _ in ids) params = tuple(ids) - with self._lock, self.conn: + with self._writable(): # Check which IDs actually exist existing = self.conn.execute( f"SELECT id FROM {self._table_name} WHERE id IN ({placeholders})", @@ -390,11 +695,31 @@ def delete_by_ids(self, ids: Iterable[int]) -> list[int]: if existing_ids: placeholders = ",".join("?" for _ in existing_ids) + # 2.6.1 aux tables. FK enforcement may be off, so cascade + # explicitly to avoid orphan rows. + self.conn.execute( + f"DELETE FROM {self._table_name}_pending_vectors " + f"WHERE doc_id IN ({placeholders})", + tuple(existing_ids), + ) + self.conn.execute( + f"DELETE FROM {self._table_name}_edges " + f"WHERE src_id IN ({placeholders}) " + f"OR dst_id IN ({placeholders})", + tuple(existing_ids) * 2, + ) + self.conn.execute( + f"DELETE FROM {self._table_name}_ttl " + f"WHERE doc_id IN ({placeholders})", + tuple(existing_ids), + ) self.conn.execute( f"DELETE FROM {self._table_name} WHERE id IN ({placeholders})", tuple(existing_ids), ) self._delete_fts_rows(existing_ids) + for doc_id in existing_ids: + self.append_event_in_tx("delete", doc_id=int(doc_id)) _logger.debug("Deleted %d documents", len(existing_ids)) return existing_ids @@ -436,9 +761,7 @@ def list_all_ids(self) -> list[int]: discipline alone. """ with self._lock: - rows = self.conn.execute( - f"SELECT id FROM {self._table_name}" - ).fetchall() + rows = self.conn.execute(f"SELECT id FROM {self._table_name}").fetchall() return [row[0] for row in rows] def get_embeddings_by_ids(self, ids: Sequence[int]) -> dict[int, Any]: @@ -497,7 +820,11 @@ def get_documents_and_embeddings_by_ids( result: dict[int, tuple[str, dict[str, Any], np.ndarray | None]] = {} for row_id, text, meta_json, emb_blob in rows: meta = json.loads(meta_json) if meta_json else {} - emb = np.frombuffer(emb_blob, dtype=np.float32) if emb_blob is not None else None + emb = ( + np.frombuffer(emb_blob, dtype=np.float32) + if emb_blob is not None + else None + ) result[row_id] = (text, meta, emb) return result @@ -618,49 +945,146 @@ def keyword_search( return [(int(row[0]), float(row[1])) for row in rows] def build_filter_clause( - self, filter_dict: dict[str, Any] | None, metadata_column: str = "metadata" + self, + filter_dict: dict[str, Any] | None, + metadata_column: str = "metadata", ) -> tuple[str, list[Any]]: """ Build SQL WHERE clause from metadata filter dictionary. + Accepts the full grammar from utils.validate_filter: + - scalar equality: {"k": v} + - list IN: {"k": [v1, v2]} + - Mongo-style operator dicts: + {"k": {"$gt": x, "$lte": y}}, {"k": {"$between": [lo, hi]}}, + {"k": {"$in": [...]}}, {"k": {"$exists": True}} + - tuple shorthand: {"k": (">", x)} (normalized internally) + + Numeric operators wrap the JSON value in CAST(... AS REAL) so + SQLite uses index-friendly numeric comparisons even though the + column itself is JSON text. + Args: - filter_dict: Metadata key-value pairs to filter by - metadata_column: Name of JSON metadata column + filter_dict: Filter dictionary in the grammar above. + metadata_column: Column expression holding the JSON document. Returns: - Tuple of (where_clause, parameters) for SQL query - - Raises: - ValueError: If filter keys are not strings or values are unsupported types + Tuple of (where_clause, parameters). where_clause is empty + string or starts with "AND (" for direct interpolation. """ if not filter_dict: return "", [] - # Validate filter structure before processing validate_filter(filter_dict) + normalized = normalize_filter(filter_dict) or {} - clauses = [] + clauses: list[str] = [] params: list[Any] = [] - for key, value in filter_dict.items(): + for key, value in normalized.items(): json_path = f"$.{key}" - if isinstance(value, (int, float)): - clauses.append(f"json_extract({metadata_column}, ?) = ?") - params.extend([json_path, value]) - elif isinstance(value, str): - # Use exact equality for string filters - clauses.append(f"json_extract({metadata_column}, ?) = ?") + text_extract = f"json_extract({metadata_column}, ?)" + num_extract = f"CAST({text_extract} AS REAL)" + type_extract = f"json_type({metadata_column}, ?)" + + if isinstance(value, dict): + self._build_operator_clauses( + json_path, + text_extract, + num_extract, + type_extract, + value, + clauses, + params, + ) + continue + + if isinstance(value, bool): + # JSON encodes bool as 0/1 via json_extract; normalize. + clauses.append(f"{text_extract} = ?") + params.extend([json_path, 1 if value else 0]) + elif isinstance(value, (int, float, str)): + clauses.append(f"{text_extract} = ?") params.extend([json_path, value]) elif isinstance(value, list): placeholders = ",".join("?" for _ in value) - clauses.append( - f"json_extract({metadata_column}, ?) IN ({placeholders})" - ) - params.extend([json_path] + value) + clauses.append(f"{text_extract} IN ({placeholders})") + params.append(json_path) + params.extend(value) else: raise ValueError(f"Unsupported filter value type for {key}") + where = " AND ".join(clauses) return f"AND ({where})" if where else "", params + def _build_operator_clauses( + self, + json_path: str, + text_extract: str, + num_extract: str, + type_extract: str, + op_dict: dict[str, Any], + clauses: list[str], + params: list[Any], + ) -> None: + """Compile a single key's operator dict into WHERE clause fragments.""" + + def numeric(sql_op: str, arg: Any) -> None: + # Gate the comparison on json_type so non-numeric values aren't + # coerced to 0.0 by CAST(... AS REAL); without this guard a row + # with `score: "oops"` would match `{"$lt": 1}` because the + # cast silently produces 0.0. Mirrors `_eval_operator_dict`, + # which rejects non-numeric `meta_value` for these ops. + clauses.append( + f"({type_extract} IN ('integer', 'real') AND {num_extract} {sql_op} ?)" + ) + params.extend([json_path, json_path, arg]) + + for op, arg in op_dict.items(): + if op == "$eq": + clauses.append(f"{text_extract} = ?") + params.extend([json_path, _coerce_scalar(arg)]) + elif op == "$ne": + # IS NOT for null-safety + difference for present values. + clauses.append(f"({text_extract} IS NULL OR {text_extract} != ?)") + params.extend([json_path, json_path, _coerce_scalar(arg)]) + elif op == "$gt": + numeric(">", arg) + elif op == "$gte": + numeric(">=", arg) + elif op == "$lt": + numeric("<", arg) + elif op == "$lte": + numeric("<=", arg) + elif op == "$in": + placeholders = ",".join("?" for _ in arg) + clauses.append(f"{text_extract} IN ({placeholders})") + params.append(json_path) + params.extend(arg) + elif op == "$nin": + placeholders = ",".join("?" for _ in arg) + clauses.append( + f"({text_extract} IS NULL OR " + f"{text_extract} NOT IN ({placeholders}))" + ) + params.append(json_path) + params.append(json_path) + params.extend(arg) + elif op == "$exists": + if arg: + clauses.append(f"{text_extract} IS NOT NULL") + else: + clauses.append(f"{text_extract} IS NULL") + params.append(json_path) + elif op == "$between": + lo, hi = arg + clauses.append( + f"({type_extract} IN ('integer', 'real') " + f"AND {num_extract} BETWEEN ? AND ?)" + ) + params.extend([json_path, json_path, lo, hi]) + else: + raise ValueError(f"Unsupported operator '{op}'") + def count(self) -> int: """Return total number of documents.""" with self._lock: @@ -734,7 +1158,7 @@ def update_metadata_batch(self, updates: list[tuple[int, dict[str, Any]]]) -> in if not updates: return 0 - with self._lock, self.conn: + with self._writable(): updated = 0 # Batch into chunks of 500 for performance for batch in _batched(updates, 500): @@ -747,7 +1171,9 @@ def update_metadata_batch(self, updates: list[tuple[int, dict[str, Any]]]) -> in ids, ).fetchall() - current_meta_map = {r[0]: (json.loads(r[1]) if r[1] else {}) for r in rows} + current_meta_map = { + r[0]: (json.loads(r[1]) if r[1] else {}) for r in rows + } # Prepare updates update_data = [] @@ -766,56 +1192,770 @@ def update_metadata_batch(self, updates: list[tuple[int, dict[str, Any]]]) -> in return updated - def check_legacy_sqlite_vec(self, vec_table_name: str) -> bool: - """ - Check if legacy sqlite-vec tables exist (for migration). + def increment_metadata(self, doc_id: int, deltas: dict[str, float | int]) -> int: + """Atomically increment numeric metadata counters (gap 4). + + Single UPDATE statement applies every delta via chained json_set, + so multi-writer races are resolved by SQLite under WAL — no + read-modify-write window in Python. Numeric values only; missing + keys treat the prior value as 0. Args: - vec_table_name: Expected name of the old vec0 virtual table + doc_id: Target document id. + deltas: Mapping of counter name -> numeric delta. Keys must + match `^[A-Za-z_][A-Za-z0-9_]*$` (validated). Values must + be int or float (bool/strings/None rejected). Returns: - True if legacy sqlite-vec data exists - """ - try: - with self._lock: - row = self.conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name=?", - (vec_table_name,), - ).fetchone() - return row is not None - except Exception: - return False - - def get_legacy_vectors(self, vec_table_name: str) -> list[tuple[int, bytes]]: + 1 if the row exists and was updated, 0 otherwise. """ - Extract vectors from legacy sqlite-vec table for migration. + return self.increment_metadata_many([(doc_id, deltas)]) + + def increment_metadata_many( + self, updates: list[tuple[int, dict[str, float | int]]] + ) -> int: + """Batch variant of increment_metadata. + + Each (doc_id, deltas) tuple becomes one UPDATE; all run inside + the same transaction so partial application on crash is + impossible. Args: - vec_table_name: Name of the old vec0 virtual table + updates: List of (doc_id, deltas) pairs. Returns: - List of (rowid, embedding_blob) tuples + Total number of rows updated (sum of UPDATE rowcounts). """ - _validate_table_name(vec_table_name) - try: - with self._lock: - rows = self.conn.execute( - f"SELECT rowid, embedding FROM {vec_table_name}" - ).fetchall() - return [(int(r[0]), r[1]) for r in rows] - except Exception as e: - _logger.warning("Failed to read legacy vectors: %s", e) - return [] + if not updates: + return 0 - def drop_legacy_vec_table(self, vec_table_name: str) -> None: - """Drop legacy sqlite-vec table after migration.""" - _validate_table_name(vec_table_name) - try: - with self._lock, self.conn: - self.conn.execute(f"DROP TABLE IF EXISTS {vec_table_name}") - _logger.info("Dropped legacy sqlite-vec table: %s", vec_table_name) - except Exception as e: - _logger.warning("Failed to drop legacy table %s: %s", vec_table_name, e) + # Pre-validate all keys/values up front so we can fail before + # opening a transaction. Keeps the SQL build loop branch-free. + for doc_id, deltas in updates: + if not isinstance(doc_id, int) or isinstance(doc_id, bool): + raise TypeError( + f"increment_metadata: doc_id must be int, got " + f"{type(doc_id).__name__}" + ) + if not deltas: + raise ValueError("increment_metadata: deltas must be a non-empty dict") + for key, val in deltas.items(): + _validate_identifier(key, "metadata counter key") + if isinstance(val, bool) or not isinstance(val, (int, float)): + raise TypeError( + f"increment_metadata: delta for '{key}' must be " + f"int or float, got {type(val).__name__}" + ) + if isinstance(val, float) and ( + val != val or val in (float("inf"), float("-inf")) + ): + raise ValueError( + f"increment_metadata: delta for '{key}' must be finite" + ) + + total = 0 + with self._writable(): + for doc_id, deltas in updates: + sql, params = self._build_increment_sql(doc_id, deltas) + cursor = self.conn.execute(sql, params) + if cursor.rowcount: + total += cursor.rowcount + self.append_event_in_tx( + "counter", + doc_id=doc_id, + payload={"deltas": dict(deltas)}, + ) + return total + + def _build_increment_sql( + self, doc_id: int, deltas: dict[str, float | int] + ) -> tuple[str, tuple[Any, ...]]: + """Compile a chained json_set UPDATE that adds all deltas atomically. + + Builds `json_set(json_set(... base ..., '$.k1', cur1+?), '$.k2', cur2+?)` + from the inside out. Key names are inlined after _validate_identifier + rejects anything outside the safe alphabet; deltas are bound params. + """ + expr = f"COALESCE({self._table_name}.metadata, '{{}}')" + params: list[Any] = [] + for key, val in deltas.items(): + cur = ( + f"COALESCE(CAST(json_extract({self._table_name}.metadata, " + f"'$.{key}') AS REAL), 0)" + ) + expr = f"json_set({expr}, '$.{key}', {cur} + ?)" + params.append(val) + params.append(doc_id) + sql = f"UPDATE {self._table_name} SET metadata = {expr} WHERE id = ?" + return sql, tuple(params) + + def get_metadata_counter( + self, doc_id: int, key: str, default: float | int = 0 + ) -> float | int | None: + """Read a single numeric counter value from metadata. + + Returns the stored number, or `default` if the key is missing + or non-numeric. Returns None if the row itself doesn't exist + (so callers can distinguish missing-row from missing-key). + """ + _validate_identifier(key, "metadata counter key") + with self._lock: + row = self.conn.execute( + f"SELECT json_extract(metadata, '$.{key}') " + f"FROM {self._table_name} WHERE id = ?", + (doc_id,), + ).fetchone() + if row is None: + return None + value = row[0] + if value is None: + return default + if isinstance(value, (int, float)): + return value + return default + + # --- Pending vector buffer (gap 1 / gap 6) ----------------------------- + + def upsert_pending_vector( + self, doc_id: int, embedding_bytes: bytes, source: str | None = None + ) -> int: + """Buffer a vector update; flush() promotes it to the HNSW index. + + Replaces any existing pending entry for the same doc_id (last-write- + wins per id) so duplicate calls before flush coalesce. Returns the + rowcount (1 if a row was written, 0 if the doc_id row is missing + and FK enforcement is on). + """ + return self.upsert_pending_vectors_many([(doc_id, embedding_bytes, source)]) + + def upsert_pending_vectors_many( + self, + rows: Sequence[tuple[int, bytes, str | None]], + ) -> int: + """Bulk variant of upsert_pending_vector; one transaction.""" + if not rows: + return 0 + with self._writable(): + cur = self.conn.executemany( + f""" + INSERT INTO {self._table_name}_pending_vectors + (doc_id, embedding, source, enqueued_at) + VALUES (?, ?, ?, unixepoch('subsec')) + ON CONFLICT(doc_id) DO UPDATE SET + embedding = excluded.embedding, + source = excluded.source, + enqueued_at = excluded.enqueued_at + """, + rows, + ) + self.append_event_in_tx( + "pending_enqueue", + payload={"count": len(rows), "ids": [int(r[0]) for r in rows[:50]]}, + ) + return cur.rowcount or 0 + + def list_pending_vectors( + self, *, limit: int | None = None + ) -> list[tuple[int, bytes, str | None]]: + """Read pending vector rows in enqueue order. + + Returns list of (doc_id, embedding_bytes, source). Rows whose + doc_id no longer exists in the main table are skipped; this is + defense-in-depth in case FK enforcement is off. + """ + sql = ( + f"SELECT p.doc_id, p.embedding, p.source FROM " + f"{self._table_name}_pending_vectors p " + f"JOIN {self._table_name} t ON t.id = p.doc_id " + f"ORDER BY p.enqueued_at ASC, p.doc_id ASC" + ) + params: tuple[Any, ...] = () + if limit is not None: + sql += " LIMIT ?" + params = (limit,) + with self._lock: + rows = self.conn.execute(sql, params).fetchall() + return [(int(r[0]), bytes(r[1]), r[2]) for r in rows] + + def delete_pending_vectors(self, doc_ids: Sequence[int]) -> int: + """Drop pending rows for the given doc_ids (post-flush cleanup).""" + if not doc_ids: + return 0 + placeholders = ",".join("?" for _ in doc_ids) + sql = ( + f"DELETE FROM {self._table_name}_pending_vectors " + f"WHERE doc_id IN ({placeholders})" + ) + with self._writable(): + cur = self.conn.execute(sql, tuple(doc_ids)) + if cur.rowcount: + self.append_event_in_tx( + "pending_flush", + payload={ + "count": int(cur.rowcount), + "ids": [int(i) for i in list(doc_ids)[:50]], + }, + ) + return cur.rowcount or 0 + + def count_pending_vectors(self) -> int: + """Number of buffered vector updates.""" + with self._lock: + row = self.conn.execute( + f"SELECT COUNT(*) FROM {self._table_name}_pending_vectors" + ).fetchone() + return int(row[0]) if row else 0 + + def get_pending_vector(self, doc_id: int) -> bytes | None: + """Return buffered vector bytes for a doc, or None if not pending.""" + with self._lock: + row = self.conn.execute( + f"SELECT embedding FROM {self._table_name}_pending_vectors " + f"WHERE doc_id = ?", + (doc_id,), + ).fetchone() + return bytes(row[0]) if row is not None else None + + # --- Edges (gap 3) ----------------------------------------------------- + + def add_edge( + self, + src_id: int, + dst_id: int, + *, + kind: str = "", + weight: float = 0.0, + bonus: float = 0.0, + hits: int = 0, + metadata: dict[str, Any] | None = None, + ) -> int: + """Insert an edge; raises if (src, dst, kind) already exists.""" + if not isinstance(kind, str): + raise TypeError(f"edge kind must be str, got {type(kind).__name__}") + _check_finite_edge_field(weight, "weight") + _check_finite_edge_field(bonus, "bonus") + with self._writable(): + cur = self.conn.execute( + f""" + INSERT INTO {self._table_name}_edges + (src_id, dst_id, kind, weight, hits, bonus, + last_touch, metadata) + VALUES (?, ?, ?, ?, ?, ?, unixepoch('subsec'), ?) + """, + ( + int(src_id), + int(dst_id), + kind, + float(weight), + int(hits), + float(bonus), + json.dumps(metadata) if metadata is not None else None, + ), + ) + self.append_event_in_tx( + "edge_add", + payload={ + "src": int(src_id), + "dst": int(dst_id), + "kind": kind, + "weight": float(weight), + }, + ) + return cur.rowcount or 0 + + def upsert_edge( + self, + src_id: int, + dst_id: int, + *, + kind: str = "", + weight: float | None = None, + bonus: float | None = None, + hits: int | None = None, + metadata: dict[str, Any] | None = None, + ) -> int: + """Insert or replace fields on an existing edge. + + Fields left as None preserve their existing value (or use the + column default on insert). + """ + if not isinstance(kind, str): + raise TypeError(f"edge kind must be str, got {type(kind).__name__}") + _check_finite_edge_field(weight, "weight") + _check_finite_edge_field(bonus, "bonus") + meta_json = json.dumps(metadata) if metadata is not None else None + with self._writable(): + cur = self.conn.execute( + f""" + INSERT INTO {self._table_name}_edges + (src_id, dst_id, kind, weight, hits, bonus, + last_touch, metadata) + VALUES (?, ?, ?, + COALESCE(?, 0.0), COALESCE(?, 0), COALESCE(?, 0.0), + unixepoch('subsec'), ?) + ON CONFLICT(src_id, dst_id, kind) DO UPDATE SET + weight = COALESCE(?, weight), + hits = COALESCE(?, hits), + bonus = COALESCE(?, bonus), + metadata = COALESCE(?, metadata), + last_touch = unixepoch('subsec') + """, + ( + int(src_id), + int(dst_id), + kind, + weight, + hits, + bonus, + meta_json, + weight, + hits, + bonus, + meta_json, + ), + ) + self.append_event_in_tx( + "edge_upsert", + payload={"src": int(src_id), "dst": int(dst_id), "kind": kind}, + ) + return cur.rowcount or 0 + + def update_edge( + self, + src_id: int, + dst_id: int, + *, + kind: str = "", + weight: float | None = None, + bonus: float | None = None, + hits: int | None = None, + metadata: dict[str, Any] | None = None, + dweight: float = 0.0, + dbonus: float = 0.0, + dhits: int = 0, + ) -> int: + """Modify an edge. + + Absolute values (weight/bonus/hits/metadata) replace the column + when not None. Deltas (dweight/dbonus/dhits) are applied + atomically via SQL `col = col + ?`. Both can be combined: an + absolute set runs *before* the delta on the same statement. + Non-existent edges are not auto-created (use upsert_edge for that). + """ + if not isinstance(kind, str): + raise TypeError(f"edge kind must be str, got {type(kind).__name__}") + _check_finite_edge_field(weight, "weight") + _check_finite_edge_field(bonus, "bonus") + _check_finite_edge_field(dweight, "dweight") + _check_finite_edge_field(dbonus, "dbonus") + meta_json = json.dumps(metadata) if metadata is not None else None + with self._writable(): + cur = self.conn.execute( + f""" + UPDATE {self._table_name}_edges SET + weight = COALESCE(?, weight) + ?, + bonus = COALESCE(?, bonus) + ?, + hits = COALESCE(?, hits) + ?, + metadata = COALESCE(?, metadata), + last_touch = unixepoch('subsec') + WHERE src_id = ? AND dst_id = ? AND kind = ? + """, + ( + weight, + float(dweight), + bonus, + float(dbonus), + hits, + int(dhits), + meta_json, + int(src_id), + int(dst_id), + kind, + ), + ) + if cur.rowcount: + self.append_event_in_tx( + "edge_update", + payload={ + "src": int(src_id), + "dst": int(dst_id), + "kind": kind, + "dweight": float(dweight), + "dhits": int(dhits), + }, + ) + return cur.rowcount or 0 + + def delete_edge(self, src_id: int, dst_id: int, *, kind: str = "") -> int: + """Delete a single edge. Returns 1 if removed, else 0.""" + with self._writable(): + cur = self.conn.execute( + f""" + DELETE FROM {self._table_name}_edges + WHERE src_id = ? AND dst_id = ? AND kind = ? + """, + (int(src_id), int(dst_id), kind), + ) + if cur.rowcount: + self.append_event_in_tx( + "edge_delete", + payload={"src": int(src_id), "dst": int(dst_id), "kind": kind}, + ) + return cur.rowcount or 0 + + def get_edges( + self, + *, + src_id: int | None = None, + dst_id: int | None = None, + kind: str | None = None, + filter: dict[str, Any] | None = None, + limit: int | None = None, + ) -> list[tuple[int, int, str, float, int, float, float, dict | None]]: + """Read edges. Filter applies via build_filter_clause to numeric + columns (weight/bonus/hits/last_touch) or to the JSON metadata.""" + clauses: list[str] = [] + params: list[Any] = [] + if src_id is not None: + clauses.append("src_id = ?") + params.append(int(src_id)) + if dst_id is not None: + clauses.append("dst_id = ?") + params.append(int(dst_id)) + if kind is not None: + clauses.append("kind = ?") + params.append(kind) + + edge_filter, meta_filter = _split_edge_filter(filter) + if edge_filter: + clauses.append(_compile_edge_column_filter(edge_filter, params)) + if meta_filter: + extra_clause, extra_params = self.build_filter_clause( + meta_filter, metadata_column="metadata" + ) + if extra_clause: + # build_filter_clause returns "AND (...)"; strip the "AND". + clauses.append(extra_clause[4:].strip()) + params.extend(extra_params) + + where = " AND ".join(clauses) if clauses else "1=1" + sql = ( + f"SELECT src_id, dst_id, kind, weight, hits, bonus, " + f"last_touch, metadata " + f"FROM {self._table_name}_edges WHERE {where} " + f"ORDER BY last_touch DESC, src_id, dst_id" + ) + if limit is not None: + sql += " LIMIT ?" + params.append(int(limit)) + with self._lock: + rows = self.conn.execute(sql, tuple(params)).fetchall() + result = [] + for r in rows: + meta = json.loads(r[7]) if r[7] else None + result.append( + ( + int(r[0]), + int(r[1]), + str(r[2]), + float(r[3]), + int(r[4]), + float(r[5]), + float(r[6]), + meta, + ) + ) + return result + + # --- Change feed (gap 7) ----------------------------------------------- + + def append_event( + self, + kind: str, + *, + doc_id: int | None = None, + payload: dict[str, Any] | None = None, + ) -> int: + """Append a single event to the change feed. + + Returns the assigned `seq`. Caller is responsible for placing + this inside an outer transaction if multiple writes must atomic. + """ + if not isinstance(kind, str) or not kind: + raise ValueError("event kind must be a non-empty string") + with self._writable(): + cur = self.conn.execute( + f""" + INSERT INTO {self._table_name}_events + (ts, kind, doc_id, payload) + VALUES (unixepoch('subsec'), ?, ?, ?) + """, + ( + kind, + int(doc_id) if doc_id is not None else None, + json.dumps(payload) if payload is not None else None, + ), + ) + return int(cur.lastrowid or 0) + + def append_event_in_tx( + self, + kind: str, + *, + doc_id: int | None = None, + payload: dict[str, Any] | None = None, + ) -> None: + """Append without opening a transaction (caller already has one).""" + if not isinstance(kind, str) or not kind: + raise ValueError("event kind must be a non-empty string") + self.conn.execute( + f""" + INSERT INTO {self._table_name}_events + (ts, kind, doc_id, payload) + VALUES (unixepoch('subsec'), ?, ?, ?) + """, + ( + kind, + int(doc_id) if doc_id is not None else None, + json.dumps(payload) if payload is not None else None, + ), + ) + + def last_event_seq(self) -> int: + """Return the highest assigned sequence number, or 0 if none.""" + with self._lock: + row = self.conn.execute( + f"SELECT MAX(seq) FROM {self._table_name}_events" + ).fetchone() + return int(row[0]) if row and row[0] is not None else 0 + + def read_events( + self, + *, + since: int = 0, + kind: str | None = None, + limit: int | None = None, + ) -> list[tuple[int, float, str, int | None, dict | None]]: + """Return events with seq > `since`, optionally filtered by kind.""" + clauses = ["seq > ?"] + params: list[Any] = [int(since)] + if kind is not None: + clauses.append("kind = ?") + params.append(kind) + sql = ( + f"SELECT seq, ts, kind, doc_id, payload " + f"FROM {self._table_name}_events " + f"WHERE {' AND '.join(clauses)} ORDER BY seq ASC" + ) + if limit is not None: + sql += " LIMIT ?" + params.append(int(limit)) + with self._lock: + rows = self.conn.execute(sql, tuple(params)).fetchall() + result = [] + for r in rows: + payload = json.loads(r[4]) if r[4] else None + doc_id = int(r[3]) if r[3] is not None else None + result.append((int(r[0]), float(r[1]), str(r[2]), doc_id, payload)) + return result + + def prune_events(self, *, before_seq: int) -> int: + """Delete events with seq < `before_seq`. Returns count deleted.""" + with self._writable(): + cur = self.conn.execute( + f"DELETE FROM {self._table_name}_events WHERE seq < ?", + (int(before_seq),), + ) + return cur.rowcount or 0 + + # --- TTL / expiry (gap 8) --------------------------------------------- + + def set_ttl( + self, + doc_id: int, + expires_at: float, + *, + on_expire: str = "delete", + ) -> int: + """Set or replace the TTL entry for a document. + + `expires_at` is a unix timestamp (seconds). `on_expire` is either + "delete" (sweep removes the row) or "callback" (sweep returns the + id but does not delete; caller acts on it). + """ + if on_expire not in ("delete", "callback"): + raise ValueError( + f"on_expire must be 'delete' or 'callback', got {on_expire!r}" + ) + with self._writable(): + self.conn.execute( + f""" + INSERT INTO {self._table_name}_ttl + (doc_id, expires_at, on_expire) + VALUES (?, ?, ?) + ON CONFLICT(doc_id) DO UPDATE SET + expires_at = excluded.expires_at, + on_expire = excluded.on_expire + """, + (int(doc_id), float(expires_at), on_expire), + ) + self.append_event_in_tx( + "ttl_set", + doc_id=doc_id, + payload={"expires_at": float(expires_at), "on_expire": on_expire}, + ) + return 1 + + def clear_ttl(self, doc_id: int) -> int: + """Remove a TTL entry. Returns 1 if removed, 0 if missing.""" + with self._writable(): + cur = self.conn.execute( + f"DELETE FROM {self._table_name}_ttl WHERE doc_id = ?", + (int(doc_id),), + ) + return cur.rowcount or 0 + + def list_expired_ttl( + self, *, now: float | None = None, limit: int | None = None + ) -> list[tuple[int, float, str]]: + """Return TTL entries whose expires_at <= now.""" + cutoff = float(now) if now is not None else None + if cutoff is None: + sql = ( + f"SELECT doc_id, expires_at, on_expire " + f"FROM {self._table_name}_ttl " + f"WHERE expires_at <= unixepoch('subsec') " + f"ORDER BY expires_at ASC" + ) + params: tuple[Any, ...] = () + else: + sql = ( + f"SELECT doc_id, expires_at, on_expire " + f"FROM {self._table_name}_ttl " + f"WHERE expires_at <= ? " + f"ORDER BY expires_at ASC" + ) + params = (cutoff,) + if limit is not None: + sql += " LIMIT ?" + params = params + (int(limit),) + with self._lock: + rows = self.conn.execute(sql, params).fetchall() + return [(int(r[0]), float(r[1]), str(r[2])) for r in rows] + + def sweep_ttl( + self, + *, + now: float | None = None, + limit: int = 1000, + ) -> tuple[list[int], list[int]]: + """Apply due TTL entries atomically. + + For each expired entry: + - on_expire == "delete": deletes from the main table (and the + 2.6.1 aux tables explicitly, since FK enforcement may be + off), then drops the TTL row. + - on_expire == "callback": leaves the doc in place and just + drops the TTL row. + + The expiration check and the TTL-row deletion happen as a single + atomic ``DELETE … RETURNING`` inside one write transaction — + otherwise a concurrent ``set_ttl`` extending the deadline (or + ``clear_ttl``) between a separate SELECT and DELETE would let + this method delete docs whose TTL was just renewed. + + Returns (deleted_ids, callback_ids). Both lists are empty when + there's nothing to do. + """ + with self._writable(): + if now is None: + claim_sql = ( + f"DELETE FROM {self._table_name}_ttl " + f"WHERE doc_id IN (" + f" SELECT doc_id FROM {self._table_name}_ttl " + f" WHERE expires_at <= unixepoch('subsec') " + f" ORDER BY expires_at ASC LIMIT ?" + f") " + f"RETURNING doc_id, on_expire" + ) + claim_params: tuple[Any, ...] = (int(limit),) + else: + claim_sql = ( + f"DELETE FROM {self._table_name}_ttl " + f"WHERE doc_id IN (" + f" SELECT doc_id FROM {self._table_name}_ttl " + f" WHERE expires_at <= ? " + f" ORDER BY expires_at ASC LIMIT ?" + f") " + f"RETURNING doc_id, on_expire" + ) + claim_params = (float(now), int(limit)) + + claimed = self.conn.execute(claim_sql, claim_params).fetchall() + if not claimed: + return [], [] + + delete_ids = [int(r[0]) for r in claimed if r[1] == "delete"] + callback_ids = [int(r[0]) for r in claimed if r[1] == "callback"] + + if delete_ids: + placeholders = ",".join("?" for _ in delete_ids) + # Children first (FK pragma may be off). The TTL row is + # already gone above via RETURNING. + self.conn.execute( + f"DELETE FROM {self._table_name}_pending_vectors " + f"WHERE doc_id IN ({placeholders})", + tuple(delete_ids), + ) + self.conn.execute( + f"DELETE FROM {self._table_name}_edges " + f"WHERE src_id IN ({placeholders}) " + f"OR dst_id IN ({placeholders})", + tuple(delete_ids) * 2, + ) + # Main row. + self.conn.execute( + f"DELETE FROM {self._table_name} WHERE id IN ({placeholders})", + tuple(delete_ids), + ) + if self._fts_enabled: + self.conn.execute( + f"DELETE FROM {self._fts_table_name} " + f"WHERE rowid IN ({placeholders})", + tuple(delete_ids), + ) + + for doc_id in delete_ids + callback_ids: + self.append_event_in_tx("ttl_expire", doc_id=doc_id) + return delete_ids, callback_ids + + def prune_edges( + self, + *, + kind: str | None = None, + max_weight: float | None = None, + idle_before: float | None = None, + ) -> int: + """Bulk-delete edges by threshold. Returns number deleted.""" + clauses: list[str] = [] + params: list[Any] = [] + if kind is not None: + clauses.append("kind = ?") + params.append(kind) + if max_weight is not None: + clauses.append("weight <= ?") + params.append(float(max_weight)) + if idle_before is not None: + clauses.append("last_touch <= ?") + params.append(float(idle_before)) + if not clauses: + raise ValueError( + "prune_edges: at least one of max_weight/idle_before/kind required" + ) + sql = f"DELETE FROM {self._table_name}_edges WHERE " + " AND ".join(clauses) + with self._writable(): + cur = self.conn.execute(sql, tuple(params)) + return cur.rowcount or 0 # ------------------------------------------------------------------ # # Hierarchical Relationships @@ -891,7 +2031,9 @@ def get_descendants( # int() coercion makes the previous f-string safe today, but the # parameter form is one less line away from injection on a future # refactor. - effective_depth = int(max_depth) if max_depth is not None else constants.MAX_HIERARCHY_DEPTH + effective_depth = ( + int(max_depth) if max_depth is not None else constants.MAX_HIERARCHY_DEPTH + ) sql = f""" WITH RECURSIVE descendants(id, text, metadata, depth) AS ( @@ -935,7 +2077,9 @@ def get_ancestors( # Apply safety cap to prevent infinite recursion from cycles. Bind # the depth as a parameter (see get_descendants for rationale). - effective_depth = int(max_depth) if max_depth is not None else constants.MAX_HIERARCHY_DEPTH + effective_depth = ( + int(max_depth) if max_depth is not None else constants.MAX_HIERARCHY_DEPTH + ) sql = f""" WITH RECURSIVE ancestors(id, text, metadata, parent_id, depth) AS ( @@ -979,7 +2123,7 @@ 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. - with self._lock, self.conn: + with self._writable(): if parent_id is not None: if parent_id == doc_id: raise ValueError("A document cannot be its own parent") @@ -1005,7 +2149,7 @@ def _ensure_cluster_table(self) -> None: if self._cluster_table_ready: return cluster_table = self._cluster_table_name - with self._lock, self.conn: + with self._writable(): # Re-check inside the lock so concurrent first-callers don't # both run the DDL. The CREATE TABLE IF NOT EXISTS is itself # idempotent, but doing the work twice defeats the early-exit. @@ -1048,7 +2192,7 @@ def save_cluster_state( meta_json = json.dumps(metadata) if metadata else None - with self._lock, self.conn: + with self._writable(): self.conn.execute( f""" INSERT OR REPLACE INTO {cluster_table} @@ -1116,7 +2260,7 @@ def delete_cluster_state(self, name: str) -> bool: self._ensure_cluster_table() cluster_table = self._cluster_table_name - with self._lock, self.conn: + with self._writable(): cursor = self.conn.execute( f"DELETE FROM {cluster_table} WHERE name = ?", (name,) ) diff --git a/src/simplevecdb/engine/search.py b/src/simplevecdb/engine/search.py index f74ff59..813d4ba 100755 --- a/src/simplevecdb/engine/search.py +++ b/src/simplevecdb/engine/search.py @@ -12,7 +12,7 @@ from collections.abc import Sequence from ..types import Document, DistanceStrategy -from ..utils import validate_filter +from ..utils import validate_filter, normalize_filter from .. import constants if TYPE_CHECKING: @@ -99,7 +99,9 @@ def similarity_search( results: list[tuple[Document, float]] = [] while len(results) < k and multiplier <= max_multiplier: - fetch_k = min(k * multiplier, index_size) if index_size > 0 else k * 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 ) @@ -127,7 +129,9 @@ def similarity_search( 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 @@ -212,7 +216,9 @@ def similarity_search_batch( if filter and 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 @@ -398,9 +404,7 @@ def max_marginal_relevance_search( 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() - ) + 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) @@ -501,21 +505,78 @@ def _resolve_query_vector(self, query: str | Sequence[float]) -> np.ndarray: return np.asarray(query, dtype=np.float32) def _matches_filter(self, metadata: dict[str, Any], filter: dict[str, Any]) -> bool: - """Check if metadata matches all filter criteria.""" - for key, value in filter.items(): + """Check if metadata matches all filter criteria. + + Mirrors catalog.build_filter_clause grammar so post-filter Python + evaluation produces the same results as a SQL pre-filter would. + """ + normalized = normalize_filter(filter) or {} + for key, value in normalized.items(): meta_value = metadata.get(key) + if isinstance(value, dict): + if not _eval_operator_dict(meta_value, value, key in metadata): + return False + continue + if isinstance(value, list): - # List filter: meta_value must be in the list if meta_value not in value: return False - elif isinstance(value, str): - # String filter: exact match (consistent with SQL build_filter_clause) + elif isinstance(value, bool): + # Direct comparison; metadata stores Python bools. + if meta_value != value: + return False + elif isinstance(value, (int, float, str)): if meta_value != value: return False else: - # Exact match for int/float if meta_value != value: return False return True + + +def _eval_operator_dict( + meta_value: Any, op_dict: dict[str, Any], key_present: bool +) -> bool: + """Evaluate an operator dict against a single metadata value. + + Matches catalog.build_filter_clause semantics: + - $ne / $nin treat missing keys as not-equal (i.e. they pass). + - numeric operators on missing/non-numeric values fail. + - $exists checks dict membership, not value truthiness. + """ + for op, arg in op_dict.items(): + if op == "$exists": + if bool(arg) != key_present: + return False + continue + + if op == "$eq": + if meta_value != arg: + return False + elif op == "$ne": + if key_present and meta_value == arg: + return False + elif op == "$in": + if meta_value not in arg: + return False + elif op == "$nin": + if key_present and meta_value in arg: + return False + elif op in ("$gt", "$gte", "$lt", "$lte", "$between"): + if not isinstance(meta_value, (int, float)) or isinstance(meta_value, bool): + return False + if op == "$gt" and not (meta_value > arg): + return False + if op == "$gte" and not (meta_value >= arg): + return False + if op == "$lt" and not (meta_value < arg): + return False + if op == "$lte" and not (meta_value <= arg): + return False + if op == "$between": + lo, hi = arg + if not (lo <= meta_value <= hi): + return False + return True diff --git a/src/simplevecdb/engine/usearch_index.py b/src/simplevecdb/engine/usearch_index.py index 3f821bb..38b39f5 100755 --- a/src/simplevecdb/engine/usearch_index.py +++ b/src/simplevecdb/engine/usearch_index.py @@ -376,9 +376,7 @@ def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: with self._write_lock: # 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_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) @@ -434,8 +432,9 @@ def save(self) -> None: finally: os.close(dir_fd) except OSError as exc: - _logger.debug("Directory fsync on %s skipped: %s", - self._path.parent, exc) + _logger.debug( + "Directory fsync on %s skipped: %s", self._path.parent, exc + ) except Exception: # Clean up the temp file on failure; the original index # at self._path is untouched. @@ -488,9 +487,7 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[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 - ) + existing_mask = np.array([int(k) in self._index for k in keys], dtype=bool) if not existing_mask.any(): _logger.warning( @@ -510,9 +507,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] = np.asarray(self._index[existing_keys], dtype=np.float32) return result def __del__(self) -> None: diff --git a/src/simplevecdb/integrations/llamaindex.py b/src/simplevecdb/integrations/llamaindex.py index 09510de..d285dee 100755 --- a/src/simplevecdb/integrations/llamaindex.py +++ b/src/simplevecdb/integrations/llamaindex.py @@ -58,9 +58,7 @@ def __init__( try: self._warn_if_legacy_collection() except Exception: - _logger.debug( - "Legacy-collection probe failed", exc_info=True - ) + _logger.debug("Legacy-collection probe failed", exc_info=True) def _warn_if_legacy_collection(self) -> None: """Probe one row; warn if it lacks the v2.6 node_id metadata stamp.""" diff --git a/src/simplevecdb/logging.py b/src/simplevecdb/logging.py index 5ee43a9..2fb63b3 100755 --- a/src/simplevecdb/logging.py +++ b/src/simplevecdb/logging.py @@ -179,6 +179,3 @@ def log_operation( exc_info=True, ) raise - - - diff --git a/src/simplevecdb/types.py b/src/simplevecdb/types.py index 968246e..7b3947d 100755 --- a/src/simplevecdb/types.py +++ b/src/simplevecdb/types.py @@ -53,41 +53,6 @@ class Quantization(StrEnum): BIT = "bit" -class MigrationRequiredError(Exception): - """Raised when a v1.x database needs migration to v2.0 usearch backend. - - This error is raised when opening a database that contains sqlite-vec - data that needs to be migrated to the usearch backend. - - Attributes: - path: Path to the database file - collections: List of collection names that need migration - total_vectors: Total number of vectors to migrate - migration_info: Full migration info dict from check_migration() - """ - - def __init__( - self, - path: str, - collections: list[str], - total_vectors: int, - migration_info: dict, - ): - self.path = path - self.collections = collections - self.total_vectors = total_vectors - self.migration_info = migration_info - - msg = ( - f"Database '{path}' requires migration from sqlite-vec to usearch.\n" - f"Collections: {', '.join(collections)} ({total_vectors} vectors total)\n\n" - f"To migrate automatically, open with: VectorDB('{path}', auto_migrate=True)\n" - f"Or check migration details first: VectorDB.check_migration('{path}')\n\n" - f"⚠️ BACKUP YOUR DATABASE BEFORE MIGRATING: cp {path} {path}.backup" - ) - super().__init__(msg) - - @dataclasses.dataclass class ClusterResult: """Result of a clustering operation.""" @@ -123,3 +88,63 @@ def metrics(self) -> dict[str, float | None]: ClusterTagCallback = Callable[[list[str]], str] + + +@dataclasses.dataclass(frozen=True, slots=True) +class Edge: + """Weighted directed edge between two documents (gap 3). + + Attributes: + src_id: Source document id. + dst_id: Destination document id. + kind: Optional edge type label (default ""). Same (src, dst, kind) + triple is unique; distinct kinds coexist between the same pair. + weight: Numeric weight (e.g. similarity, plasticity). + bonus: Secondary weight (free for caller — e.g. priority bias). + hits: Counter of edge traversals/reinforcements. + last_touch: Unix timestamp (seconds) of the most recent write. + metadata: Optional extra JSON metadata. + """ + + src_id: int + dst_id: int + kind: str = "" + weight: float = 0.0 + bonus: float = 0.0 + hits: int = 0 + last_touch: float = 0.0 + metadata: dict | None = None + + +@dataclasses.dataclass(frozen=True, slots=True) +class Event: + """Change-feed entry (gap 7). + + Attributes: + seq: Monotonic sequence number assigned by SQLite. + ts: Unix timestamp (seconds) at append time. + kind: Event kind (insert/update/delete/edge/counter/ttl/flush/...). + doc_id: Optional document id this event refers to. + payload: Optional JSON-decoded payload dict. + """ + + seq: int + ts: float + kind: str + doc_id: int | None = None + payload: dict | None = None + + +@dataclasses.dataclass(frozen=True, slots=True) +class TTLEntry: + """A pending expiry hook (gap 8). + + Attributes: + doc_id: Target document id. + expires_at: Unix timestamp (seconds) at which the entry expires. + on_expire: Action when the entry expires ("delete" or "callback"). + """ + + doc_id: int + expires_at: float + on_expire: str = "delete" diff --git a/src/simplevecdb/utils.py b/src/simplevecdb/utils.py index 4cbb344..890c936 100755 --- a/src/simplevecdb/utils.py +++ b/src/simplevecdb/utils.py @@ -268,64 +268,238 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any: return decorator +# Range/numeric filter operators (gap 5). Mongo-style operator dicts: +# {"score": {"$gt": 0.5, "$lte": 0.9}} +# {"tag": {"$in": ["a", "b"]}} +# {"flag": {"$exists": True}} +# Plus tuple shorthand normalized to the same operator dicts: +# {"score": (">", 0.5)} -> {"$gt": 0.5} +# {"score": ("range", 0.5, 0.9)} -> {"$between": [0.5, 0.9]} +_FILTER_OPERATORS: frozenset[str] = frozenset( + { + "$eq", + "$ne", + "$gt", + "$gte", + "$lt", + "$lte", + "$in", + "$nin", + "$exists", + "$between", + } +) + +_TUPLE_OP_MAP: dict[str, str] = { + "==": "$eq", + "eq": "$eq", + "!=": "$ne", + "ne": "$ne", + ">": "$gt", + "gt": "$gt", + ">=": "$gte", + "gte": "$gte", + "<": "$lt", + "lt": "$lt", + "<=": "$lte", + "lte": "$lte", + "in": "$in", + "nin": "$nin", + "exists": "$exists", + "range": "$between", + "between": "$between", +} + + +def _is_finite_number(x: Any) -> bool: + if not isinstance(x, (int, float)) or isinstance(x, bool): + return False + if isinstance(x, float) and (x != x or x in (float("inf"), float("-inf"))): + return False + return True + + +def _normalize_filter_value(key: str, value: Any) -> Any: + """Normalize tuple shorthand into operator dicts; pass through others.""" + if isinstance(value, tuple): + if not value: + raise ValueError(f"Filter tuple for '{key}' must not be empty") + op_raw = value[0] + if not isinstance(op_raw, str): + raise ValueError( + f"Filter tuple operator for '{key}' must be a string, " + f"got {type(op_raw).__name__}: {op_raw!r}" + ) + op = _TUPLE_OP_MAP.get(op_raw) + if op is None: + raise ValueError( + f"Unknown tuple operator '{op_raw}' for '{key}'. " + f"Valid: {sorted(set(_TUPLE_OP_MAP))}" + ) + rest = list(value[1:]) + if op == "$between": + if len(rest) != 2: + raise ValueError( + f"'{key}' range/between expects exactly 2 args, got {len(rest)}" + ) + return {op: rest} + if op in ("$in", "$nin"): + arg = rest[0] if len(rest) == 1 else rest + if not isinstance(arg, list): + arg = list(arg) if isinstance(arg, (tuple, set)) else [arg] + return {op: arg} + if len(rest) != 1: + raise ValueError( + f"'{key}' operator '{op_raw}' expects exactly 1 arg, got {len(rest)}" + ) + return {op: rest[0]} + return value + + +def normalize_filter( + filter_dict: dict[str, Any] | None, +) -> dict[str, Any] | None: + """Convert tuple shorthand to operator dicts; otherwise return as-is. + + Pure: callers can rely on the result not aliasing the input for keys + that needed conversion. + """ + if not filter_dict: + return filter_dict + return {k: _normalize_filter_value(k, v) for k, v in filter_dict.items()} + + def validate_filter(filter_dict: dict[str, Any] | None) -> None: """ Validate metadata filter structure before SQL generation. - Ensures filter keys are strings and values are supported types. - Call this before building SQL WHERE clauses to provide clear error - messages for invalid filters. + Accepts: + - scalar equality: {"category": "tech"} + - list IN: {"tag": ["a", "b"]} + - operator dicts: {"score": {"$gt": 0.5, "$lte": 0.9}} + - tuple shorthand: {"score": (">", 0.5)} (normalized internally) Args: filter_dict: Metadata filter dictionary to validate. Raises: - ValueError: If filter keys are not strings or values are unsupported types. + ValueError: If keys are not strings, operators unknown, or values + are unsupported types/non-finite. Example: >>> validate_filter({"category": "tech", "score": 0.95}) # OK - >>> validate_filter({123: "value"}) # Raises ValueError + >>> validate_filter({"score": {"$gt": 0.5}}) # OK + >>> validate_filter({"score": (">", 0.5)}) # OK + >>> validate_filter({123: "value"}) # ValueError """ if filter_dict is None: return - for key, value in filter_dict.items(): + for key, raw_value in filter_dict.items(): if not isinstance(key, str): raise ValueError( f"Filter keys must be strings, got {type(key).__name__}: {key!r}" ) + # 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) + + if isinstance(value, dict): + _validate_operator_dict(key, value) + continue + + if isinstance(value, bool): + # bool is a subclass of int; allow as exact equality. + continue if not isinstance(value, (int, float, str, list)): raise ValueError( - f"Filter value for '{key}' must be int, float, str, or list, " - f"got {type(value).__name__}: {value!r}" + f"Filter value for '{key}' must be int, float, str, list, " + f"or operator dict, got {type(value).__name__}: {value!r}" ) - if isinstance(value, float) and ( - value != value or value == float("inf") or value == float("-inf") - ): + if isinstance(value, float) and not _is_finite_number(value): + raise ValueError(f"Filter value for '{key}' must be finite, got {value!r}") + if isinstance(value, list): + _validate_filter_list(key, value) + + +# SQLite's SQLITE_MAX_VARIABLE_NUMBER defaults to 999 in builds <3.32 and +# 32766 from 3.32 onward. Cap at the universally-safe minimum so $in/$nin +# never blow past the limit at runtime regardless of which SQLite the host +# python is linked against. +_FILTER_LIST_MAX_LEN = 999 + + +def _validate_filter_list(key: str, value: list[Any]) -> None: + if not value: + raise ValueError(f"Filter list for '{key}' must not be empty") + if len(value) > _FILTER_LIST_MAX_LEN: + raise ValueError( + f"Filter list for '{key}' has {len(value)} items; max " + f"{_FILTER_LIST_MAX_LEN} (SQLite parameter limit)" + ) + for i, item in enumerate(value): + if isinstance(item, bool): + continue + if not isinstance(item, (int, float, str)): raise ValueError( - f"Filter value for '{key}' must be finite, got {value!r}" + f"Filter list items for '{key}' must be int, float, or str, " + f"got {type(item).__name__} at index {i}: {item!r}" ) - if isinstance(value, list): - if not value: + if isinstance(item, float) and not _is_finite_number(item): + raise ValueError( + f"Filter list item for '{key}' at index {i} must be finite, " + f"got {item!r}" + ) + + +def _validate_operator_dict(key: str, op_dict: dict[str, Any]) -> None: + if not op_dict: + raise ValueError(f"Operator dict for '{key}' must not be empty") + for op, arg in op_dict.items(): + if op not in _FILTER_OPERATORS: + raise ValueError( + f"Unknown operator '{op}' for '{key}'. " + f"Valid: {sorted(_FILTER_OPERATORS)}" + ) + if op in ("$gt", "$gte", "$lt", "$lte"): + if not _is_finite_number(arg): + raise ValueError(f"'{key}' {op} expects a finite number, got {arg!r}") + elif op in ("$eq", "$ne"): + if isinstance(arg, bool): + continue + if arg is None or isinstance(arg, str): + continue + if isinstance(arg, (int, float)) and _is_finite_number(arg): + continue + raise ValueError( + f"'{key}' {op} expects scalar (str/number/bool/None), " + f"got {type(arg).__name__}: {arg!r}" + ) + elif op in ("$in", "$nin"): + if not isinstance(arg, list): raise ValueError( - f"Filter list for '{key}' must not be empty" + f"'{key}' {op} expects a list, got {type(arg).__name__}" + ) + _validate_filter_list(key, arg) + elif op == "$exists": + if not isinstance(arg, bool): + raise ValueError( + f"'{key}' $exists expects a bool, got {type(arg).__name__}" + ) + elif op == "$between": + if not isinstance(arg, (list, tuple)) or len(arg) != 2: + raise ValueError(f"'{key}' $between expects [lo, hi], got {arg!r}") + lo, hi = arg + if not (_is_finite_number(lo) and _is_finite_number(hi)): + raise ValueError( + f"'{key}' $between bounds must be finite numbers, got {arg!r}" + ) + if lo > hi: + # SQL `BETWEEN lo AND hi` evaluates to false when lo > hi, + # silently returning no rows. Surface the mistake instead. + raise ValueError( + f"'{key}' $between requires lo <= hi, got [{lo}, {hi}]" ) - for i, item in enumerate(value): - if not isinstance(item, (int, float, str)): - raise ValueError( - f"Filter list items for '{key}' must be int, float, or str, " - f"got {type(item).__name__} at index {i}: {item!r}" - ) - if isinstance(item, float) and ( - item != item - or item == float("inf") - or item == float("-inf") - ): - raise ValueError( - f"Filter list item for '{key}' at index {i} must be finite, " - f"got {item!r}" - ) - @contextmanager diff --git a/tests/integration/test_rag.py b/tests/integration/test_rag.py index 088ab77..44819cd 100755 --- a/tests/integration/test_rag.py +++ b/tests/integration/test_rag.py @@ -8,6 +8,7 @@ _ollama_available = False try: from ollama import Client as OllamaClient + _ollama_available = True except ImportError: OllamaClient = Mock() # type: ignore @@ -19,9 +20,10 @@ def test_rag_end_to_end(populated_db: VectorDB, monkeypatch): # Mock embed_texts to return a 4D vector matching populated_db mock_embed = Mock(return_value=[[0.1, 0.1, 0.1, 0.1]]) - + # We need to mock the module import import sys + mock_module = Mock() mock_module.embed_texts = mock_embed monkeypatch.setitem(sys.modules, "simplevecdb.embeddings.models", mock_module) @@ -35,7 +37,9 @@ def mock_generate(prompt) -> dict[str, str]: # Simple RAG chain (real code would use langchain/llama_index) query = "What color is grape?" - contexts = populated_db.collection("default").similarity_search(query, k=2) # embed query in real + contexts = populated_db.collection("default").similarity_search( + query, k=2 + ) # embed query in real context_str = "\n".join(doc.page_content for doc, _ in contexts) prompt = f"Context: {context_str}\nQuestion: {query}" @@ -48,9 +52,7 @@ def mock_generate(prompt) -> dict[str, str]: # Real Ollama test — runs only when a local Ollama server has the # `qwen3.5:0.8b` model pulled. Skipped in CI (no Ollama daemon, no model) # and skipped locally when the daemon is unreachable. -@pytest.mark.skipif( - not _ollama_available, reason="Ollama not installed" -) +@pytest.mark.skipif(not _ollama_available, reason="Ollama not installed") @pytest.mark.skipif( bool(os.environ.get("CI")), reason="CI environments do not run a local Ollama server", @@ -70,11 +72,11 @@ def test_rag_with_ollama(populated_db): # But populated_db fixture has 4D vectors, which won't match real embeddings. # So this test is conceptually flawed unless we use a real DB. # We'll just fix the syntax error for now. - - # Mocking embedding for the sake of the test structure, + + # Mocking embedding for the sake of the test structure, # assuming we had a real embedding function available. - query_emb = [0.1, 0.1, 0.1, 0.1] - + query_emb = [0.1, 0.1, 0.1, 0.1] + contexts = populated_db.collection("default").similarity_search(query_emb, k=2) context = "\n".join(d.page_content for d, _ in contexts) response = client.generate( diff --git a/tests/integration/test_server.py b/tests/integration/test_server.py index 2b24e2e..2f2f8f4 100755 --- a/tests/integration/test_server.py +++ b/tests/integration/test_server.py @@ -93,7 +93,9 @@ def test_run_server(): 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.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" diff --git a/tests/unit/core/test_core_additional_coverage.py b/tests/unit/core/test_core_additional_coverage.py index 73bdf86..1b94044 100755 --- a/tests/unit/core/test_core_additional_coverage.py +++ b/tests/unit/core/test_core_additional_coverage.py @@ -197,16 +197,19 @@ def test_filter_with_numeric_value(self, tmp_path): db.close() def test_filter_with_unsupported_type_raises(self, tmp_path): - """Filter with unsupported type raises ValueError.""" + """Operator dicts must use a recognized $-prefixed operator (gap 5).""" db = VectorDB(str(tmp_path / "unsupported_filter.db")) collection = db.collection("test") collection.add_texts( ["doc"], embeddings=[[1.0, 0.0]], ) - # dict as filter value is not supported - with pytest.raises(ValueError, match="must be int, float, str, or list"): + # Plain string-keyed nested dicts are rejected as unknown operators. + with pytest.raises(ValueError, match="Unknown operator"): collection._catalog.build_filter_clause({"key": {"nested": "dict"}}) + # Bytes values are also unsupported. + with pytest.raises(ValueError, match="must be int, float, str"): + collection._catalog.build_filter_clause({"key": b"bytes"}) db.close() def test_keyword_search_empty_query(self, tmp_path): diff --git a/tests/unit/core/test_filters.py b/tests/unit/core/test_filters.py index 1b472b3..3d5a618 100755 --- a/tests/unit/core/test_filters.py +++ b/tests/unit/core/test_filters.py @@ -26,12 +26,13 @@ def test_build_filter_clause_in_list(): def test_build_filter_clause_unsupported_type(): - """Test build_filter_clause with unsupported value type.""" + """Operator dicts require $-prefixed operators (gap 5).""" db = VectorDB(":memory:") collection = db.collection("default") - filter_dict = {"key": {"nested": "dict"}} # Dict is not supported - with pytest.raises(ValueError, match="must be int, float, str, or list"): - collection._catalog.build_filter_clause(filter_dict) + with pytest.raises(ValueError, match="Unknown operator"): + collection._catalog.build_filter_clause({"key": {"nested": "dict"}}) + with pytest.raises(ValueError, match="must be int, float, str"): + collection._catalog.build_filter_clause({"key": b"bytes"}) def test_filter_advanced(): diff --git a/tests/unit/core/test_missing_coverage.py b/tests/unit/core/test_missing_coverage.py index 86c5bd6..5adeb69 100644 --- a/tests/unit/core/test_missing_coverage.py +++ b/tests/unit/core/test_missing_coverage.py @@ -2,7 +2,6 @@ from __future__ import annotations -import sqlite3 import sys from pathlib import Path from unittest.mock import MagicMock, patch @@ -12,7 +11,7 @@ from simplevecdb import VectorDB from simplevecdb.core import get_optimal_batch_size -from simplevecdb.types import ClusterResult, MigrationRequiredError +from simplevecdb.types import ClusterResult # ------------------------------------------------------------------ # @@ -128,9 +127,7 @@ def test_rebuild_index_no_embeddings_raises(self, tmp_path): # Insert a doc but clear embeddings from catalog collection.add_texts(["test"], embeddings=[[0.1, 0.2]]) # Wipe the embeddings column - db.conn.execute( - f"UPDATE {collection._table_name} SET embedding = NULL" - ) + db.conn.execute(f"UPDATE {collection._table_name} SET embedding = NULL") db.conn.commit() with pytest.raises(RuntimeError, match="No embeddings found"): collection.rebuild_index() @@ -164,110 +161,6 @@ def test_del_calls_close(self): assert db._closed is True -# ------------------------------------------------------------------ # -# MigrationRequiredError in VectorDB init (lines 1439-1440) -# ------------------------------------------------------------------ # - - -class TestMigrationRequired: - def test_migration_required_error_raised(self, tmp_path): - """When legacy data exists and auto_migrate=False, raise MigrationRequiredError.""" - db_path = tmp_path / "legacy.db" - # Create a database with a legacy vec_index table - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vec_index (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - # Insert a fake vector - fake_vec = np.array([0.1, 0.2, 0.3], dtype=np.float32).tobytes() - conn.execute("INSERT INTO vec_index VALUES (1, ?)", (fake_vec,)) - conn.commit() - conn.close() - - with pytest.raises(MigrationRequiredError): - VectorDB(str(db_path), auto_migrate=False) - - -# ------------------------------------------------------------------ # -# check_migration edge cases (lines 1668-1670, 1703-1716, 1722-1739, 1745) -# ------------------------------------------------------------------ # - - -class TestCheckMigration: - def test_nonexistent_path(self): - result = VectorDB.check_migration("/nonexistent/path/db.sqlite") - assert result["needs_migration"] is False - assert result["collections"] == [] - assert result["total_vectors"] == 0 - - def test_corrupted_database(self, tmp_path): - """DatabaseError should return no-migration-needed result.""" - db_path = tmp_path / "corrupt.db" - db_path.write_bytes(b"not a sqlite database at all") - result = VectorDB.check_migration(str(db_path)) - assert result["needs_migration"] is False - - def test_no_legacy_tables(self, tmp_path): - db_path = tmp_path / "clean.db" - conn = sqlite3.connect(str(db_path)) - conn.execute("CREATE TABLE unrelated (id INTEGER)") - conn.commit() - conn.close() - result = VectorDB.check_migration(str(db_path)) - assert result["needs_migration"] is False - assert result["collections"] == [] - - def test_legacy_default_collection(self, tmp_path): - """Detect legacy vec_index table with data.""" - db_path = tmp_path / "legacy_default.db" - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vec_index (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - fake_vec = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32).tobytes() - conn.execute("INSERT INTO vec_index VALUES (1, ?)", (fake_vec,)) - conn.execute("INSERT INTO vec_index VALUES (2, ?)", (fake_vec,)) - conn.commit() - conn.close() - - result = VectorDB.check_migration(str(db_path)) - assert result["needs_migration"] is True - assert "default" in result["collections"] - assert result["total_vectors"] == 2 - assert result["estimated_size_mb"] >= 0 - assert "ROLLBACK" in result["rollback_notes"] - - def test_legacy_named_collections(self, tmp_path): - """Detect legacy vectors_{name} tables.""" - db_path = tmp_path / "legacy_named.db" - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vectors_products (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - fake_vec = np.array([0.5, 0.6], dtype=np.float32).tobytes() - conn.execute("INSERT INTO vectors_products VALUES (1, ?)", (fake_vec,)) - conn.commit() - conn.close() - - result = VectorDB.check_migration(str(db_path)) - assert result["needs_migration"] is True - assert "products" in result["collections"] - assert result["total_vectors"] == 1 - - def test_empty_legacy_table_ignored(self, tmp_path): - """Legacy tables with 0 rows are not flagged.""" - db_path = tmp_path / "empty_legacy.db" - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vec_index (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - conn.commit() - conn.close() - - result = VectorDB.check_migration(str(db_path)) - assert result["needs_migration"] is False - - # ------------------------------------------------------------------ # # Cross-collection search: parallel failure handling (lines 1552-1553) # ------------------------------------------------------------------ # @@ -285,6 +178,7 @@ def test_parallel_search_handles_failure(self): # Make c2's search raise def failing_search(*args, **kwargs): raise RuntimeError("Simulated failure") + c2.similarity_search = failing_search results = db.search_collections([0.1, 0.2], k=5, parallel=True) @@ -310,8 +204,7 @@ def callback(progress): progress_reports.append(progress) items = [ - (f"doc{i}", {"i": i}, [float(i) * 0.1, float(i) * 0.2]) - for i in range(5) + (f"doc{i}", {"i": i}, [float(i) * 0.1, float(i) * 0.2]) for i in range(5) ] # Consume the generator @@ -357,79 +250,6 @@ def test_process_streaming_batch_auto_embed(self, tmp_path): db.close() -# ------------------------------------------------------------------ # -# _migrate_from_sqlite_vec_if_needed (lines 276-312) -# ------------------------------------------------------------------ # - - -class TestMigrateFromSqliteVec: - def test_migration_skipped_when_no_legacy_table(self): - """Migration is a no-op when no legacy table exists.""" - db = VectorDB(":memory:") - collection = db.collection("default") - # No crash, no error - migration silently skips - assert collection.count() == 0 - - def test_migration_with_empty_legacy_data(self, tmp_path): - """Migration handles empty legacy table gracefully.""" - db_path = tmp_path / "empty_legacy.db" - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vec_index (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - conn.commit() - conn.close() - - db = VectorDB(str(db_path), auto_migrate=True) - collection = db.collection("default") - # Should handle empty table without error - assert collection.count() == 0 - db.close() - - def test_migration_with_legacy_data(self, tmp_path): - """Migration transfers vectors from legacy table to usearch.""" - db_path = tmp_path / "migrate.db" - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vec_index (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - vec1 = np.array([0.1, 0.2, 0.3], dtype=np.float32).tobytes() - vec2 = np.array([0.4, 0.5, 0.6], dtype=np.float32).tobytes() - conn.execute("INSERT INTO vec_index VALUES (1, ?)", (vec1,)) - conn.execute("INSERT INTO vec_index VALUES (2, ?)", (vec2,)) - conn.commit() - conn.close() - - db = VectorDB(str(db_path), auto_migrate=True) - collection = db.collection("default") - # Vectors should be migrated to usearch index - assert collection._index.size == 2 - db.close() - - def test_migration_failure_raises_runtime_error(self, tmp_path): - """Migration failure raises RuntimeError with context.""" - db_path = tmp_path / "fail_migrate.db" - conn = sqlite3.connect(str(db_path)) - conn.execute( - "CREATE TABLE vec_index (rowid INTEGER PRIMARY KEY, embedding BLOB)" - ) - # Insert invalid blob data - conn.execute("INSERT INTO vec_index VALUES (1, ?)", (b"invalid",)) - conn.commit() - conn.close() - - db = VectorDB(str(db_path), auto_migrate=True) - # Migration may fail when trying to deserialize invalid blob - # The exact error depends on numpy's frombuffer behavior - try: - db.collection("default") - # If it doesn't raise, the data was somehow processable - except RuntimeError as e: - assert "Failed to migrate" in str(e) - finally: - db.close() - - # ------------------------------------------------------------------ # # _resolve_index_path with encrypted index (line 261) # ------------------------------------------------------------------ # diff --git a/tests/unit/core/test_v25_correctness.py b/tests/unit/core/test_v25_correctness.py index 2fe92a1..65fc5d8 100644 --- a/tests/unit/core/test_v25_correctness.py +++ b/tests/unit/core/test_v25_correctness.py @@ -64,7 +64,9 @@ def test_fts_subtables_excluded(self): 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}" + assert fts_leaked == [], ( + f"FTS tables leaked into list_collections: {fts_leaked}" + ) def test_sorted_output(self): db = VectorDB(":memory:") diff --git a/tests/unit/core/test_v25_features.py b/tests/unit/core/test_v25_features.py index 34f42e9..b5eb90e 100644 --- a/tests/unit/core/test_v25_features.py +++ b/tests/unit/core/test_v25_features.py @@ -69,11 +69,15 @@ def test_filter_returns_fewer_than_k_when_insufficient(self): 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)] + 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"}) + results = col.similarity_search( + _rand_embedding(42), k=10, filter={"color": "red"} + ) assert len(results) == 3 def test_no_filter_match_returns_empty(self): @@ -161,7 +165,9 @@ 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) + 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)) diff --git a/tests/unit/core/test_v25_robustness.py b/tests/unit/core/test_v25_robustness.py index cdb22b3..ec7983d 100644 --- a/tests/unit/core/test_v25_robustness.py +++ b/tests/unit/core/test_v25_robustness.py @@ -35,7 +35,9 @@ 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_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 @@ -52,7 +54,9 @@ 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_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 @@ -60,8 +64,10 @@ async def flaky(): raise sqlite3.OperationalError("database is locked") return "done" - with patch("asyncio.sleep", new_callable=AsyncMock) as mock_asleep, \ - patch("time.sleep") as mock_tsleep: + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_asleep, + patch("time.sleep") as mock_tsleep, + ): result = await flaky() assert result == "done" @@ -73,7 +79,9 @@ 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_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 @@ -87,7 +95,9 @@ async def bad(): 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_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") @@ -283,10 +293,12 @@ def test_load_or_create_mmap_decision(self, tmp_path: Path): 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): + 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 @@ -296,10 +308,12 @@ def test_load_or_create_mmap_decision(self, tmp_path: Path): 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): + 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 diff --git a/tests/unit/embeddings/test_repo_id_validation.py b/tests/unit/embeddings/test_repo_id_validation.py index ca94e0c..cf1e8c0 100644 --- a/tests/unit/embeddings/test_repo_id_validation.py +++ b/tests/unit/embeddings/test_repo_id_validation.py @@ -80,10 +80,12 @@ class TestLoadModelEnforcesTrustRemoteCodeFalse: """``load_model`` must always pass trust_remote_code=False.""" def test_trust_remote_code_forced_off(self): - with patch("simplevecdb.embeddings.models._load_snapshot_download") as snap, \ - patch( - "simplevecdb.embeddings.models._load_sentence_transformer_cls" - ) as st_cls: + with ( + patch("simplevecdb.embeddings.models._load_snapshot_download") as snap, + patch( + "simplevecdb.embeddings.models._load_sentence_transformer_cls" + ) as st_cls, + ): snap.return_value = lambda **kw: "/tmp/fake-model-path" # noqa: ARG005 st_cls.return_value = lambda *args, **kwargs: kwargs diff --git a/tests/unit/embeddings/test_server.py b/tests/unit/embeddings/test_server.py index dc209be..7c324a5 100755 --- a/tests/unit/embeddings/test_server.py +++ b/tests/unit/embeddings/test_server.py @@ -122,7 +122,9 @@ def test_server_run_with_args(): 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.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" @@ -136,7 +138,9 @@ def test_server_run_default_config(): 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.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 @@ -157,7 +161,9 @@ def test_server_run_cli_args(): 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.uvicorn.Server", return_value=mock_server + ): with patch.object( sys, "argv", ["script", "--host", "192.168.1.1", "--port", "7000"] ): diff --git a/tests/unit/embeddings/test_v25_enhancements.py b/tests/unit/embeddings/test_v25_enhancements.py index 6b63b82..f83e169 100644 --- a/tests/unit/embeddings/test_v25_enhancements.py +++ b/tests/unit/embeddings/test_v25_enhancements.py @@ -76,7 +76,9 @@ 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): + 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) @@ -141,9 +143,7 @@ def test_no_warmup_skips_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: + 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 diff --git a/tests/unit/engine/test_v26_quantization_clustering.py b/tests/unit/engine/test_v26_quantization_clustering.py index af1b042..f6176cd 100644 --- a/tests/unit/engine/test_v26_quantization_clustering.py +++ b/tests/unit/engine/test_v26_quantization_clustering.py @@ -58,6 +58,7 @@ def test_unit_norm_vector_accepted(self): def test_just_over_unit_warns_and_clips(self): # Reset the module-level latch so this test sees the warning. from simplevecdb.engine import quantization as q + q._INT8_RANGE_WARNED = False strat = QuantizationStrategy(Quantization.INT8) v = np.array([1.5, 0.0, 0.0, 0.0], dtype=np.float32) @@ -69,6 +70,7 @@ def test_just_over_unit_warns_and_clips(self): def test_just_under_negative_unit_warns_and_clips(self): from simplevecdb.engine import quantization as q + q._INT8_RANGE_WARNED = False strat = QuantizationStrategy(Quantization.INT8) v = np.array([-2.0, 0.0, 0.0, 0.0], dtype=np.float32) @@ -81,6 +83,7 @@ def test_warning_only_emitted_once_per_process(self): # Latch behavior: a second call with a still-out-of-range vector # must not re-warn after the first warn has fired. from simplevecdb.engine import quantization as q + q._INT8_RANGE_WARNED = False strat = QuantizationStrategy(Quantization.INT8) v = np.array([1.5, 0.0, 0.0, 0.0], dtype=np.float32) @@ -88,11 +91,13 @@ def test_warning_only_emitted_once_per_process(self): strat.serialize(v) # Second call must not raise the warning again. import warnings + with warnings.catch_warnings(record=True) as record: warnings.simplefilter("always") strat.serialize(v) int8_warnings = [ - w for w in record + w + for w in record if issubclass(w.category, DeprecationWarning) and "INT8 quantization" in str(w.message) ] @@ -113,6 +118,7 @@ def test_empty_vector_does_not_raise(self): def test_warning_message_includes_max_abs(self): from simplevecdb.engine import quantization as q + q._INT8_RANGE_WARNED = False strat = QuantizationStrategy(Quantization.INT8) v = np.array([3.7, -0.5, 0.2], dtype=np.float32) diff --git a/tests/unit/integrations/test_langchain_coverage.py b/tests/unit/integrations/test_langchain_coverage.py index 6c09d62..d25d1ec 100755 --- a/tests/unit/integrations/test_langchain_coverage.py +++ b/tests/unit/integrations/test_langchain_coverage.py @@ -64,7 +64,9 @@ def test_langchain_similarity_search_with_score_returns_scores(tmp_path): assert doc.page_content == "content" assert score == 0.25 mock_embedding.embed_query.assert_called_once_with("query") - mock_col.similarity_search.assert_called_once_with(query=[0.5] * 3, k=1, filter=None) + mock_col.similarity_search.assert_called_once_with( + query=[0.5] * 3, k=1, filter=None + ) def test_langchain_mmr_requires_embedding(tmp_path): diff --git a/tests/unit/integrations/test_llamaindex_v26.py b/tests/unit/integrations/test_llamaindex_v26.py index aab3730..0947f11 100644 --- a/tests/unit/integrations/test_llamaindex_v26.py +++ b/tests/unit/integrations/test_llamaindex_v26.py @@ -166,9 +166,7 @@ class TestDeleteNodesFilters: def test_delete_nodes_with_filters_raises(self, tmp_path): store = _make_store(tmp_path, "filters") - filters = MetadataFilters( - filters=[ExactMatchFilter(key="source", value="x")] - ) + filters = MetadataFilters(filters=[ExactMatchFilter(key="source", value="x")]) with pytest.raises(NotImplementedError, match="filters"): store.delete_nodes(filters=filters) diff --git a/tests/unit/test_async_coverage.py b/tests/unit/test_async_coverage.py index 9f94bd9..2dcf01f 100644 --- a/tests/unit/test_async_coverage.py +++ b/tests/unit/test_async_coverage.py @@ -76,7 +76,9 @@ async def test_async_auto_tag(sample_texts, sample_embeddings): await collection.add_texts(texts=texts, embeddings=emb.tolist()) cluster_result = await collection.cluster(n_clusters=3, random_state=42) - tags = await collection.auto_tag(cluster_result, method="keywords", n_keywords=3) + tags = await collection.auto_tag( + cluster_result, method="keywords", n_keywords=3 + ) assert isinstance(tags, dict) assert len(tags) > 0 diff --git a/tests/unit/test_catalog_coverage.py b/tests/unit/test_catalog_coverage.py index 8e00e38..573501f 100644 --- a/tests/unit/test_catalog_coverage.py +++ b/tests/unit/test_catalog_coverage.py @@ -204,8 +204,8 @@ def test_empty_filter_returns_empty(self, catalog): assert params == [] def test_unsupported_filter_type_raises(self, catalog): - """Line 525: unsupported value type raises ValueError.""" - with pytest.raises(ValueError, match="must be int, float, str, or list"): + """Unsupported scalar types still raise ValueError (gap 5 grammar).""" + with pytest.raises(ValueError, match="must be int, float, str"): catalog.build_filter_clause({"key": object()}) @@ -233,66 +233,6 @@ def test_get_all_docs_with_filter(self, catalog): assert meta["category"] == "a" -class TestLegacyVec: - """Cover lines 630-631, 643-650, 654-659.""" - - def test_check_legacy_returns_false_on_exception(self, catalog): - """Lines 630-631: exception during check -> False.""" - # Replace conn with a mock that raises on execute - mock_conn = MagicMock() - mock_conn.execute.side_effect = sqlite3.OperationalError("boom") - catalog.conn = mock_conn - result = catalog.check_legacy_sqlite_vec("old_vec_table") - assert result is False - - def test_check_legacy_returns_false_no_table(self, catalog): - """Line 629: table doesn't exist -> False.""" - result = catalog.check_legacy_sqlite_vec("nonexistent_vec") - assert result is False - - def test_get_legacy_vectors_failure(self, catalog): - """Lines 648-650: get_legacy_vectors returns empty on error.""" - result = catalog.get_legacy_vectors("nonexistent_table") - assert result == [] - - def test_get_legacy_vectors_success(self, catalog): - """Lines 643-647: get_legacy_vectors reads from table.""" - catalog.conn.execute( - "CREATE TABLE old_vec (embedding BLOB)" - ) - catalog.conn.execute( - "INSERT INTO old_vec (rowid, embedding) VALUES (1, ?)", - (b"\x00\x01\x02\x03",), - ) - catalog.conn.commit() - - result = catalog.get_legacy_vectors("old_vec") - assert len(result) == 1 - assert result[0][0] == 1 - assert result[0][1] == b"\x00\x01\x02\x03" - - def test_drop_legacy_vec_table(self, catalog): - """Lines 654-659: drop legacy table.""" - catalog.conn.execute("CREATE TABLE old_vec2 (embedding BLOB)") - catalog.conn.commit() - - catalog.drop_legacy_vec_table("old_vec2") - - # Table should be gone - row = catalog.conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='old_vec2'" - ).fetchone() - assert row is None - - def test_drop_legacy_vec_table_failure_logged(self, catalog): - """Line 659: failure during drop is logged, not raised.""" - # Replace conn with a mock that raises on execute - mock_conn = MagicMock() - mock_conn.execute.side_effect = sqlite3.OperationalError("locked") - catalog.conn = mock_conn - catalog.drop_legacy_vec_table("some_table") # should not raise - - class TestClusterStateOperations: """Cover line 768 (list_cluster_states).""" diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index cf80467..f3bc43e 100755 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -361,9 +361,7 @@ def test_normalize_l2(): assert np.allclose(normalize_l2(zero_vec), zero_vec) -@pytest.mark.skipif( - not _has_langchain, reason="langchain-core not installed" -) +@pytest.mark.skipif(not _has_langchain, reason="langchain-core not installed") def test_as_langchain(empty_db): """Test LangChain integration factory method.""" lc_store = empty_db.as_langchain() @@ -373,9 +371,7 @@ def test_as_langchain(empty_db): assert isinstance(lc_store, SimpleVecDBVectorStore) -@pytest.mark.skipif( - not _has_llamaindex, reason="llama-index not installed" -) +@pytest.mark.skipif(not _has_llamaindex, reason="llama-index not installed") def test_as_llama_index(empty_db): """Test LlamaIndex integration factory method.""" li_store = empty_db.as_llama_index() @@ -551,39 +547,6 @@ def test_rebuild_index_empty_collection(tmp_path): db.close() -def test_check_migration_no_legacy(tmp_path): - """Test check_migration() on a fresh v2.0 database.""" - db_path = str(tmp_path / "fresh.db") - - # Create a new database (no legacy data) - db = VectorDB(db_path) - collection = db.collection("default") - collection.add_texts(["test"], embeddings=[[0.1] * 64]) - db.close() - - # Check migration status - info = VectorDB.check_migration(db_path) - - assert info["needs_migration"] is False - assert info["collections"] == [] - assert info["total_vectors"] == 0 - - -def test_check_migration_nonexistent(): - """Test check_migration() on nonexistent file.""" - info = VectorDB.check_migration("/nonexistent/path.db") - - assert info["needs_migration"] is False - assert info["collections"] == [] - - -def test_check_migration_memory(): - """Test check_migration() on :memory: database.""" - info = VectorDB.check_migration(":memory:") - - assert info["needs_migration"] is False - - def test_adaptive_search_uses_exact_for_small_collections(tmp_path): """Test that search uses brute-force (exact) for small collections.""" from simplevecdb import constants diff --git a/tests/unit/test_encryption_coverage.py b/tests/unit/test_encryption_coverage.py index d5ef3f9..08a90a7 100644 --- a/tests/unit/test_encryption_coverage.py +++ b/tests/unit/test_encryption_coverage.py @@ -71,7 +71,9 @@ def test_cipher_version_none_raises_encryption_error(self, tmp_path: Path): mock_sqlcipher.connect.return_value = mock_conn mock_pkg = MagicMock(dbapi2=mock_sqlcipher) - with patch.dict("sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher}): + with patch.dict( + "sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher} + ): with pytest.raises(EncryptionError, match="not active"): create_encrypted_connection(tmp_path / "test.db", "passphrase") @@ -100,7 +102,9 @@ def side_effect(*args, **kwargs): mock_sqlcipher.connect.return_value = mock_conn mock_pkg = MagicMock(dbapi2=mock_sqlcipher) - with patch.dict("sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher}): + with patch.dict( + "sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher} + ): with pytest.raises(EncryptionError, match="wrong key"): create_encrypted_connection(tmp_path / "test.db", "passphrase") @@ -127,7 +131,9 @@ def side_effect(*args, **kwargs): mock_sqlcipher.connect.return_value = mock_conn mock_pkg = MagicMock(dbapi2=mock_sqlcipher) - with patch.dict("sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher}): + with patch.dict( + "sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher} + ): with pytest.raises(EncryptionError, match="Failed to verify"): create_encrypted_connection(tmp_path / "test.db", "passphrase") @@ -137,7 +143,9 @@ def test_outer_generic_exception_wraps(self, tmp_path: Path): mock_sqlcipher.connect.side_effect = OSError("disk full") mock_pkg = MagicMock(dbapi2=mock_sqlcipher) - with patch.dict("sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher}): + with patch.dict( + "sys.modules", {"sqlcipher3": mock_pkg, "sqlcipher3.dbapi2": mock_sqlcipher} + ): with pytest.raises(EncryptionError, match="Failed to create encrypted"): create_encrypted_connection(tmp_path / "test.db", "passphrase") @@ -165,7 +173,10 @@ def test_encrypt_file_import_error(self, tmp_path: Path): input_file = tmp_path / "plain.bin" input_file.write_bytes(b"data") - with patch.dict("sys.modules", {"cryptography": None, "cryptography.hazmat.primitives.ciphers.aead": None}): + with patch.dict( + "sys.modules", + {"cryptography": None, "cryptography.hazmat.primitives.ciphers.aead": None}, + ): with patch( "builtins.__import__", side_effect=_make_import_blocker("cryptography"), @@ -191,7 +202,10 @@ def test_decrypt_file_import_error(self, tmp_path: Path): enc_file = tmp_path / "enc.bin" enc_file.write_bytes(os.urandom(100)) - with patch.dict("sys.modules", {"cryptography": None, "cryptography.hazmat.primitives.ciphers.aead": None}): + with patch.dict( + "sys.modules", + {"cryptography": None, "cryptography.hazmat.primitives.ciphers.aead": None}, + ): with patch( "builtins.__import__", side_effect=_make_import_blocker("cryptography"), @@ -287,7 +301,9 @@ def test_decrypt_index_non_usearch_suffix(self, tmp_path: Path): def _make_import_blocker(blocked_module: str): """Create an __import__ side_effect that blocks a specific module.""" - real_import = __builtins__.__import__ if hasattr(__builtins__, '__import__') else __import__ + real_import = ( + __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + ) def blocker(name, *args, **kwargs): if name == blocked_module or name.startswith(blocked_module + "."): diff --git a/tests/unit/test_encryption_v1_format.py b/tests/unit/test_encryption_v1_format.py index ccd67d4..f8ed009 100644 --- a/tests/unit/test_encryption_v1_format.py +++ b/tests/unit/test_encryption_v1_format.py @@ -132,9 +132,7 @@ def test_decrypt_v0_legacy_blob(self, tmp_path: Path): decrypt_file(enc, out, TEST_KEY) assert out.read_bytes() == original - def test_decrypt_v0_blob_starting_with_sv_bytes_still_works( - self, tmp_path: Path - ): + def test_decrypt_v0_blob_starting_with_sv_bytes_still_works(self, tmp_path: Path): # A v0 nonce that *happens* to start with 'SV' but whose 3rd byte is # not the version sentinel must still decrypt as v0. Decrypt logic # only strips the header when *both* the magic bytes AND the version diff --git a/tests/unit/test_error_handling.py b/tests/unit/test_error_handling.py index 028288b..0da73ba 100755 --- a/tests/unit/test_error_handling.py +++ b/tests/unit/test_error_handling.py @@ -194,12 +194,20 @@ def test_invalid_key_type(self): validate_filter({None: "value"}) # type: ignore def test_invalid_value_type(self): - """Unsupported value types raise ValueError.""" - with pytest.raises(ValueError, match="must be int, float, str, or list"): + """Unsupported value types raise ValueError. + + Operator dicts now accept $-prefixed operators (gap 5); plain + nested dicts are rejected as unknown operators. Tuple shorthand + is accepted but the first element must be a recognized op string. + """ + with pytest.raises(ValueError, match="Unknown operator"): validate_filter({"key": {"nested": "dict"}}) - with pytest.raises(ValueError, match="must be int, float, str, or list"): - validate_filter({"key": (1, 2, 3)}) # tuple not allowed + with pytest.raises(ValueError, match="operator.*must be a string"): + validate_filter({"key": (1, 2, 3)}) # numeric op invalid + + with pytest.raises(ValueError, match="must be int, float, str"): + validate_filter({"key": b"bytes"}) def test_invalid_list_item_type(self): """List items must be int, float, or str.""" @@ -209,6 +217,15 @@ def test_invalid_list_item_type(self): with pytest.raises(ValueError, match="list items.*must be int, float, or str"): validate_filter({"tags": [[1, 2, 3]]}) # nested list + def test_filter_list_length_capped(self): + """Lists exceeding the SQLite parameter limit raise.""" + # Just under the cap is fine; just over raises. + validate_filter({"ids": list(range(999))}) + with pytest.raises(ValueError, match="SQLite parameter limit"): + validate_filter({"ids": list(range(1000))}) + with pytest.raises(ValueError, match="SQLite parameter limit"): + validate_filter({"ids": {"$in": list(range(1000))}}) + # ============================================================================ # Logging tests @@ -327,8 +344,8 @@ def test_invalid_filter_raises_on_keyword_search(self): embeddings=[[0.1, 0.2, 0.3]], ) - # Search with invalid filter - with pytest.raises(ValueError, match="must be int, float, str, or list"): + # Search with invalid filter (unknown operator name). + with pytest.raises(ValueError, match="Unknown operator"): collection.keyword_search( query="hello", k=5, @@ -425,55 +442,3 @@ def writer(collection, writer_id): db1.close() db2.close() - - -# ============================================================================ -# Migration error tests -# ============================================================================ - - -class TestMigrationRequiredError: - """Tests for MigrationRequiredError blocking behavior.""" - - def test_migration_error_attributes(self): - """MigrationRequiredError has correct attributes.""" - from simplevecdb import MigrationRequiredError - - error = MigrationRequiredError( - path="/path/to/db.db", - collections=["default", "docs"], - total_vectors=1000, - migration_info={"needs_migration": True}, - ) - - assert error.path == "/path/to/db.db" - assert error.collections == ["default", "docs"] - assert error.total_vectors == 1000 - assert error.migration_info["needs_migration"] is True - assert "1000 vectors" in str(error) - assert "auto_migrate=True" in str(error) - - def test_new_db_no_migration_error(self): - """New databases don't raise MigrationRequiredError.""" - # auto_migrate=False should be fine for new databases - db = VectorDB(":memory:") # Default is auto_migrate=False - collection = db.collection("test") - collection.add_texts(["hello"], embeddings=[[0.1, 0.2, 0.3]]) - assert collection.count() == 1 - db.close() - - def test_check_migration_new_db(self, tmp_path): - """check_migration returns no migration for new databases.""" - db_path = str(tmp_path / "new.db") - - # Create a new v2.0 database - db = VectorDB(db_path, auto_migrate=True) - collection = db.collection("test") - collection.add_texts(["hello"], embeddings=[[0.1, 0.2, 0.3]]) - db.close() - - # Check migration - should be empty - info = VectorDB.check_migration(db_path) - assert info["needs_migration"] is False - assert info["collections"] == [] - assert info["total_vectors"] == 0 diff --git a/tests/unit/test_multi_collection.py b/tests/unit/test_multi_collection.py index b3eefd2..5aca7b9 100755 --- a/tests/unit/test_multi_collection.py +++ b/tests/unit/test_multi_collection.py @@ -1,54 +1,56 @@ - import pytest from simplevecdb import VectorDB, Quantization + def test_multi_collection_basic(): db = VectorDB(":memory:") - + # Create two collections c1 = db.collection("c1", quantization=Quantization.FLOAT) c2 = db.collection("c2", quantization=Quantization.FLOAT) - + # Add data c1.add_texts(["doc1"], embeddings=[[0.1, 0.2]]) c2.add_texts(["doc2"], embeddings=[[0.9, 0.8]]) - + # Search c1 res1 = c1.similarity_search([0.1, 0.2], k=1) assert len(res1) == 1 assert res1[0][0].page_content == "doc1" - + # Search c2 res2 = c2.similarity_search([0.9, 0.8], k=1) assert len(res2) == 1 assert res2[0][0].page_content == "doc2" - + # Ensure isolation res1_cross = c1.similarity_search([0.9, 0.8], k=1) # Should still find doc1 because it's the only doc in c1, but distance should be large assert res1_cross[0][0].page_content == "doc1" - + # Check isolation by count (if we had more docs) # Let's add another doc to c1 c1.add_texts(["doc1b"], embeddings=[[0.11, 0.21]]) assert len(c1.similarity_search([0.1, 0.2], k=10)) == 2 assert len(c2.similarity_search([0.9, 0.8], k=10)) == 1 + def test_collection_persistence(tmp_path): db_path = tmp_path / "test.db" db = VectorDB(db_path) - + c1 = db.collection("users") c1.add_texts(["alice"], embeddings=[[1.0, 0.0]]) - + db.close() - + # Reopen db2 = VectorDB(db_path) c1_reopened = db2.collection("users") res = c1_reopened.similarity_search([1.0, 0.0], k=1) assert res[0][0].page_content == "alice" + def test_invalid_collection_name(): db = VectorDB(":memory:") with pytest.raises(ValueError): @@ -56,11 +58,12 @@ def test_invalid_collection_name(): with pytest.raises(ValueError): db.collection("drop table users;") + def test_default_collection_compat(): db = VectorDB(":memory:") # Explicitly use collection("default") as add_texts is removed from VectorDB db.collection("default").add_texts(["default_doc"], embeddings=[[0.5, 0.5]]) - + # Access via default collection explicitly def_col = db.collection("default") res = def_col.similarity_search([0.5, 0.5], k=1) diff --git a/tests/unit/test_search.py b/tests/unit/test_search.py index a5be8be..66de872 100755 --- a/tests/unit/test_search.py +++ b/tests/unit/test_search.py @@ -34,12 +34,12 @@ def test_similarity_search_basic(db): # Query vector must match the dimension of stored vectors (3D). query = [0.95, 0.95, 0.95] results = db.collection("default").similarity_search(query, k=2) - + assert len(results) == 2 # "grape" ([0.85, 0.85, 0.85]) and "orange" ([0.9, 0.9, 0.9]) are closest to [0.95, 0.95, 0.95] # Note: The order depends on exact distance calculations. # Both are very close to the query direction (1,1,1). - + # Verify that results are returned with scores. assert 0 <= results[0][1] < 0.1 assert results[0][1] <= results[1][1] @@ -48,8 +48,10 @@ def test_similarity_search_basic(db): def test_similarity_search_filter(db): """Test similarity search with metadata filtering.""" # Query with 3D vector matching the database dimension. - results = db.collection("default").similarity_search([0.95] * 3, k=4, filter={"likes": [10, 15]}) - + results = db.collection("default").similarity_search( + [0.95] * 3, k=4, filter={"likes": [10, 15]} + ) + assert len(results) == 2 # Should match "apple" (10) and "orange" (15) found_texts = {r[0].page_content for r in results} assert found_texts == {"apple", "orange"} @@ -61,7 +63,7 @@ def test_recall_gold_standard(populated_db): Uses 'populated_db' fixture from conftest.py which has 4D vectors. """ query = np.array([0.95, 0.95, 0.95, 0.95]) - + # Reconstruct embeddings from the fixture logic for ground truth calculation all_embs = np.array( [ @@ -71,23 +73,28 @@ def test_recall_gold_standard(populated_db): [0.85, 0.85, 0.85, 0.85], ] ) - + # Normalize for cosine similarity comparison all_embs = all_embs / np.linalg.norm(all_embs, axis=1, keepdims=True) query_norm = query / np.linalg.norm(query) - + # Compute ground truth similarities sims = np.dot(all_embs, query_norm) - + # Dynamically determine expected top-k based on numpy calculation # This integrates 'sims' to ensure the test validates against the actual math top_k_indices = np.argsort(-sims)[:2] - all_texts = ["apple is red", "banana is yellow", "orange is orange", "grape is purple"] + all_texts = [ + "apple is red", + "banana is yellow", + "orange is orange", + "grape is purple", + ] expected = [all_texts[i] for i in top_k_indices] - + results = populated_db.collection("default").similarity_search(query, k=2) result_texts = [r[0].page_content for r in results] - + # Calculate recall intersection = set(result_texts) & set(expected) recall = len(intersection) / 2 @@ -99,13 +106,13 @@ def test_quantization_search(quant_db): # Generate random 128D vectors embs = np.random.randn(10, 128).astype(np.float32) embs /= np.linalg.norm(embs, axis=1, keepdims=True) - + collection = quant_db.collection("default") collection.add_texts(["t"] * 10, embeddings=embs.tolist()) - + # Search with one of the inserted vectors results = collection.similarity_search(embs[0], k=1) - + # Expect the vector to find itself with very low distance assert results[0][1] < 0.05 @@ -116,12 +123,12 @@ def test_mmr_diversity(): collection = db.collection("default") # A and B are very similar. C is orthogonal. # Query is close to A and B. - + texts = ["A", "B", "C"] embeddings = [ - [1.0, 0.0, 0.0], # A + [1.0, 0.0, 0.0], # A [0.99, 0.01, 0.0], # B (very close to A) - [0.0, 1.0, 0.0], # C (orthogonal) + [0.0, 1.0, 0.0], # C (orthogonal) ] collection.add_texts(texts, embeddings=embeddings) @@ -143,7 +150,9 @@ def test_delete_by_ids(): """Test that deleting items removes them from search results.""" db = VectorDB(":memory:") collection = db.collection("default") - ids = collection.add_texts(["a", "b"], embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + ids = collection.add_texts( + ["a", "b"], embeddings=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]] + ) assert len(ids) == 2 # Delete the first item diff --git a/tests/unit/test_search_missing_coverage.py b/tests/unit/test_search_missing_coverage.py index e4c0dad..55ab27b 100644 --- a/tests/unit/test_search_missing_coverage.py +++ b/tests/unit/test_search_missing_coverage.py @@ -26,7 +26,9 @@ def db_3d(tmp_path): {"cat": "A", "score": 30}, {"cat": "B", "score": 40}, ] - col.add_texts(["doc1", "doc2", "doc3", "doc4"], embeddings=embeddings, metadatas=metadatas) + col.add_texts( + ["doc1", "doc2", "doc3", "doc4"], embeddings=embeddings, metadatas=metadatas + ) return db @@ -178,9 +180,7 @@ def test_mmr_candidates_fewer_than_k(self, db_3d): """Line 376: when candidates <= k, return all candidates directly.""" col = db_3d.collection("test") # Request k=10 but only 4 docs exist, so candidates <= k - results = col.max_marginal_relevance_search( - [1.0, 0.0, 0.0], k=10, fetch_k=4 - ) + results = col.max_marginal_relevance_search([1.0, 0.0, 0.0], k=10, fetch_k=4) assert len(results) <= 4 assert len(results) > 0 diff --git a/tests/unit/test_usearch_index_missing_coverage.py b/tests/unit/test_usearch_index_missing_coverage.py index 93ab8f1..e039a38 100644 --- a/tests/unit/test_usearch_index_missing_coverage.py +++ b/tests/unit/test_usearch_index_missing_coverage.py @@ -17,7 +17,9 @@ class TestUnpackBits: def test_unpack_bits_roundtrip(self): from simplevecdb.engine.usearch_index import _pack_bits, _unpack_bits - vectors = np.array([[1.0, -0.5, 0.3, -0.8, 0.0, 0.1, -0.2, 0.9]], dtype=np.float32) + vectors = np.array( + [[1.0, -0.5, 0.3, -0.8, 0.0, 0.1, -0.2, 0.9]], dtype=np.float32 + ) packed = _pack_bits(vectors) unpacked = _unpack_bits(packed, ndim=8) # Positive -> 1 -> +1, negative/zero -> 0 -> -1 @@ -27,10 +29,13 @@ def test_unpack_bits_roundtrip(self): def test_unpack_bits_multi_row(self): from simplevecdb.engine.usearch_index import _pack_bits, _unpack_bits - vectors = np.array([ - [1.0, -1.0, 1.0, -1.0], - [-1.0, 1.0, -1.0, 1.0], - ], dtype=np.float32) + vectors = np.array( + [ + [1.0, -1.0, 1.0, -1.0], + [-1.0, 1.0, -1.0, 1.0], + ], + dtype=np.float32, + ) packed = _pack_bits(vectors) unpacked = _unpack_bits(packed, ndim=4) assert unpacked.shape == (2, 4) @@ -58,6 +63,7 @@ def test_large_index_uses_mmap(self, tmp_path): # Patch the constants module that _load_or_create imports import simplevecdb.constants as real_constants + original_threshold = real_constants.USEARCH_MMAP_THRESHOLD try: real_constants.USEARCH_MMAP_THRESHOLD = 0 # Everything is "large" diff --git a/tests/unit/test_v26_1_features.py b/tests/unit/test_v26_1_features.py new file mode 100644 index 0000000..af6734a --- /dev/null +++ b/tests/unit/test_v26_1_features.py @@ -0,0 +1,655 @@ +"""Tests for SimpleVecDB 2.6.1 catalog extensions (gaps 1-10).""" + +from __future__ import annotations + +import sqlite3 +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pytest + +from simplevecdb import VectorDB + + +@pytest.fixture +def db_with_docs(tmp_path): + db = VectorDB(str(tmp_path / "v261.db")) + c = db.collection("default", store_embeddings=True) + embs = np.random.rand(5, 4).astype(np.float32) + c.add_texts( + ["a", "b", "c", "d", "e"], + embeddings=embs, + metadatas=[ + {"i": i, "score": float(i), "tag": "x" if i % 2 == 0 else "y"} + for i in range(5) + ], + ) + yield db, c, embs + db.close() + + +# ---------------------------- gap 1: update_embedding --------------------- + + +class TestUpdateEmbedding: + def test_buffers_without_hnsw_change(self, db_with_docs): + db, c, embs = db_with_docs + before_size = c._index.size + c.update_embedding(1, np.zeros(4, dtype=np.float32)) + assert c.pending.size() == 1 + # HNSW size unchanged until flush. + assert c._index.size == before_size + + def test_flush_promotes_to_hnsw(self, db_with_docs): + db, c, embs = db_with_docs + new = np.ones(4, dtype=np.float32) + c.update_embedding(1, new) + assert c.pending.flush() == 1 + assert c.pending.size() == 0 + # The doc closest to all-ones is now id=1. + results = c.similarity_search(new, k=1) + assert results[0][0].metadata["i"] == 0 # may match by direction + + def test_dim_mismatch_raises(self, db_with_docs): + db, c, _ = db_with_docs + with pytest.raises(ValueError, match="!= index dim"): + c.update_embedding(1, np.zeros(7, dtype=np.float32)) + + def test_blend_toward_combines_with_pending(self, db_with_docs): + db, c, _ = db_with_docs + centroid = np.full(4, 0.5, dtype=np.float32) + n = c.pending.blend_toward([1, 2], centroid, alpha=0.5) + assert n == 2 + assert c.pending.size() == 2 + + +# ---------------------------- gap 2: transaction -------------------------- + + +class TestTransaction: + def test_db_transaction_commits_on_success(self, db_with_docs): + db, c, _ = db_with_docs + with db.transaction() as tx: + tx["default"].increment_metadata(1, {"hits": 1}) + tx["default"].edges.add_edge(1, 2, weight=0.7) + assert c.counters.get(1, "hits") == 1 + assert len(c.edges.get_edges(src=1)) == 1 + + def test_db_transaction_rolls_back_on_error(self, db_with_docs): + db, c, _ = db_with_docs + c.increment_metadata(1, {"hits": 1}) + baseline = c.counters.get(1, "hits") + with pytest.raises(RuntimeError): + with db.transaction() as tx: + tx["default"].increment_metadata(1, {"hits": 5}) + tx["default"].edges.add_edge(2, 3, weight=0.3) + raise RuntimeError("boom") + # SQL state preserved. + assert c.counters.get(1, "hits") == baseline + assert len(c.edges.get_edges(src=2)) == 0 + + def test_collection_tx_yields_collection(self, db_with_docs): + db, c, _ = db_with_docs + with c.tx() as ctx: + assert ctx is c + ctx.increment_metadata(1, {"hits": 3}) + assert c.counters.get(1, "hits") == 3 + + def test_outer_commit_failure_propagates(self, db_with_docs): + """Outer commit failures must surface, not be swallowed.""" + db, _c, _ = db_with_docs + real_conn = db.conn + calls = {"n": 0} + + class FlakyConn: + def __getattr__(self, name): + return getattr(real_conn, name) + + def commit(self): + calls["n"] += 1 + raise sqlite3.OperationalError("disk full") + + db.conn = FlakyConn() # type: ignore[assignment] + try: + with pytest.raises(sqlite3.OperationalError, match="disk full"): + with db.transaction() as tx: + tx["default"].increment_metadata(1, {"hits": 1}) + finally: + db.conn = real_conn + assert calls["n"] == 1 + + +# ---------------------------- gap 3: edges -------------------------------- + + +class TestEdges: + def test_crud(self, db_with_docs): + db, c, _ = db_with_docs + assert c.edges.add_edge(1, 2, kind="similar", weight=0.7) == 1 + edges = c.edges.get_edges(src=1) + assert len(edges) == 1 + assert edges[0].weight == pytest.approx(0.7) + # Update with absolute set. + c.edges.update_edge(1, 2, kind="similar", weight=0.9) + assert c.edges.get_edges(src=1)[0].weight == pytest.approx(0.9) + # Delete. + assert c.edges.delete_edge(1, 2, kind="similar") == 1 + assert c.edges.get_edges(src=1) == [] + + def test_atomic_delta_concurrency(self, db_with_docs): + db, c, _ = db_with_docs + c.edges.add_edge(1, 2, kind="x", weight=0.0) + + def bump(_): + c.edges.update_edge(1, 2, kind="x", dweight=0.01, dhits=1) + + with ThreadPoolExecutor(max_workers=8) as ex: + list(ex.map(bump, range(200))) + e = c.edges.get_edges(src=1, dst=2, kind="x")[0] + assert e.hits == 200 + assert e.weight == pytest.approx(2.0, abs=1e-6) + + def test_range_filter_on_weight(self, db_with_docs): + db, c, _ = db_with_docs + c.edges.add_edge(1, 2, weight=0.05) + c.edges.add_edge(1, 3, weight=0.5) + c.edges.add_edge(1, 4, weight=0.95) + low = c.edges.get_edges(filter={"weight": {"$lt": 0.1}}) + assert {e.dst_id for e in low} == {2} + between = c.edges.get_edges(filter={"weight": {"$between": [0.1, 0.9]}}) + assert {e.dst_id for e in between} == {3} + + def test_prune(self, db_with_docs): + db, c, _ = db_with_docs + c.edges.add_edge(1, 2, weight=0.05) + c.edges.add_edge(1, 3, weight=0.5) + n = c.edges.prune(max_weight=0.1) + assert n == 1 + + +# ---------------------------- gap 4: counters ----------------------------- + + +class TestCounters: + def test_dict_increment(self, db_with_docs): + db, c, _ = db_with_docs + c.increment_metadata(1, {"retrieval_count": 1, "drift_total": 0.02}) + c.increment_metadata(1, {"retrieval_count": 1, "drift_total": 0.05}) + assert c.counters.get(1, "retrieval_count") == 2 + assert c.counters.get(1, "drift_total") == pytest.approx(0.07) + + def test_concurrent_increments_are_atomic(self, db_with_docs): + db, c, _ = db_with_docs + + def bump(_): + c.increment_metadata(2, {"hits": 1}) + + with ThreadPoolExecutor(max_workers=8) as ex: + list(ex.map(bump, range(800))) + assert c.counters.get(2, "hits") == 800 + + def test_invalid_key_rejected(self, db_with_docs): + db, c, _ = db_with_docs + with pytest.raises(ValueError, match="metadata counter key"): + c.increment_metadata(1, {"no spaces": 1}) + + def test_non_numeric_rejected(self, db_with_docs): + db, c, _ = db_with_docs + with pytest.raises(TypeError): + c.increment_metadata(1, {"x": "string"}) + + +# ---------------------------- gap 5: range filters ------------------------ + + +class TestRangeFilters: + def test_mongo_operator_dict(self, db_with_docs): + db, c, embs = db_with_docs + r = c.similarity_search( + embs[0], k=10, filter={"score": {"$gt": 1.5, "$lt": 4.0}} + ) + scores = sorted(d.metadata["score"] for d, _ in r) + assert scores == [2.0, 3.0] + + def test_tuple_shorthand(self, db_with_docs): + db, c, embs = db_with_docs + r = c.similarity_search(embs[0], k=10, filter={"score": ("range", 1.5, 4.0)}) + scores = sorted(d.metadata["score"] for d, _ in r) + assert scores == [2.0, 3.0, 4.0] + + def test_in_and_nin(self, db_with_docs): + db, c, embs = db_with_docs + r = c.similarity_search(embs[0], k=10, filter={"tag": {"$in": ["x"]}}) + assert all(d.metadata["tag"] == "x" for d, _ in r) + r = c.similarity_search(embs[0], k=10, filter={"tag": {"$nin": ["x"]}}) + assert all(d.metadata["tag"] != "x" for d, _ in r) + + def test_unknown_operator_raises(self, db_with_docs): + db, c, embs = db_with_docs + with pytest.raises(ValueError, match="Unknown operator"): + c.similarity_search(embs[0], k=1, filter={"score": {"$bogus": 1}}) + + +# ---------------------------- gap 7: change feed -------------------------- + + +class TestEvents: + def test_mutation_appends_event(self, db_with_docs): + db, c, _ = db_with_docs + before = c.events.last_seq() + c.increment_metadata(1, {"hits": 1}) + c.edges.add_edge(1, 2, weight=0.5) + after = c.events.last_seq() + assert after >= before + 2 + + def test_read_filters_by_kind(self, db_with_docs): + db, c, _ = db_with_docs + c.edges.add_edge(1, 2, weight=0.5) + c.edges.delete_edge(1, 2) + adds = c.events.read(kind="edge_add") + dels = c.events.read(kind="edge_delete") + assert len(adds) >= 1 and len(dels) >= 1 + + +# ---------------------------- gap 8: TTL ---------------------------------- + + +class TestTTL: + def test_sweep_deletes_expired(self, db_with_docs): + db, c, _ = db_with_docs + c.ttl.set(1, seconds=-10) + deleted, callback_ids = c.ttl.sweep() + assert 1 in deleted + assert callback_ids == [] + # Doc removed. + assert c._catalog.count() == 4 + + def test_callback_keeps_doc(self, db_with_docs): + db, c, _ = db_with_docs + c.ttl.set(2, seconds=-5, on_expire="callback") + deleted, callback_ids = c.ttl.sweep() + assert deleted == [] + assert 2 in callback_ids + + def test_background_sweep(self, db_with_docs): + db, c, _ = db_with_docs + c.ttl.set(3, seconds=0.2) + c.ttl.start_background(interval=0.1) + try: + time.sleep(0.6) + finally: + c.ttl.stop_background() + # 3 should have been swept. + ids_left = {row[0] for row in c._catalog.get_all_docs_with_text()} + assert 3 not in ids_left + + +# ---------------------------- gap 9: maintenance -------------------------- + + +class TestMaintenance: + def test_threshold_triggers_rebuild(self, db_with_docs): + db, c, _ = db_with_docs + c.update_embedding(1, np.zeros(4, dtype=np.float32)) + c.pending.flush() + ran = c.maintenance.rebuild_if_needed(max_pending=1) + assert ran is True + # Subsequent call doesn't rebuild again until threshold re-passed. + assert c.maintenance.rebuild_if_needed(max_pending=1) is False + + +# ---------------------------- coverage gap fillers ------------------------ + + +class TestCountersGet: + def test_get_returns_stored_value(self, db_with_docs): + _, c, _ = db_with_docs + c.counters.increment(1, {"hits": 3}) + assert c.counters.get(1, "hits") == 3 + + def test_get_returns_default_for_missing_key(self, db_with_docs): + _, c, _ = db_with_docs + assert c.counters.get(1, "never_set", default=42) == 42 + + def test_get_returns_none_for_missing_row(self, db_with_docs): + _, c, _ = db_with_docs + assert c.counters.get(99999, "hits") is None + + +class TestExistsOperator: + def test_exists_true_matches_present_key(self, db_with_docs): + _, c, embs = db_with_docs + results = c.similarity_search( + embs[0], k=10, filter={"score": {"$exists": True}} + ) + # Every seeded doc has a "score" metadata field. + assert len(results) == 5 + + def test_exists_false_matches_absent_key(self, db_with_docs): + _, c, embs = db_with_docs + results = c.similarity_search( + embs[0], k=10, filter={"missing": {"$exists": False}} + ) + assert len(results) == 5 + + +class TestEventsObservability: + def test_last_seq_grows_with_appends(self, db_with_docs): + _, c, _ = db_with_docs + before = c.events.last_seq() + c.events.append("manual", payload={"k": 1}) + c.events.append("manual", payload={"k": 2}) + assert c.events.last_seq() == before + 2 + + def test_prune_drops_old_events(self, db_with_docs): + _, c, _ = db_with_docs + for i in range(5): + c.events.append("noise", payload={"i": i}) + cutoff = c.events.last_seq() + c.events.append("keep", payload={"i": "kept"}) + # prune deletes seq < before_seq, so cutoff+1 covers all noise rows. + removed = c.events.prune(before_seq=cutoff + 1) + assert removed >= 5 + kinds = [e.kind for e in c.events.read()] + assert "keep" in kinds + assert "noise" not in kinds + + def test_subscribe_yields_new_events(self, db_with_docs): + _, c, _ = db_with_docs + start = c.events.last_seq() + c.events.append("first", payload={"n": 1}) + c.events.append("second", payload={"n": 2}) + + gen = c.events.subscribe(since=start, poll_interval=0.001, batch=10) + try: + seen = [next(gen), next(gen)] + finally: + gen.close() + kinds = [e.kind for e in seen] + assert "first" in kinds and "second" in kinds + + +# ---------------------------- fortification guards ------------------------ + + +class TestUpdateEmbeddingFortification: + def test_nan_vector_rejected(self, db_with_docs): + _, c, _ = db_with_docs + bad = np.array([1.0, float("nan"), 0.0, 0.0], dtype=np.float32) + with pytest.raises(ValueError, match="finite"): + c.update_embedding(1, bad) + + def test_inf_vector_rejected(self, db_with_docs): + _, c, _ = db_with_docs + bad = np.array([float("inf"), 0.0, 0.0, 0.0], dtype=np.float32) + with pytest.raises(ValueError, match="finite"): + c.update_embedding(1, bad) + + def test_2d_vector_rejected(self, db_with_docs): + _, c, _ = db_with_docs + bad = np.zeros((1, 4), dtype=np.float32) + with pytest.raises(ValueError, match="1-D"): + c.update_embedding(1, bad) + + +class TestBetweenOperator: + def test_between_inclusive_range_matches(self, db_with_docs): + _, c, embs = db_with_docs + results = c.similarity_search( + embs[0], k=10, filter={"score": {"$between": [1.0, 3.0]}} + ) + scores = sorted(r[0].metadata["score"] for r in results) + assert scores == [1.0, 2.0, 3.0] + + def test_between_tuple_shorthand(self, db_with_docs): + _, c, embs = db_with_docs + results = c.similarity_search( + embs[0], k=10, filter={"score": ("range", 0.0, 1.0)} + ) + scores = sorted(r[0].metadata["score"] for r in results) + assert scores == [0.0, 1.0] + + def test_between_lo_greater_than_hi_rejected(self, db_with_docs): + _, c, embs = db_with_docs + with pytest.raises(ValueError, match="lo <= hi"): + c.similarity_search( + embs[0], k=10, filter={"score": {"$between": [5.0, 1.0]}} + ) + + def test_between_non_finite_rejected(self, db_with_docs): + _, c, embs = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.similarity_search( + embs[0], k=10, filter={"score": {"$between": [0.0, float("inf")]}} + ) + + def test_between_wrong_arity_rejected(self, db_with_docs): + _, c, embs = db_with_docs + with pytest.raises(ValueError): + c.similarity_search(embs[0], k=10, filter={"score": {"$between": [1.0]}}) + + +class TestEdgeFortification: + def test_add_edge_nan_weight_rejected(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.edges.add_edge(1, 2, weight=float("nan")) + + def test_add_edge_inf_bonus_rejected(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.edges.add_edge(1, 2, bonus=float("inf")) + + def test_upsert_edge_nan_weight_rejected(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.edges.upsert(1, 2, weight=float("nan")) + + def test_update_edge_nan_dweight_rejected(self, db_with_docs): + _, c, _ = db_with_docs + c.edges.add_edge(1, 2, weight=0.5) + with pytest.raises(ValueError, match="finite"): + c.edges.update_edge(1, 2, dweight=float("nan")) + + def test_update_edge_inf_dbonus_rejected(self, db_with_docs): + _, c, _ = db_with_docs + c.edges.add_edge(1, 2, weight=0.5) + with pytest.raises(ValueError, match="finite"): + c.edges.update_edge(1, 2, dbonus=float("inf")) + + +class TestTTLFortification: + def test_set_requires_one_of_seconds_or_expires_at(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="expires_at or seconds"): + c.ttl.set(1) + + def test_set_rejects_both_seconds_and_expires_at(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="exactly one"): + c.ttl.set(1, expires_at=time.time() + 5, seconds=10) + + def test_set_rejects_invalid_on_expire(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="on_expire"): + c.ttl.set(1, seconds=5, on_expire="bogus") + + def test_set_rejects_nan_seconds(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.ttl.set(1, seconds=float("nan")) + + def test_set_rejects_inf_seconds(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.ttl.set(1, seconds=float("inf")) + + def test_set_rejects_nan_expires_at(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="finite"): + c.ttl.set(1, expires_at=float("nan")) + + def test_clear_returns_zero_for_missing(self, db_with_docs): + _, c, _ = db_with_docs + assert c.ttl.clear(99999) == 0 + + def test_clear_returns_one_after_set(self, db_with_docs): + _, c, _ = db_with_docs + c.ttl.set(1, seconds=60) + assert c.ttl.clear(1) == 1 + + def test_start_background_rejects_zero_interval(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="positive finite"): + c.ttl.start_background(interval=0) + + def test_start_background_rejects_negative_interval(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="positive finite"): + c.ttl.start_background(interval=-1.0) + + def test_start_background_rejects_nan_interval(self, db_with_docs): + _, c, _ = db_with_docs + with pytest.raises(ValueError, match="positive finite"): + c.ttl.start_background(interval=float("nan")) + + def test_start_background_idempotent_and_stops_cleanly(self, db_with_docs): + _, c, _ = db_with_docs + c.ttl.start_background(interval=60.0) + # Idempotent — second call is a no-op. + c.ttl.start_background(interval=60.0) + c.ttl.stop_background() + # After clean stop the next start spawns fresh. + c.ttl.start_background(interval=60.0) + c.ttl.stop_background() + + +# ----- regression: numeric filter type guard -------------------------------- + + +class TestNumericFilterTypeGuard: + """SQL numeric ops must reject non-numeric JSON values. + + Without ``json_type`` gating, ``CAST(json_extract(...) AS REAL)`` + coerces strings/null/objects to 0.0, so ``{"score": "oops"}`` would + spuriously match ``{"score": {"$lt": 1}}``. The python-side + ``_matches_filter`` already filters these out, so SQL and Python + were disagreeing. + """ + + def _build_collection(self, tmp_path): + db = VectorDB(str(tmp_path / "guard.db")) + c = db.collection("default") + c.add_texts( + ["numeric-half", "numeric-five", "string-oops", "missing-key"], + embeddings=np.random.rand(4, 4).astype(np.float32), + metadatas=[ + {"score": 0.5}, + {"score": 5.0}, + {"score": "oops"}, # would CAST to 0.0 without the guard + {"other": "no score key"}, + ], + ) + return db, c + + def test_lt_skips_string_value(self, tmp_path): + db, c = self._build_collection(tmp_path) + try: + docs = c.get_documents(filter_dict={"score": {"$lt": 1.0}}) + texts = {t for _, t, _ in docs} + # Only the numeric 0.5 row passes; "oops" must NOT coerce to 0.0. + assert texts == {"numeric-half"} + finally: + db.close() + + def test_between_skips_string_value(self, tmp_path): + db, c = self._build_collection(tmp_path) + try: + docs = c.get_documents(filter_dict={"score": {"$between": (-1, 1)}}) + texts = {t for _, t, _ in docs} + assert texts == {"numeric-half"} + finally: + db.close() + + def test_gt_skips_string_and_missing(self, tmp_path): + db, c = self._build_collection(tmp_path) + try: + docs = c.get_documents(filter_dict={"score": {"$gt": 0.0}}) + texts = {t for _, t, _ in docs} + # Both numeric rows pass; string and missing are excluded. + assert texts == {"numeric-half", "numeric-five"} + finally: + db.close() + + def test_sql_and_python_agree_on_string_value(self, tmp_path): + """SQL pre-filter (get_documents) and Python post-filter + (similarity_search) must produce the same set of rows.""" + db, c = self._build_collection(tmp_path) + try: + sql_texts = { + text + for _, text, _ in c.get_documents(filter_dict={"score": {"$lt": 1.0}}) + } + # similarity_search applies _matches_filter post-fetch + hits = c.similarity_search([0.0] * 4, k=10, filter={"score": {"$lt": 1.0}}) + py_texts = {doc.page_content for doc, _ in hits} + assert sql_texts == py_texts == {"numeric-half"} + finally: + db.close() + + +# ----- regression: TTL sweep atomicity -------------------------------------- + + +class TestTTLSweepAtomic: + """sweep_ttl must claim expired rows atomically. + + Before the fix, ``sweep_ttl`` ran a SELECT and then a DELETE in two + separate steps; a concurrent ``set_ttl`` extending the deadline (or + ``clear_ttl``) between those two steps would still see the doc + deleted off the stale read. Post-fix it uses + ``DELETE … RETURNING`` inside one write transaction. + """ + + def test_basic_sweep_still_works(self, db_with_docs): + _, c, _ = db_with_docs + c.ttl.set(1, expires_at=time.time() - 1, on_expire="delete") + c.ttl.set(2, expires_at=time.time() - 1, on_expire="callback") + deleted, callbacks = c.ttl.sweep() + assert deleted == [1] + assert callbacks == [2] + # TTL rows for both must be cleared (RETURNING removed them). + assert c.ttl.sweep() == ([], []) + + def test_extension_between_logical_check_and_delete(self, db_with_docs): + """If a TTL is renewed before the sweep transaction runs, the + atomic ``DELETE … RETURNING`` must observe the new value and + skip the row — even when the test races by snapshotting the + cutoff first. + """ + _, c, _ = db_with_docs + c.ttl.set(3, expires_at=time.time() - 5, on_expire="delete") + # Capture a "what would expire as of cutoff_t0" snapshot, then + # extend the TTL before sweeping with that same cutoff. + cutoff_t0 = time.time() + c.ttl.set(3, expires_at=time.time() + 3600, on_expire="delete") + deleted, callbacks = c.ttl.sweep(now=cutoff_t0) + # The renewed TTL is now > cutoff_t0 → must NOT be swept. + assert deleted == [] + assert callbacks == [] + # Doc itself still present. + rows = c.get_documents() + assert any(doc_id == 3 for doc_id, _, _ in rows) + + def test_clear_between_logical_check_and_delete(self, db_with_docs): + _, c, _ = db_with_docs + c.ttl.set(4, expires_at=time.time() - 1, on_expire="delete") + cutoff_t0 = time.time() + c.ttl.clear(4) + # Cleared TTL means there's no row to claim → no delete. + deleted, callbacks = c.ttl.sweep(now=cutoff_t0) + assert deleted == [] + assert callbacks == [] + rows = c.get_documents() + assert any(doc_id == 4 for doc_id, _, _ in rows) diff --git a/tests/unit/test_v26_encryption_review_pass_3.py b/tests/unit/test_v26_encryption_review_pass_3.py index 558b240..32f01a6 100644 --- a/tests/unit/test_v26_encryption_review_pass_3.py +++ b/tests/unit/test_v26_encryption_review_pass_3.py @@ -60,8 +60,8 @@ def test_two_encryptions_use_distinct_nonces(self, tmp_path, random_key): # Skip the 3-byte v1 header; nonce follows immediately. header_len = len(_ENC_MAGIC) + 1 - nonce_a = out_a.read_bytes()[header_len:header_len + AES_NONCE_SIZE] - nonce_b = out_b.read_bytes()[header_len:header_len + AES_NONCE_SIZE] + nonce_a = out_a.read_bytes()[header_len : header_len + AES_NONCE_SIZE] + nonce_b = out_b.read_bytes()[header_len : header_len + AES_NONCE_SIZE] assert len(nonce_a) == AES_NONCE_SIZE assert len(nonce_b) == AES_NONCE_SIZE @@ -69,9 +69,7 @@ def test_two_encryptions_use_distinct_nonces(self, tmp_path, random_key): class TestWrongKeyDoesNotCreateOutput: - def test_decrypt_with_wrong_key_never_writes_output( - self, tmp_path, random_key - ): + def test_decrypt_with_wrong_key_never_writes_output(self, tmp_path, random_key): plaintext = b"sensitive contents" src = tmp_path / "p.bin" src.write_bytes(plaintext) @@ -85,15 +83,11 @@ def test_decrypt_with_wrong_key_never_writes_output( decrypt_file(encrypted, output, wrong_key) # Authentication failure must short-circuit before any write. - assert not output.exists(), ( - "Output file must not exist after wrong-key decrypt" - ) + assert not output.exists(), "Output file must not exist after wrong-key decrypt" class TestHeaderAADBinding: - def test_tampering_with_magic_byte_fails_authentication( - self, tmp_path, random_key - ): + def test_tampering_with_magic_byte_fails_authentication(self, tmp_path, random_key): plaintext = b"tamper test" src = tmp_path / "p.bin" src.write_bytes(plaintext) @@ -141,9 +135,7 @@ def test_existing_salt_sidecar_is_not_overwritten(self, tmp_path): result = _resolve_salt(resource, create_if_missing=True) - assert result == existing, ( - "Existing sidecar must be preserved, not overwritten" - ) + assert result == existing, "Existing sidecar must be preserved, not overwritten" assert salt_path.read_bytes() == existing def test_concurrent_creators_converge_on_one_salt(self, tmp_path): diff --git a/tests/unit/test_v26_misc.py b/tests/unit/test_v26_misc.py index d6a42a4..604c1fb 100644 --- a/tests/unit/test_v26_misc.py +++ b/tests/unit/test_v26_misc.py @@ -59,7 +59,8 @@ def test_distinct_docs_with_same_text_kept_separate(self, db_with_dup_text): # results — pre-2.6.0 they collapsed into one, with whichever # metadata happened to come last winning. sources = sorted( - r[0].metadata.get("source") for r in results + r[0].metadata.get("source") + for r in results if r[0].page_content.startswith("the quick brown fox") ) assert sources == ["A", "B"], ( @@ -157,9 +158,7 @@ def test_reload_attaches_null_handler(self): importlib.reload(svc_logging) root = logging.getLogger("simplevecdb") - null_handlers = [ - h for h in root.handlers if isinstance(h, logging.NullHandler) - ] + null_handlers = [h for h in root.handlers if isinstance(h, logging.NullHandler)] assert len(null_handlers) == 1, ( "simplevecdb root logger must have exactly one NullHandler " "after the logging module is imported" diff --git a/tests/unit/test_v26_review_pass_3.py b/tests/unit/test_v26_review_pass_3.py index a77df2f..4398ff0 100644 --- a/tests/unit/test_v26_review_pass_3.py +++ b/tests/unit/test_v26_review_pass_3.py @@ -171,12 +171,8 @@ def test_filter_does_not_inflate_vector_rrf_score(self, tmp_path): try: # 10 documents; only the last carries category=keep. texts = [f"doc number {i}" for i in range(10)] - metas = [{"category": "drop"} for _ in range(9)] + [ - {"category": "keep"} - ] - embs = [ - [float(i), float(10 - i), 0.0, 0.0] for i in range(10) - ] + metas = [{"category": "drop"} for _ in range(9)] + [{"category": "keep"}] + embs = [[float(i), float(10 - i), 0.0, 0.0] for i in range(10)] col.add_texts(texts, metadatas=metas, embeddings=embs) # Hybrid search with a filter that drops the top 9 vector hits. @@ -218,9 +214,7 @@ def test_two_docs_same_text_different_ids_both_appear(self, tmp_path): query_vector=[0.5, 0.5, 0.0, 0.0], k=5, ) - variants = { - doc.metadata.get("variant") for doc, _ in results - } + variants = {doc.metadata.get("variant") for doc, _ in results} assert variants == {"a", "b"}, ( f"Both documents should surface; got variants={variants}" ) diff --git a/tests/unit/test_v26_review_pass_4.py b/tests/unit/test_v26_review_pass_4.py index 90c5802..ec9ef30 100644 --- a/tests/unit/test_v26_review_pass_4.py +++ b/tests/unit/test_v26_review_pass_4.py @@ -107,9 +107,7 @@ class TestLegacyPassphraseDBStillOpens: def sqlcipher_module(self): return pytest.importorskip("sqlcipher3") - def test_pre_2_6_passphrase_db_opens_under_2_6( - self, tmp_path, sqlcipher_module - ): + def test_pre_2_6_passphrase_db_opens_under_2_6(self, tmp_path, sqlcipher_module): # Step 1: create a SQLCipher DB the pre-2.6 way — passphrase # PRAGMA, SQLCipher does its own internal KDF, no sidecar. db_path = tmp_path / "legacy.db" @@ -132,9 +130,7 @@ def test_pre_2_6_passphrase_db_opens_under_2_6( conn = create_encrypted_connection(db_path, passphrase) try: - row = conn.execute( - "SELECT val FROM secrets WHERE id=1" - ).fetchone() + row = conn.execute("SELECT val FROM secrets WHERE id=1").fetchone() assert row is not None assert row[0] == "hello" finally: @@ -148,9 +144,7 @@ def test_pre_2_6_passphrase_db_opens_under_2_6( "KDF; a sidecar was incorrectly created." ) - def test_post_2_6_db_still_uses_raw_key_path( - self, tmp_path, sqlcipher_module - ): + def test_post_2_6_db_still_uses_raw_key_path(self, tmp_path, sqlcipher_module): """Brand-new DBs written under 2.6+ must use the new raw-key path with a sidecar; reopening must still work.""" from simplevecdb.encryption import create_encrypted_connection diff --git a/uv.lock b/uv.lock index fe7356c..bd1eaac 100755 --- a/uv.lock +++ b/uv.lock @@ -4705,7 +4705,7 @@ wheels = [ [[package]] name = "simplevecdb" -version = "2.6.0" +version = "2.6.1" source = { editable = "." } dependencies = [ { name = "cryptography" }, @@ -4715,7 +4715,6 @@ dependencies = [ { name = "python-dotenv" }, { name = "scikit-learn" }, { name = "sqlcipher3-binary" }, - { name = "sqlite-vec" }, { name = "usearch" }, ] @@ -4774,7 +4773,6 @@ requires-dist = [ { name = "scikit-learn", specifier = ">=1.3.0" }, { name = "sentence-transformers", marker = "extra == 'server'", specifier = ">=5.0" }, { name = "sqlcipher3-binary", specifier = ">=0.5.0" }, - { name = "sqlite-vec", specifier = ">=0.1.6" }, { name = "usearch", specifier = ">=2.16.3" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = ">=0.30" }, ] @@ -4900,18 +4898,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/ce/fbe8da71656eefd5fd4db70cb6d7b5911ba34279985ee3c11476a1f8550d/sqlcipher3_binary-0.6.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f9bccb9e942a04bbf920714940c8f9e00134cd655b11ba4c6f1306bcc3504d6c", size = 3246791, upload-time = "2025-12-31T16:40:41.484Z" }, ] -[[package]] -name = "sqlite-vec" -version = "0.1.9" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/68/85/9fad0045d8e7c8df3e0fa5a56c630e8e15ad6e5ca2e6106fceb666aa6638/sqlite_vec-0.1.9-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:1b62a7f0a060d9475575d4e599bbf94a13d85af896bc1ce86ee80d1b5b48e5fb", size = 131171, upload-time = "2026-03-31T08:02:31.717Z" }, - { url = "https://files.pythonhosted.org/packages/a4/3d/3677e0cd2f92e5ebc43cd29fbf565b75582bff1ccfa0b8327c7508e1084f/sqlite_vec-0.1.9-py3-none-macosx_11_0_arm64.whl", hash = "sha256:1d52e30513bae4cc9778ddbf6145610434081be4c3afe57cd877893bad9f6b6c", size = 165434, upload-time = "2026-03-31T08:02:32.712Z" }, - { url = "https://files.pythonhosted.org/packages/00/d4/f2b936d3bdc38eadcbd2a87875815db36430fab0363182ba5d12cd8e0b51/sqlite_vec-0.1.9-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e921e592f24a5f9a18f590b6ddd530eb637e2d474e3b1972f9bbeb773aa3cb9", size = 160076, upload-time = "2026-03-31T08:02:33.796Z" }, - { url = "https://files.pythonhosted.org/packages/6f/ad/6afd073b0f817b3e03f9e37ad626ae341805891f23c74b5292818f49ac63/sqlite_vec-0.1.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:1515727990b49e79bcaf75fdee2ffc7d461f8b66905013231251f1c8938e7786", size = 163388, upload-time = "2026-03-31T08:02:34.888Z" }, - { url = "https://files.pythonhosted.org/packages/42/89/81b2907cda14e566b9bf215e2ad82fc9b349edf07d2010756ffdb902f328/sqlite_vec-0.1.9-py3-none-win_amd64.whl", hash = "sha256:4a28dc12fa4b53d7b1dced22da2488fade444e96b5d16fd2d698cd670675cf32", size = 292804, upload-time = "2026-03-31T08:02:36.035Z" }, -] - [[package]] name = "stack-data" version = "0.6.3"