From 71103560d7f72071a20d9ac14ebc636de49af7d2 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Fri, 16 Jan 2026 03:46:34 -0600 Subject: [PATCH 1/5] feat: add cross-collection search API Add list_collections() and search_collections() methods to VectorDB for searching across multiple collections with unified, ranked results. - Score normalization: 1/(1+distance) for fair comparison across metrics - Parallel search with ThreadPoolExecutor for >1 collections - Dimension validation across searched collections - Async versions in AsyncVectorDB - 20 comprehensive unit tests - Documentation and examples --- README.md | 2 +- docs/api/core.md | 55 ++++++ docs/examples.md | 80 ++++++++ src/simplevecdb/async_core.py | 32 ++++ src/simplevecdb/core.py | 121 ++++++++++++ tests/unit/test_cross_collection_search.py | 212 +++++++++++++++++++++ 6 files changed, 501 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_cross_collection_search.py diff --git a/README.md b/README.md index d4d1941..2c382a5 100644 --- a/README.md +++ b/README.md @@ -369,7 +369,7 @@ pip install torch --index-url https://download.pytorch.org/whl/cu118 - [x] SQLCipher encryption (at-rest data protection) - [x] Streaming insert API for large-scale ingestion - [x] Hierarchical document relationships (parent/child) -- [ ] Cross-collection search +- [x] Cross-collection search - [ ] Vector clustering and auto-tagging Vote on features or propose new ones in [GitHub Discussions](https://github.com/coderdayton/simplevecdb/discussions). diff --git a/docs/api/core.md b/docs/api/core.md index e93f9f1..b1f83c9 100644 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -8,6 +8,8 @@ The main database class for managing vector collections. options: members: - collection + - list_collections + - search_collections - vacuum - close - check_migration @@ -145,3 +147,56 @@ results = collection.similarity_search( | `get_descendants(doc_id, max_depth)` | All nested children recursively | | `get_ancestors(doc_id)` | Path from document to root | | `set_parent(doc_id, parent_id)` | Move document to new parent (or None to orphan) | + +### Cross-Collection Search + +Search across multiple collections with unified, ranked results: + +```python +from simplevecdb import VectorDB + +db = VectorDB("app.db") + +# Initialize collections +users = db.collection("users") +products = db.collection("products") +docs = db.collection("docs") + +# Add data to each collection +users.add_texts(["Alice likes hiking"], embeddings=[[0.1]*384]) +products.add_texts(["Hiking boots", "Trail map"], embeddings=[[0.2]*384, [0.15]*384]) +docs.add_texts(["Mountain hiking guide"], embeddings=[[0.12]*384]) + +# List initialized collections +print(db.list_collections()) # ['users', 'products', 'docs'] + +# Search across ALL collections +results = db.search_collections([0.1]*384, k=5) +for doc, score, collection_name in results: + print(f"[{collection_name}] {doc.page_content} (score: {score:.3f})") + +# Search specific collections only +results = db.search_collections( + [0.1]*384, + collections=["users", "products"], # Exclude 'docs' + k=3 +) + +# With metadata filtering (applies to all collections) +results = db.search_collections( + [0.1]*384, + k=10, + filter={"category": "outdoor"} +) + +# Disable score normalization (returns inverted distances) +results = db.search_collections([0.1]*384, normalize_scores=False) + +# Sequential search (disable parallelism) +results = db.search_collections([0.1]*384, parallel=False) +``` + +| Method | Description | +|--------|-------------| +| `list_collections()` | Names of all initialized collections | +| `search_collections(query, collections, k, filter, normalize_scores, parallel)` | Search across multiple collections with merged results | diff --git a/docs/examples.md b/docs/examples.md index a9b46b9..285e497 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -249,6 +249,86 @@ async def main(): results = asyncio.run(main()) ``` +### Cross-Collection Search (v2.2+) + +Search across multiple collections with unified ranking: + +```python +from simplevecdb import VectorDB + +db = VectorDB("multi_tenant.db") + +# Create domain-specific collections +users = db.collection("users") +products = db.collection("products") +support = db.collection("support_tickets") + +# 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] +) + +# 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 +) + +# Apply metadata filter across all searched collections +results = db.search_collections( + query, + k=10, + filter={"category": "software"} +) + +# List available collections +print(db.list_collections()) # ['users', 'products', 'support_tickets'] +``` + +**Async cross-collection search:** + +```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()) +``` + ## Benchmark Scripts ### Backend Benchmark diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index 8a57068..f1e0866 100644 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -285,6 +285,38 @@ def collection( ) return self._collections[cache_key] + def list_collections(self) -> list[str]: + """Return names of all initialized collections.""" + return self._db.list_collections() + + async def search_collections( + self, + query: Sequence[float], + collections: list[str] | None = None, + k: int = 10, + filter: dict[str, Any] | None = None, + *, + normalize_scores: bool = True, + parallel: bool = True, + ) -> list[tuple[Document, float, str]]: + """ + Search across multiple collections with merged, ranked results. + + 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, + ), + ) + async def vacuum(self, checkpoint_wal: bool = True) -> None: """ Reclaim disk space by rebuilding the database file. diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 06a3065..579f370 100644 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -1084,6 +1084,127 @@ def __init__( migration_info=migration_info, ) + def list_collections(self) -> list[str]: + """ + Return names of all initialized collections. + + Only returns collections that have been accessed via `collection()` in this + session. Does not scan the database for collections created in previous sessions. + + Returns: + List of collection names currently cached in this VectorDB instance. + + Example: + >>> db = VectorDB("app.db") + >>> db.collection("users") + >>> db.collection("products") + >>> db.list_collections() + ['users', 'products'] + """ + return list(self._collections.keys()) + + def search_collections( + self, + query: Sequence[float], + collections: list[str] | None = None, + k: int = 10, + filter: dict[str, Any] | None = None, + *, + normalize_scores: bool = True, + parallel: bool = True, + ) -> list[tuple[Document, float, str]]: + """ + Search across multiple collections with merged, ranked results. + + Performs similarity search on each collection and merges results using + score normalization for fair comparison across distance metrics. + + Args: + query: Query vector (must match dimension of all searched collections). + collections: List of collection names to search. None searches all + initialized collections (from list_collections()). + k: Number of top results to return after merging. + filter: Optional metadata filter applied to all collections. + normalize_scores: If True, convert distances to similarity scores + in [0, 1] range using `1 / (1 + distance)`. Enables fair + comparison across COSINE [0,2] and L2 [0,∞) metrics. + parallel: If True, search collections concurrently using ThreadPoolExecutor. + + Returns: + List of (Document, similarity_score, collection_name) tuples, + sorted by descending similarity score (highest first). + + Raises: + ValueError: If no collections specified and none initialized, + or if collections have mismatched dimensions. + KeyError: If a specified collection name doesn't exist. + + Example: + >>> db = VectorDB("app.db") + >>> db.collection("users").add_texts(["alice"], embeddings=[[0.1]*384]) + >>> db.collection("products").add_texts(["widget"], embeddings=[[0.2]*384]) + >>> results = db.search_collections([0.15]*384, k=2) + >>> for doc, score, coll in results: + ... print(f"{coll}: {doc.page_content} ({score:.3f})") + """ + target_names = ( + collections if collections is not None else self.list_collections() + ) + + if not target_names: + return [] + + # Resolve and validate collections + targets: list[VectorCollection] = [] + dims: set[int | None] = set() + for name in target_names: + if name not in self._collections: + raise KeyError( + f"Collection '{name}' not initialized. Call db.collection('{name}') first." + ) + coll = self._collections[name] + targets.append(coll) + dims.add(coll._dim) + + # Check dimension consistency (ignore None for empty collections) + dims.discard(None) + if len(dims) > 1: + raise ValueError( + f"Dimension mismatch across collections: {dims}. " + "All searched collections must have the same embedding dimension." + ) + + # Search function for each collection + def _search_one(coll: VectorCollection) -> list[tuple[Document, float, str]]: + results = coll.similarity_search(query, k=k, filter=filter) + return [(doc, dist, coll.name) for doc, dist in results] + + # Execute searches + all_results: list[tuple[Document, float, str]] = [] + if parallel and len(targets) > 1: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=min(len(targets), 8)) as executor: + futures = [executor.submit(_search_one, coll) for coll in targets] + for future in futures: + all_results.extend(future.result()) + else: + for coll in targets: + all_results.extend(_search_one(coll)) + + # Normalize scores: similarity = 1 / (1 + distance) + if normalize_scores: + all_results = [ + (doc, 1.0 / (1.0 + dist), name) for doc, dist, name in all_results + ] + else: + # Invert for sorting (lower distance = higher rank) + all_results = [(doc, -dist, name) for doc, dist, name in all_results] + + # Sort by score descending and take top k + all_results.sort(key=lambda x: x[1], reverse=True) + return all_results[:k] + def collection( self, name: str = "default", diff --git a/tests/unit/test_cross_collection_search.py b/tests/unit/test_cross_collection_search.py new file mode 100644 index 0000000..87e377e --- /dev/null +++ b/tests/unit/test_cross_collection_search.py @@ -0,0 +1,212 @@ +"""Tests for cross-collection search functionality.""" + +import pytest +from simplevecdb import VectorDB, Quantization + + +class TestListCollections: + def test_empty_db_returns_empty_list(self): + db = VectorDB(":memory:") + assert db.list_collections() == [] + + def test_returns_initialized_collections(self): + db = VectorDB(":memory:") + db.collection("users") + db.collection("products") + db.collection("orders") + + result = db.list_collections() + assert set(result) == {"users", "products", "orders"} + + def test_includes_default_collection(self): + db = VectorDB(":memory:") + db.collection("default") + db.collection("other") + + result = db.list_collections() + assert "default" in result + assert "other" in result + + +class TestSearchCollections: + @pytest.fixture + def db_with_collections(self): + db = VectorDB(":memory:") + c1 = db.collection("c1", quantization=Quantization.FLOAT) + c2 = db.collection("c2", quantization=Quantization.FLOAT) + c3 = db.collection("c3", quantization=Quantization.FLOAT) + + c1.add_texts(["doc1_a", "doc1_b"], embeddings=[[0.1, 0.2], [0.15, 0.25]]) + c2.add_texts(["doc2_a", "doc2_b"], embeddings=[[0.9, 0.8], [0.85, 0.75]]) + c3.add_texts(["doc3_a"], embeddings=[[0.5, 0.5]]) + + return db + + def test_search_all_collections(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.1, 0.2], k=10) + + assert len(results) == 5 + doc, score, coll_name = results[0] + assert doc.page_content == "doc1_a" + assert coll_name == "c1" + assert 0 < score <= 1 + + def test_search_specific_collections(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.9, 0.8], collections=["c2"], k=5) + + assert len(results) == 2 + for doc, score, coll_name in results: + assert coll_name == "c2" + + def test_search_multiple_specific_collections(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.5, 0.5], collections=["c1", "c3"], k=10) + + collection_names = {r[2] for r in results} + assert collection_names == {"c1", "c3"} + + def test_k_limits_results(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.5, 0.5], k=2) + + assert len(results) == 2 + + def test_empty_collections_list_returns_empty(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.5, 0.5], collections=[]) + + assert results == [] + + def test_no_initialized_collections_returns_empty(self): + db = VectorDB(":memory:") + results = db.search_collections([0.5, 0.5]) + + assert results == [] + + def test_nonexistent_collection_raises_keyerror(self, db_with_collections): + db = db_with_collections + with pytest.raises(KeyError, match="not initialized"): + db.search_collections([0.5, 0.5], collections=["nonexistent"]) + + def test_normalized_scores_in_zero_one_range(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.1, 0.2], k=10, normalize_scores=True) + + for doc, score, coll_name in results: + assert 0 < score <= 1, f"Score {score} not in (0, 1]" + + def test_results_sorted_by_score_descending(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.1, 0.2], k=10) + + scores = [r[1] for r in results] + assert scores == sorted(scores, reverse=True) + + def test_filter_applies_across_collections(self): + db = VectorDB(":memory:") + c1 = db.collection("c1", quantization=Quantization.FLOAT) + c2 = db.collection("c2", quantization=Quantization.FLOAT) + + c1.add_texts( + ["match1", "nomatch1"], + embeddings=[[0.1, 0.2], [0.11, 0.21]], + metadatas=[{"category": "A"}, {"category": "B"}], + ) + c2.add_texts( + ["match2", "nomatch2"], + embeddings=[[0.9, 0.8], [0.91, 0.81]], + metadatas=[{"category": "A"}, {"category": "B"}], + ) + + results = db.search_collections([0.5, 0.5], k=10, filter={"category": "A"}) + + assert len(results) == 2 + for doc, score, coll_name in results: + assert "match" in doc.page_content + + def test_parallel_vs_sequential_same_results(self, db_with_collections): + db = db_with_collections + query = [0.3, 0.4] + + parallel_results = db.search_collections(query, k=5, parallel=True) + sequential_results = db.search_collections(query, k=5, parallel=False) + + parallel_docs = {r[0].page_content for r in parallel_results} + sequential_docs = {r[0].page_content for r in sequential_results} + assert parallel_docs == sequential_docs + + def test_unnormalized_scores(self, db_with_collections): + db = db_with_collections + results = db.search_collections([0.1, 0.2], k=10, normalize_scores=False) + + scores = [r[1] for r in results] + assert scores == sorted(scores, reverse=True) + + def test_single_collection_same_as_direct_search(self): + db = VectorDB(":memory:") + coll = db.collection("test", quantization=Quantization.FLOAT) + coll.add_texts(["a", "b", "c"], embeddings=[[0.1, 0.2], [0.5, 0.5], [0.9, 0.8]]) + + query = [0.5, 0.5] + direct_results = coll.similarity_search(query, k=3) + cross_results = db.search_collections(query, collections=["test"], k=3) + + direct_docs = [r[0].page_content for r in direct_results] + cross_docs = [r[0].page_content for r in cross_results] + assert direct_docs == cross_docs + + +class TestDimensionValidation: + def test_mismatched_dimensions_raises_valueerror(self): + db = VectorDB(":memory:") + c1 = db.collection("c1", quantization=Quantization.FLOAT) + c2 = db.collection("c2", quantization=Quantization.FLOAT) + + c1.add_texts(["doc1"], embeddings=[[0.1, 0.2]]) + c2.add_texts(["doc2"], embeddings=[[0.1, 0.2, 0.3]]) + + with pytest.raises(ValueError, match="Dimension mismatch"): + db.search_collections([0.1, 0.2], collections=["c1", "c2"]) + + def test_empty_collection_dimension_ignored(self): + db = VectorDB(":memory:") + c1 = db.collection("c1", quantization=Quantization.FLOAT) + db.collection("c2", quantization=Quantization.FLOAT) + + c1.add_texts(["doc1"], embeddings=[[0.1, 0.2]]) + + results = db.search_collections([0.1, 0.2], collections=["c1", "c2"]) + assert len(results) == 1 + assert results[0][0].page_content == "doc1" + + +class TestAsyncCrossCollectionSearch: + @pytest.fixture + def async_db(self): + from simplevecdb.async_core import AsyncVectorDB + + return AsyncVectorDB(":memory:") + + @pytest.mark.asyncio + async def test_list_collections(self, async_db): + async_db.collection("users") + async_db.collection("products") + + result = async_db.list_collections() + assert set(result) == {"users", "products"} + + @pytest.mark.asyncio + async def test_search_collections(self, async_db): + c1 = async_db.collection("c1", quantization=Quantization.FLOAT) + c2 = async_db.collection("c2", quantization=Quantization.FLOAT) + + await c1.add_texts(["doc1"], embeddings=[[0.1, 0.2]]) + await c2.add_texts(["doc2"], embeddings=[[0.9, 0.8]]) + + results = await async_db.search_collections([0.5, 0.5], k=10) + + assert len(results) == 2 + docs = {r[0].page_content for r in results} + assert docs == {"doc1", "doc2"} From 91a6b4b494c0b088cf642c7a277aee7e3b975524 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Sat, 17 Jan 2026 13:26:28 -0600 Subject: [PATCH 2/5] fix: address clustering assignment and dependency messages --- README.md | 37 +- docs/CHANGELOG.md | 101 ++++ docs/api/core.md | 124 +++++ docs/api/types.md | 202 ++++++++ docs/guides/clustering.md | 403 ++++++++++++++++ docs/index.md | 36 +- mkdocs.yml | 3 + pyproject.toml | 41 +- src/simplevecdb/async_core.py | 179 ++++++++ src/simplevecdb/core.py | 371 ++++++++++++++- src/simplevecdb/encryption.py | 7 +- src/simplevecdb/engine/catalog.py | 195 ++++++++ src/simplevecdb/engine/clustering.py | 248 ++++++++++ src/simplevecdb/engine/usearch_index.py | 41 ++ src/simplevecdb/types.py | 42 +- tests/unit/test_async.py | 86 ++++ tests/unit/test_clustering.py | 585 ++++++++++++++++++++++++ uv.lock | 132 +++++- 18 files changed, 2783 insertions(+), 50 deletions(-) create mode 100644 docs/api/types.md create mode 100644 docs/guides/clustering.md create mode 100644 src/simplevecdb/engine/clustering.py create mode 100644 tests/unit/test_clustering.py diff --git a/README.md b/README.md index 2c382a5..a6cef2f 100644 --- a/README.md +++ b/README.md @@ -46,16 +46,20 @@ SimpleVecDB brings **Chroma-like simplicity** to a single **SQLite file**. Built ## Installation ```bash -# Core library only (lightweight, 50MB) +# Standard installation (includes clustering, encryption) pip install simplevecdb -# With local embeddings server + HuggingFace models (500MB+) +# With local embeddings server (adds 500MB+ models) pip install "simplevecdb[server]" - -# With encryption support (SQLCipher) -pip install "simplevecdb[encryption]" ``` +**What's included by default:** +- Vector search with HNSW indexing +- Clustering (K-means, MiniBatch K-means, HDBSCAN) +- Encryption (SQLCipher AES-256) +- Async support +- LangChain & LlamaIndex integrations + **Verify Installation:** ```bash @@ -282,6 +286,23 @@ parent = collection.get_parent(child_ids[0]) descendants = collection.get_descendants(parent_ids[0]) ``` +### 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 | @@ -301,6 +322,8 @@ descendants = collection.get_descendants(parent_ids[0]) | **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+) | ## Performance Benchmarks @@ -370,7 +393,9 @@ pip install torch --index-url https://download.pytorch.org/whl/cu118 - [x] Streaming insert API for large-scale ingestion - [x] Hierarchical document relationships (parent/child) - [x] Cross-collection search -- [ ] Vector clustering and auto-tagging +- [x] Vector clustering and auto-tagging (v2.2) +- [ ] Incremental clustering (online learning) +- [ ] Cluster visualization exports Vote on features or propose new ones in [GitHub Discussions](https://github.com/coderdayton/simplevecdb/discussions). diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 167a92d..4581f2a 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -5,6 +5,107 @@ 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.2.0] - 2026-01-17 + +### Added + +- **Vector Clustering & Auto-Tagging** - Discover natural groupings in embeddings + - `VectorCollection.cluster()` - Cluster documents by semantic similarity + - **K-means**: Classic centroid-based clustering for balanced clusters + - **MiniBatch K-means**: Scalable variant for large datasets (default) + - **HDBSCAN**: Density-based clustering that auto-discovers cluster count + - `VectorCollection.auto_tag()` - Generate descriptive tags for clusters + - TF-IDF method (default): Extract keywords with highest TF-IDF scores + - Frequency method: Extract most common words per cluster + - Custom callback: Implement custom tagging logic (e.g., LLM-based) + - `VectorCollection.assign_cluster_metadata()` - Persist cluster IDs to document metadata + - `VectorCollection.get_cluster_members()` - Retrieve all documents in a cluster + +- **Cluster Quality Metrics** - Evaluate clustering results + - `ClusterResult.inertia` - Sum of squared distances to centroids (K-means only, lower is better) + - `ClusterResult.silhouette_score` - Cluster separation metric (-1 to 1, higher is better) + - `ClusterResult.metrics()` - Get all metrics as dictionary + +- **Cluster Persistence** - Save and reuse cluster configurations + - `VectorCollection.save_cluster()` - Save cluster centroids and metadata to database + - `VectorCollection.load_cluster()` - Load saved cluster configuration + - `VectorCollection.list_clusters()` - List all saved cluster configurations + - `VectorCollection.delete_cluster()` - Delete a saved cluster configuration + - `VectorCollection.assign_to_cluster()` - Assign new documents to saved clusters without re-clustering + +- **Async Clustering Support** - Full async/await parity for all clustering operations + - `AsyncVectorCollection.cluster()`, `auto_tag()`, `assign_cluster_metadata()`, `get_cluster_members()` + - `AsyncVectorCollection.save_cluster()`, `load_cluster()`, `list_clusters()`, `delete_cluster()`, `assign_to_cluster()` + +- **New Dependencies** - Now included in standard installation + - `scikit-learn>=1.3.0` - K-means, MiniBatch K-means, silhouette score + - `hdbscan>=0.8.33` - Density-based clustering + - `sqlcipher3-binary>=0.5.0` - Encryption support (previously optional) + - `cryptography>=41.0` - Encryption utilities (previously optional) + +- **Documentation** + - New comprehensive clustering guide: `docs/guides/clustering.md` + - Algorithm comparison and selection guide + - Quality metrics interpretation + - Cluster persistence workflows + - Use cases: product categorization, topic discovery, customer segmentation, duplicate detection + - Best practices and troubleshooting + - New types reference: `docs/api/types.md` + - Complete `ClusterResult` API documentation + - `Document`, `DistanceStrategy`, `Quantization`, `ClusterAlgorithm` reference + - Updated README.md and docs/index.md with clustering sections + - Enhanced `docs/api/core.md` with clustering examples + +### Changed + +- **pyproject.toml**: Updated `scikit-learn` minimum version from `1.0` to `1.3.0` for improved clustering stability + +### Testing + +- Added 26 clustering tests in `tests/unit/test_clustering.py`: + - 16 core clustering tests (algorithms, auto-tagging, metadata persistence, edge cases) + - 4 cluster metrics tests (inertia, silhouette, metrics method) + - 6 cluster persistence tests (save/load/list/delete/assign) +- Added 3 async clustering tests in `tests/unit/test_async.py` +- Total test count: 305 (up from 292) + +### Installation + +Clustering and encryption are now included by default: + +```bash +pip install simplevecdb +``` + +No extra installation steps required! + +### Example + +```python +from simplevecdb import VectorDB + +db = VectorDB("products.db") +collection = db.collection("items") + +# Cluster documents +result = collection.cluster(n_clusters=5, algorithm="minibatch_kmeans") + +# Generate tags and persist +tags = collection.auto_tag(result, method="tfidf", n_keywords=3) +collection.assign_cluster_metadata(result, tags) + +# Save for fast assignment of new documents +collection.save_cluster("categories", result, metadata={"tags": tags}) + +# Later: assign new documents without re-clustering +new_ids = collection.add_texts(new_texts, embeddings=new_embeddings) +collection.assign_to_cluster("categories", new_ids) + +# Evaluate quality +print(f"Silhouette Score: {result.silhouette_score:.2f}") # 0.62 +print(f"Inertia: {result.inertia:.2f}") # 1523.45 +``` + ## [2.0.0] - 2025-12-23 ### Breaking Changes diff --git a/docs/api/core.md b/docs/api/core.md index b1f83c9..3d12ca1 100644 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -36,6 +36,15 @@ A named collection of vectors within a database. - get_descendants - get_ancestors - set_parent + - cluster + - auto_tag + - assign_cluster_metadata + - get_cluster_members + - save_cluster + - load_cluster + - list_clusters + - delete_cluster + - assign_to_cluster ## Quick Reference @@ -200,3 +209,118 @@ results = db.search_collections([0.1]*384, parallel=False) |--------|-------------| | `list_collections()` | Names of all initialized collections | | `search_collections(query, collections, k, filter, normalize_scores, parallel)` | Search across multiple collections with merged results | + +### Clustering & Auto-Tagging + +Group similar documents and generate descriptive tags: + +```python +from simplevecdb import VectorDB + +db = VectorDB("app.db") +collection = db.collection("docs") + +# Add documents with embeddings +collection.add_texts(texts, embeddings=embeddings) + +# Cluster documents into groups +result = collection.cluster( + n_clusters=5, + algorithm="minibatch_kmeans", # or "kmeans", "hdbscan" + random_state=42 +) +print(result.summary()) # {0: 42, 1: 38, 2: 15, 3: 3, 4: 2} + +# Generate keyword tags for each cluster +tags = collection.auto_tag(result, n_keywords=5) +# {0: 'machine learning, neural network, deep', 1: 'database, sql, query', ...} + +# Persist cluster assignments to metadata +collection.assign_cluster_metadata(result, tags) + +# Query documents by cluster +ml_docs = collection.get_cluster_members(0) +db_docs = collection.similarity_search(query, filter={"cluster": 1}) + +# Custom tagging callback +def summarize_cluster(texts: list[str]) -> str: + return f"Group of {len(texts)} docs about {texts[0][:20]}..." + +custom_tags = collection.auto_tag(result, method="custom", custom_callback=summarize_cluster) +``` + +| Method | Description | +|--------|-------------| +| `cluster(n_clusters, algorithm, filter, sample_size)` | Cluster documents by embedding similarity | +| `auto_tag(result, method, n_keywords, custom_callback)` | Generate descriptive tags for clusters | +| `assign_cluster_metadata(result, tags, metadata_key)` | Persist cluster IDs to document metadata | +| `get_cluster_members(cluster_id, metadata_key)` | Retrieve all documents in a cluster | +| `save_cluster(name, result, metadata)` | Save cluster centroids for later assignment | +| `load_cluster(name)` | Load saved cluster configuration | +| `list_clusters()` | List all saved cluster configurations | +| `delete_cluster(name)` | Delete a saved cluster configuration | +| `assign_to_cluster(name, doc_ids, metadata_key)` | Assign documents to saved clusters | + +**Algorithms:** + +| Algorithm | Best For | Requires n_clusters | +|-----------|----------|-------------------| +| `minibatch_kmeans` | Large datasets (default) | Yes | +| `kmeans` | Small datasets, precise centroids | Yes | +| `hdbscan` | Unknown cluster count, density-based | No | + +Clustering is included in the standard installation (no extras needed). + +### Cluster Metrics + +Access clustering quality metrics to evaluate results: + +```python +result = collection.cluster(n_clusters=5, random_state=42) + +# Inertia (K-means only): sum of squared distances to centroids +# Lower is better; indicates tighter clusters +print(f"Inertia: {result.inertia}") + +# Silhouette score: measure of cluster separation (-1 to 1) +# Higher is better; >0.5 indicates good clustering +print(f"Silhouette: {result.silhouette_score}") + +# Get all metrics as dict +metrics = result.metrics() +# {'inertia': 1523.45, 'silhouette_score': 0.62} +``` + +### Cluster Persistence + +Save cluster configurations for fast assignment of new documents: + +```python +# 1. Cluster your documents +result = collection.cluster(n_clusters=5, random_state=42) +tags = collection.auto_tag(result) + +# 2. Save cluster state (centroids + metadata) +collection.save_cluster( + "product_categories", + result, + metadata={"tags": tags, "version": 1} +) + +# 3. Later: assign new documents without re-clustering +new_ids = collection.add_texts(new_texts, embeddings=new_embeddings) +collection.assign_to_cluster("product_categories", new_ids) + +# List saved clusters +clusters = collection.list_clusters() +# [{'name': 'product_categories', 'n_clusters': 5, 'algorithm': 'minibatch_kmeans', ...}] + +# Load cluster for inspection +saved = collection.load_cluster("product_categories") +if saved: + result, meta = saved + print(f"Loaded {result.n_clusters} clusters, tags: {meta['tags']}") + +# Delete when no longer needed +collection.delete_cluster("product_categories") +``` diff --git a/docs/api/types.md b/docs/api/types.md new file mode 100644 index 0000000..c97f3d6 --- /dev/null +++ b/docs/api/types.md @@ -0,0 +1,202 @@ +# Types Reference + +Core type definitions and data classes used throughout SimpleVecDB. + +## Document + +::: simplevecdb.types.Document + options: + show_root_heading: true + show_source: false + +Represents a document with text content, embeddings, and metadata. + +**Example:** +```python +from simplevecdb.types import Document + +doc = Document( + id=1, + page_content="Paris is the capital of France.", + metadata={"category": "geography", "verified": True}, + embedding=[0.1, 0.2, 0.3, ...] +) +``` + +## ClusterResult + +::: simplevecdb.types.ClusterResult + options: + show_root_heading: true + show_source: false + members: + - n_clusters + - labels + - centroids + - algorithm + - inertia + - silhouette_score + - summary + - metrics + +Result of clustering operation with quality metrics. + +**Example:** +```python +result = collection.cluster(n_clusters=5) + +print(f"Clusters: {result.n_clusters}") +print(f"Algorithm: {result.algorithm}") +print(f"Silhouette: {result.silhouette_score:.2f}") +print(f"Inertia: {result.inertia:.2f}") + +# Cluster size distribution +print(result.summary()) +# {0: 42, 1: 38, 2: 15, 3: 3, 4: 2} + +# All metrics +metrics = result.metrics() +# {'inertia': 1523.45, 'silhouette_score': 0.62} +``` + +### Fields + +| Field | Type | Description | +|-------|------|-------------| +| `n_clusters` | `int` | Number of clusters discovered | +| `labels` | `np.ndarray` | Cluster ID for each document (shape: `[n_docs]`) | +| `centroids` | `np.ndarray \| None` | Cluster centroids (shape: `[n_clusters, dim]`). `None` for HDBSCAN. | +| `algorithm` | `ClusterAlgorithm` | Algorithm used: `"kmeans"`, `"minibatch_kmeans"`, or `"hdbscan"` | +| `inertia` | `float \| None` | Sum of squared distances to centroids (K-means only, lower is better) | +| `silhouette_score` | `float \| None` | Cluster separation metric (-1 to 1, higher is better) | + +### Methods + +#### `summary() -> dict[int, int]` + +Returns cluster size distribution. + +```python +result = collection.cluster(n_clusters=3) +sizes = result.summary() +# {0: 50, 1: 30, 2: 20} +``` + +#### `metrics() -> dict[str, float | None]` + +Returns all quality metrics as a dictionary. + +```python +metrics = result.metrics() +# {'inertia': 1523.45, 'silhouette_score': 0.62} +``` + +## Enums + +### DistanceStrategy + +::: simplevecdb.types.DistanceStrategy + options: + show_root_heading: true + show_source: false + +Distance metrics for vector similarity. + +| Value | Description | Use Case | +|-------|-------------|----------| +| `COSINE` | Cosine similarity (default) | Text embeddings, normalized vectors | +| `EUCLIDEAN` | L2 distance | Image embeddings, spatial data | +| `INNER_PRODUCT` | Dot product | Pre-normalized embeddings | + +**Example:** +```python +from simplevecdb import VectorDB, DistanceStrategy + +collection = db.collection("docs", distance_strategy=DistanceStrategy.COSINE) +``` + +### Quantization + +::: simplevecdb.types.Quantization + options: + show_root_heading: true + show_source: false + +Vector compression strategies. + +| Value | Precision | Compression | Speed | Use Case | +|-------|-----------|-------------|-------|----------| +| `FLOAT` | 32-bit | 1x | Baseline | High precision required | +| `FLOAT16` | 16-bit | 2x | Fast | Recommended default | +| `INT8` | 8-bit | 4x | Faster | Large collections | +| `BIT` | 1-bit | 32x | Fastest | Massive scale, approximate search | + +**Example:** +```python +from simplevecdb import VectorDB, Quantization + +# 2x memory savings, minimal precision loss +collection = db.collection("docs", quantization=Quantization.FLOAT16) + +# 32x compression for massive scale +collection = db.collection("docs", quantization=Quantization.BIT) +``` + +### ClusterAlgorithm + +::: simplevecdb.types.ClusterAlgorithm + options: + show_root_heading: true + show_source: false + +Clustering algorithms. + +| Value | Description | Requires n_clusters | Provides Centroids | +|-------|-------------|--------------------|--------------------| +| `kmeans` | Classic K-means | Yes | Yes | +| `minibatch_kmeans` | Scalable K-means (default) | Yes | Yes | +| `hdbscan` | Density-based clustering | No | No | + +**Example:** +```python +# Auto-discover cluster count +result = collection.cluster(algorithm="hdbscan", min_cluster_size=10) + +# Fixed cluster count with centroids +result = collection.cluster(n_clusters=5, algorithm="minibatch_kmeans") +``` + +## Type Aliases + +### EmbeddingVector + +```python +EmbeddingVector = list[float] | np.ndarray +``` + +Represents a single embedding vector. + +### MetadataDict + +```python +MetadataDict = dict[str, Any] +``` + +Document metadata with arbitrary key-value pairs. + +### FilterDict + +```python +FilterDict = dict[str, Any] +``` + +Metadata filter for search operations. Supports equality matching. + +**Example:** +```python +results = collection.similarity_search( + query_vector, + k=10, + filter={"category": "tech", "verified": True} +) +``` diff --git a/docs/guides/clustering.md b/docs/guides/clustering.md new file mode 100644 index 0000000..c7aa907 --- /dev/null +++ b/docs/guides/clustering.md @@ -0,0 +1,403 @@ +# Vector Clustering Guide + +Vector clustering discovers natural groupings in your embeddings, enabling automatic categorization, topic discovery, and semantic organization of documents. + +## Installation + +Clustering is included in the standard installation: + +```bash +pip install simplevecdb +``` + +**Dependencies included:** +- `scikit-learn>=1.3.0` — K-means and MiniBatch K-means algorithms +- `hdbscan>=0.8.33` — Density-based clustering + +No extra installation steps required! + +## Quick Start + +```python +from simplevecdb import VectorDB + +db = VectorDB("products.db") +collection = db.collection("items") + +# Add documents with embeddings +collection.add_texts(texts=descriptions, embeddings=embeddings) + +# Cluster into 5 groups +result = collection.cluster(n_clusters=5) + +# Auto-generate descriptive tags +tags = collection.auto_tag(result, method="tfidf", n_keywords=3) +# {'0': ['electronics', 'wireless', 'bluetooth'], '1': ['clothing', ...], ...} + +# Persist cluster IDs to document metadata +collection.assign_cluster_metadata(result, tags) + +# Retrieve all documents in cluster 0 +docs = collection.get_cluster_members(0) +``` + +## Algorithms + +### K-Means (`kmeans`) + +Classic centroid-based clustering. Best for balanced, spherical clusters. + +```python +result = collection.cluster( + n_clusters=5, + algorithm="kmeans", + random_state=42 # Reproducible results +) +``` + +**Pros:** +- Fast and deterministic +- Works well with balanced clusters +- Provides cluster centroids for assignment + +**Cons:** +- Requires specifying `n_clusters` upfront +- Sensitive to outliers +- Assumes spherical cluster shapes + +**Best for:** Product categorization, customer segmentation, content organization + +### MiniBatch K-Means (`minibatch_kmeans`, default) + +Scalable variant of K-means using mini-batches. 3-10x faster on large datasets. + +```python +result = collection.cluster( + n_clusters=10, + algorithm="minibatch_kmeans", + sample_size=5000, # Use subset for speed + random_state=42 +) +``` + +**Pros:** +- Scales to millions of documents +- Memory-efficient +- Nearly identical quality to K-means + +**Cons:** +- Slightly less stable than K-means +- Still requires `n_clusters` + +**Best for:** Large-scale document clustering, real-time categorization + +### HDBSCAN (`hdbscan`) + +Density-based clustering that automatically discovers cluster count and handles noise. + +```python +result = collection.cluster( + algorithm="hdbscan", + min_cluster_size=10 # Minimum documents per cluster +) +``` + +**Pros:** +- Automatically determines optimal cluster count +- Handles noise (assigns label `-1` to outliers) +- Discovers non-spherical clusters + +**Cons:** +- Slower than K-means variants +- No centroids (cannot assign new documents) +- Requires tuning `min_cluster_size` + +**Best for:** Exploratory analysis, topic discovery, anomaly detection + +## Cluster Quality Metrics + +Evaluate clustering quality with built-in metrics: + +```python +result = collection.cluster(n_clusters=5) + +# Silhouette Score: -1 to 1 (higher is better) +# Measures how well-separated clusters are +print(f"Silhouette: {result.silhouette_score:.2f}") +# > 0.7: Strong clustering +# 0.5-0.7: Reasonable clustering +# < 0.5: Weak clustering + +# Inertia: Sum of squared distances to centroids (K-means only) +# Lower is better (indicates tighter clusters) +print(f"Inertia: {result.inertia:.2f}") + +# Get all metrics as dict +metrics = result.metrics() +# {'inertia': 1523.45, 'silhouette_score': 0.62} +``` + +## Auto-Tagging + +Generate human-readable labels for clusters: + +### TF-IDF Method (default) + +Extracts keywords with highest TF-IDF scores per cluster. + +```python +tags = collection.auto_tag(result, method="tfidf", n_keywords=5) +# {'0': ['machine', 'learning', 'neural', 'network', 'deep'], ...} +``` + +**Best for:** Text documents with distinct vocabulary per cluster + +### Frequency Method + +Extracts most common words per cluster. + +```python +tags = collection.auto_tag(result, method="frequency", n_keywords=3) +``` + +**Best for:** Short documents, social media posts + +### Custom Callback + +Implement custom tagging logic: + +```python +def custom_tagger(cluster_id: int, texts: list[str]) -> list[str]: + # Your logic here (e.g., LLM-based summarization) + return ["tag1", "tag2", "tag3"] + +tags = collection.auto_tag(result, custom_callback=custom_tagger) +``` + +## Cluster Persistence + +Save cluster configurations for fast assignment of new documents without re-clustering. + +### Save Cluster State + +```python +result = collection.cluster(n_clusters=5) +tags = collection.auto_tag(result) + +collection.save_cluster( + "product_categories", + result, + metadata={"tags": tags, "version": 1, "created_at": "2026-01-17"} +) +``` + +### Load Cluster State + +```python +loaded = collection.load_cluster("product_categories") +if loaded: + result, metadata = loaded + print(f"Loaded {result.n_clusters} clusters") + print(f"Tags: {metadata['tags']}") +``` + +### Assign New Documents + +```python +# Add new documents +new_ids = collection.add_texts(new_texts, embeddings=new_embeddings) + +# Assign to nearest cluster centroids +assigned_count = collection.assign_to_cluster("product_categories", new_ids) +print(f"Assigned {assigned_count} documents") + +# Retrieve assigned documents +docs = collection.get_cluster_members(0) +``` + +### List and Delete + +```python +# List all saved clusters +clusters = collection.list_clusters() +for c in clusters: + print(f"{c['name']}: {c['n_clusters']} clusters, {c['algorithm']}") + +# Delete when no longer needed +collection.delete_cluster("product_categories") +``` + +## Filtering and Sampling + +Cluster subsets of your collection: + +```python +# Cluster only verified documents +result = collection.cluster( + n_clusters=3, + filter={"verified": True} +) + +# Use random sample for speed (large collections) +result = collection.cluster( + n_clusters=10, + sample_size=10000 # Cluster 10k random documents +) +``` + +## Async Support + +All clustering methods have async equivalents: + +```python +from simplevecdb import AsyncVectorDB + +async with AsyncVectorDB("products.db") as db: + collection = db.collection("items") + + result = await collection.cluster(n_clusters=5) + tags = await collection.auto_tag(result) + await collection.save_cluster("categories", result) + + new_ids = await collection.add_texts(texts, embeddings=embeddings) + await collection.assign_to_cluster("categories", new_ids) +``` + +## Use Cases + +### Product Categorization + +```python +# Cluster products by description embeddings +result = collection.cluster(n_clusters=20, algorithm="minibatch_kmeans") +tags = collection.auto_tag(result, n_keywords=5) +collection.assign_cluster_metadata(result, tags) + +# Save for new products +collection.save_cluster("product_taxonomy", result, metadata={"tags": tags}) +``` + +### Topic Discovery + +```python +# Let HDBSCAN discover natural topics +result = collection.cluster(algorithm="hdbscan", min_cluster_size=50) +tags = collection.auto_tag(result, method="tfidf", n_keywords=10) + +# Analyze cluster sizes +for cluster_id in range(result.n_clusters): + docs = collection.get_cluster_members(cluster_id) + print(f"Topic {cluster_id}: {len(docs)} docs - {tags[str(cluster_id)]}") +``` + +### Customer Segmentation + +```python +# Cluster customer profiles +result = collection.cluster(n_clusters=8, random_state=42) +collection.assign_cluster_metadata(result) + +# Target marketing campaigns per segment +segment_0_customers = collection.get_cluster_members(0) +``` + +### Duplicate Detection + +```python +# Use high cluster count to find near-duplicates +result = collection.cluster(n_clusters=1000, algorithm="minibatch_kmeans") +collection.assign_cluster_metadata(result) + +# Find potential duplicates in same cluster +for cluster_id in range(result.n_clusters): + docs = collection.get_cluster_members(cluster_id) + if len(docs) > 1: + print(f"Potential duplicates in cluster {cluster_id}: {len(docs)} docs") +``` + +## Best Practices + +### Choosing Cluster Count + +1. **Elbow Method**: Plot inertia vs. `n_clusters`, look for "elbow" +2. **Silhouette Analysis**: Maximize silhouette score +3. **Domain Knowledge**: Use business requirements (e.g., 10 product categories) +4. **HDBSCAN**: Let algorithm decide + +```python +# Elbow method +inertias = [] +for k in range(2, 20): + result = collection.cluster(n_clusters=k) + inertias.append(result.inertia) +# Plot and find elbow +``` + +### Performance Optimization + +```python +# Large collections: use sampling + MiniBatch K-means +result = collection.cluster( + n_clusters=50, + algorithm="minibatch_kmeans", + sample_size=50000 +) + +# Small collections: use K-means for stability +result = collection.cluster( + n_clusters=5, + algorithm="kmeans", + random_state=42 +) +``` + +### Reproducibility + +Always set `random_state` for deterministic results: + +```python +result = collection.cluster(n_clusters=5, random_state=42) +``` + +### Metadata Organization + +Use consistent metadata keys: + +```python +collection.assign_cluster_metadata(result, tags, metadata_key="category") +collection.assign_cluster_metadata(result, tags, metadata_key="topic") +``` + +## API Reference + +See [VectorCollection API](../api/core.md#clustering--auto-tagging) for complete method signatures and parameters. + +## Troubleshooting + +**ValueError: n_clusters must be >= 2** + +K-means requires at least 2 clusters. Use HDBSCAN for single-cluster detection. + +**Silhouette score is None** + +Occurs when clustering produces only 1 cluster or all documents in separate clusters. Adjust `n_clusters` or `min_cluster_size`. + +**HDBSCAN assigns all documents to noise (cluster -1)** + +Decrease `min_cluster_size` or increase document count. HDBSCAN needs sufficient density. + +**Slow clustering on large collections** + +Use `sample_size` parameter or switch to `minibatch_kmeans`: + +```python +result = collection.cluster(n_clusters=10, sample_size=10000) +``` + +**Cluster assignments change between runs** + +Set `random_state` for reproducibility: + +```python +result = collection.cluster(n_clusters=5, random_state=42) +``` diff --git a/docs/index.md b/docs/index.md index 108fa06..bdef15c 100644 --- a/docs/index.md +++ b/docs/index.md @@ -46,16 +46,20 @@ SimpleVecDB brings **Chroma-like simplicity** to a single **SQLite file**. Built ## Installation ```bash -# Core library only (lightweight, 50MB) +# Standard installation (includes clustering, encryption) pip install simplevecdb -# With local embeddings server + HuggingFace models (500MB+) +# With local embeddings server (adds 500MB+ models) pip install "simplevecdb[server]" - -# With encryption support (SQLCipher) -pip install "simplevecdb[encryption]" ``` +**What's included by default:** +- Vector search with HNSW indexing +- Clustering (K-means, MiniBatch K-means, HDBSCAN) +- Encryption (SQLCipher AES-256) +- Async support +- LangChain & LlamaIndex integrations + **Verify Installation:** ```bash @@ -247,6 +251,23 @@ child_ids = collection.add_texts( children = collection.get_children(parent_ids[0]) ``` +### 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) +``` + +See [Clustering Guide](guides/clustering.md) for algorithms, metrics, and use cases. + ## Feature Matrix | Feature | Status | Description | @@ -266,6 +287,7 @@ children = collection.get_children(parent_ids[0]) | **Built-in Encryption** | ✅ | SQLCipher AES-256 at-rest encryption via `[encryption]` | | **Streaming Insert** | ✅ | Memory-efficient large-scale ingestion with progress | | **Document Hierarchies** | ✅ | Parent/child relationships for chunked docs | +| **Vector Clustering** | ✅ | K-means, MiniBatch K-means, HDBSCAN with auto-tagging | ## Performance Benchmarks @@ -339,8 +361,8 @@ pip install torch --index-url https://download.pytorch.org/whl/cu118 - [x] SQLCipher encryption (at-rest data protection) - [x] Streaming insert API for large-scale ingestion - [x] Hierarchical document relationships (parent/child) -- [ ] Cross-collection search -- [ ] Vector clustering and auto-tagging +- [x] Cross-collection search +- [x] Vector clustering and auto-tagging (v2.2) Vote on features or propose new ones in [GitHub Discussions](https://github.com/coderdayton/simplevecdb/discussions). diff --git a/mkdocs.yml b/mkdocs.yml index 4fbbf48..79e2ccc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -39,6 +39,8 @@ nav: - Changelog: CHANGELOG.md - Contributing: CONTRIBUTING.md - Setup: ENV_SETUP.md + - Guides: + - Clustering: guides/clustering.md - API: - VectorDB: api/core.md - Async API: api/async.md @@ -46,6 +48,7 @@ nav: - Encryption: api/encryption.md - Integrations: api/integrations.md - Configuration: api/config.md + - Types: api/types.md - Engine: - Catalog: api/engine/catalog.md - Search: api/engine/search.md diff --git a/pyproject.toml b/pyproject.toml index ba02fcb..bf2ec7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,12 +8,31 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ - "usearch>=2.12", - "numpy>=2.0", - "python-dotenv>=1.2.1", - "psutil>=5.9.0", + "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 + "cryptography>=41.0", # Encryption utilities ] +[project.optional-dependencies] +server = [ + "fastapi>=0.115", + "uvicorn[standard]>=0.30", + "sentence-transformers>=5.0", # 500MB+ models +] +dev = [ + "pytest>=7.0", + "pytest-asyncio>=0.21", + "pytest-cov>=4.0", + "black>=23.0", + "ruff>=0.1.0", + "mypy>=1.0", +] +examples = ["ollama"] + [dependency-groups] dev = [ "pytest>=8.0", @@ -39,20 +58,10 @@ dev = [ "pymdown-extensions>=10.0", "sqlcipher3-binary>=0.5.0", "cryptography>=41.0", + "scikit-learn>=1.3.0", + "hdbscan>=0.8.33", ] -[project.optional-dependencies] -server = [ - "fastapi>=0.115", - "uvicorn[standard]>=0.30", - "sentence-transformers>=5.0", -] -encryption = [ - "sqlcipher3-binary>=0.5.0", - "cryptography>=41.0", -] -examples = ["ollama"] - [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index f1e0866..8c5fcc1 100644 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -213,6 +213,185 @@ async def remove_texts( lambda: self._collection.remove_texts(texts, filter), ) + # ───────────────────────────────────────────────────────────────────────── + # Clustering Methods (Async) + # ───────────────────────────────────────────────────────────────────────── + + async def cluster( + self, + n_clusters: int | None = None, + algorithm: str = "minibatch_kmeans", + *, + filter: dict[str, Any] | None = None, + sample_size: int | None = None, + min_cluster_size: int = 5, + random_state: int | None = None, + ) -> Any: + """ + Cluster documents by their embeddings (async). + + See VectorCollection.cluster for full documentation. + """ + + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._executor, + lambda: self._collection.cluster( + n_clusters, + algorithm, # type: ignore[arg-type] + filter=filter, + sample_size=sample_size, + min_cluster_size=min_cluster_size, + random_state=random_state, + ), + ) + + async def auto_tag( + self, + cluster_result: Any, + *, + method: str = "keywords", + n_keywords: int = 5, + custom_callback: Any = None, + ) -> dict[int, str]: + """ + Generate descriptive tags for clusters (async). + + 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, + ), + ) + + async def assign_cluster_metadata( + self, + cluster_result: Any, + tags: dict[int, str] | None = None, + *, + metadata_key: str = "cluster", + tag_key: str = "cluster_tag", + ) -> int: + """ + Persist cluster assignments to metadata (async). + + 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, + ), + ) + + async def get_cluster_members( + self, + cluster_id: int, + *, + metadata_key: str = "cluster", + ) -> list[Document]: + """ + Get all documents in a cluster (async). + + 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 + ), + ) + + async def save_cluster( + self, + name: str, + cluster_result: Any, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """ + Save cluster state for later assignment (async). + + 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 + ), + ) + + async def load_cluster( + self, + name: str, + ) -> tuple[Any, dict[str, Any]] | None: + """ + Load saved cluster state (async). + + 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), + ) + + async def list_clusters(self) -> list[dict[str, Any]]: + """ + List all saved cluster configurations (async). + + See VectorCollection.list_clusters for full documentation. + """ + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._executor, + self._collection.list_clusters, + ) + + async def delete_cluster(self, name: str) -> bool: + """ + Delete a saved cluster configuration (async). + + 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), + ) + + async def assign_to_cluster( + self, + name: str, + doc_ids: list[int], + *, + metadata_key: str = "cluster", + ) -> int: + """ + Assign documents to a saved cluster (async). + + 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 + ), + ) + class AsyncVectorDB: """ diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 579f370..37da53e 100644 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -28,12 +28,15 @@ MigrationRequiredError, StreamingProgress, ProgressCallback, + ClusterResult, + ClusterTagCallback, ) from .utils import _import_optional from .engine.quantization import QuantizationStrategy from .engine.search import SearchEngine from .engine.catalog import CatalogManager from .engine.usearch_index import UsearchIndex +from .engine.clustering import ClusterEngine, ClusterAlgorithm from . import constants from .encryption import ( create_encrypted_connection, @@ -984,6 +987,366 @@ def set_parent(self, doc_id: int, parent_id: int | None) -> bool: """ return self._catalog.set_parent(doc_id, parent_id) + # ───────────────────────────────────────────────────────────────────────── + # Clustering Methods + # ───────────────────────────────────────────────────────────────────────── + + def cluster( + self, + n_clusters: int | None = None, + algorithm: ClusterAlgorithm = "minibatch_kmeans", + *, + filter: dict[str, Any] | None = None, + sample_size: int | None = None, + min_cluster_size: int = 5, + random_state: int | None = None, + ) -> ClusterResult: + """ + Cluster documents in the collection by their embeddings. + + Requires scikit-learn and hdbscan (included in the standard install). + + Args: + n_clusters: Number of clusters (required for kmeans/minibatch_kmeans). + algorithm: Clustering algorithm - 'kmeans', 'minibatch_kmeans', or 'hdbscan'. + filter: Optional metadata filter to cluster a subset of documents. + sample_size: If set, cluster a random sample and assign rest to nearest centroid. + min_cluster_size: Minimum cluster size (HDBSCAN only). + random_state: Random seed for reproducibility. + + Returns: + ClusterResult with labels, centroids, and doc_ids. + + Raises: + ImportError: If scikit-learn or hdbscan (for HDBSCAN) not installed. + ValueError: If n_clusters required but not provided. + + Example: + >>> result = collection.cluster(n_clusters=5) + >>> print(result.summary()) # {0: 42, 1: 38, 2: 20, ...} + """ + engine = ClusterEngine() + + doc_ids = list(self._index.keys()) + if not doc_ids: + return ClusterResult( + labels=np.array([], dtype=np.int32), + centroids=None, + doc_ids=[], + n_clusters=0, + algorithm=algorithm, + ) + + if filter: + filtered_ids = set( + self._catalog.find_ids_by_filter( + filter, self._catalog.build_filter_clause + ) + ) + doc_ids = [d for d in doc_ids if d in filtered_ids] + + vectors = self._index.get(np.array(doc_ids, dtype=np.uint64)) + + effective_n_clusters = n_clusters + if n_clusters is not None and algorithm in ("kmeans", "minibatch_kmeans"): + effective_n_clusters = min(n_clusters, len(doc_ids)) + + if sample_size and sample_size < len(doc_ids): + rng = np.random.default_rng(random_state) + sample_indices = rng.choice(len(doc_ids), sample_size, replace=False) + sample_ids = [doc_ids[i] for i in sample_indices] + sample_vectors = vectors[sample_indices] + + result = engine.cluster_vectors( + sample_vectors, + sample_ids, + algorithm=algorithm, + n_clusters=effective_n_clusters, + min_cluster_size=min_cluster_size, + random_state=random_state, + ) + + if result.centroids is not None: + remaining_mask = np.ones(len(doc_ids), dtype=bool) + remaining_mask[sample_indices] = False + remaining_ids = [doc_ids[i] for i, m in enumerate(remaining_mask) if m] + remaining_vectors = vectors[remaining_mask] + + remaining_labels = engine.assign_to_nearest_centroid( + remaining_vectors, result.centroids + ) + + all_ids = sample_ids + remaining_ids + all_labels = np.concatenate([result.labels, remaining_labels]) + + order = np.argsort(all_ids) + return ClusterResult( + labels=all_labels[order], + centroids=result.centroids, + doc_ids=[all_ids[i] for i in order], + n_clusters=result.n_clusters, + algorithm=algorithm, + ) + return result + + return engine.cluster_vectors( + vectors, + doc_ids, + algorithm=algorithm, + n_clusters=effective_n_clusters, + min_cluster_size=min_cluster_size, + random_state=random_state, + ) + + def auto_tag( + self, + cluster_result: ClusterResult, + *, + method: str = "keywords", + n_keywords: int = 5, + custom_callback: ClusterTagCallback | None = None, + ) -> dict[int, str]: + """ + Generate descriptive tags for each cluster. + + Args: + cluster_result: Result from cluster() method. + method: Tagging method - 'keywords' (TF-IDF) or 'custom'. + n_keywords: Number of keywords per cluster (for 'keywords' method). + custom_callback: Custom function (texts: list[str]) -> str for 'custom' method. + + Returns: + Dict mapping cluster_id -> tag string. + + Example: + >>> result = collection.cluster(n_clusters=3) + >>> tags = collection.auto_tag(result) + >>> print(tags) # {0: 'machine learning, neural', 1: 'database, sql', ...} + """ + docs = self._catalog.get_documents_by_ids(cluster_result.doc_ids) + + cluster_texts: dict[int, list[str]] = {} + for doc_id, label in zip(cluster_result.doc_ids, cluster_result.labels): + label_int = int(label) + if label_int not in cluster_texts: + cluster_texts[label_int] = [] + if doc_id in docs: + cluster_texts[label_int].append(docs[doc_id][0]) + + if method == "custom" and custom_callback: + return { + cluster_id: custom_callback(texts) + for cluster_id, texts in cluster_texts.items() + } + + engine = ClusterEngine() + return engine.generate_keywords(cluster_texts, n_keywords) + + def assign_cluster_metadata( + self, + cluster_result: ClusterResult, + tags: dict[int, str] | None = None, + *, + metadata_key: str = "cluster", + tag_key: str = "cluster_tag", + ) -> int: + """ + Persist cluster assignments to document metadata. + + After calling this, you can filter by cluster: filter={"cluster": 2} + + Args: + cluster_result: Result from cluster() method. + tags: Optional cluster tags from auto_tag(). If provided, also sets tag_key. + metadata_key: Metadata key for cluster ID (default: "cluster"). + tag_key: Metadata key for cluster tag (default: "cluster_tag"). + + Returns: + Number of documents updated. + + Example: + >>> result = collection.cluster(n_clusters=5) + >>> tags = collection.auto_tag(result) + >>> collection.assign_cluster_metadata(result, tags) + >>> # Now filter by cluster + >>> docs = collection.similarity_search(query, filter={"cluster": 2}) + """ + updates: list[tuple[int, dict[str, Any]]] = [] + for doc_id, label in zip(cluster_result.doc_ids, cluster_result.labels): + meta: dict[str, Any] = {metadata_key: int(label)} + if tags and int(label) in tags: + meta[tag_key] = tags[int(label)] + updates.append((doc_id, meta)) + + return self._catalog.update_metadata_batch(updates) + + def get_cluster_members( + self, + cluster_id: int, + *, + metadata_key: str = "cluster", + ) -> list[Document]: + """ + Get all documents in a cluster (requires assign_cluster_metadata first). + + Args: + cluster_id: Cluster ID to retrieve. + metadata_key: Metadata key where cluster is stored (default: "cluster"). + + Returns: + List of Documents in the cluster. + """ + rows = self._catalog.get_all_docs_with_text( + filter_dict={metadata_key: cluster_id}, + filter_builder=self._catalog.build_filter_clause, + ) + return [Document(page_content=text, metadata=meta) for _, text, meta in rows] + + def save_cluster( + self, + name: str, + cluster_result: ClusterResult, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """ + Save cluster state for later reuse without re-clustering. + + Persists centroids and algorithm info so new documents can be assigned + to existing clusters using assign_to_cluster(). + + Args: + name: Unique name for this cluster configuration. + cluster_result: Result from cluster() method. + metadata: Optional additional metadata (tags, metrics, etc.). + + Example: + >>> result = collection.cluster(n_clusters=5) + >>> tags = collection.auto_tag(result) + >>> collection.save_cluster("product_categories", result, metadata={"tags": tags}) + """ + centroids_bytes = None + if cluster_result.centroids is not None: + centroids_bytes = cluster_result.centroids.tobytes() + + self._catalog.save_cluster_state( + name=name, + algorithm=cluster_result.algorithm, + n_clusters=cluster_result.n_clusters, + centroids=centroids_bytes, + metadata=metadata, + ) + + def load_cluster(self, name: str) -> tuple[ClusterResult, dict[str, Any]] | None: + """ + Load a saved cluster configuration. + + Args: + name: Name of the saved cluster configuration. + + Returns: + Tuple of (ClusterResult with centroids, metadata dict) or None if not found. + + Example: + >>> saved = collection.load_cluster("product_categories") + >>> if saved: + ... result, meta = saved + ... print(f"Loaded {result.n_clusters} clusters") + """ + state = self._catalog.load_cluster_state(name) + if state is None: + return None + + algorithm, n_clusters, centroids_bytes, metadata = state + + centroids = None + if centroids_bytes is not None: + dim = self._dim + if dim: + centroids = np.frombuffer(centroids_bytes, dtype=np.float32).reshape( + n_clusters, dim + ) + + result = ClusterResult( + labels=np.array([], dtype=np.int32), + centroids=centroids, + doc_ids=[], + n_clusters=n_clusters, + algorithm=algorithm, + ) + return result, metadata + + def list_clusters(self) -> list[dict[str, Any]]: + """List all saved cluster configurations.""" + return self._catalog.list_cluster_states() + + def delete_cluster(self, name: str) -> bool: + """Delete a saved cluster configuration.""" + return self._catalog.delete_cluster_state(name) + + def assign_to_cluster( + self, + name: str, + doc_ids: list[int] | None = None, + *, + metadata_key: str = "cluster", + ) -> int: + """ + Assign documents to clusters using saved centroids. + + Fast assignment without re-clustering - uses nearest centroid matching. + Useful for assigning newly added documents to existing cluster structure. + + Args: + name: Name of saved cluster configuration (from save_cluster). + doc_ids: Document IDs to assign. If None, assigns all unassigned docs. + metadata_key: Metadata key to store cluster assignment. + + Returns: + Number of documents assigned. + + Raises: + ValueError: If cluster not found or has no centroids (HDBSCAN). + + Example: + >>> # Add new documents + >>> new_ids = collection.add_texts(new_texts, embeddings=new_embs) + >>> # Assign to existing clusters + >>> collection.assign_to_cluster("product_categories", new_ids) + """ + saved = self.load_cluster(name) + if saved is None: + raise ValueError(f"Cluster '{name}' not found") + + result, _ = saved + if result.centroids is None: + raise ValueError( + f"Cluster '{name}' has no centroids (HDBSCAN clusters cannot be used for assignment)" + ) + + if doc_ids is None: + all_ids = list(self._index.keys()) + # Get all documents to check for metadata key existence + all_docs = self._catalog.get_all_docs_with_text() + assigned_ids = { + doc_id for doc_id, _, meta in all_docs if metadata_key in meta + } + doc_ids = [d for d in all_ids if d not in assigned_ids] + + if not doc_ids: + return 0 + + vectors = self._index.get(np.array(doc_ids, dtype=np.uint64)) + + engine = ClusterEngine() + labels = engine.assign_to_nearest_centroid(vectors, result.centroids) + + updates = [ + (doc_id, {metadata_key: int(label)}) + for doc_id, label in zip(doc_ids, labels) + ] + return self._catalog.update_metadata_batch(updates) + def count(self) -> int: """Return the number of documents in the collection.""" return self._catalog.count() @@ -1009,7 +1372,6 @@ class VectorDB: Encryption (optional): - SQLite encrypted via SQLCipher (transparent page-level AES-256) - Index files encrypted via AES-256-GCM (at-rest only, zero runtime overhead) - - Install with: pip install simplevecdb[encryption] """ def __init__( @@ -1028,8 +1390,7 @@ def __init__( distance_strategy: Default distance metric for similarity search. quantization: Default vector compression strategy. encryption_key: Optional passphrase or 32-byte key for at-rest encryption. - Requires simplevecdb[encryption] extras. Encrypts both SQLite - (via SQLCipher) and usearch index files (via AES-256-GCM). + 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. @@ -1037,8 +1398,8 @@ def __init__( 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 - simplevecdb[encryption] not installed. + EncryptionUnavailableError: If encryption_key provided but encryption + dependencies are missing. EncryptionError: If encrypted database cannot be opened (wrong key). ValueError: If encryption_key used with ":memory:" database. """ diff --git a/src/simplevecdb/encryption.py b/src/simplevecdb/encryption.py index 264eb93..484fd50 100644 --- a/src/simplevecdb/encryption.py +++ b/src/simplevecdb/encryption.py @@ -19,7 +19,8 @@ db = VectorDB("secure.db", encryption_key=os.urandom(32)) Requirements: - pip install simplevecdb[encryption] + Included in the standard install. If missing, reinstall: + pip install --force-reinstall simplevecdb """ from __future__ import annotations @@ -57,8 +58,8 @@ class EncryptionUnavailableError(ImportError): def __init__(self) -> None: super().__init__( - "Encryption requires additional dependencies. " - "Install with: pip install simplevecdb[encryption]" + "Encryption requires sqlcipher3-binary and cryptography. " + "Reinstall simplevecdb: pip install --force-reinstall simplevecdb" ) diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 8a74547..f905156 100644 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -487,6 +487,87 @@ def count(self) -> int: row = self.conn.execute(f"SELECT COUNT(*) FROM {self._table_name}").fetchone() return row[0] if row else 0 + def get_all_docs_with_text( + self, + filter_dict: dict[str, Any] | None = None, + filter_builder: Callable[[dict[str, Any], str], tuple[str, list[Any]]] + | None = None, + ) -> list[tuple[int, str, dict[str, Any]]]: + """ + Get all documents with their text content. + + Args: + filter_dict: Optional metadata filter + filter_builder: Function to build filter clause + + Returns: + List of (doc_id, text, metadata) tuples + """ + filter_clause = "" + filter_params: list[Any] = [] + if filter_dict and filter_builder: + filter_clause, filter_params = filter_builder(filter_dict, "metadata") + + sql = f""" + SELECT id, text, metadata FROM {self._table_name} + WHERE 1=1 {filter_clause} + ORDER BY id + """ + rows = self.conn.execute(sql, tuple(filter_params)).fetchall() + result = [] + for row_id, text, meta_json in rows: + meta = json.loads(meta_json) if meta_json else {} + result.append((int(row_id), text, meta)) + return result + + def update_metadata_batch(self, updates: list[tuple[int, dict[str, Any]]]) -> int: + """ + Update metadata for multiple documents in a single transaction. + + Merges new metadata with existing metadata (shallow merge). + + Args: + updates: List of (doc_id, metadata_updates) tuples + + Returns: + Number of documents updated + """ + if not updates: + return 0 + + updated = 0 + # Batch into chunks of 500 for performance + for i in range(0, len(updates), 500): + batch = updates[i : i + 500] + ids = [u[0] for u in batch] + + # Fetch all existing metadata in one query + placeholders = ",".join(["?"] * len(ids)) + rows = self.conn.execute( + f"SELECT id, metadata FROM {self._table_name} WHERE id IN ({placeholders})", + ids, + ).fetchall() + + current_meta_map = {r[0]: (json.loads(r[1]) if r[1] else {}) for r in rows} + + # Prepare updates + update_data = [] + for doc_id, meta_updates in batch: + if doc_id in current_meta_map: + meta = current_meta_map[doc_id] + meta.update(meta_updates) + update_data.append((json.dumps(meta), doc_id)) + updated += 1 + + if update_data: + self.conn.executemany( + f"UPDATE {self._table_name} SET metadata = ? WHERE id = ?", + update_data, + ) + + self.conn.commit() + return updated + def check_legacy_sqlite_vec(self, vec_table_name: str) -> bool: """ Check if legacy sqlite-vec tables exist (for migration). @@ -706,3 +787,117 @@ def set_parent(self, doc_id: int, parent_id: int | None) -> bool: (parent_id, doc_id), ) return cursor.rowcount > 0 + + # ------------------------------------------------------------------ # + # Cluster State Persistence + # ------------------------------------------------------------------ # + + def _ensure_cluster_table(self) -> None: + """Create cluster state table if it doesn't exist.""" + cluster_table = f"{self._table_name}_clusters" + self.conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {cluster_table} ( + name TEXT PRIMARY KEY, + algorithm TEXT NOT NULL, + n_clusters INTEGER NOT NULL, + centroids BLOB, + created_at TEXT DEFAULT CURRENT_TIMESTAMP, + metadata TEXT + ) + """ + ) + self.conn.commit() + + def save_cluster_state( + self, + name: str, + algorithm: str, + n_clusters: int, + centroids: bytes | None, + metadata: dict[str, Any] | None = None, + ) -> None: + """ + Save cluster state for later reuse. + + Args: + name: Unique name for this cluster configuration + algorithm: Algorithm used (kmeans, minibatch_kmeans, hdbscan) + n_clusters: Number of clusters + centroids: Serialized centroid array (numpy bytes) + metadata: Additional metadata (inertia, silhouette, etc.) + """ + self._ensure_cluster_table() + cluster_table = f"{self._table_name}_clusters" + + meta_json = json.dumps(metadata) if metadata else None + + self.conn.execute( + f""" + INSERT OR REPLACE INTO {cluster_table} + (name, algorithm, n_clusters, centroids, metadata, created_at) + VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP) + """, + (name, algorithm, n_clusters, centroids, meta_json), + ) + self.conn.commit() + + def load_cluster_state( + self, name: str + ) -> tuple[str, int, bytes | None, dict[str, Any]] | None: + """ + Load saved cluster state. + + Args: + name: Name of the cluster configuration + + Returns: + Tuple of (algorithm, n_clusters, centroids_bytes, metadata) or None + """ + self._ensure_cluster_table() + cluster_table = f"{self._table_name}_clusters" + + row = self.conn.execute( + f"SELECT algorithm, n_clusters, centroids, metadata FROM {cluster_table} WHERE name = ?", + (name,), + ).fetchone() + + if not row: + return None + + algorithm, n_clusters, centroids, meta_json = row + metadata = json.loads(meta_json) if meta_json else {} + return (algorithm, n_clusters, centroids, metadata) + + def list_cluster_states(self) -> list[dict[str, Any]]: + """List all saved cluster configurations.""" + self._ensure_cluster_table() + cluster_table = f"{self._table_name}_clusters" + + rows = self.conn.execute( + f"SELECT name, algorithm, n_clusters, created_at, metadata FROM {cluster_table}" + ).fetchall() + + result = [] + for name, algorithm, n_clusters, created_at, meta_json in rows: + result.append( + { + "name": name, + "algorithm": algorithm, + "n_clusters": n_clusters, + "created_at": created_at, + "metadata": json.loads(meta_json) if meta_json else {}, + } + ) + return result + + def delete_cluster_state(self, name: str) -> bool: + """Delete a saved cluster configuration.""" + self._ensure_cluster_table() + cluster_table = f"{self._table_name}_clusters" + + cursor = self.conn.execute( + f"DELETE FROM {cluster_table} WHERE name = ?", (name,) + ) + self.conn.commit() + return cursor.rowcount > 0 diff --git a/src/simplevecdb/engine/clustering.py b/src/simplevecdb/engine/clustering.py new file mode 100644 index 0000000..ef4f6d5 --- /dev/null +++ b/src/simplevecdb/engine/clustering.py @@ -0,0 +1,248 @@ +"""Vector clustering engine for SimpleVecDB.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Literal + +import numpy as np + +from ..utils import _import_optional + +if TYPE_CHECKING: + from ..types import ClusterResult + +_logger = logging.getLogger("simplevecdb.engine.clustering") + +ClusterAlgorithm = Literal["kmeans", "minibatch_kmeans", "hdbscan"] + + +class ClusterEngine: + """Handles vector clustering and tag generation.""" + + def cluster_vectors( + self, + vectors: np.ndarray, + doc_ids: list[int], + algorithm: ClusterAlgorithm = "minibatch_kmeans", + n_clusters: int | None = None, + *, + min_cluster_size: int = 5, + random_state: int | None = None, + ) -> ClusterResult: + """ + Cluster vectors using the specified algorithm. + + Args: + vectors: 2D array of shape (n_samples, n_features) + doc_ids: Document IDs corresponding to each vector + algorithm: Clustering algorithm to use + n_clusters: Number of clusters (required for kmeans variants) + min_cluster_size: Minimum cluster size (HDBSCAN only) + random_state: Random seed for reproducibility + + Returns: + ClusterResult with labels, centroids, and metadata + """ + from ..types import ClusterResult + + if len(vectors) == 0: + return ClusterResult( + labels=np.array([], dtype=np.int32), + centroids=None, + doc_ids=[], + n_clusters=0, + algorithm=algorithm, + ) + + if algorithm == "hdbscan": + labels, centroids, inertia = self._hdbscan(vectors, min_cluster_size) + elif algorithm == "minibatch_kmeans": + if n_clusters is None: + raise ValueError("n_clusters required for minibatch_kmeans") + labels, centroids, inertia = self._minibatch_kmeans( + vectors, n_clusters, random_state + ) + elif algorithm == "kmeans": + if n_clusters is None: + raise ValueError("n_clusters required for kmeans") + labels, centroids, inertia = self._kmeans(vectors, n_clusters, random_state) + else: + raise ValueError(f"Unknown algorithm: {algorithm}") + + n_clusters_found = len(set(labels)) - (1 if -1 in labels else 0) + + silhouette = self._compute_silhouette(vectors, labels, n_clusters_found) + + return ClusterResult( + labels=labels, + centroids=centroids, + doc_ids=list(doc_ids), + n_clusters=n_clusters_found, + algorithm=algorithm, + inertia=inertia, + silhouette_score=silhouette, + ) + + def _kmeans( + self, + vectors: np.ndarray, + n_clusters: int, + random_state: int | None, + ) -> tuple[np.ndarray, np.ndarray, float]: + """Standard K-means clustering.""" + sklearn_cluster = _import_optional("sklearn.cluster") + if sklearn_cluster is None: + raise ImportError( + "scikit-learn required for clustering. Install with: " + "pip install --force-reinstall simplevecdb" + ) + + kmeans = sklearn_cluster.KMeans( + n_clusters=n_clusters, + random_state=random_state, + n_init="auto", + ) + labels = kmeans.fit_predict(vectors) + return ( + labels.astype(np.int32), + kmeans.cluster_centers_.astype(np.float32), + float(kmeans.inertia_), + ) + + def _minibatch_kmeans( + self, + vectors: np.ndarray, + n_clusters: int, + random_state: int | None, + ) -> tuple[np.ndarray, np.ndarray, float]: + """Mini-batch K-means for large datasets.""" + sklearn_cluster = _import_optional("sklearn.cluster") + if sklearn_cluster is None: + raise ImportError( + "scikit-learn required for clustering. Install with: " + "pip install --force-reinstall simplevecdb" + ) + + batch_size = min(1024, len(vectors)) + kmeans = sklearn_cluster.MiniBatchKMeans( + n_clusters=n_clusters, + random_state=random_state, + batch_size=batch_size, + n_init="auto", + ) + labels = kmeans.fit_predict(vectors) + return ( + labels.astype(np.int32), + kmeans.cluster_centers_.astype(np.float32), + float(kmeans.inertia_), + ) + + def _hdbscan( + self, + vectors: np.ndarray, + min_cluster_size: int, + ) -> tuple[np.ndarray, None, None]: + """HDBSCAN density-based clustering (discovers natural clusters).""" + hdbscan_mod = _import_optional("hdbscan") + if hdbscan_mod is None: + raise ImportError( + "hdbscan required for density-based clustering. Install with: " + "pip install --force-reinstall simplevecdb" + ) + + clusterer = hdbscan_mod.HDBSCAN( + min_cluster_size=min_cluster_size, + metric="euclidean", + ) + labels = clusterer.fit_predict(vectors) + # HDBSCAN has no centroids or inertia + return labels.astype(np.int32), None, None + + def _compute_silhouette( + self, + vectors: np.ndarray, + labels: np.ndarray, + n_clusters: int, + ) -> float | None: + """Compute silhouette score; returns None if invalid clustering or sklearn unavailable.""" + if n_clusters < 2: + return None + + sklearn_metrics = _import_optional("sklearn.metrics") + if sklearn_metrics is None: + return None + + mask = labels >= 0 + valid_vectors = vectors[mask] + valid_labels = labels[mask] + n_valid = len(valid_vectors) + n_unique = len(set(valid_labels)) + + if n_valid < 2 or n_unique < 2 or n_unique >= n_valid: + return None + + return float(sklearn_metrics.silhouette_score(valid_vectors, valid_labels)) + + def generate_keywords( + self, + cluster_texts: dict[int, list[str]], + n_keywords: int = 5, + ) -> dict[int, str]: + """ + Generate keyword tags for clusters using TF-IDF. + + Args: + cluster_texts: Mapping of cluster_id -> list of document texts + n_keywords: Number of keywords per cluster + + Returns: + Mapping of cluster_id -> comma-separated keywords + """ + sklearn_text = _import_optional("sklearn.feature_extraction.text") + if sklearn_text is None: + raise ImportError( + "scikit-learn required for keyword extraction. Install with: " + "pip install --force-reinstall simplevecdb" + ) + + tags: dict[int, str] = {} + + for cluster_id, texts in cluster_texts.items(): + if cluster_id == -1: + tags[cluster_id] = "outliers" + continue + + if not texts: + tags[cluster_id] = f"cluster_{cluster_id}" + continue + + try: + vectorizer = sklearn_text.TfidfVectorizer( + max_features=100, + stop_words="english", + ngram_range=(1, 2), + min_df=1, + max_df=0.95, + ) + tfidf = vectorizer.fit_transform(texts) + feature_names = vectorizer.get_feature_names_out() + + scores = np.asarray(tfidf.sum(axis=0)).flatten() + top_indices = scores.argsort()[-n_keywords:][::-1] + keywords = [feature_names[i] for i in top_indices] + + tags[cluster_id] = ", ".join(keywords) + except ValueError: + tags[cluster_id] = f"cluster_{cluster_id}" + + return tags + + def assign_to_nearest_centroid( + self, + vectors: np.ndarray, + centroids: np.ndarray, + ) -> np.ndarray: + """Assign vectors to nearest centroid (for out-of-sample assignment).""" + distances = np.linalg.norm(vectors[:, np.newaxis] - centroids, axis=2) + return np.argmin(distances, axis=1).astype(np.int32) diff --git a/src/simplevecdb/engine/usearch_index.py b/src/simplevecdb/engine/usearch_index.py index 0f9543a..8bb86e2 100644 --- a/src/simplevecdb/engine/usearch_index.py +++ b/src/simplevecdb/engine/usearch_index.py @@ -414,6 +414,47 @@ def __len__(self) -> int: def __contains__(self, key: int) -> bool: return self.contains(key) + def keys(self) -> list[int]: + """Return all keys in the index.""" + if self._index is None: + return [] + return [int(k) for k in self._index.keys] + + def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: + """ + Retrieve vectors by their keys. + + Args: + keys: Array of keys to retrieve + + Returns: + Array of vectors, shape (len(keys), ndim). Missing keys return zeros. + """ + if self._index is None or len(keys) == 0: + return np.array([], dtype=np.float32).reshape(0, self._ndim or 1) + + keys = np.asarray(keys, dtype=np.uint64) + vectors = [] + missing_logged = False + for key in keys: + try: + vec = self._index.get(int(key)) + vectors.append(np.asarray(vec, dtype=np.float32)) + except KeyError: + if not missing_logged: + _logger.warning( + "Index missing key(s); returning zero vectors. " + "Index may be out of sync with catalog." + ) + missing_logged = True + vectors.append(np.zeros(self._ndim or 1, dtype=np.float32)) + + return ( + np.stack(vectors) + if vectors + else np.array([], dtype=np.float32).reshape(0, self._ndim or 1) + ) + def __del__(self) -> None: try: self.close() diff --git a/src/simplevecdb/types.py b/src/simplevecdb/types.py index c7d5f0c..968246e 100644 --- a/src/simplevecdb/types.py +++ b/src/simplevecdb/types.py @@ -3,7 +3,10 @@ import dataclasses from dataclasses import field from enum import Enum -from typing import Callable, TypedDict +from typing import Callable, TypedDict, TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np class StrEnum(str, Enum): @@ -83,3 +86,40 @@ def __init__( f"⚠️ BACKUP YOUR DATABASE BEFORE MIGRATING: cp {path} {path}.backup" ) super().__init__(msg) + + +@dataclasses.dataclass +class ClusterResult: + """Result of a clustering operation.""" + + labels: np.ndarray + centroids: np.ndarray | None + doc_ids: list[int] + n_clusters: int + algorithm: str + inertia: float | None = None + silhouette_score: float | None = None + + def get_cluster_doc_ids(self, cluster_id: int) -> list[int]: + """Get document IDs belonging to a specific cluster.""" + return [ + doc_id + for doc_id, label in zip(self.doc_ids, self.labels) + if label == cluster_id + ] + + def summary(self) -> dict[int, int]: + """Return cluster_id -> member count mapping.""" + from collections import Counter + + return dict(Counter(int(label) for label in self.labels)) + + def metrics(self) -> dict[str, float | None]: + """Return clustering quality metrics.""" + return { + "inertia": self.inertia, + "silhouette_score": self.silhouette_score, + } + + +ClusterTagCallback = Callable[[list[str]], str] diff --git a/tests/unit/test_async.py b/tests/unit/test_async.py index b899503..fc48a18 100644 --- a/tests/unit/test_async.py +++ b/tests/unit/test_async.py @@ -259,3 +259,89 @@ async def test_async_similarity_search_batch(sample_texts, sample_embeddings): doc, score = query_results[0] assert doc.page_content == sample_texts[i] assert score < 0.01 + + +@pytest.mark.asyncio +async def test_async_cluster_persistence_roundtrip(): + """Test save_cluster and load_cluster async methods.""" + async with AsyncVectorDB(":memory:") as db: + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(20, 384).astype(np.float32) + embeddings /= np.linalg.norm(embeddings, axis=1, keepdims=True) + + await collection.add_texts( + texts=[f"doc_{i}" for i in range(20)], + embeddings=embeddings.tolist(), + ) + + result = await collection.cluster(n_clusters=3, random_state=42) + await collection.save_cluster("test_cluster", result, metadata={"v": 1}) + + loaded = await collection.load_cluster("test_cluster") + assert loaded is not None + loaded_result, loaded_meta = loaded + assert loaded_result.n_clusters == 3 + assert loaded_meta == {"v": 1} + + +@pytest.mark.asyncio +async def test_async_list_and_delete_clusters(): + """Test list_clusters and delete_cluster async methods.""" + async with AsyncVectorDB(":memory:") as db: + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, 384).astype(np.float32) + embeddings /= np.linalg.norm(embeddings, axis=1, keepdims=True) + + await collection.add_texts( + texts=[f"doc_{i}" for i in range(10)], + embeddings=embeddings.tolist(), + ) + + result = await collection.cluster(n_clusters=2, random_state=42) + await collection.save_cluster("cluster_a", result) + await collection.save_cluster("cluster_b", result) + + clusters = await collection.list_clusters() + assert len(clusters) == 2 + names = {c["name"] for c in clusters} + assert names == {"cluster_a", "cluster_b"} + + deleted = await collection.delete_cluster("cluster_a") + assert deleted is True + + clusters = await collection.list_clusters() + assert len(clusters) == 1 + assert clusters[0]["name"] == "cluster_b" + + +@pytest.mark.asyncio +async def test_async_assign_to_cluster(): + """Test assign_to_cluster async method.""" + async with AsyncVectorDB(":memory:") as db: + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, 384).astype(np.float32) + embeddings /= np.linalg.norm(embeddings, axis=1, keepdims=True) + + await collection.add_texts( + texts=[f"doc_{i}" for i in range(10)], + embeddings=embeddings.tolist(), + ) + + result = await collection.cluster(n_clusters=2, random_state=42) + await collection.save_cluster("saved", result) + + new_emb = np.random.randn(3, 384).astype(np.float32) + new_emb /= np.linalg.norm(new_emb, axis=1, keepdims=True) + new_ids = await collection.add_texts( + texts=["new_a", "new_b", "new_c"], + embeddings=new_emb.tolist(), + ) + + assigned = await collection.assign_to_cluster("saved", new_ids) + assert assigned == 3 diff --git a/tests/unit/test_clustering.py b/tests/unit/test_clustering.py new file mode 100644 index 0000000..f77fc06 --- /dev/null +++ b/tests/unit/test_clustering.py @@ -0,0 +1,585 @@ +"""Tests for clustering and auto-tagging functionality.""" + +from __future__ import annotations + +import numpy as np +import pytest +from pathlib import Path +from typing import Literal + +from simplevecdb import VectorDB +from simplevecdb.types import ClusterResult + + +sklearn = pytest.importorskip("sklearn", reason="sklearn required for clustering tests") + + +class TestClustering: + @pytest.fixture + def db_path(self, tmp_path: Path) -> Path: + return tmp_path / "test_clustering.db" + + @pytest.fixture + def dim(self) -> int: + return 32 + + def make_clustered_embeddings( + self, n_per_cluster: int, n_clusters: int, dim: int + ) -> tuple[list[str], np.ndarray]: + np.random.seed(42) + texts = [] + embeddings = [] + for c in range(n_clusters): + for i in range(n_per_cluster): + texts.append(f"cluster_{c}_doc_{i}") + emb = np.random.randn(dim).astype(np.float32) + emb[c] += 10.0 + embeddings.append(emb) + return texts, np.array(embeddings) + + def test_cluster_basic(self, db_path: Path, dim: int): + """Basic clustering returns ClusterResult.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(5, 3, dim) + collection.add_texts(texts, embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=3, random_state=42) + + assert isinstance(result, ClusterResult) + assert result.n_clusters == 3 + assert len(result.labels) == 15 + assert len(result.doc_ids) == 15 + assert result.centroids is not None + assert result.centroids.shape == (3, dim) + assert result.algorithm == "minibatch_kmeans" + + db.close() + + def test_cluster_empty_collection(self, db_path: Path): + """Clustering empty collection returns empty result.""" + db = VectorDB(db_path) + collection = db.collection("test") + + result = collection.cluster(n_clusters=3) + + assert result.n_clusters == 0 + assert len(result.labels) == 0 + assert len(result.doc_ids) == 0 + + db.close() + + def test_cluster_algorithms(self, db_path: Path, dim: int): + """Test different clustering algorithms.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(5, 3, dim) + collection.add_texts(texts, embeddings=embeddings.tolist()) + + algorithms: tuple[Literal["kmeans", "minibatch_kmeans"], ...] = ( + "kmeans", + "minibatch_kmeans", + ) + for algo in algorithms: + result = collection.cluster(n_clusters=3, algorithm=algo, random_state=42) + assert result.n_clusters == 3 + assert result.algorithm == algo + + db.close() + + def test_cluster_with_filter(self, db_path: Path, dim: int): + """Clustering respects metadata filter.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(5, 2, dim) + metadatas = [{"group": "A"}] * 5 + [{"group": "B"}] * 5 + collection.add_texts(texts, embeddings=embeddings.tolist(), metadatas=metadatas) + + result = collection.cluster( + n_clusters=2, filter={"group": "A"}, random_state=42 + ) + + assert len(result.doc_ids) == 5 + + db.close() + + def test_cluster_with_sample_size(self, db_path: Path, dim: int): + """Clustering with sample_size clusters sample and assigns rest.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(10, 3, dim) + collection.add_texts(texts, embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=3, sample_size=9, random_state=42) + + assert len(result.doc_ids) == 30 + assert result.n_clusters == 3 + + db.close() + + def test_cluster_result_summary(self, db_path: Path, dim: int): + """ClusterResult.summary() returns cluster counts.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(5, 3, dim) + collection.add_texts(texts, embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=3, random_state=42) + summary = result.summary() + + assert isinstance(summary, dict) + assert sum(summary.values()) == 15 + + db.close() + + def test_cluster_result_get_cluster_doc_ids(self, db_path: Path, dim: int): + """ClusterResult.get_cluster_doc_ids() returns docs in cluster.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts, embeddings = self.make_clustered_embeddings(5, 3, dim) + collection.add_texts(texts, embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=3, random_state=42) + + all_from_clusters = [] + for cluster_id in range(3): + cluster_docs = result.get_cluster_doc_ids(cluster_id) + all_from_clusters.extend(cluster_docs) + + assert set(all_from_clusters) == set(result.doc_ids) + + db.close() + + +class TestAutoTag: + @pytest.fixture + def db_path(self, tmp_path: Path) -> Path: + return tmp_path / "test_autotag.db" + + def test_auto_tag_keywords(self, db_path: Path): + """auto_tag generates TF-IDF keywords.""" + db = VectorDB(db_path) + collection = db.collection("test") + + texts = [ + "machine learning neural network", + "deep learning artificial intelligence", + "database sql query optimization", + "sql index performance tuning", + ] + np.random.seed(42) + embeddings = np.random.randn(4, 32).astype(np.float32) + embeddings[:2, 0] += 10.0 + embeddings[2:, 1] += 10.0 + + collection.add_texts(texts, embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=2, random_state=42) + tags = collection.auto_tag(result, n_keywords=3) + + assert isinstance(tags, dict) + assert len(tags) == 2 + + db.close() + + def test_auto_tag_custom_callback(self, db_path: Path): + """auto_tag with custom callback.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(4, 32).astype(np.float32) + collection.add_texts(["a", "b", "c", "d"], embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=2, random_state=42) + + def custom_tag(texts: list[str]) -> str: + return f"custom_{len(texts)}" + + tags = collection.auto_tag(result, method="custom", custom_callback=custom_tag) + + for tag in tags.values(): + assert tag.startswith("custom_") + + db.close() + + +class TestClusterMetadataPersistence: + @pytest.fixture + def db_path(self, tmp_path: Path) -> Path: + return tmp_path / "test_cluster_meta.db" + + @pytest.fixture + def dim(self) -> int: + return 32 + + def test_assign_cluster_metadata(self, db_path: Path, dim: int): + """assign_cluster_metadata persists cluster IDs.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(6, dim).astype(np.float32) + embeddings[:3, 0] += 10.0 + embeddings[3:, 1] += 10.0 + + collection.add_texts( + [f"doc_{i}" for i in range(6)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + updated = collection.assign_cluster_metadata(result) + + assert updated == 6 + + db.close() + + def test_assign_cluster_metadata_with_tags(self, db_path: Path, dim: int): + """assign_cluster_metadata persists tags too.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(4, dim).astype(np.float32) + collection.add_texts( + ["ml doc", "ml text", "db doc", "db text"], + embeddings=embeddings.tolist(), + ) + + result = collection.cluster(n_clusters=2, random_state=42) + tags = {0: "machine_learning", 1: "database"} + collection.assign_cluster_metadata(result, tags) + + members = collection.get_cluster_members(0) + for doc in members: + assert "cluster" in doc.metadata + assert "cluster_tag" in doc.metadata + + db.close() + + def test_get_cluster_members(self, db_path: Path, dim: int): + """get_cluster_members retrieves docs by cluster.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(6, dim).astype(np.float32) + embeddings[:3, 0] += 10.0 + embeddings[3:, 1] += 10.0 + + collection.add_texts( + [f"doc_{i}" for i in range(6)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + collection.assign_cluster_metadata(result) + + summary = result.summary() + for cluster_id, count in summary.items(): + members = collection.get_cluster_members(cluster_id) + assert len(members) == count + + db.close() + + def test_custom_metadata_keys(self, db_path: Path, dim: int): + """Custom metadata keys work.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(4, dim).astype(np.float32) + collection.add_texts(["a", "b", "c", "d"], embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=2, random_state=42) + tags = {0: "first", 1: "second"} + collection.assign_cluster_metadata( + result, tags, metadata_key="my_cluster", tag_key="my_tag" + ) + + members = collection.get_cluster_members(0, metadata_key="my_cluster") + for doc in members: + assert "my_cluster" in doc.metadata + assert "my_tag" in doc.metadata + + db.close() + + +class TestClusteringEdgeCases: + @pytest.fixture + def db_path(self, tmp_path: Path) -> Path: + return tmp_path / "test_cluster_edge.db" + + def test_cluster_requires_n_clusters_for_kmeans(self, db_path: Path): + """kmeans algorithms require n_clusters.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(5, 32).astype(np.float32) + collection.add_texts(["a", "b", "c", "d", "e"], embeddings=embeddings.tolist()) + + with pytest.raises(ValueError, match="n_clusters required"): + collection.cluster(algorithm="kmeans") + + db.close() + + def test_cluster_single_document(self, db_path: Path): + """Clustering single doc works.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + collection.add_texts(["only"], embeddings=[np.random.randn(32).tolist()]) + + result = collection.cluster(n_clusters=1, random_state=42) + + assert result.n_clusters == 1 + assert len(result.doc_ids) == 1 + + db.close() + + def test_more_clusters_than_docs(self, db_path: Path): + """Requesting more clusters than docs gives fewer clusters.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(3, 32).astype(np.float32) + collection.add_texts(["a", "b", "c"], embeddings=embeddings.tolist()) + + result = collection.cluster(n_clusters=10, random_state=42) + + assert result.n_clusters <= 3 + + db.close() + + +class TestClusterMetrics: + @pytest.fixture + def db_path(self, tmp_path: Path) -> Path: + return tmp_path / "test_cluster_metrics.db" + + @pytest.fixture + def dim(self) -> int: + return 32 + + def test_cluster_result_has_inertia(self, db_path: Path, dim: int): + """K-means clustering populates inertia metric.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(20, dim).astype(np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(20)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=3, algorithm="kmeans", random_state=42) + + assert result.inertia is not None + assert result.inertia > 0 + + db.close() + + def test_cluster_result_has_silhouette(self, db_path: Path, dim: int): + """Clustering with enough samples populates silhouette score.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(20, dim).astype(np.float32) + embeddings[:10, 0] += 10.0 + embeddings[10:, 1] += 10.0 + + collection.add_texts( + [f"doc_{i}" for i in range(20)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + + assert result.silhouette_score is not None + assert -1 <= result.silhouette_score <= 1 + + db.close() + + def test_cluster_result_metrics_method(self, db_path: Path, dim: int): + """ClusterResult.metrics() returns dict with available metrics.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(20, dim).astype(np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(20)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=3, random_state=42) + metrics = result.metrics() + + assert isinstance(metrics, dict) + assert "inertia" in metrics + assert "silhouette_score" in metrics + assert metrics["inertia"] is not None + assert metrics["inertia"] > 0 + + db.close() + + def test_silhouette_none_for_single_cluster(self, db_path: Path, dim: int): + """Silhouette is None when only one cluster.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(5, dim).astype(np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(5)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=1, random_state=42) + + assert result.silhouette_score is None + + db.close() + + +class TestClusterPersistence: + @pytest.fixture + def db_path(self, tmp_path: Path) -> Path: + return tmp_path / "test_cluster_persist.db" + + @pytest.fixture + def dim(self) -> int: + return 32 + + def test_save_and_load_cluster(self, db_path: Path, dim: int): + """save_cluster and load_cluster round-trip correctly.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, dim).astype(np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(10)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=3, random_state=42) + collection.save_cluster("test_cluster", result, metadata={"version": 1}) + + loaded = collection.load_cluster("test_cluster") + + assert loaded is not None + loaded_result, loaded_meta = loaded + assert loaded_result.n_clusters == 3 + assert loaded_result.algorithm == result.algorithm + assert loaded_result.centroids is not None + assert result.centroids is not None + assert loaded_result.centroids.shape == result.centroids.shape + assert loaded_meta == {"version": 1} + + db.close() + + def test_list_clusters(self, db_path: Path, dim: int): + """list_clusters returns saved cluster info.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, dim).astype(np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(10)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + collection.save_cluster("cluster_a", result) + collection.save_cluster("cluster_b", result, metadata={"tag": "test"}) + + clusters = collection.list_clusters() + + assert len(clusters) == 2 + names = {c["name"] for c in clusters} + assert names == {"cluster_a", "cluster_b"} + + db.close() + + def test_delete_cluster(self, db_path: Path, dim: int): + """delete_cluster removes saved cluster.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, dim).astype(np.float32) + collection.add_texts( + [f"doc_{i}" for i in range(10)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + collection.save_cluster("to_delete", result) + + assert collection.load_cluster("to_delete") is not None + + deleted = collection.delete_cluster("to_delete") + + assert deleted is True + assert collection.load_cluster("to_delete") is None + + db.close() + + def test_load_nonexistent_cluster_returns_none(self, db_path: Path): + """load_cluster returns None for unknown cluster.""" + db = VectorDB(db_path) + collection = db.collection("test") + + assert collection.load_cluster("nonexistent") is None + + db.close() + + def test_assign_to_cluster(self, db_path: Path, dim: int): + """assign_to_cluster assigns docs using saved centroids.""" + db = VectorDB(db_path) + collection = db.collection("test") + + np.random.seed(42) + embeddings = np.random.randn(10, dim).astype(np.float32) + embeddings[:5, 0] += 10.0 + embeddings[5:, 1] += 10.0 + + collection.add_texts( + [f"doc_{i}" for i in range(10)], embeddings=embeddings.tolist() + ) + + result = collection.cluster(n_clusters=2, random_state=42) + collection.save_cluster("saved", result) + + new_embs = np.random.randn(4, dim).astype(np.float32) + new_embs[:2, 0] += 10.0 + new_embs[2:, 1] += 10.0 + new_ids = collection.add_texts( + ["new_a", "new_b", "new_c", "new_d"], embeddings=new_embs.tolist() + ) + + assigned = collection.assign_to_cluster("saved", new_ids) + + assert assigned == 4 + + cluster_0 = collection.get_cluster_members(0) + cluster_1 = collection.get_cluster_members(1) + all_assigned = len(cluster_0) + len(cluster_1) + assert all_assigned >= 4 + + db.close() + + def test_assign_to_cluster_raises_for_unknown(self, db_path: Path): + """assign_to_cluster raises ValueError for unknown cluster.""" + db = VectorDB(db_path) + collection = db.collection("test") + + with pytest.raises(ValueError, match="not found"): + collection.assign_to_cluster("nonexistent", [1, 2, 3]) + + db.close() diff --git a/uv.lock b/uv.lock index c3c4fea..b2e1ece 100644 --- a/uv.lock +++ b/uv.lock @@ -366,6 +366,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/fe/3aed5d0be4d404d12d36ab97e2f1791424d9ca39c2f754a6285d59a3b01d/beautifulsoup4-4.14.2-py3-none-any.whl", hash = "sha256:5ef6fa3a8cbece8488d66985560f97ed091e22bbc4e9c2338508a9d5de6d4515", size = 106392, upload-time = "2025-09-29T10:05:43.771Z" }, ] +[[package]] +name = "black" +version = "25.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "mypy-extensions" }, + { name = "packaging" }, + { name = "pathspec" }, + { name = "platformdirs" }, + { name = "pytokens" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c4/d9/07b458a3f1c525ac392b5edc6b191ff140b596f9d77092429417a54e249d/black-25.12.0.tar.gz", hash = "sha256:8d3dd9cea14bff7ddc0eb243c811cdb1a011ebb4800a5f0335a01a68654796a7", size = 659264, upload-time = "2025-12-08T01:40:52.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/d5/8d3145999d380e5d09bb00b0f7024bf0a8ccb5c07b5648e9295f02ec1d98/black-25.12.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:f85ba1ad15d446756b4ab5f3044731bf68b777f8f9ac9cdabd2425b97cd9c4e8", size = 1895720, upload-time = "2025-12-08T01:46:58.197Z" }, + { url = "https://files.pythonhosted.org/packages/06/97/7acc85c4add41098f4f076b21e3e4e383ad6ed0a3da26b2c89627241fc11/black-25.12.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:546eecfe9a3a6b46f9d69d8a642585a6eaf348bcbbc4d87a19635570e02d9f4a", size = 1727193, upload-time = "2025-12-08T01:52:26.674Z" }, + { url = "https://files.pythonhosted.org/packages/24/f0/fdf0eb8ba907ddeb62255227d29d349e8256ef03558fbcadfbc26ecfe3b2/black-25.12.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:17dcc893da8d73d8f74a596f64b7c98ef5239c2cd2b053c0f25912c4494bf9ea", size = 1774506, upload-time = "2025-12-08T01:46:25.721Z" }, + { url = "https://files.pythonhosted.org/packages/e4/f5/9203a78efe00d13336786b133c6180a9303d46908a9aa72d1104ca214222/black-25.12.0-cp310-cp310-win_amd64.whl", hash = "sha256:09524b0e6af8ba7a3ffabdfc7a9922fb9adef60fed008c7cd2fc01f3048e6e6f", size = 1416085, upload-time = "2025-12-08T01:46:06.073Z" }, + { url = "https://files.pythonhosted.org/packages/ba/cc/7a6090e6b081c3316282c05c546e76affdce7bf7a3b7d2c3a2a69438bd01/black-25.12.0-cp310-cp310-win_arm64.whl", hash = "sha256:b162653ed89eb942758efeb29d5e333ca5bb90e5130216f8369857db5955a7da", size = 1226038, upload-time = "2025-12-08T01:45:29.388Z" }, + { url = "https://files.pythonhosted.org/packages/60/ad/7ac0d0e1e0612788dbc48e62aef8a8e8feffac7eb3d787db4e43b8462fa8/black-25.12.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d0cfa263e85caea2cff57d8f917f9f51adae8e20b610e2b23de35b5b11ce691a", size = 1877003, upload-time = "2025-12-08T01:43:29.967Z" }, + { url = "https://files.pythonhosted.org/packages/e8/dd/a237e9f565f3617a88b49284b59cbca2a4f56ebe68676c1aad0ce36a54a7/black-25.12.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1a2f578ae20c19c50a382286ba78bfbeafdf788579b053d8e4980afb079ab9be", size = 1712639, upload-time = "2025-12-08T01:52:46.756Z" }, + { url = "https://files.pythonhosted.org/packages/12/80/e187079df1ea4c12a0c63282ddd8b81d5107db6d642f7d7b75a6bcd6fc21/black-25.12.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d3e1b65634b0e471d07ff86ec338819e2ef860689859ef4501ab7ac290431f9b", size = 1758143, upload-time = "2025-12-08T01:45:29.137Z" }, + { url = "https://files.pythonhosted.org/packages/93/b5/3096ccee4f29dc2c3aac57274326c4d2d929a77e629f695f544e159bfae4/black-25.12.0-cp311-cp311-win_amd64.whl", hash = "sha256:a3fa71e3b8dd9f7c6ac4d818345237dfb4175ed3bf37cd5a581dbc4c034f1ec5", size = 1420698, upload-time = "2025-12-08T01:45:53.379Z" }, + { url = "https://files.pythonhosted.org/packages/7e/39/f81c0ffbc25ffbe61c7d0385bf277e62ffc3e52f5ee668d7369d9854fadf/black-25.12.0-cp311-cp311-win_arm64.whl", hash = "sha256:51e267458f7e650afed8445dc7edb3187143003d52a1b710c7321aef22aa9655", size = 1229317, upload-time = "2025-12-08T01:46:35.606Z" }, + { url = "https://files.pythonhosted.org/packages/d1/bd/26083f805115db17fda9877b3c7321d08c647df39d0df4c4ca8f8450593e/black-25.12.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:31f96b7c98c1ddaeb07dc0f56c652e25bdedaac76d5b68a059d998b57c55594a", size = 1924178, upload-time = "2025-12-08T01:49:51.048Z" }, + { url = "https://files.pythonhosted.org/packages/89/6b/ea00d6651561e2bdd9231c4177f4f2ae19cc13a0b0574f47602a7519b6ca/black-25.12.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:05dd459a19e218078a1f98178c13f861fe6a9a5f88fc969ca4d9b49eb1809783", size = 1742643, upload-time = "2025-12-08T01:49:59.09Z" }, + { url = "https://files.pythonhosted.org/packages/6d/f3/360fa4182e36e9875fabcf3a9717db9d27a8d11870f21cff97725c54f35b/black-25.12.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c1f68c5eff61f226934be6b5b80296cf6939e5d2f0c2f7d543ea08b204bfaf59", size = 1800158, upload-time = "2025-12-08T01:44:27.301Z" }, + { url = "https://files.pythonhosted.org/packages/f8/08/2c64830cb6616278067e040acca21d4f79727b23077633953081c9445d61/black-25.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:274f940c147ddab4442d316b27f9e332ca586d39c85ecf59ebdea82cc9ee8892", size = 1426197, upload-time = "2025-12-08T01:45:51.198Z" }, + { url = "https://files.pythonhosted.org/packages/d4/60/a93f55fd9b9816b7432cf6842f0e3000fdd5b7869492a04b9011a133ee37/black-25.12.0-cp312-cp312-win_arm64.whl", hash = "sha256:169506ba91ef21e2e0591563deda7f00030cb466e747c4b09cb0a9dae5db2f43", size = 1237266, upload-time = "2025-12-08T01:45:10.556Z" }, + { url = "https://files.pythonhosted.org/packages/c8/52/c551e36bc95495d2aa1a37d50566267aa47608c81a53f91daa809e03293f/black-25.12.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a05ddeb656534c3e27a05a29196c962877c83fa5503db89e68857d1161ad08a5", size = 1923809, upload-time = "2025-12-08T01:46:55.126Z" }, + { url = "https://files.pythonhosted.org/packages/a0/f7/aac9b014140ee56d247e707af8db0aae2e9efc28d4a8aba92d0abd7ae9d1/black-25.12.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9ec77439ef3e34896995503865a85732c94396edcc739f302c5673a2315e1e7f", size = 1742384, upload-time = "2025-12-08T01:49:37.022Z" }, + { url = "https://files.pythonhosted.org/packages/74/98/38aaa018b2ab06a863974c12b14a6266badc192b20603a81b738c47e902e/black-25.12.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e509c858adf63aa61d908061b52e580c40eae0dfa72415fa47ac01b12e29baf", size = 1798761, upload-time = "2025-12-08T01:46:05.386Z" }, + { url = "https://files.pythonhosted.org/packages/16/3a/a8ac542125f61574a3f015b521ca83b47321ed19bb63fe6d7560f348bfe1/black-25.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:252678f07f5bac4ff0d0e9b261fbb029fa530cfa206d0a636a34ab445ef8ca9d", size = 1429180, upload-time = "2025-12-08T01:45:34.903Z" }, + { url = "https://files.pythonhosted.org/packages/e6/2d/bdc466a3db9145e946762d52cd55b1385509d9f9004fec1c97bdc8debbfb/black-25.12.0-cp313-cp313-win_arm64.whl", hash = "sha256:bc5b1c09fe3c931ddd20ee548511c64ebf964ada7e6f0763d443947fd1c603ce", size = 1239350, upload-time = "2025-12-08T01:46:09.458Z" }, + { url = "https://files.pythonhosted.org/packages/35/46/1d8f2542210c502e2ae1060b2e09e47af6a5e5963cb78e22ec1a11170b28/black-25.12.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:0a0953b134f9335c2434864a643c842c44fba562155c738a2a37a4d61f00cad5", size = 1917015, upload-time = "2025-12-08T01:53:27.987Z" }, + { url = "https://files.pythonhosted.org/packages/41/37/68accadf977672beb8e2c64e080f568c74159c1aaa6414b4cd2aef2d7906/black-25.12.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2355bbb6c3b76062870942d8cc450d4f8ac71f9c93c40122762c8784df49543f", size = 1741830, upload-time = "2025-12-08T01:54:36.861Z" }, + { url = "https://files.pythonhosted.org/packages/ac/76/03608a9d8f0faad47a3af3a3c8c53af3367f6c0dd2d23a84710456c7ac56/black-25.12.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9678bd991cc793e81d19aeeae57966ee02909877cb65838ccffef24c3ebac08f", size = 1791450, upload-time = "2025-12-08T01:44:52.581Z" }, + { url = "https://files.pythonhosted.org/packages/06/99/b2a4bd7dfaea7964974f947e1c76d6886d65fe5d24f687df2d85406b2609/black-25.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:97596189949a8aad13ad12fcbb4ae89330039b96ad6742e6f6b45e75ad5cfd83", size = 1452042, upload-time = "2025-12-08T01:46:13.188Z" }, + { url = "https://files.pythonhosted.org/packages/b2/7c/d9825de75ae5dd7795d007681b752275ea85a1c5d83269b4b9c754c2aaab/black-25.12.0-cp314-cp314-win_arm64.whl", hash = "sha256:778285d9ea197f34704e3791ea9404cd6d07595745907dd2ce3da7a13627b29b", size = 1267446, upload-time = "2025-12-08T01:46:14.497Z" }, + { url = "https://files.pythonhosted.org/packages/68/11/21331aed19145a952ad28fca2756a1433ee9308079bd03bd898e903a2e53/black-25.12.0-py3-none-any.whl", hash = "sha256:48ceb36c16dbc84062740049eef990bb2ce07598272e673c17d1a7720c71c828", size = 206191, upload-time = "2025-12-08T01:40:50.963Z" }, +] + [[package]] name = "bleach" version = "6.3.0" @@ -1199,6 +1243,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, ] +[[package]] +name = "hdbscan" +version = "0.8.41" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "joblib" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "scikit-learn" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0c/22/32a66dd4ce72145ec1b792c794b98897be467bdf18c10aa8b48275530b11/hdbscan-0.8.41.tar.gz", hash = "sha256:e41e823e5bb21ff2173f252d226266b1dda82bdbba5d89106eafb251429dff3d", size = 7091384, upload-time = "2025-12-12T15:48:30.807Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/7a/8bc50300f7c1240284b8a69d6c69c59ba37e6349bc9ca097760d3efae077/hdbscan-0.8.41-cp310-cp310-macosx_13_0_x86_64.whl", hash = "sha256:0589ea22e225e4ed6fae8b0a6ac6d18c0aff15165b8ddea7561962e1006b7e63", size = 755274, upload-time = "2025-12-12T15:49:22.72Z" }, + { url = "https://files.pythonhosted.org/packages/10/d9/cf2dc6c14ff2a85f2f48a5c3e034df3690b655f5dc09e2e7db6bc140e0ce/hdbscan-0.8.41-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3970e33b7b370cdca5d0fc5d31171c4b5d588590ec6f04b83e08f743099ba950", size = 4162517, upload-time = "2025-12-12T15:48:27.787Z" }, + { url = "https://files.pythonhosted.org/packages/63/f1/37daced2420b5edaea6fb91875fa4c08e1ef72eaf6f84bc8d4442f45c1dc/hdbscan-0.8.41-cp310-cp310-win_amd64.whl", hash = "sha256:7a689386170d91d1bd9386665b521e4ae66b6a78e0b7e34265ea5b1aa1eb165f", size = 687021, upload-time = "2025-12-12T15:48:57.521Z" }, + { url = "https://files.pythonhosted.org/packages/70/58/5c1cbfac6dd5fd4310da17b09950c68dba1a7c1bdd267eb31468b140bd3a/hdbscan-0.8.41-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:be948fc76d0035d93d309920f14ab6a0185580a63c5a63b05739f08b45dc6c03", size = 1394443, upload-time = "2025-12-12T15:49:56.098Z" }, + { url = "https://files.pythonhosted.org/packages/ef/f1/9a17849751488049003a6af08b270eac1e0135d1d29dfd006bcc4edcca00/hdbscan-0.8.41-cp311-cp311-win_amd64.whl", hash = "sha256:0af3e3bab1eb6b07ea497afc4d2db1b58122974fb052bd21f0ea4b42fcf8d535", size = 687100, upload-time = "2025-12-12T15:49:57.879Z" }, + { url = "https://files.pythonhosted.org/packages/f6/6b/b589c0e903e00108c62f324e99840ab050f1da344fab9cf143ce8ebf1d38/hdbscan-0.8.41-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:07ae4c44098449bd9de12145ad17e92ef699754a43988cbea2dd1a95b89bf142", size = 1393282, upload-time = "2025-12-12T15:52:38.472Z" }, + { url = "https://files.pythonhosted.org/packages/59/ab/6314e52aee546cc14b74fbb575b8713eeec2255880ec99d6490838306d3a/hdbscan-0.8.41-cp312-cp312-win_amd64.whl", hash = "sha256:dce39272d2d4f1dde50dde9cc428cadb84ed16326de872b01761f7ec4f690419", size = 671752, upload-time = "2025-12-12T15:51:33.213Z" }, + { url = "https://files.pythonhosted.org/packages/90/51/0befb66e11c5989b7ec419da2bc652023d30113d4bf4df09cf42a42494d8/hdbscan-0.8.41-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c2f0111395bd1beba1c095edf6b123ec529c8ebd7f4ccd02aaabd6b016454de", size = 1385747, upload-time = "2025-12-12T15:49:25.195Z" }, + { url = "https://files.pythonhosted.org/packages/37/a6/e208ef8bb6e9e97e4b274951160a5bb754f20a79d5673563222d79b00461/hdbscan-0.8.41-cp313-cp313-win_amd64.whl", hash = "sha256:e90f6b9e2fcc94f9ac09d537f8b414191d1a837d62a355edd78e12820b63f0e2", size = 671718, upload-time = "2025-12-12T15:51:39.379Z" }, +] + [[package]] name = "hf-xet" version = "1.2.0" @@ -3781,6 +3850,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/e5/fecf13f06e5e5f67e8837d777d1bc43fac0ed2b77a676804df5c34744727/python_json_logger-4.0.0-py3-none-any.whl", hash = "sha256:af09c9daf6a813aa4cc7180395f50f2a9e5fa056034c9953aec92e381c5ba1e2", size = 15548, upload-time = "2025-10-06T04:15:17.553Z" }, ] +[[package]] +name = "pytokens" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/8d/a762be14dae1c3bf280202ba3172020b2b0b4c537f94427435f19c413b72/pytokens-0.3.0.tar.gz", hash = "sha256:2f932b14ed08de5fcf0b391ace2642f858f1394c0857202959000b68ed7a458a", size = 17644, upload-time = "2025-11-05T13:36:35.34Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/25/d9db8be44e205a124f6c98bc0324b2bb149b7431c53877fc6d1038dddaf5/pytokens-0.3.0-py3-none-any.whl", hash = "sha256:95b2b5eaf832e469d141a378872480ede3f251a5a5041b8ec6e581d3ac71bbf3", size = 12195, upload-time = "2025-11-05T13:36:33.183Z" }, +] + [[package]] name = "pytz" version = "2025.2" @@ -4530,17 +4608,24 @@ name = "simplevecdb" version = "2.1.0" source = { editable = "." } dependencies = [ + { name = "cryptography" }, + { name = "hdbscan" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "psutil" }, - { name = "python-dotenv" }, + { name = "scikit-learn" }, + { name = "sqlcipher3-binary" }, + { name = "sqlite-vec" }, { name = "usearch" }, ] [package.optional-dependencies] -encryption = [ - { name = "cryptography" }, - { name = "sqlcipher3-binary" }, +dev = [ + { name = "black" }, + { name = "mypy" }, + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "pytest-cov" }, + { name = "ruff" }, ] examples = [ { name = "ollama" }, @@ -4555,6 +4640,7 @@ server = [ dev = [ { name = "cryptography" }, { name = "fastapi" }, + { name = "hdbscan" }, { name = "jupyterlab" }, { name = "langchain-core" }, { name = "langchain-openai" }, @@ -4572,6 +4658,7 @@ dev = [ { name = "pytest-asyncio" }, { name = "pytest-cov" }, { name = "ruff" }, + { name = "scikit-learn" }, { name = "sentence-transformers" }, { name = "sqlcipher3-binary" }, { name = "types-psutil" }, @@ -4580,23 +4667,31 @@ dev = [ [package.metadata] requires-dist = [ - { name = "cryptography", marker = "extra == 'encryption'", specifier = ">=41.0" }, + { name = "black", marker = "extra == 'dev'", specifier = ">=23.0" }, + { name = "cryptography", specifier = ">=41.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = ">=0.115" }, - { name = "numpy", specifier = ">=2.0" }, + { name = "hdbscan", specifier = ">=0.8.33" }, + { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.0" }, + { name = "numpy", specifier = ">=1.24" }, { name = "ollama", marker = "extra == 'examples'" }, - { name = "psutil", specifier = ">=5.9.0" }, - { name = "python-dotenv", specifier = ">=1.2.1" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.21" }, + { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.1.0" }, + { name = "scikit-learn", specifier = ">=1.3.0" }, { name = "sentence-transformers", marker = "extra == 'server'", specifier = ">=5.0" }, - { name = "sqlcipher3-binary", marker = "extra == 'encryption'", specifier = ">=0.5.0" }, - { name = "usearch", specifier = ">=2.12" }, + { 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" }, ] -provides-extras = ["server", "encryption", "examples"] +provides-extras = ["server", "dev", "examples"] [package.metadata.requires-dev] dev = [ { name = "cryptography", specifier = ">=41.0" }, { name = "fastapi", specifier = ">=0.115" }, + { name = "hdbscan", specifier = ">=0.8.33" }, { name = "jupyterlab" }, { name = "langchain-core", specifier = ">=1.0.7" }, { name = "langchain-openai", specifier = ">=1.0.3" }, @@ -4614,6 +4709,7 @@ dev = [ { name = "pytest-asyncio", specifier = ">=0.23.0" }, { name = "pytest-cov" }, { name = "ruff", specifier = ">=0.1.0" }, + { name = "scikit-learn", specifier = ">=1.3.0" }, { name = "sentence-transformers", specifier = ">=5.0" }, { name = "sqlcipher3-binary", specifier = ">=0.5.0" }, { name = "types-psutil", specifier = ">=5.9.0" }, @@ -4824,6 +4920,18 @@ 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.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/ed/aabc328f29ee6814033d008ec43e44f2c595447d9cccd5f2aabe60df2933/sqlite_vec-0.1.6-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:77491bcaa6d496f2acb5cc0d0ff0b8964434f141523c121e313f9a7d8088dee3", size = 164075, upload-time = "2024-11-20T16:40:29.847Z" }, + { url = "https://files.pythonhosted.org/packages/a7/57/05604e509a129b22e303758bfa062c19afb020557d5e19b008c64016704e/sqlite_vec-0.1.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fdca35f7ee3243668a055255d4dee4dea7eed5a06da8cad409f89facf4595361", size = 165242, upload-time = "2024-11-20T16:40:31.206Z" }, + { url = "https://files.pythonhosted.org/packages/f2/48/dbb2cc4e5bad88c89c7bb296e2d0a8df58aab9edc75853728c361eefc24f/sqlite_vec-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b0519d9cd96164cd2e08e8eed225197f9cd2f0be82cb04567692a0a4be02da3", size = 103704, upload-time = "2024-11-20T16:40:33.729Z" }, + { url = "https://files.pythonhosted.org/packages/80/76/97f33b1a2446f6ae55e59b33869bed4eafaf59b7f4c662c8d9491b6a714a/sqlite_vec-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:823b0493add80d7fe82ab0fe25df7c0703f4752941aee1c7b2b02cec9656cb24", size = 151556, upload-time = "2024-11-20T16:40:35.387Z" }, + { url = "https://files.pythonhosted.org/packages/6a/98/e8bc58b178266eae2fcf4c9c7a8303a8d41164d781b32d71097924a6bebe/sqlite_vec-0.1.6-py3-none-win_amd64.whl", hash = "sha256:c65bcfd90fa2f41f9000052bcb8bb75d38240b2dae49225389eca6c3136d3f0c", size = 281540, upload-time = "2024-11-20T16:40:37.296Z" }, +] + [[package]] name = "stack-data" version = "0.6.3" From c12000ca45f798b658ef631c6f4d2afced378a55 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Sat, 17 Jan 2026 13:36:30 -0600 Subject: [PATCH 3/5] docs: fix clustering anchor links --- docs/api/core.md | 2 ++ docs/api/types.md | 7 +------ docs/guides/clustering.md | 2 +- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/docs/api/core.md b/docs/api/core.md index 3d12ca1..cc7a83c 100644 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -210,6 +210,8 @@ results = db.search_collections([0.1]*384, parallel=False) | `list_collections()` | Names of all initialized collections | | `search_collections(query, collections, k, filter, normalize_scores, parallel)` | Search across multiple collections with merged results | + + ### Clustering & Auto-Tagging Group similar documents and generate descriptive tags: diff --git a/docs/api/types.md b/docs/api/types.md index c97f3d6..bd8dcf7 100644 --- a/docs/api/types.md +++ b/docs/api/types.md @@ -144,12 +144,7 @@ collection = db.collection("docs", quantization=Quantization.BIT) ### ClusterAlgorithm -::: simplevecdb.types.ClusterAlgorithm - options: - show_root_heading: true - show_source: false - -Clustering algorithms. +Clustering algorithms (string literals accepted by `VectorCollection.cluster`). | Value | Description | Requires n_clusters | Provides Centroids | |-------|-------------|--------------------|--------------------| diff --git a/docs/guides/clustering.md b/docs/guides/clustering.md index c7aa907..8342f7b 100644 --- a/docs/guides/clustering.md +++ b/docs/guides/clustering.md @@ -370,7 +370,7 @@ collection.assign_cluster_metadata(result, tags, metadata_key="topic") ## API Reference -See [VectorCollection API](../api/core.md#clustering--auto-tagging) for complete method signatures and parameters. +See [VectorCollection API](../api/core.md#clustering-auto-tagging) for complete method signatures and parameters. ## Troubleshooting From 75f0c2fb1eddd2cee4edaf494c44c6e7c9be6c79 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Sat, 17 Jan 2026 13:38:19 -0600 Subject: [PATCH 4/5] chore: bump version to 2.2.0 --- pyproject.toml | 2 +- src/simplevecdb/__init__.py | 2 +- uv.lock | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bf2ec7e..37e6ee5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "simplevecdb" -version = "2.1.0" +version = "2.2.0" description = "Dead-simple local vector database powered by usearch HNSW." authors = [{ name = "Dayton Dunbar", email = "coderdayton14@gmail.com" }] license = { text = "MIT" } diff --git a/src/simplevecdb/__init__.py b/src/simplevecdb/__init__.py index 0a61b98..57fd179 100644 --- a/src/simplevecdb/__init__.py +++ b/src/simplevecdb/__init__.py @@ -17,7 +17,7 @@ from .utils import DatabaseLockedError, retry_on_lock, validate_filter from .encryption import EncryptionError, EncryptionUnavailableError -__version__ = "2.1.0" +__version__ = "2.2.0" __all__ = [ # Core classes "VectorDB", diff --git a/uv.lock b/uv.lock index b2e1ece..81bd97c 100644 --- a/uv.lock +++ b/uv.lock @@ -4605,7 +4605,7 @@ wheels = [ [[package]] name = "simplevecdb" -version = "2.1.0" +version = "2.2.0" source = { editable = "." } dependencies = [ { name = "cryptography" }, From 5b2d20982dbdca9198f1f1dd7263e5773bdd03d5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 17 Jan 2026 19:45:17 +0000 Subject: [PATCH 5/5] Initial plan