diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index 68a7d5d..f3b9e55 100644 --- a/.github/workflows/quality.yml +++ b/.github/workflows/quality.yml @@ -7,8 +7,11 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: {python-version: "3.12"} - - run: pip install -e '.[dev,cli]' - - run: ruff check --select E9,F63,F7,F82 vecstream tests benchmarks/ann_benchmark.py - - run: mypy --follow-imports=skip vecstream/hnsw_index.py vecstream/vector_store.py vecstream/binary_store.py + - run: pip install numpy hypothesis mypy pytest ruff + - run: ruff check --select E,F,I,B vecstream tests benchmarks/ann_benchmark.py + - run: ruff format --check vecstream tests benchmarks/ann_benchmark.py + - run: mypy vecstream - run: pytest -q - - run: python benchmarks/ann_benchmark.py --vectors 500 --queries 10 --dim 32 --ef-search 10 20 + - run: python -m benchmarks.ann_benchmark --experiment frontier --vectors 300 --queries 10 --dim 32 --ef-search 10 20 + - run: python -m benchmarks.ann_benchmark --experiment insertion-order --vectors 300 --queries 10 --dim 32 --ef-search 20 --order-seeds 1 2 + - run: python -m benchmarks.ann_benchmark --experiment deletion --vectors 300 --queries 10 --dim 32 --ef-search 20 diff --git a/.gitignore b/.gitignore index 9876d6e..623c70d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,63 +1,49 @@ -# Python +# Python bytecode and native extensions __pycache__/ *.py[cod] -*$py.class *.so + +# Build output .Python build/ -develop-eggs/ dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg -# Virtual Environment -venv/ +# Virtual environments and local configuration .venv/ +venv/ env/ ENV/ .env -# IDE -.idea/ -.vscode/ -*.swp -*.swo -.DS_Store - -# Testing +# Test, type-check, and lint caches .coverage htmlcov/ +.hypothesis/ +.mypy_cache/ .pytest_cache/ +.ruff_cache/ .tox/ -# Project specific -*.vec -*.idx +# Editors and operating systems +.DS_Store +.idea/ +.vscode/ +*.swp +*.swo + +# VecStream local stores *.db -benchmarks/results/ -!benchmarks/results/ -benchmarks/results/* -!benchmarks/results/smoke-1000x64-seed42.json +*.idx +*.vec +.vecstream/ +vector_db/ +test_env/ +test_vectors/ tests/data/ -.vecstream/ -.vscode -vector_db -test_vectors -test_env -.vscode -vecstream/__pycache__/ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class +# Benchmark runs are generated by default. These three measured Stage Two +# experiments are intentional, reviewable repository artifacts. +benchmarks/results/* +!benchmarks/results/frontier-1000x64.json +!benchmarks/results/insertion-order-1000x64.json +!benchmarks/results/deletion-1000x64.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 29538c8..6b569b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,19 +2,21 @@ All notable changes to VecStream 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). +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). ## [Unreleased] ### Added -- Nightly maintenance (#27) +- Exact-search recall oracle and three reproducible ANN experiments. +- Comprehensive HNSW invariant validation and property-based mutation tests. +- Versioned, crash-safe vector checkpoints with explicit corruption errors. -### Fixed -- fix: create PR for changelog updates (#26) +### Changed +- Correct HNSW level generation, diversity selection, reciprocal pruning, and stable float32 normalization. +- Reduced the supported architecture to `Collection`, `BinaryVectorStore`, and `HNSWIndex`. -### Fixed -- Fixed metadata filtering in HNSW search by increasing the number of candidates from 100 to 1000 to ensure enough matches are found after filtering +### Removed +- Legacy manager/query layers, JSON store, client/server, CLI, embedding dependencies, examples, and obsolete benchmark artifacts. ## [0.3.4] - 2024-03-30 diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index d4129d6..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1,7 +0,0 @@ -include LICENSE -include README.md -include requirements.txt -recursive-include docs * -recursive-include vecstream *.py -recursive-exclude * __pycache__ -recursive-exclude * *.py[cod] \ No newline at end of file diff --git a/README.md b/README.md index dc92c06..4a44fd8 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ The project is deliberately single-machine and small enough that the important a - generation-based local checkpoints without pickle - recovery to the previous complete checkpoint and explicit corruption errors - property tests for randomized graph mutations -- a reproducible `ef_search` recall/latency benchmark +- reproducible frontier, insertion-order, and deletion experiments ## Architecture @@ -50,17 +50,21 @@ This is a committed smoke measurement, not a large-scale performance claim: 1,00 | ef_search | recall@10 | p50 ms | p99 ms | QPS | |---:|---:|---:|---:|---:| -| 10 | 0.758 | 0.538 | 1.092 | 1,731 | -| 20 | 0.910 | 0.826 | 1.552 | 1,136 | -| 40 | 0.988 | 1.241 | 2.298 | 753 | -| 80 | 0.996 | 1.652 | 2.808 | 567 | -| 160 | 1.000 | 1.907 | 2.694 | 505 | +| 10 | 0.768 | 0.522 | 0.666 | 1,902 | +| 20 | 0.908 | 0.746 | 0.857 | 1,327 | +| 40 | 0.990 | 1.106 | 1.189 | 903 | +| 80 | 0.996 | 1.475 | 1.652 | 674 | +| 160 | 1.000 | 1.777 | 1.861 | 562 | -Build time was 7.77 s (129 vectors/s). The result demonstrates the expected frontier: greater search effort recovers more exact neighbors while reducing throughput. Full raw output is in [`benchmarks/results/smoke-1000x64-seed42.json`](benchmarks/results/smoke-1000x64-seed42.json). Results are hardware- and dataset-specific. +Build time was 7.37 s (136 vectors/s). The result demonstrates the expected frontier: greater search effort recovers more exact neighbors while reducing throughput. Full raw output is in [`benchmarks/results/frontier-1000x64.json`](benchmarks/results/frontier-1000x64.json). Results are hardware- and dataset-specific. + +At `ef_search=20`, five deterministic insertion orders produced recall@10 from 0.904 to 0.922 (mean 0.9116, population standard deviation 0.0061). This small experiment demonstrates topology sensitivity; it does not establish statistical significance. After cumulative physical deletion of 20%, recall was 0.916 versus a baseline of 0.908; rebuilding the remaining vectors raised it to 0.954. This dataset did not show monotonic deletion degradation, so none is claimed. The rebuild result still shows that an unrepaired graph and a fresh graph are measurably different. These files record the base Git revision and `working_tree_dirty=true`, because Stage Two measurements were produced before its changes were committed. ```bash -python benchmarks/ann_benchmark.py --vectors 10000 --dim 128 --queries 100 \ - --ef-search 10 20 40 80 160 --seed 42 +python -m benchmarks.ann_benchmark --experiment frontier --vectors 10000 \ + --dim 128 --queries 100 --ef-search 10 20 40 80 160 +python -m benchmarks.ann_benchmark --experiment insertion-order --ef-search 20 +python -m benchmarks.ann_benchmark --experiment deletion --ef-search 20 ``` Use 100,000 vectors and dimensions 128/384/768 for full experiments; those runs are deliberately excluded from CI. @@ -74,12 +78,14 @@ Each mutation writes an immutable `gen-*/` directory containing `manifest.json`, ```bash python -m venv .venv . .venv/bin/activate -pip install -e '.[dev,cli]' -ruff check vecstream tests benchmarks/ann_benchmark.py +python -m pip install numpy hypothesis mypy pytest ruff +ruff check --select E,F,I,B vecstream tests benchmarks/ann_benchmark.py mypy vecstream pytest -q ``` +VecStream is run directly from the repository and is not published as a Python package. + See [HNSW design](docs/design.md), [limitations](docs/limitations.md), [performance notes](docs/performance.md), and the [repository audit](docs/audit.md). ## Scope and limitations diff --git a/__init__.py b/__init__.py deleted file mode 100644 index 0519ecb..0000000 --- a/__init__.py +++ /dev/null @@ -1 +0,0 @@ - \ No newline at end of file diff --git a/benchmarks/README.md b/benchmarks/README.md index ddbef40..6778fbe 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,12 +1,16 @@ # ANN benchmarks -`ann_benchmark.py` measures VecStream HNSW results against vectorized exact cosine top-k. It reports environment and configuration, build throughput, recall@k, latency percentiles, and QPS while sweeping `ef_search`. +`ann_benchmark.py` measures HNSW against vectorized exact cosine top-k. It provides three deliberately narrow experiments: + +- `frontier`: sweep `ef_search` and report recall/latency. +- `insertion-order`: index one fixed dataset under independently seeded shuffles. +- `deletion`: measure baseline, cumulative 5%/10%/20% physical deletion, and rebuild. ```bash -python benchmarks/ann_benchmark.py --vectors 10000 --dim 128 --queries 100 \ - --ef-search 10 20 40 80 160 --seed 42 --output run.json +python -m benchmarks.ann_benchmark --experiment frontier --vectors 10000 \ + --dim 128 --queries 100 --ef-search 10 20 40 80 160 --output run.json +python -m benchmarks.ann_benchmark --experiment insertion-order --ef-search 20 +python -m benchmarks.ann_benchmark --experiment deletion --ef-search 20 ``` -The committed `results/smoke-1000x64-seed42.json` is a small measured regression artifact. It must not be generalized to other hardware, dimensions, distributions, or dataset sizes. Full 10k/100k experiments are manual because this pure-Python implementation builds too slowly for routine CI at those sizes. - -Historical graph images in this directory predate the exact-recall harness and are retained only as project history; they are not evidence for current performance claims. +Dataset, graph-level, insertion-order, and deletion RNG seeds are separate. Each JSON result records the Git revision, Python/NumPy versions, platform, architecture, configuration, and seeds. The three committed 1,000 × 64 results are small measured examples, not claims about other workloads. Full 10k/100k experiments remain manual. diff --git a/benchmarks/ann_benchmark.py b/benchmarks/ann_benchmark.py index ec98654..1cecdda 100644 --- a/benchmarks/ann_benchmark.py +++ b/benchmarks/ann_benchmark.py @@ -1,8 +1,4 @@ -"""Reproducible exact-vs-HNSW recall/latency benchmark. - -CI example: python benchmarks/ann_benchmark.py --vectors 1000 --queries 25 --dim 64 -Full runs are intentionally opt-in; use --vectors 100000 and desired dimensions. -""" +"""Three reproducible ANN experiments: frontier, insertion order, deletion.""" from __future__ import annotations @@ -10,6 +6,7 @@ import json import platform import statistics +import subprocess import sys import time from pathlib import Path @@ -19,57 +16,180 @@ from vecstream import HNSWIndex, VectorStore -def percentile(values: list[float], p: float) -> float: - return float(np.percentile(np.asarray(values), p)) +def environment() -> dict: + try: + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], capture_output=True, check=True, text=True + ).stdout.strip() + working_tree_dirty = bool( + subprocess.run( + ["git", "status", "--porcelain"], + capture_output=True, + check=True, + text=True, + ).stdout.strip() + ) + except (OSError, subprocess.CalledProcessError): + revision = "unknown" + working_tree_dirty = None + return { + "vecstream_revision": revision, + "working_tree_dirty": working_tree_dirty, + "platform": platform.platform(), + "architecture": platform.machine(), + "processor": platform.processor(), + "python": sys.version.split()[0], + "numpy": np.__version__, + } + + +def dataset(count: int, dimension: int, queries: int, seed: int) -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(seed) + return ( + rng.normal(size=(count, dimension)).astype(np.float32), + rng.normal(size=(queries, dimension)).astype(np.float32), + ) + +def exact_truth(vectors: dict[str, np.ndarray], queries: np.ndarray, k: int) -> list[set[str]]: + store = VectorStore() + for item_id, vector in vectors.items(): + store.add_vector(item_id, vector) + return [ + {item_id for item_id, _ in store.exact_search(query, k=k, threshold=-1.0)} + for query in queries + ] -def run(args: argparse.Namespace) -> dict: - rng = np.random.default_rng(args.seed) - vectors = rng.normal(size=(args.vectors, args.dim)).astype(np.float32) - queries = rng.normal(size=(args.queries, args.dim)).astype(np.float32) - exact = VectorStore() - index = HNSWIndex(args.dim, M=args.M, ef_construction=args.ef_construction, seed=args.seed) + +def evaluate( + index: HNSWIndex, + queries: np.ndarray, + truth: list[set[str]], + k: int, + ef_search: int, + warmups: int, +) -> dict: + for query in queries[:warmups]: + index.search(query, k=k, ef_search=ef_search) + latencies: list[float] = [] + recalls: list[float] = [] + for query, expected in zip(queries, truth, strict=True): + started = time.perf_counter_ns() + actual = {item_id for item_id, _ in index.search(query, k=k, ef_search=ef_search)} + latencies.append((time.perf_counter_ns() - started) / 1e6) + recalls.append(len(expected & actual) / min(k, len(expected))) + return { + f"recall@{k}": statistics.mean(recalls), + "p50_ms": float(np.percentile(latencies, 50)), + "p99_ms": float(np.percentile(latencies, 99)), + "qps": 1000.0 / statistics.mean(latencies), + } + + +def build( + vectors: np.ndarray, order: np.ndarray, args: argparse.Namespace, graph_seed: int +) -> tuple[HNSWIndex, float]: + index = HNSWIndex(args.dim, M=args.M, ef_construction=args.ef_construction, seed=graph_seed) started = time.perf_counter() - for i, vector in enumerate(vectors): - item_id = str(i) - exact.add_vector(item_id, vector) - index.add_item(item_id, vector) - build_seconds = time.perf_counter() - started + for position in order: + index.add_item(str(int(position)), vectors[position]) + elapsed = time.perf_counter() - started index.validate() - truth = [{item_id for item_id, _ in exact.exact_search(q, k=args.k, threshold=-1.0)} for q in queries] + return index, elapsed + + +def frontier(args: argparse.Namespace, vectors: np.ndarray, queries: np.ndarray) -> dict: + live = {str(i): vector for i, vector in enumerate(vectors)} + truth = exact_truth(live, queries, args.k) + index, build_seconds = build(vectors, np.arange(len(vectors)), args, args.graph_seed) + return { + "build": {"seconds": build_seconds, "vectors_per_second": len(vectors) / build_seconds}, + "search": [ + {"ef_search": ef, **evaluate(index, queries, truth, args.k, ef, args.warmups)} + for ef in args.ef_search + ], + } + + +def insertion_order(args: argparse.Namespace, vectors: np.ndarray, queries: np.ndarray) -> dict: + truth = exact_truth({str(i): vector for i, vector in enumerate(vectors)}, queries, args.k) rows = [] - for ef in args.ef_search: - for query in queries[: args.warmups]: - index.search(query, k=args.k, ef_search=ef) - latencies, recalls = [], [] - for query, expected in zip(queries, truth): - started = time.perf_counter_ns() - actual = {item_id for item_id, _ in index.search(query, k=args.k, ef_search=ef)} - latencies.append((time.perf_counter_ns() - started) / 1e6) - recalls.append(len(expected & actual) / args.k) - rows.append({ - "ef_search": ef, - f"recall@{args.k}": statistics.mean(recalls), - "p50_ms": percentile(latencies, 50), - "p95_ms": percentile(latencies, 95), - "p99_ms": percentile(latencies, 99), - "qps": 1000.0 / statistics.mean(latencies), - }) + ef = args.ef_search[0] + for order_seed in args.order_seeds: + order = np.random.default_rng(order_seed).permutation(len(vectors)) + index, build_seconds = build(vectors, order, args, args.graph_seed) + rows.append( + { + "insertion_order_seed": order_seed, + "build_seconds": build_seconds, + **evaluate(index, queries, truth, args.k, ef, args.warmups), + } + ) + recalls = [row[f"recall@{args.k}"] for row in rows] return { - "environment": {"platform": platform.platform(), "processor": platform.processor(), - "python": sys.version.split()[0], "numpy": np.__version__}, - "configuration": { - key: (str(value) if isinstance(value, Path) else value) - for key, value in vars(args).items() - if key != "output" + "ef_search": ef, + "runs": rows, + "summary": { + "mean_recall": statistics.mean(recalls), + "recall_stddev": statistics.pstdev(recalls), + "min_recall": min(recalls), + "max_recall": max(recalls), }, - "build": {"seconds": build_seconds, "vectors_per_second": args.vectors / build_seconds}, - "search": rows, + } + + +def deletion(args: argparse.Namespace, vectors: np.ndarray, queries: np.ndarray) -> dict: + index, build_seconds = build(vectors, np.arange(len(vectors)), args, args.graph_seed) + live = {str(i): vector for i, vector in enumerate(vectors)} + deletion_order = np.random.default_rng(args.deletion_seed).permutation(len(vectors)) + rows = [] + prior_target = 0 + for fraction in [0.0, 0.05, 0.10, 0.20]: + target = int(len(vectors) * fraction) + for position in deletion_order[prior_target:target]: + item_id = str(int(position)) + index.remove_item(item_id) + del live[item_id] + prior_target = target + index.validate() + truth = exact_truth(live, queries, args.k) + rows.append( + { + "deleted_fraction": fraction, + "remaining_vectors": len(live), + **evaluate(index, queries, truth, args.k, args.ef_search[0], args.warmups), + } + ) + rebuilt = HNSWIndex( + args.dim, M=args.M, ef_construction=args.ef_construction, seed=args.graph_seed + ) + started = time.perf_counter() + for item_id, vector in live.items(): + rebuilt.add_item(item_id, vector) + rebuild_seconds = time.perf_counter() - started + rebuilt.validate() + truth = exact_truth(live, queries, args.k) + rows.append( + { + "deleted_fraction": 0.20, + "remaining_vectors": len(live), + "state": "rebuilt", + **evaluate(rebuilt, queries, truth, args.k, args.ef_search[0], args.warmups), + } + ) + return { + "initial_build_seconds": build_seconds, + "rebuild_seconds": rebuild_seconds, + "ef_search": args.ef_search[0], + "stages": rows, } def main() -> None: parser = argparse.ArgumentParser() + parser.add_argument( + "--experiment", choices=["frontier", "insertion-order", "deletion"], default="frontier" + ) parser.add_argument("--vectors", type=int, default=10_000) parser.add_argument("--queries", type=int, default=100) parser.add_argument("--dim", type=int, default=128) @@ -78,10 +198,22 @@ def main() -> None: parser.add_argument("--ef-construction", type=int, default=100) parser.add_argument("--ef-search", type=int, nargs="+", default=[10, 20, 40, 80, 160]) parser.add_argument("--warmups", type=int, default=5) - parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--dataset-seed", type=int, default=42) + parser.add_argument("--graph-seed", type=int, default=43) + parser.add_argument("--order-seeds", type=int, nargs="+", default=[1, 2, 3, 4, 5]) + parser.add_argument("--deletion-seed", type=int, default=44) parser.add_argument("--output", type=Path) args = parser.parse_args() - result = run(args) + vectors, queries = dataset(args.vectors, args.dim, args.queries, args.dataset_seed) + runner = {"frontier": frontier, "insertion-order": insertion_order, "deletion": deletion}[ + args.experiment + ] + result = { + "environment": environment(), + "experiment": args.experiment, + "configuration": {key: value for key, value in vars(args).items() if key != "output"}, + "results": runner(args, vectors, queries), + } rendered = json.dumps(result, indent=2) print(rendered) if args.output: diff --git a/benchmarks/graphs/accuracy_vs_speed.png b/benchmarks/graphs/accuracy_vs_speed.png deleted file mode 100644 index 9791926..0000000 Binary files a/benchmarks/graphs/accuracy_vs_speed.png and /dev/null differ diff --git a/benchmarks/graphs/insertion_performance.png b/benchmarks/graphs/insertion_performance.png deleted file mode 100644 index e224f2a..0000000 Binary files a/benchmarks/graphs/insertion_performance.png and /dev/null differ diff --git a/benchmarks/graphs/latency_distribution.png b/benchmarks/graphs/latency_distribution.png deleted file mode 100644 index 262e470..0000000 Binary files a/benchmarks/graphs/latency_distribution.png and /dev/null differ diff --git a/benchmarks/graphs/memory_usage.png b/benchmarks/graphs/memory_usage.png deleted file mode 100644 index 05d1cce..0000000 Binary files a/benchmarks/graphs/memory_usage.png and /dev/null differ diff --git a/benchmarks/graphs/query_performance.png b/benchmarks/graphs/query_performance.png deleted file mode 100644 index a413e11..0000000 Binary files a/benchmarks/graphs/query_performance.png and /dev/null differ diff --git a/benchmarks/results/deletion-1000x64.json b/benchmarks/results/deletion-1000x64.json new file mode 100644 index 0000000..918c1b4 --- /dev/null +++ b/benchmarks/results/deletion-1000x64.json @@ -0,0 +1,83 @@ +{ + "environment": { + "vecstream_revision": "ea285da89d72a3e9d05e96ac7d760ccb03cdcf4a", + "working_tree_dirty": true, + "platform": "macOS-27.0-arm64-arm-64bit-Mach-O", + "architecture": "arm64", + "processor": "arm", + "python": "3.14.3", + "numpy": "2.5.3" + }, + "experiment": "deletion", + "configuration": { + "experiment": "deletion", + "vectors": 1000, + "queries": 50, + "dim": 64, + "k": 10, + "M": 16, + "ef_construction": 100, + "ef_search": [ + 20 + ], + "warmups": 5, + "dataset_seed": 42, + "graph_seed": 43, + "order_seeds": [ + 1, + 2, + 3, + 4, + 5 + ], + "deletion_seed": 44 + }, + "results": { + "initial_build_seconds": 7.59150420900005, + "rebuild_seconds": 5.429771624998466, + "ef_search": 20, + "stages": [ + { + "deleted_fraction": 0.0, + "remaining_vectors": 1000, + "recall@10": 0.908, + "p50_ms": 0.7918335, + "p99_ms": 0.8779598299999999, + "qps": 1266.2045695296988 + }, + { + "deleted_fraction": 0.05, + "remaining_vectors": 950, + "recall@10": 0.916, + "p50_ms": 0.718459, + "p99_ms": 0.81786067, + "qps": 1381.205892743672 + }, + { + "deleted_fraction": 0.1, + "remaining_vectors": 900, + "recall@10": 0.916, + "p50_ms": 0.7032499999999999, + "p99_ms": 0.770887, + "qps": 1410.7041351942103 + }, + { + "deleted_fraction": 0.2, + "remaining_vectors": 800, + "recall@10": 0.916, + "p50_ms": 0.6229374999999999, + "p99_ms": 0.71083683, + "qps": 1591.5921660942295 + }, + { + "deleted_fraction": 0.2, + "remaining_vectors": 800, + "state": "rebuilt", + "recall@10": 0.954, + "p50_ms": 0.7115, + "p99_ms": 0.77522308, + "qps": 1408.3086832370143 + } + ] + } +} diff --git a/benchmarks/results/frontier-1000x64.json b/benchmarks/results/frontier-1000x64.json new file mode 100644 index 0000000..cc0b7b2 --- /dev/null +++ b/benchmarks/results/frontier-1000x64.json @@ -0,0 +1,82 @@ +{ + "environment": { + "vecstream_revision": "ea285da89d72a3e9d05e96ac7d760ccb03cdcf4a", + "working_tree_dirty": true, + "platform": "macOS-27.0-arm64-arm-64bit-Mach-O", + "architecture": "arm64", + "processor": "arm", + "python": "3.14.3", + "numpy": "2.5.3" + }, + "experiment": "frontier", + "configuration": { + "experiment": "frontier", + "vectors": 1000, + "queries": 50, + "dim": 64, + "k": 10, + "M": 16, + "ef_construction": 100, + "ef_search": [ + 10, + 20, + 40, + 80, + 160 + ], + "warmups": 5, + "dataset_seed": 42, + "graph_seed": 43, + "order_seeds": [ + 1, + 2, + 3, + 4, + 5 + ], + "deletion_seed": 44 + }, + "results": { + "build": { + "seconds": 7.368972875001418, + "vectors_per_second": 135.7041228082696 + }, + "search": [ + { + "ef_search": 10, + "recall@10": 0.768, + "p50_ms": 0.5220625, + "p99_ms": 0.6656896599999998, + "qps": 1901.921267610745 + }, + { + "ef_search": 20, + "recall@10": 0.908, + "p50_ms": 0.7463960000000001, + "p99_ms": 0.8570514899999998, + "qps": 1327.4483829660974 + }, + { + "ef_search": 40, + "recall@10": 0.99, + "p50_ms": 1.106354, + "p99_ms": 1.18859933, + "qps": 903.4688523963364 + }, + { + "ef_search": 80, + "recall@10": 0.996, + "p50_ms": 1.4753545, + "p99_ms": 1.6519151699999997, + "qps": 673.8146814758736 + }, + { + "ef_search": 160, + "recall@10": 1.0, + "p50_ms": 1.7770415, + "p99_ms": 1.86097267, + "qps": 561.7590671761492 + } + ] + } +} diff --git a/benchmarks/results/insertion-order-1000x64.json b/benchmarks/results/insertion-order-1000x64.json new file mode 100644 index 0000000..efd3ce6 --- /dev/null +++ b/benchmarks/results/insertion-order-1000x64.json @@ -0,0 +1,86 @@ +{ + "environment": { + "vecstream_revision": "ea285da89d72a3e9d05e96ac7d760ccb03cdcf4a", + "working_tree_dirty": true, + "platform": "macOS-27.0-arm64-arm-64bit-Mach-O", + "architecture": "arm64", + "processor": "arm", + "python": "3.14.3", + "numpy": "2.5.3" + }, + "experiment": "insertion-order", + "configuration": { + "experiment": "insertion-order", + "vectors": 1000, + "queries": 50, + "dim": 64, + "k": 10, + "M": 16, + "ef_construction": 100, + "ef_search": [ + 20 + ], + "warmups": 5, + "dataset_seed": 42, + "graph_seed": 43, + "order_seeds": [ + 1, + 2, + 3, + 4, + 5 + ], + "deletion_seed": 44 + }, + "results": { + "ef_search": 20, + "runs": [ + { + "insertion_order_seed": 1, + "build_seconds": 7.245805540998845, + "recall@10": 0.914, + "p50_ms": 0.738146, + "p99_ms": 0.85547358, + "qps": 1350.7580764999675 + }, + { + "insertion_order_seed": 2, + "build_seconds": 7.160134833000484, + "recall@10": 0.91, + "p50_ms": 0.763146, + "p99_ms": 0.81895866, + "qps": 1316.2195657459965 + }, + { + "insertion_order_seed": 3, + "build_seconds": 7.097335499998735, + "recall@10": 0.922, + "p50_ms": 0.7316245, + "p99_ms": 1.101236249999999, + "qps": 1337.7733853623724 + }, + { + "insertion_order_seed": 4, + "build_seconds": 7.23738725000112, + "recall@10": 0.904, + "p50_ms": 0.7155210000000001, + "p99_ms": 0.78817041, + "qps": 1393.906409889777 + }, + { + "insertion_order_seed": 5, + "build_seconds": 7.005199167000683, + "recall@10": 0.908, + "p50_ms": 0.708187, + "p99_ms": 0.80132342, + "qps": 1398.718499705444 + } + ], + "summary": { + "mean_recall": 0.9116000000000001, + "recall_stddev": 0.006118823416311347, + "min_recall": 0.904, + "max_recall": 0.922 + } + } +} diff --git a/benchmarks/results/smoke-1000x64-seed42.json b/benchmarks/results/smoke-1000x64-seed42.json deleted file mode 100644 index dc5efe5..0000000 --- a/benchmarks/results/smoke-1000x64-seed42.json +++ /dev/null @@ -1,71 +0,0 @@ -{ - "environment": { - "platform": "macOS-27.0-arm64-arm-64bit-Mach-O", - "processor": "arm", - "python": "3.14.3", - "numpy": "2.5.3" - }, - "configuration": { - "vectors": 1000, - "queries": 50, - "dim": 64, - "k": 10, - "M": 16, - "ef_construction": 100, - "ef_search": [ - 10, - 20, - 40, - 80, - 160 - ], - "warmups": 5, - "seed": 42 - }, - "build": { - "seconds": 7.771016916000008, - "vectors_per_second": 128.68328698925703 - }, - "search": [ - { - "ef_search": 10, - "recall@10": 0.758, - "p50_ms": 0.538146, - "p95_ms": 0.9373293499999998, - "p99_ms": 1.0919939199999997, - "qps": 1731.4193254514971 - }, - { - "ef_search": 20, - "recall@10": 0.91, - "p50_ms": 0.8258125000000001, - "p95_ms": 1.3074291999999998, - "p99_ms": 1.5520239899999992, - "qps": 1135.9571186184144 - }, - { - "ef_search": 40, - "recall@10": 0.988, - "p50_ms": 1.241104, - "p95_ms": 1.8408544999999998, - "p99_ms": 2.2981681699999985, - "qps": 752.7536292322537 - }, - { - "ef_search": 80, - "recall@10": 0.996, - "p50_ms": 1.651875, - "p95_ms": 2.24563475, - "p99_ms": 2.8083543299999985, - "qps": 566.532924106207 - }, - { - "ef_search": 160, - "recall@10": 1.0, - "p50_ms": 1.9066670000000001, - "p95_ms": 2.51787285, - "p99_ms": 2.6939151699999995, - "qps": 505.18766605329137 - } - ] -} diff --git a/benchmarks/test_cli_stress.py b/benchmarks/test_cli_stress.py deleted file mode 100644 index 5b6ad0d..0000000 --- a/benchmarks/test_cli_stress.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Stress tests for CLI operations under high load.""" - -import tempfile -import os -import numpy as np -from click.testing import CliRunner -from vecstream.cli import main as cli_main -from rich.console import Console -from rich.table import Table -from rich import box - -console = Console() - -def print_benchmark_results(benchmark): - """Print benchmark results in a nice table format.""" - if not benchmark or not hasattr(benchmark, 'stats'): - console.print("[red]No benchmark results available[/red]") - return - - stats = benchmark.stats - - # Create the table - table = Table( - title="🚀 VecStream CLI Performance Results", - box=box.ROUNDED, - header_style="bold magenta", - title_style="bold cyan", - border_style="blue" - ) - - # Add columns - table.add_column("Metric", style="cyan") - table.add_column("Value", justify="right", style="green") - - # Add rows - table.add_row("Min Time", f"{stats.min * 1000:.2f} ms") - table.add_row("Max Time", f"{stats.max * 1000:.2f} ms") - table.add_row("Mean Time", f"{stats.mean * 1000:.2f} ms") - table.add_row("Std Dev", f"{stats.stddev * 1000:.2f} ms") - table.add_row("Operations/sec", f"{1 / stats.mean:.2f}") - table.add_row("Rounds", str(stats.rounds)) - - # Print the table - console.print() - console.print(table) - console.print() - -def generate_large_batch_data(num_vectors=1000, dimension=384): - """Generate a large batch of test vectors.""" - return { - f"stress_vec_{i}": np.random.random(dimension).astype(np.float32).tolist() - for i in range(num_vectors) - } - -def test_cli_bulk_add_stress(benchmark): - """Stress test adding many vectors through CLI.""" - runner = CliRunner() - vectors = generate_large_batch_data(num_vectors=1000) - - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "stress_test.vec") - - def bulk_add(): - for vec_id, vector in list(vectors.items())[:100]: # Test with 100 vectors per batch - result = runner.invoke(cli_main, [ - 'add', - '--id', vec_id, - '--text', f"Test vector {vec_id}", # Using text instead of raw vector - '--db-path', db_path - ]) - assert result.exit_code == 0 - - try: - result = benchmark(bulk_add) - console.rule("[bold red]Bulk Add Stress Test") - print_benchmark_results(result) - except Exception as e: - console.print(f"[red]Error in bulk add test: {str(e)}[/red]") - -def test_cli_concurrent_search_stress(benchmark): - """Stress test concurrent search operations.""" - runner = CliRunner() - - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "stress_test.vec") - - # First populate the database - for i in range(500): # Add 500 vectors for search - runner.invoke(cli_main, [ - 'add', - '--id', f'stress_vec_{i}', - '--text', f'Test vector {i}', - '--db-path', db_path - ]) - - def concurrent_search(): - # Simulate concurrent searches - for _ in range(10): # 10 concurrent searches - result = runner.invoke(cli_main, [ - 'search', - '--text', 'Test vector', # Search by text - '--k', '10', - '--db-path', db_path - ]) - assert result.exit_code == 0 - - try: - result = benchmark(concurrent_search) - console.rule("[bold red]Concurrent Search Stress Test") - print_benchmark_results(result) - except Exception as e: - console.print(f"[red]Error in concurrent search test: {str(e)}[/red]") - -def test_cli_mixed_operations_stress(benchmark): - """Stress test mixed operations (add, search, delete).""" - runner = CliRunner() - - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "stress_test.vec") - - # Initial population - for i in range(50): - runner.invoke(cli_main, [ - 'add', - '--id', f'stress_vec_{i}', - '--text', f'Test vector {i}', - '--db-path', db_path - ]) - - def mixed_operations(): - # Add operation - vec_id = f"stress_vec_{np.random.randint(1000)}" - runner.invoke(cli_main, [ - 'add', - '--id', vec_id, - '--text', f'Test vector {vec_id}', - '--db-path', db_path - ]) - - # Search operation - runner.invoke(cli_main, [ - 'search', - '--text', 'Test vector', - '--k', '5', - '--db-path', db_path - ]) - - # Delete operation - runner.invoke(cli_main, [ - 'remove', # Changed from delete to remove - '--id', vec_id, - '--db-path', db_path - ]) - - try: - result = benchmark(mixed_operations) - console.rule("[bold red]Mixed Operations Stress Test") - print_benchmark_results(result) - except Exception as e: - console.print(f"[red]Error in mixed operations test: {str(e)}[/red]") - -def test_cli_large_batch_processing(benchmark): - """Stress test processing large batches of vectors.""" - runner = CliRunner() - - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "stress_test.vec") - - def batch_process(): - # Generate and process a large batch of vectors - batch_size = 100 - - # Add vectors - for i in range(batch_size): - runner.invoke(cli_main, [ - 'add', - '--id', f'stress_vec_{i}', - '--text', f'Test vector {i}', - '--db-path', db_path - ]) - - # Perform batch search - runner.invoke(cli_main, [ - 'search', - '--text', 'Test vector', - '--k', str(batch_size // 2), - '--db-path', db_path - ]) - - try: - result = benchmark(batch_process) - console.rule("[bold red]Large Batch Processing Test") - print_benchmark_results(result) - except Exception as e: - console.print(f"[red]Error in batch processing test: {str(e)}[/red]") \ No newline at end of file diff --git a/benchmarks/test_vector_operations.py b/benchmarks/test_vector_operations.py deleted file mode 100644 index 029b38d..0000000 --- a/benchmarks/test_vector_operations.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Benchmark tests for vector operations.""" - -import numpy as np -from vecstream import VectorStore - -def generate_random_vectors(num_vectors, dimension): - """Generate random vectors for testing.""" - return { - f"vec_{i}": np.random.random(dimension).astype(np.float32) - for i in range(num_vectors) - } - -def test_add_vector_benchmark(benchmark): - """Benchmark vector addition performance.""" - store = VectorStore() - vector = np.random.random(384).astype(np.float32) - - def add_vector(): - store.add_vector("test_vec", vector) - - benchmark(add_vector) - -def test_search_similar_benchmark(benchmark): - """Benchmark similarity search performance.""" - store = VectorStore() - dimension = 384 - num_vectors = 10000 - - # Add test vectors - vectors = generate_random_vectors(num_vectors, dimension) - for vec_id, vector in vectors.items(): - store.add_vector(vec_id, vector) - - # Create query vector - query = np.random.random(dimension).astype(np.float32) - - def search_similar(): - store.search_similar(query, k=10) - - benchmark(search_similar) - -def test_bulk_add_benchmark(benchmark): - """Benchmark bulk vector addition performance.""" - store = VectorStore() - dimension = 384 - num_vectors = 1000 - vectors = generate_random_vectors(num_vectors, dimension) - - def bulk_add(): - for vec_id, vector in vectors.items(): - store.add_vector(vec_id, vector) - - benchmark(bulk_add) - -def test_vector_normalization_benchmark(benchmark): - """Benchmark vector normalization performance.""" - dimension = 384 - num_vectors = 1000 - vectors = list(generate_random_vectors(num_vectors, dimension).values()) - - def normalize_vectors(): - norms = np.linalg.norm(vectors, axis=1, keepdims=True) - normalized = vectors / norms - return normalized - - benchmark(normalize_vectors) \ No newline at end of file diff --git a/docs/advanced_usage.md b/docs/advanced_usage.md deleted file mode 100644 index 4455493..0000000 --- a/docs/advanced_usage.md +++ /dev/null @@ -1,287 +0,0 @@ -# Advanced Usage Guide - -## Custom Embedding Models - -VecStream supports any embedding model that produces numerical vectors. Here's how to use a custom model: - -```python -from transformers import AutoModel, AutoTokenizer -from vecstream import VectorStore - -class CustomEmbedding: - def __init__(self): - self.model = AutoModel.from_pretrained("your-model-name") - self.tokenizer = AutoTokenizer.from_pretrained("your-model-name") - - def encode(self, text): - tokens = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True) - outputs = self.model(**tokens) - return outputs.last_hidden_state.mean(dim=1).squeeze().tolist() - -# Use custom embeddings -embedder = CustomEmbedding() -store = VectorStore() - -# Add vectors -text = "Example text" -vector = embedder.encode(text) -store.add_vector("doc1", vector) -``` - -## Batch Operations - -For better performance when adding many vectors: - -```python -from vecstream import VectorStore, IndexManager -import numpy as np - -store = VectorStore() -index = IndexManager(store) - -# Generate random vectors -vectors = [(f"vec_{i}", np.random.randn(384)) for i in range(10000)] - -# Batch add vectors -for vec_id, vector in vectors: - store.add_vector(vec_id, vector.tolist()) - -# Create index once after adding all vectors -index.create_index() -``` - -## Persistent Storage with Automatic Saving - -```python -from vecstream import PersistentVectorStore -import time - -class AutoSaveStore(PersistentVectorStore): - def __init__(self, filepath, save_interval=300): # 5 minutes - super().__init__(filepath) - self.last_save = time.time() - self.save_interval = save_interval - - def add_vector(self, id, vector): - super().add_vector(id, vector) - if time.time() - self.last_save > self.save_interval: - self.save() - self.last_save = time.time() - -# Use auto-saving store -store = AutoSaveStore("vectors.vec", save_interval=60) # Save every minute -``` - -## Client-Server with Load Balancing - -```python -from vecstream import VectorDBServer, VectorDBClient -import random - -class LoadBalancedClient: - def __init__(self, servers): - self.servers = servers - self.clients = [VectorDBClient(host, port) for host, port in servers] - - def add_vector(self, id, vector): - # Round-robin distribution - client_idx = hash(id) % len(self.clients) - return self.clients[client_idx].add_vector(id, vector) - - def search_similar(self, query, k=5): - # Query all servers and merge results - all_results = [] - for client in self.clients: - results = client.search_similar(query, k=k) - all_results.extend(results) - - # Sort by similarity and return top k - return sorted(all_results, key=lambda x: x[1], reverse=True)[:k] - -# Usage -servers = [ - ("localhost", 5000), - ("localhost", 5001), - ("localhost", 5002) -] -client = LoadBalancedClient(servers) -``` - -## Custom Index Strategies - -```python -from vecstream import IndexManager -import faiss - -class FaissIndexManager(IndexManager): - def create_index(self): - dimension = next(iter(self.store.vectors.values())).shape[0] - self.index = faiss.IndexFlatL2(dimension) - - # Add existing vectors - vectors = list(self.store.vectors.values()) - if vectors: - self.index.add(np.array(vectors)) - - def update_index(self): - self.create_index() # Recreate index for simplicity - - def search(self, query, k): - distances, indices = self.index.search(np.array([query]), k) - vector_ids = list(self.store.vectors.keys()) - return [(vector_ids[i], 1/(1 + d)) for i, d in zip(indices[0], distances[0])] - -# Usage -store = VectorStore() -index = FaissIndexManager(store) -``` - -## Thread-Safe Vector Store - -```python -from vecstream import VectorStore -from threading import Lock - -class ThreadSafeVectorStore(VectorStore): - def __init__(self): - super().__init__() - self.lock = Lock() - - def add_vector(self, id, vector): - with self.lock: - super().add_vector(id, vector) - - def get_vector(self, id): - with self.lock: - return super().get_vector(id) - - def remove_vector(self, id): - with self.lock: - super().remove_vector(id) - - def search_similar(self, query, k=5): - with self.lock: - return super().search_similar(query, k) - -# Usage -store = ThreadSafeVectorStore() -``` - -## Performance Monitoring - -```python -from vecstream import VectorStore -import time -import logging - -class MonitoredVectorStore(VectorStore): - def __init__(self): - super().__init__() - self.logger = logging.getLogger("vecstream") - self.metrics = { - "add_count": 0, - "search_count": 0, - "total_add_time": 0, - "total_search_time": 0 - } - - def add_vector(self, id, vector): - start = time.time() - super().add_vector(id, vector) - duration = time.time() - start - - self.metrics["add_count"] += 1 - self.metrics["total_add_time"] += duration - self.logger.info(f"Added vector {id} in {duration:.4f}s") - - def search_similar(self, query, k=5): - start = time.time() - results = super().search_similar(query, k) - duration = time.time() - start - - self.metrics["search_count"] += 1 - self.metrics["total_search_time"] += duration - self.logger.info(f"Search completed in {duration:.4f}s") - - return results - - def get_metrics(self): - metrics = self.metrics.copy() - metrics["avg_add_time"] = ( - metrics["total_add_time"] / metrics["add_count"] - if metrics["add_count"] > 0 else 0 - ) - metrics["avg_search_time"] = ( - metrics["total_search_time"] / metrics["search_count"] - if metrics["search_count"] > 0 else 0 - ) - return metrics - -# Usage -store = MonitoredVectorStore() -``` - -## Error Recovery - -```python -from vecstream import PersistentVectorStore -import json -import os - -class RecoverableStore(PersistentVectorStore): - def __init__(self, filepath): - super().__init__(filepath) - self.backup_path = filepath + ".backup" - self.transaction_log = filepath + ".log" - - def add_vector(self, id, vector): - # Log operation before executing - with open(self.transaction_log, "a") as f: - json.dump({"op": "add", "id": id, "vector": vector}, f) - f.write("\n") - - super().add_vector(id, vector) - - def save(self): - # Create backup of current state - if os.path.exists(self.filepath): - os.rename(self.filepath, self.backup_path) - - try: - super().save() - # Clear transaction log after successful save - os.remove(self.transaction_log) - except Exception as e: - # Restore from backup if save fails - if os.path.exists(self.backup_path): - os.rename(self.backup_path, self.filepath) - raise e - -# Usage -store = RecoverableStore("vectors.vec") -``` - -## Best Practices Summary - -1. **Memory Management** - - Use batch operations for large datasets - - Implement periodic saving for persistent storage - - Monitor memory usage and implement cleanup strategies - -2. **Performance** - - Create indices after batch additions - - Use appropriate batch sizes - - Implement caching for frequent queries - - Consider distributed storage for large datasets - -3. **Reliability** - - Implement proper error handling - - Use transaction logs for critical operations - - Create backup strategies - - Monitor system performance - -4. **Scalability** - - Use load balancing for distributed setups - - Implement proper thread safety - - Consider sharding for large datasets - - Monitor and optimize resource usage diff --git a/docs/api_reference.md b/docs/api_reference.md deleted file mode 100644 index 24c61a3..0000000 --- a/docs/api_reference.md +++ /dev/null @@ -1,182 +0,0 @@ -# VecStream API Reference - -## Core Classes - -### VectorStore - -The base class for vector storage and retrieval. - -```python -from vecstream import VectorStore - -store = VectorStore() -``` - -#### Methods - -- `add_vector(id: str, vector: List[float]) -> None` - - Add a vector to the store with the given ID - - Raises `ValueError` if vector dimensions don't match - -- `get_vector(id: str) -> List[float]` - - Retrieve a vector by ID - - Raises `KeyError` if ID not found - -- `remove_vector(id: str) -> None` - - Remove a vector from the store - - Raises `KeyError` if ID not found - -- `search_similar(query: List[float], k: int = 5, threshold: float = 0.0) -> List[Tuple[str, float]]` - - Find k most similar vectors to the query - - Returns list of (id, similarity) tuples - -### PersistentVectorStore - -Extends VectorStore with persistent storage capabilities. - -```python -from vecstream import PersistentVectorStore - -store = PersistentVectorStore("vectors.vec") -``` - -#### Additional Methods - -- `save() -> None` - - Save the current state to disk - -- `load() -> None` - - Load the state from disk - -### IndexManager - -Manages vector indexing for efficient similarity search. - -```python -from vecstream import IndexManager - -index = IndexManager(vector_store) -``` - -#### Methods - -- `create_index() -> None` - - Create or recreate the search index - -- `update_index() -> None` - - Update the index with new vectors - -- `search(query: List[float], k: int) -> List[Tuple[str, float]]` - - Search the index for similar vectors - -### QueryEngine - -High-level interface for vector queries. - -```python -from vecstream import QueryEngine - -engine = QueryEngine(index_manager) -``` - -#### Methods - -- `search(query: List[float], k: int = 5, threshold: float = 0.0) -> List[Tuple[str, float]]` - - Search for similar vectors with optional filtering - -### VectorDBServer - -Server component for distributed vector storage. - -```python -from vecstream import VectorDBServer - -server = VectorDBServer(host="localhost", port=5000) -server.start() -``` - -#### Methods - -- `start() -> None` - - Start the server - -- `stop() -> None` - - Stop the server - -### VectorDBClient - -Client component for connecting to a VectorDBServer. - -```python -from vecstream import VectorDBClient - -client = VectorDBClient(host="localhost", port=5000) -``` - -#### Methods - -- `add_vector(id: str, vector: List[float]) -> None` - - Add a vector through the server - -- `search_similar(query: List[float], k: int = 5) -> List[Tuple[str, float]]` - - Search for similar vectors through the server - -## CLI Commands - -### Add Vector -```bash -vecstream add TEXT ID [--model MODEL_NAME] -``` - -### Search Vectors -```bash -vecstream search TEXT [--k K] [--threshold THRESHOLD] [--model MODEL_NAME] -``` - -### Get Vector -```bash -vecstream get ID -``` - -### Remove Vector -```bash -vecstream remove ID -``` - -### Database Info -```bash -vecstream info -``` - -### Clear Database -```bash -vecstream clear -``` - -## Error Handling - -The library uses the following exception types: - -- `VectorDimensionError`: When vector dimensions don't match -- `VectorNotFoundError`: When a vector ID is not found -- `IndexNotInitializedError`: When trying to search without an index -- `ConnectionError`: When client-server communication fails - -## Best Practices - -1. **Memory Management** - - Use PersistentVectorStore for large datasets - - Call save() periodically when adding many vectors - -2. **Performance Optimization** - - Create indices before searching - - Batch vector additions when possible - - Use appropriate k values for search - -3. **Error Handling** - - Always catch specific exceptions - - Implement proper cleanup in client-server mode - -4. **Thread Safety** - - VectorStore operations are not thread-safe by default - - Use appropriate synchronization in multi-threaded environments \ No newline at end of file diff --git a/docs/audit.md b/docs/audit.md index 71b7c9b..d01318b 100644 --- a/docs/audit.md +++ b/docs/audit.md @@ -12,12 +12,19 @@ The initial repository had a useful narrow core but its claims exceeded its evid - Exact search looped in Python and repeatedly normalized vectors; it is now a vectorized oracle. - Persistence pickled a dictionary, overwrote files, and suppressed every exception. It now uses versioned immutable generations and explicit errors. - There was no invariant checker, deterministic seed, statistical level test, property test, or ANN recall benchmark. -- Packaging made large ML libraries core dependencies. Core now requires only NumPy; CLI and text embedding are extras. +- Packaging made large ML libraries core dependencies. Core now requires only NumPy. - Unsupported speed, compression, optimization, and scalability claims were removed. +## Stage Two cleanup + +- Removed `IndexManager`, `QueryEngine`, `PersistentVectorStore`, client/server, CLI, text-embedding dependencies, and their compatibility exports. +- Removed obsolete examples, stress scripts, benchmark code, generated plots, and historical results. +- The only supported path is now `Collection` composed from `BinaryVectorStore` and `HNSWIndex`; `VectorStore` supplies exact-search ground truth. + +## Stage Two evidence + +- Stage Two added measured insertion-order variance and physical-deletion/rebuild experiments using exact ground truth. + ## Remaining findings -- `CollectionManager`, `IndexManager`, `QueryEngine`, CLI, and legacy client/server overlap. The primary path is `Collection` + store + HNSW; compatibility deletion requires a breaking release. -- The legacy client/server API mismatch was repaired, but the protocol remains a compatibility layer rather than the project's focus. -- HNSW graph persistence, tombstone compaction, generation cleanup, adversarial datasets, insertion-order variance, and memory accounting remain open. -- Old benchmark images/results predate the recall-aware harness and are not current evidence. +- HNSW graph persistence, tombstone compaction, generation cleanup, adversarial datasets, and memory accounting remain open. diff --git a/docs/design.md b/docs/design.md index 6f3f40b..5edeba0 100644 --- a/docs/design.md +++ b/docs/design.md @@ -20,6 +20,6 @@ Querying greedily descends upper layers, then performs a best-first level-0 sear Graph vectors are normalized at insertion, making cosine distance `1 - dot(a, b)`. Zero vectors normalize to zero and therefore have similarity zero. -`validate()` checks entry-point, node/level membership, edge target, reciprocal-edge, degree, and finite-value invariants. Level-0 degree is at most `2M`; upper-layer degree is at most `M`. +`validate()` checks that the entry point occupies the maximum level; IDs agree across raw vectors, normalized vectors, and levels; nodes occupy exactly their allowed levels; edges have existing reciprocal targets; self-edges are absent; degree limits hold; and raw/normalized vectors have valid shapes, finite values, and normalized norms of zero or one. Level-0 degree is at most `2M`; upper-layer degree is at most `M`. The graph requires vector memory `O(Nd)` and adjacency memory approximately `O(NM)`. Expected search behavior is logarithmic on well-formed data, but HNSW has no worst-case logarithmic guarantee and quality depends on parameters, distribution, and insertion order. diff --git a/docs/limitations.md b/docs/limitations.md index b88da80..a7f4564 100644 --- a/docs/limitations.md +++ b/docs/limitations.md @@ -3,9 +3,9 @@ - The HNSW graph is in memory and is rebuilt from collection vectors at startup. Graph serialization remains future work; startup includes construction cost. - Search is approximate. Quality is recall against exact top-k, not generic “accuracy.” - Metadata predicates are post-filtered. A matching vector outside the ANN candidate set can be missed; results are not exact filtered top-k. -- Deletion removes a node and incident edges. Returned IDs are correct, but navigability may degrade. Tombstones plus compaction are not yet implemented. +- Deletion removes a node and incident edges. Returned IDs remain correct, but links are not repaired. In the committed Gaussian experiment, recall did not decline monotonically through 20% deletion; rebuilding improved recall from 0.916 to 0.954. - Objects provide no internal locking. Concurrent read-only searches on an immutable index do not mutate state, but mutation requires exclusive ownership. There is no process-safe writer protocol. - Checkpoints protect vector-store mutation visibility and can recover one prior generation. They are not multi-collection transactions; old generations are not yet garbage-collected. - Python heap and adjacency traversal constrain throughput. Threads are not claimed to scale CPU-bound query work. - The engine is single-machine, with no replication, authentication, distributed query, or transactional multi-client semantics. -- HNSW quality varies with distribution and insertion order. Clustered, duplicate, concentrated, and shuffled-order studies remain to be published. +- HNSW quality varies with distribution and insertion order. The committed five-order experiment measured recall@10 from 0.904 to 0.922. Clustered and concentrated distribution studies remain unmeasured. diff --git a/docs/performance.md b/docs/performance.md index 4c90adb..898372d 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -1,8 +1,12 @@ # Performance methodology and current evidence -The harness creates one seeded dataset, builds exact and HNSW stores, computes exact top-k once, then sweeps `ef_search`. Each configuration performs warmups followed by measured single-query searches. It records the environment, dataset, seed, HNSW parameters, build throughput, recall@k, latency percentiles, and QPS. +The harness creates one seeded dataset, computes exact top-k, and runs three experiments: an `ef_search` frontier, independently seeded insertion-order shuffles, and cumulative physical deletion followed by rebuild. Dataset and topology/order/deletion RNG streams are separate. Each configuration performs warmups followed by measured single-query searches and records the Git revision and environment. -The committed 1,000 × 64 smoke run shows recall@10 rising from 0.758 to 1.000 as `ef_search` rises from 10 to 160, while QPS falls from 1,731 to 505. This is a regression signal, not a 100k-scale claim. +The 1,000 × 64 frontier shows recall@10 rising from 0.768 to 1.000 as `ef_search` rises from 10 to 160, while QPS falls from 1,902 to 562. The committed files identify the base revision and explicitly record that the working tree contained the Stage Two changes. + +Five insertion orders at `ef_search=20` ranged from 0.904 to 0.922 recall@10 (mean 0.9116, population standard deviation 0.0061). This demonstrates insertion-order sensitivity on one small Gaussian dataset, not statistical significance. + +Physical deletion produced recall@10 of 0.908, 0.916, 0.916, and 0.916 at 0%, 5%, 10%, and 20% deleted. A rebuild at 20% yielded 0.954. Deletion did not monotonically reduce recall here. The honest conclusion is that physical deletion changes graph quality unpredictably and rebuilding produced a better graph for this run. ## Measurement-driven optimization diff --git a/examples/example_advanced_features.py b/examples/example_advanced_features.py deleted file mode 100644 index 3ea4c40..0000000 --- a/examples/example_advanced_features.py +++ /dev/null @@ -1,227 +0,0 @@ -#!/usr/bin/env python -""" -Advanced features example for VecStream, demonstrating: -1. HNSW indexing for fast similarity search -2. Collections/namespaces for organizing vectors -3. Metadata filtering for fine-grained search -""" - -import time -import numpy as np -import os -from pathlib import Path -from rich.console import Console -from rich.table import Table - -from vecstream.collections import CollectionManager -from vecstream.hnsw_index import HNSWIndex - -# Create a console for pretty output -console = Console() - -# Create a temporary directory for this example -temp_dir = os.path.join(Path.home(), ".vecstream_example") -os.makedirs(temp_dir, exist_ok=True) - -console.print("\n[bold cyan]VecStream Advanced Features Example[/bold cyan]") -console.print("[dim]This example demonstrates HNSW indexing, collections, and metadata filtering[/dim]\n") - -# Initialize the collection manager -manager = CollectionManager(temp_dir, use_hnsw=True) - -# Create collections for different vector types -science_collection = manager.create_collection("science") -tech_collection = manager.create_collection("technology") -art_collection = manager.create_collection("art") - -console.print("[green]Created three collections: science, technology, art[/green]") - -# Generate some random document vectors and metadata -def create_random_vectors(n=1000, dim=128): - """Create random normalized vectors.""" - vectors = np.random.randn(n, dim) - # Normalize - norms = np.linalg.norm(vectors, axis=1, keepdims=True) - vectors = vectors / norms - return vectors - -# Generate vectors -console.print("[yellow]Generating random vectors...[/yellow]") -n_vectors = 10000 -dim = 128 -all_vectors = create_random_vectors(n_vectors, dim) - -# Split vectors between collections -science_vectors = all_vectors[:3000] -tech_vectors = all_vectors[3000:7000] -art_vectors = all_vectors[7000:] - -# Add vectors to collections with appropriate metadata -console.print("[yellow]Adding vectors to collections with metadata...[/yellow]") - -# Science collection -for i in range(len(science_vectors)): - category = np.random.choice(["physics", "biology", "chemistry"]) - year = np.random.randint(2000, 2024) - metadata = { - "category": category, - "year": year, - "is_peer_reviewed": bool(np.random.choice([True, False], p=[0.8, 0.2])), - "details": { - "citations": np.random.randint(0, 1000), - "journal": np.random.choice(["Nature", "Science", "PNAS", "PLoS ONE"]) - } - } - science_collection.add_vector(f"sci_{i}", science_vectors[i].tolist(), metadata) - -# Technology collection -for i in range(len(tech_vectors)): - category = np.random.choice(["ai", "web", "mobile", "cloud"]) - year = np.random.randint(2010, 2024) - metadata = { - "category": category, - "year": year, - "is_open_source": bool(np.random.choice([True, False])), - "details": { - "stars": np.random.randint(0, 50000), - "language": np.random.choice(["Python", "JavaScript", "Go", "Rust"]) - } - } - tech_collection.add_vector(f"tech_{i}", tech_vectors[i].tolist(), metadata) - -# Art collection -for i in range(len(art_vectors)): - category = np.random.choice(["painting", "sculpture", "digital", "photography"]) - year = np.random.randint(1800, 2024) - metadata = { - "category": category, - "year": year, - "is_exhibited": bool(np.random.choice([True, False])), - "details": { - "artist": np.random.choice(["Anonymous", "Smith", "Johnson", "Garcia", "Wong"]), - "medium": np.random.choice(["Oil", "Acrylic", "Digital", "Mixed"]) - } - } - art_collection.add_vector(f"art_{i}", art_vectors[i].tolist(), metadata) - -console.print("[green]Added vectors to collections with rich metadata![/green]") - -# Show collection stats -console.print("\n[bold cyan]Collection Statistics:[/bold cyan]") -table = Table() -table.add_column("Collection", style="cyan") -table.add_column("Vectors", style="green") -table.add_column("Using HNSW", style="yellow") - -for collection_name in ["science", "technology", "art"]: - stats = manager.get_collection_stats(collection_name) - table.add_row( - collection_name, - str(stats["vector_count"]), - "Yes" if stats["using_hnsw"] else "No" - ) - -console.print(table) - -# Benchmark search performance -console.print("\n[bold cyan]Search Performance Benchmark:[/bold cyan]") - -# Create a query vector (normalized random vector) -query = np.random.randn(dim) -query = query / np.linalg.norm(query) - -# Get the technology collection -tech_coll = manager.get_collection("technology") - -# Benchmark HNSW search -console.print("\n[yellow]Benchmarking HNSW search...[/yellow]") -start_time = time.time() -hnsw_results = tech_coll.search_similar(query.tolist(), k=10) -hnsw_time = time.time() - start_time -console.print(f"[green]HNSW search time: {hnsw_time:.6f} seconds[/green]") - -# Benchmark standard search (by temporarily disabling HNSW) -console.print("\n[yellow]Benchmarking standard search...[/yellow]") -# Store the HNSW index and temporarily disable it -saved_index = tech_coll.hnsw_index -tech_coll.hnsw_index = None -tech_coll.use_hnsw = False - -start_time = time.time() -standard_results = tech_coll.search_similar(query.tolist(), k=10) -standard_time = time.time() - start_time -console.print(f"[green]Standard search time: {standard_time:.6f} seconds[/green]") - -# Restore HNSW index -tech_coll.hnsw_index = saved_index -tech_coll.use_hnsw = True - -# Speed comparison -speedup = standard_time / hnsw_time if hnsw_time > 0 else float('inf') -console.print(f"[bold green]HNSW speedup: {speedup:.2f}x faster![/bold green]") - -# Demonstrate metadata filtering -console.print("\n[bold cyan]Metadata Filtering Examples:[/bold cyan]") - -# Example 1: Filter science papers by category and year -filter1 = {"category": "physics", "year": 2022} -console.print(f"\n[yellow]Searching science collection with filter: {filter1}[/yellow]") -start_time = time.time() -results1 = science_collection.search_similar( - query.tolist(), - k=5, - filter_metadata=filter1 -) -filter_time = time.time() - start_time - -# Print results -table = Table(title="Physics papers from 2022") -table.add_column("ID", style="cyan") -table.add_column("Similarity", style="green") -table.add_column("Metadata", style="white") - -for id, score in results1: - _, metadata = science_collection.get_vector_with_metadata(id) - table.add_row(id, f"{score:.4f}", str(metadata)) - -console.print(table) -console.print(f"[green]Filter search time: {filter_time:.6f} seconds[/green]") - -# Example 2: Filter technology repositories by nested property -filter2 = {"details.language": "Python", "is_open_source": True} -console.print(f"\n[yellow]Searching technology collection with nested filter: {filter2}[/yellow]") -results2 = tech_collection.search_similar( - query.tolist(), - k=5, - filter_metadata=filter2 -) - -# Print results -table = Table(title="Open source Python projects") -table.add_column("ID", style="cyan") -table.add_column("Similarity", style="green") -table.add_column("Stars", style="yellow") -table.add_column("Year", style="blue") - -for id, score in results2: - _, metadata = tech_collection.get_vector_with_metadata(id) - table.add_row( - id, - f"{score:.4f}", - str(metadata["details"]["stars"]), - str(metadata["year"]) - ) - -console.print(table) - -# Clean up -console.print("\n[yellow]Do you want to keep the example data? (y/n)[/yellow]") -response = input().lower() -if response != 'y': - import shutil - shutil.rmtree(temp_dir) - console.print("[green]Example data cleaned up.[/green]") -else: - console.print(f"[green]Example data kept at: {temp_dir}[/green]") - -console.print("\n[bold cyan]Example completed![/bold cyan]") \ No newline at end of file diff --git a/examples/example_binary_store.py b/examples/example_binary_store.py deleted file mode 100644 index 643bd4a..0000000 --- a/examples/example_binary_store.py +++ /dev/null @@ -1,147 +0,0 @@ -""" -Example script demonstrating binary vector storage with VecStream. -""" - -from vecstream.binary_store import BinaryVectorStore -import numpy as np -import os -from rich.console import Console -from rich.table import Table - -console = Console() - -def format_size(size_bytes): - """Format size in bytes to human readable format.""" - for unit in ['B', 'KB', 'MB', 'GB']: - if size_bytes < 1024: - return f"{size_bytes:.2f} {unit}" - size_bytes /= 1024 - return f"{size_bytes:.2f} TB" - -def main(): - # Create a binary store in the 'vector_db' directory - store_dir = "vector_db" - store = BinaryVectorStore(store_dir) - - # Sample documents with metadata - documents = [ - { - "text": "Machine learning is a subset of artificial intelligence", - "tags": ["AI", "ML"], - "category": "technology" - }, - { - "text": "Python is a popular programming language", - "tags": ["programming", "Python"], - "category": "software" - }, - { - "text": "Neural networks are inspired by biological brains", - "tags": ["AI", "neuroscience"], - "category": "science" - }, - { - "text": "Data science involves statistical analysis", - "tags": ["data", "statistics"], - "category": "analytics" - }, - { - "text": "Deep learning is revolutionizing AI", - "tags": ["AI", "deep learning"], - "category": "technology" - } - ] - - console.print("\n🔄 [bold cyan]Adding documents to binary store...[/bold cyan]") - # Add documents with vectors and metadata - for i, doc in enumerate(documents): - # Simulate document embedding (normally you'd use a proper embedding model) - vector = np.random.random(384).astype(np.float32).tolist() - doc_id = f"doc_{i}" - - # Add vector with metadata - store.add_vector( - id=doc_id, - vector=vector, - metadata={ - "text": doc["text"], - "tags": doc["tags"], - "category": doc["category"] - } - ) - - # Display store contents - console.print("\n📚 [bold cyan]Current documents in store:[/bold cyan]") - table = Table(show_header=True, header_style="bold magenta") - table.add_column("ID", style="dim") - table.add_column("Text") - table.add_column("Tags", style="cyan") - table.add_column("Category", style="green") - - for i in range(len(documents)): - doc_id = f"doc_{i}" - vector, metadata = store.get_vector_with_metadata(doc_id) - table.add_row( - doc_id, - metadata["text"], - ", ".join(metadata["tags"]), - metadata["category"] - ) - - console.print(table) - - # Perform similarity search - console.print("\n🔍 [bold cyan]Performing similarity search...[/bold cyan]") - query_vector = np.random.random(384).astype(np.float32).tolist() - results = store.search_similar(query_vector, k=3) - - # Display search results - console.print("\n[bold cyan]Top 3 similar documents:[/bold cyan]") - table = Table(show_header=True, header_style="bold magenta") - table.add_column("Score", style="cyan") - table.add_column("Text") - table.add_column("Category", style="green") - - for doc_id, score in results: - _, metadata = store.get_vector_with_metadata(doc_id) - table.add_row( - f"{score:.4f}", - metadata["text"], - metadata["category"] - ) - - console.print(table) - - # Display storage information - vectors_size, metadata_size = store.get_store_size() - total_size = vectors_size + metadata_size - - console.print("\n💾 [bold cyan]Storage Information:[/bold cyan]") - table = Table(show_header=True, header_style="bold magenta") - table.add_column("Metric", style="dim") - table.add_column("Size", style="cyan") - table.add_column("Percentage", style="green") - - table.add_row( - "Vectors File", - format_size(vectors_size), - f"{(vectors_size/total_size)*100:.1f}%" - ) - table.add_row( - "Metadata File", - format_size(metadata_size), - f"{(metadata_size/total_size)*100:.1f}%" - ) - table.add_row( - "Total Size", - format_size(total_size), - "100%" - ) - - console.print(table) - - # Show storage location - console.print(f"\n📁 [bold cyan]Storage Location:[/bold cyan] {os.path.abspath(store_dir)}") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/examples/example_inmemory.py b/examples/example_inmemory.py deleted file mode 100644 index 14ddc6a..0000000 --- a/examples/example_inmemory.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Example script demonstrating in-memory vector storage and search with VecStream. -""" - -from vecstream.vector_store import VectorStore -import numpy as np - -def main(): - # Initialize the in-memory vector store - vector_store = VectorStore() - - # Sample data to store (simulated document embeddings) - documents = [ - "Machine learning is a subset of artificial intelligence", - "Python is a popular programming language", - "Neural networks are inspired by biological brains", - "Data science involves statistical analysis", - "Deep learning is revolutionizing AI" - ] - - # Dictionary to store document texts - doc_texts = {} - - print("🔄 Adding documents to vector store...") - # Add documents with unique IDs - for i, doc in enumerate(documents): - # Simulate document embedding (normally you'd use a proper embedding model) - vector = np.random.random(384).astype(np.float32).tolist() - doc_id = f"doc_{i}" - vector_store.add_vector(id=doc_id, vector=vector) - doc_texts[doc_id] = doc - - print("\n📚 Current documents in store:") - for i in range(len(documents)): - doc_id = f"doc_{i}" - vector = vector_store.get_vector(doc_id) - print(f"ID: {doc_id} | Text: {doc_texts[doc_id]}") - - # Perform a similarity search - print("\n🔍 Performing similarity search...") - query_vector = np.random.random(384).astype(np.float32).tolist() - results = vector_store.search_similar(query_vector, k=3) - - print("\nTop 3 similar documents:") - for doc_id, score in results: - print(f"Score: {score:.4f} | Text: {doc_texts[doc_id]}") - - # Demonstrate removing a document - print("\n🗑️ Removing document 'doc_0'...") - vector_store.remove_vector("doc_0") - doc_texts.pop("doc_0") - - # Verify removal - print("\n📚 Remaining documents:") - for i in range(1, len(documents)): - doc_id = f"doc_{i}" - try: - vector = vector_store.get_vector(doc_id) - print(f"ID: {doc_id} | Text: {doc_texts[doc_id]}") - except KeyError: - continue - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000..77490d5 --- /dev/null +++ b/mypy.ini @@ -0,0 +1,4 @@ +[mypy] +python_version = 3.10 +packages = vecstream +ignore_missing_imports = true diff --git a/pyproject.toml b/pyproject.toml deleted file mode 100644 index 7e91765..0000000 --- a/pyproject.toml +++ /dev/null @@ -1,37 +0,0 @@ -[build-system] -requires = ["setuptools>=68"] -build-backend = "setuptools.build_meta" - -[project] -name = "vecstream" -version = "0.4.0" -description = "A compact vector-search engine with a from-scratch HNSW implementation" -readme = "README.md" -requires-python = ">=3.10" -license = {text = "MIT"} -authors = [{name = "Torin Etheridge"}] -dependencies = ["numpy>=1.24"] - -[project.optional-dependencies] -cli = ["click>=8", "rich>=13"] -dev = ["hypothesis>=6", "mypy>=1.10", "pytest>=8", "ruff>=0.6"] -text = ["sentence-transformers>=2.2"] - -[project.scripts] -vecstream = "vecstream.cli:main" - -[tool.setuptools.packages.find] -include = ["vecstream*"] - -[tool.pytest.ini_options] -testpaths = ["tests"] -addopts = "--strict-markers" - -[tool.ruff] -line-length = 100 -target-version = "py310" - -[tool.mypy] -python_version = "3.10" -packages = ["vecstream"] -ignore_missing_imports = true diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..6d9ad96 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +testpaths = tests +addopts = --strict-markers diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000..6c4c0a3 --- /dev/null +++ b/ruff.toml @@ -0,0 +1,2 @@ +line-length = 100 +target-version = "py310" diff --git a/run_tests.py b/run_tests.py deleted file mode 100644 index cf5bb2f..0000000 --- a/run_tests.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python -""" -Run all VecStream tests. -""" - -import os -import sys -import subprocess -import argparse - -def run_tests(run_benchmarks=False): - """Run all VecStream tests using pytest. - - Args: - run_benchmarks: Whether to run benchmarks in addition to tests - - Returns: - Exit code (0 if all tests pass, non-zero otherwise) - """ - # Get the current directory - current_dir = os.path.dirname(os.path.abspath(__file__)) - exit_code = 0 - - # Print header - print("\n===== Running VecStream unit tests =====\n") - - # Run unit tests - result = subprocess.run( - ["pytest", "-xvs", "tests/unit/"], - cwd=current_dir, - capture_output=False - ) - - if result.returncode != 0: - exit_code = result.returncode - - # Run integration tests - print("\n===== Running VecStream integration tests =====\n") - - result_integration = subprocess.run( - ["pytest", "-xvs", "tests/integration/"], - cwd=current_dir, - capture_output=False - ) - - if result_integration.returncode != 0: - exit_code = result_integration.returncode - - # Run benchmarks if requested - if run_benchmarks: - print("\n===== Running VecStream benchmarks =====\n") - - result_benchmarks = subprocess.run( - ["python3", "tests/benchmarks/benchmark_vector_operations.py"], - cwd=current_dir, - capture_output=False - ) - - if result_benchmarks.returncode != 0: - exit_code = result_benchmarks.returncode - - # Return exit code (0 if all tests pass, non-zero otherwise) - return exit_code - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Run VecStream tests") - parser.add_argument( - "--benchmarks", - action="store_true", - help="Run benchmarks in addition to tests" - ) - args = parser.parse_args() - - sys.exit(run_tests(run_benchmarks=args.benchmarks)) \ No newline at end of file diff --git a/sample_texts.txt b/sample_texts.txt deleted file mode 100644 index b019cf9..0000000 --- a/sample_texts.txt +++ /dev/null @@ -1,10 +0,0 @@ -The quick brown fox jumps over the lazy dog. -A fast orange fox leaps across a sleepy canine. -The weather is beautiful today. -The sun is shining brightly in the sky. -Machine learning is fascinating. -Artificial intelligence is transforming industries. -Data science helps us understand patterns. -Neural networks can learn complex relationships. -Deep learning models achieve remarkable results. -Natural language processing is advancing rapidly. \ No newline at end of file diff --git a/scripts/run_server.py b/scripts/run_server.py deleted file mode 100644 index 8c2d9c7..0000000 --- a/scripts/run_server.py +++ /dev/null @@ -1,98 +0,0 @@ -#!/usr/bin/env python3 -import socket -import threading -import argparse -import sys -import os -import signal - -# Add parent directory to path to import vector_db -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) - -from vector_db.command_interface import CommandInterface - -class VectorDBServer: - def __init__(self, host='127.0.0.1', port=6379): - self.host = host - self.port = port - self.command_interface = CommandInterface() - self.server_socket = None - self.running = False - - def start(self): - self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - try: - self.server_socket.bind((self.host, self.port)) - self.server_socket.listen(5) - self.running = True - print(f"Server running on {self.host}:{self.port}") - - while self.running: - try: - client_socket, address = self.server_socket.accept() - print(f"Connection from {address}") - client_thread = threading.Thread( - target=self.handle_client, - args=(client_socket,) - ) - client_thread.daemon = True - client_thread.start() - except KeyboardInterrupt: - break - except Exception as e: - print(f"Error accepting connection: {e}") - - finally: - if self.server_socket: - self.server_socket.close() - - def stop(self): - self.running = False - if self.server_socket: - self.server_socket.close() - print("Server stopped") - - def handle_client(self, client_socket): - try: - while True: - data = client_socket.recv(1024) - if not data: - break - - command = data.decode('utf-8').strip() - if command.lower() == "quit" or command.lower() == "exit": - break - - response = self.command_interface.execute_command(command) - client_socket.send(f"{response}\n".encode('utf-8')) - - except Exception as e: - print(f"Error handling client: {e}") - finally: - client_socket.close() - -def signal_handler(sig, frame): - print("\nShutting down server...") - if 'server' in globals(): - server.stop() - sys.exit(0) - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Vector DB Server") - parser.add_argument('--host', default='127.0.0.1', help='Host address') - parser.add_argument('--port', type=int, default=6379, help='Port number') - - args = parser.parse_args() - - signal.signal(signal.SIGINT, signal_handler) - signal.signal(signal.SIGTERM, signal_handler) - - server = VectorDBServer(args.host, args.port) - - try: - server.start() - except KeyboardInterrupt: - print("\nShutting down server...") - server.stop() diff --git a/tests/benchmarks/benchmark_vector_operations.py b/tests/benchmarks/benchmark_vector_operations.py deleted file mode 100644 index 52316d2..0000000 --- a/tests/benchmarks/benchmark_vector_operations.py +++ /dev/null @@ -1,226 +0,0 @@ -""" -Benchmark tests for VecStream vector operations. -""" -import time -import os -import numpy as np -import tempfile -import shutil -from vecstream import VectorStore -from vecstream.hnsw_index import HNSWIndex -from vecstream.collections import Collection, CollectionManager -from vecstream.persistent_store import PersistentVectorStore - -def generate_random_vectors(num_vectors, dimension): - """Generate random vectors for benchmarking.""" - vectors = [] - for i in range(num_vectors): - vec = np.random.randn(dimension) - vec = vec / np.linalg.norm(vec) # Normalize - vectors.append((f"vec_{i}", vec.tolist())) - return vectors - -def benchmark_add_vectors(num_vectors=10000, dimension=384): - """Benchmark vector addition performance.""" - store = VectorStore() - vectors = generate_random_vectors(num_vectors, dimension) - - start_time = time.time() - for vec_id, vector in vectors: - store.add_vector(vec_id, vector) - end_time = time.time() - - add_time = end_time - start_time - vectors_per_second = num_vectors / add_time - - return { - "operation": "add_vectors", - "num_vectors": num_vectors, - "dimension": dimension, - "total_time": add_time, - "vectors_per_second": vectors_per_second - } - -def benchmark_search(num_vectors=10000, dimension=384, k=10, num_queries=100): - """Benchmark similarity search performance.""" - # Setup - store = VectorStore() - vectors = generate_random_vectors(num_vectors, dimension) - for vec_id, vector in vectors: - store.add_vector(vec_id, vector) - - # Generate random queries - queries = [vec for _, vec in generate_random_vectors(num_queries, dimension)] - - # Benchmark search - start_time = time.time() - for query in queries: - store.search_similar(query, k=k) - end_time = time.time() - - search_time = end_time - start_time - queries_per_second = num_queries / search_time - - return { - "operation": "search", - "num_vectors": num_vectors, - "dimension": dimension, - "k": k, - "num_queries": num_queries, - "total_time": search_time, - "queries_per_second": queries_per_second - } - -def benchmark_hnsw_index(num_vectors=10000, dimension=128, k=10, num_queries=100, M=16, ef_construction=200): - """Benchmark HNSW index performance.""" - # Setup - hnsw_index = HNSWIndex(dim=dimension, M=M, ef_construction=ef_construction) - vectors = generate_random_vectors(num_vectors, dimension) - - # Benchmark add items - start_time = time.time() - for vec_id, vector in vectors: - hnsw_index.add_item(vec_id, np.array(vector)) - add_time = time.time() - start_time - - # Generate random queries - queries = [np.array(vec) for _, vec in generate_random_vectors(num_queries, dimension)] - - # Benchmark search - start_time = time.time() - for query in queries: - hnsw_index.search(query, k=k) - search_time = time.time() - start_time - - return { - "operation": "hnsw_index", - "num_vectors": num_vectors, - "dimension": dimension, - "M": M, - "ef_construction": ef_construction, - "k": k, - "num_queries": num_queries, - "add_time": add_time, - "search_time": search_time, - "vectors_per_second_add": num_vectors / add_time, - "queries_per_second": num_queries / search_time - } - -def benchmark_collections(num_vectors=5000, dimension=128, k=10, num_queries=50): - """Benchmark collections performance.""" - temp_dir = tempfile.mkdtemp() - try: - # Setup - collection_dir = os.path.join(temp_dir, "collections") - manager = CollectionManager(collection_dir, use_hnsw=True) - collection = manager.create_collection("benchmark_collection") - - vectors = generate_random_vectors(num_vectors, dimension) - - # Create metadata for vectors - metadata_list = [] - for i in range(num_vectors): - metadata = { - "category": f"category_{i % 5}", - "price": i % 100 + 10, - "in_stock": i % 2 == 0, - "details": { - "color": f"color_{i % 10}", - "brand": f"brand_{i % 8}" - } - } - metadata_list.append(metadata) - - # Benchmark add vectors with metadata - start_time = time.time() - for i, (vec_id, vector) in enumerate(vectors): - collection.add_vector(vec_id, vector, metadata_list[i]) - add_time = time.time() - start_time - - # Generate random queries - queries = [vec for _, vec in generate_random_vectors(num_queries, dimension)] - - # Benchmark simple search - start_time = time.time() - for query in queries: - collection.search_similar(query, k=k) - search_time = time.time() - start_time - - # Benchmark search with filter - start_time = time.time() - for query in queries: - collection.search_similar(query, k=k, filter_metadata={"category": "category_1"}) - filter_search_time = time.time() - start_time - - # Benchmark search with nested filter - start_time = time.time() - for query in queries: - collection.search_similar(query, k=k, filter_metadata={"details.brand": "brand_2"}) - nested_filter_time = time.time() - start_time - - return { - "operation": "collections", - "num_vectors": num_vectors, - "dimension": dimension, - "k": k, - "num_queries": num_queries, - "add_time": add_time, - "search_time": search_time, - "filter_search_time": filter_search_time, - "nested_filter_time": nested_filter_time, - "vectors_per_second_add": num_vectors / add_time, - "queries_per_second": num_queries / search_time, - "filter_queries_per_second": num_queries / filter_search_time, - "nested_filter_queries_per_second": num_queries / nested_filter_time - } - finally: - shutil.rmtree(temp_dir) - -def run_benchmarks(): - """Run all benchmarks and save results.""" - import json - import os - from datetime import datetime - - results = { - "timestamp": datetime.now().isoformat(), - "benchmarks": [] - } - - # Run add vectors benchmark with different sizes - for num_vectors in [1000, 10000]: - result = benchmark_add_vectors(num_vectors=num_vectors) - results["benchmarks"].append(result) - - # Run search benchmark with different sizes - for num_vectors in [1000, 10000]: - result = benchmark_search(num_vectors=num_vectors) - results["benchmarks"].append(result) - - # Run HNSW index benchmark - for num_vectors in [1000, 10000]: - result = benchmark_hnsw_index(num_vectors=num_vectors, dimension=128) - results["benchmarks"].append(result) - - # Run collections benchmark - result = benchmark_collections(num_vectors=5000, dimension=128) - results["benchmarks"].append(result) - - # Save results - os.makedirs("tests/benchmarks/results", exist_ok=True) - filename = f"benchmark_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" - filepath = os.path.join("tests/benchmarks/results", filename) - - with open(filepath, "w") as f: - json.dump(results, f, indent=2) - - return results - -if __name__ == "__main__": - results = run_benchmarks() - print("Benchmark results:") - for benchmark in results["benchmarks"]: - print(f"\n{benchmark['operation']}:") - for key, value in benchmark.items(): - if key != "operation": - print(f" {key}: {value}") \ No newline at end of file diff --git a/tests/benchmarks/results/benchmark_results_20250316_162423.json b/tests/benchmarks/results/benchmark_results_20250316_162423.json deleted file mode 100644 index 70e0aa8..0000000 --- a/tests/benchmarks/results/benchmark_results_20250316_162423.json +++ /dev/null @@ -1,53 +0,0 @@ -{ - "timestamp": "2025-03-16T16:24:04.601378", - "benchmarks": [ - { - "operation": "add_vectors", - "num_vectors": 1000, - "dimension": 384, - "total_time": 0.006134986877441406, - "vectors_per_second": 162999.5336545935 - }, - { - "operation": "add_vectors", - "num_vectors": 10000, - "dimension": 384, - "total_time": 0.057238101959228516, - "vectors_per_second": 174708.79812057948 - }, - { - "operation": "add_vectors", - "num_vectors": 100000, - "dimension": 384, - "total_time": 0.5889060497283936, - "vectors_per_second": 169806.37241903102 - }, - { - "operation": "search", - "num_vectors": 1000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 0.13065099716186523, - "queries_per_second": 765.397908720962 - }, - { - "operation": "search", - "num_vectors": 10000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 1.300077199935913, - "queries_per_second": 76.91850915078695 - }, - { - "operation": "search", - "num_vectors": 100000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 13.37088680267334, - "queries_per_second": 7.478935501870097 - } - ] -} \ No newline at end of file diff --git a/tests/benchmarks/results/benchmark_results_20250317_215701.json b/tests/benchmarks/results/benchmark_results_20250317_215701.json deleted file mode 100644 index 5e60454..0000000 --- a/tests/benchmarks/results/benchmark_results_20250317_215701.json +++ /dev/null @@ -1,78 +0,0 @@ -{ - "timestamp": "2025-03-17T21:55:28.177175", - "benchmarks": [ - { - "operation": "add_vectors", - "num_vectors": 1000, - "dimension": 384, - "total_time": 0.007476091384887695, - "vectors_per_second": 133759.73466849508 - }, - { - "operation": "add_vectors", - "num_vectors": 10000, - "dimension": 384, - "total_time": 0.0725250244140625, - "vectors_per_second": 137883.44203660847 - }, - { - "operation": "search", - "num_vectors": 1000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 0.15668606758117676, - "queries_per_second": 638.2188381120195 - }, - { - "operation": "search", - "num_vectors": 10000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 1.6238160133361816, - "queries_per_second": 61.5833315958911 - }, - { - "operation": "hnsw_index", - "num_vectors": 1000, - "dimension": 128, - "M": 16, - "ef_construction": 200, - "k": 10, - "num_queries": 100, - "add_time": 0.31320905685424805, - "search_time": 0.0154571533203125, - "vectors_per_second_add": 3192.7556950096446, - "queries_per_second": 6469.496544916091 - }, - { - "operation": "hnsw_index", - "num_vectors": 10000, - "dimension": 128, - "M": 16, - "ef_construction": 200, - "k": 10, - "num_queries": 100, - "add_time": 8.220263957977295, - "search_time": 0.0254666805267334, - "vectors_per_second_add": 1216.505948120507, - "queries_per_second": 3926.69943360015 - }, - { - "operation": "collections", - "num_vectors": 5000, - "dimension": 128, - "k": 10, - "num_queries": 50, - "add_time": 81.87265396118164, - "search_time": 0.07735610008239746, - "filter_search_time": 0.08352398872375488, - "nested_filter_time": 0.0872490406036377, - "vectors_per_second_add": 61.07045219727035, - "queries_per_second": 646.3614368710607, - "filter_queries_per_second": 598.6304146150004, - "nested_filter_queries_per_second": 573.0722040502911 - } - ] -} \ No newline at end of file diff --git a/tests/benchmarks/results/benchmark_results_20250317_215756.json b/tests/benchmarks/results/benchmark_results_20250317_215756.json deleted file mode 100644 index 06ccb19..0000000 --- a/tests/benchmarks/results/benchmark_results_20250317_215756.json +++ /dev/null @@ -1,78 +0,0 @@ -{ - "timestamp": "2025-03-17T21:56:24.623781", - "benchmarks": [ - { - "operation": "add_vectors", - "num_vectors": 1000, - "dimension": 384, - "total_time": 0.006494760513305664, - "vectors_per_second": 153970.26540875886 - }, - { - "operation": "add_vectors", - "num_vectors": 10000, - "dimension": 384, - "total_time": 0.07675909996032715, - "vectors_per_second": 130277.7130681377 - }, - { - "operation": "search", - "num_vectors": 1000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 0.16481781005859375, - "queries_per_second": 606.730546683331 - }, - { - "operation": "search", - "num_vectors": 10000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 1.7749769687652588, - "queries_per_second": 56.33875918376777 - }, - { - "operation": "hnsw_index", - "num_vectors": 1000, - "dimension": 128, - "M": 16, - "ef_construction": 200, - "k": 10, - "num_queries": 100, - "add_time": 0.6160528659820557, - "search_time": 0.01662898063659668, - "vectors_per_second_add": 1623.2373148786355, - "queries_per_second": 6013.59771746455 - }, - { - "operation": "hnsw_index", - "num_vectors": 10000, - "dimension": 128, - "M": 16, - "ef_construction": 200, - "k": 10, - "num_queries": 100, - "add_time": 7.680387735366821, - "search_time": 0.033773183822631836, - "vectors_per_second_add": 1302.0175991834078, - "queries_per_second": 2960.9290176838094 - }, - { - "operation": "collections", - "num_vectors": 5000, - "dimension": 128, - "k": 10, - "num_queries": 50, - "add_time": 81.1283688545227, - "search_time": 0.07472586631774902, - "filter_search_time": 0.07747101783752441, - "nested_filter_time": 0.07507610321044922, - "vectors_per_second_add": 61.63072265098625, - "queries_per_second": 669.1123497637378, - "filter_queries_per_second": 645.4026472823963, - "nested_filter_queries_per_second": 665.990879412624 - } - ] -} \ No newline at end of file diff --git a/tests/benchmarks/results/benchmark_results_20250317_221204.json b/tests/benchmarks/results/benchmark_results_20250317_221204.json deleted file mode 100644 index af8a71e..0000000 --- a/tests/benchmarks/results/benchmark_results_20250317_221204.json +++ /dev/null @@ -1,78 +0,0 @@ -{ - "timestamp": "2025-03-17T22:10:37.931558", - "benchmarks": [ - { - "operation": "add_vectors", - "num_vectors": 1000, - "dimension": 384, - "total_time": 0.0068950653076171875, - "vectors_per_second": 145031.25864453666 - }, - { - "operation": "add_vectors", - "num_vectors": 10000, - "dimension": 384, - "total_time": 0.07133793830871582, - "vectors_per_second": 140177.86660338956 - }, - { - "operation": "search", - "num_vectors": 1000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 0.15673303604125977, - "queries_per_second": 638.0275819685846 - }, - { - "operation": "search", - "num_vectors": 10000, - "dimension": 384, - "k": 10, - "num_queries": 100, - "total_time": 1.6680800914764404, - "queries_per_second": 59.9491598221094 - }, - { - "operation": "hnsw_index", - "num_vectors": 1000, - "dimension": 128, - "M": 16, - "ef_construction": 200, - "k": 10, - "num_queries": 100, - "add_time": 0.5421199798583984, - "search_time": 0.017783641815185547, - "vectors_per_second_add": 1844.6101179690882, - "queries_per_second": 5623.145193725701 - }, - { - "operation": "hnsw_index", - "num_vectors": 10000, - "dimension": 128, - "M": 16, - "ef_construction": 200, - "k": 10, - "num_queries": 100, - "add_time": 7.078456878662109, - "search_time": 0.023868799209594727, - "vectors_per_second_add": 1412.7372916750871, - "queries_per_second": 4189.56978614166 - }, - { - "operation": "collections", - "num_vectors": 5000, - "dimension": 128, - "k": 10, - "num_queries": 50, - "add_time": 76.09976983070374, - "search_time": 0.06311297416687012, - "filter_search_time": 0.06307291984558105, - "nested_filter_time": 0.06317496299743652, - "vectors_per_second_add": 65.70322106260387, - "queries_per_second": 792.2301342953742, - "filter_queries_per_second": 792.733238328161, - "nested_filter_queries_per_second": 791.4527785640155 - } - ] -} \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 9c269ac..0a5e768 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,30 +1,20 @@ """ Common test fixtures and configuration for VecStream tests. """ -import os -import pytest + import tempfile + import numpy as np -from vecstream import ( - VectorStore, - PersistentVectorStore, - BinaryVectorStore, - HNSWIndex, - Collection, - CollectionManager -) +import pytest + +from vecstream import BinaryVectorStore, CollectionManager, HNSWIndex, VectorStore + @pytest.fixture def vector_store(): """Fixture that provides a clean VectorStore instance.""" return VectorStore() -@pytest.fixture -def persistent_store(): - """Fixture that provides a temporary PersistentVectorStore instance.""" - with tempfile.TemporaryDirectory() as temp_dir: - store = PersistentVectorStore(os.path.join(temp_dir, "test.vec")) - yield store @pytest.fixture def binary_store(): @@ -33,6 +23,7 @@ def binary_store(): store = BinaryVectorStore(temp_dir) yield store + @pytest.fixture def sample_vectors(): """Fixture that provides sample vectors for testing.""" @@ -42,11 +33,13 @@ def sample_vectors(): ("vec3", [0.0, 0.0, 1.0]), ] + @pytest.fixture def hnsw_index(): """Fixture that provides a clean HNSW index.""" return HNSWIndex(dim=3, M=16, ef_construction=100) + @pytest.fixture def populated_hnsw_index(hnsw_index, sample_vectors): """Fixture that provides an HNSW index populated with sample vectors.""" @@ -54,19 +47,22 @@ def populated_hnsw_index(hnsw_index, sample_vectors): hnsw_index.add_item(id, np.array(vector)) return hnsw_index + @pytest.fixture def temp_directory(): """Fixture that provides a temporary directory for testing.""" with tempfile.TemporaryDirectory() as temp_dir: yield temp_dir + @pytest.fixture def collection_manager(temp_directory): """Fixture that provides a CollectionManager instance.""" return CollectionManager(temp_directory) + @pytest.fixture def test_collection(collection_manager): """Fixture that provides a test collection.""" collection = collection_manager.create_collection("test") - return collection \ No newline at end of file + return collection diff --git a/tests/integration/test_features_integration.py b/tests/integration/test_features_integration.py index 12cbccc..96210e0 100644 --- a/tests/integration/test_features_integration.py +++ b/tests/integration/test_features_integration.py @@ -7,15 +7,15 @@ """ import os -import time import shutil -import numpy as np import tempfile + +import numpy as np import pytest -from pathlib import Path +from vecstream.collections import CollectionManager from vecstream.hnsw_index import HNSWIndex -from vecstream.collections import Collection, CollectionManager + @pytest.fixture def temp_dir(): @@ -24,12 +24,13 @@ def temp_dir(): yield temp_dir shutil.rmtree(temp_dir) + def test_hnsw_indexing(temp_dir): """Test the HNSW indexing functionality.""" # Create an HNSW index directly dim = 3 hnsw_index = HNSWIndex(dim=dim, M=16, ef_construction=100) - + # Add vectors to the index vectors = [ ("vec1", np.array([1.0, 0.0, 0.0])), @@ -38,84 +39,117 @@ def test_hnsw_indexing(temp_dir): ("vec4", np.array([0.0, 1.0, 0.0])), ("vec5", np.array([0.0, 0.0, 1.0])), ] - + for id, vector in vectors: hnsw_index.add_item(id, vector) - + # Search for nearest neighbors query = np.array([0.8, 0.2, 0.0]) results = hnsw_index.search(query, k=3) - + # Verify we got reasonable results assert len(results) > 0, "Expected at least one result" for _, sim in results: assert sim >= 0.0, "Expected positive similarity" - + # Test removing a vector hnsw_index.remove_item("vec1") new_results = hnsw_index.search(query, k=3) assert "vec1" not in [r[0] for r in new_results], "vec1 should be removed from results" + def test_collections_and_filtering(temp_dir): """Test collections with metadata filtering.""" # Create a collection manager collection_dir = os.path.join(temp_dir, "collections") manager = CollectionManager(collection_dir, use_hnsw=True) - + # Create two collections images = manager.create_collection("images") texts = manager.create_collection("texts") - + # Add vectors to collections with metadata images.add_vector("img1", [1.0, 0.0, 0.0], {"format": "jpg", "size": "1024x768"}) images.add_vector("img2", [0.0, 1.0, 0.0], {"format": "png", "size": "800x600"}) - + texts.add_vector("text1", [0.0, 0.0, 1.0], {"language": "en", "word_count": 500}) texts.add_vector("text2", [0.7, 0.7, 0.0], {"language": "fr", "word_count": 300}) - + # Verify collections collections = manager.list_collections() assert "images" in collections, "Images collection not found" assert "texts" in collections, "Texts collection not found" - + # Get collection stats img_stats = manager.get_collection_stats("images") - assert img_stats["vector_count"] == 2, f"Expected 2 vectors in images collection, got {img_stats['vector_count']}" - + assert img_stats["vector_count"] == 2, ( + f"Expected 2 vectors in images collection, got {img_stats['vector_count']}" + ) + # Create a collection for testing filtering products = manager.create_collection("products") - + # Add products with various metadata - products.add_vector("prod1", [1.0, 0.0, 0.0], - {"category": "electronics", "price": 100, "in_stock": True, - "details": {"color": "black", "brand": "Sony"}}) - - products.add_vector("prod2", [0.9, 0.1, 0.0], - {"category": "electronics", "price": 200, "in_stock": False, - "details": {"color": "silver", "brand": "Apple"}}) - - products.add_vector("prod3", [0.0, 1.0, 0.0], - {"category": "clothing", "price": 50, "in_stock": True, - "details": {"color": "blue", "size": "M"}}) - - products.add_vector("prod4", [0.0, 0.0, 1.0], - {"category": "books", "price": 15, "in_stock": True, - "details": {"author": "Jane Doe", "pages": 300}}) - + products.add_vector( + "prod1", + [1.0, 0.0, 0.0], + { + "category": "electronics", + "price": 100, + "in_stock": True, + "details": {"color": "black", "brand": "Sony"}, + }, + ) + + products.add_vector( + "prod2", + [0.9, 0.1, 0.0], + { + "category": "electronics", + "price": 200, + "in_stock": False, + "details": {"color": "silver", "brand": "Apple"}, + }, + ) + + products.add_vector( + "prod3", + [0.0, 1.0, 0.0], + { + "category": "clothing", + "price": 50, + "in_stock": True, + "details": {"color": "blue", "size": "M"}, + }, + ) + + products.add_vector( + "prod4", + [0.0, 0.0, 1.0], + { + "category": "books", + "price": 15, + "in_stock": True, + "details": {"author": "Jane Doe", "pages": 300}, + }, + ) + # Test basic metadata filtering query = [0.5, 0.5, 0.0] filter1 = {"category": "electronics", "in_stock": True} - + results1 = products.search_similar(query, k=2, filter_metadata=filter1) assert len(results1) > 0, "No results found with filter" - + # Test nested metadata filtering with dot notation filter2 = {"details.color": "black"} results2 = products.search_similar(query, k=2, filter_metadata=filter2) assert len(results2) > 0, "No results found with nested filter" - + # Test combined filters filter3 = {"category": "electronics", "details.brand": "Apple"} results3 = products.search_similar(query, k=2, filter_metadata=filter3) assert len(results3) > 0, "No results found with combined filter" - assert results3[0][0] == "prod2", f"Expected prod2 as result for combined filter, got {results3[0][0]}" \ No newline at end of file + assert results3[0][0] == "prod2", ( + f"Expected prod2 as result for combined filter, got {results3[0][0]}" + ) diff --git a/tests/property/test_graph_invariants.py b/tests/property/test_graph_invariants.py index 0501256..890c771 100644 --- a/tests/property/test_graph_invariants.py +++ b/tests/property/test_graph_invariants.py @@ -1,5 +1,6 @@ import numpy as np -from hypothesis import given, settings, strategies as st +from hypothesis import given, settings +from hypothesis import strategies as st from vecstream.hnsw_index import HNSWIndex @@ -19,3 +20,57 @@ def test_random_insert_delete_sequences_preserve_invariants(dimension, count, se for i in rng.permutation(count)[: count // 2]: index.remove_item(str(i)) index.validate() + + +operation = st.tuples( + st.sampled_from(["insert", "update", "delete", "search", "validate"]), + st.integers(min_value=0, max_value=24), +) + + +@given( + dimension=st.integers(min_value=1, max_value=8), + seed=st.integers(min_value=0, max_value=2**32 - 1), + operations=st.lists(operation, min_size=1, max_size=80), +) +@settings(max_examples=35, deadline=None) +def test_random_operation_sequences(dimension, seed, operations): + rng = np.random.default_rng(seed) + index = HNSWIndex(dimension, M=4, ef_construction=16, seed=seed) + for action, numeric_id in operations: + item_id = str(numeric_id) + if action in {"insert", "update"}: + index.add_item(item_id, rng.normal(size=dimension).astype(np.float32)) + index.validate() + elif action == "delete" and item_id in index.nodes: + index.remove_item(item_id) + index.validate() + elif action == "search": + results = index.search(rng.normal(size=dimension), k=5, ef_search=10) + assert len(results) <= 5 + assert all(result_id in index.nodes for result_id, _ in results) + assert all(np.isfinite(score) for _, score in results) + assert [score for _, score in results] == sorted( + [score for _, score in results], reverse=True + ) + else: + index.validate() + + +def test_larger_ef_has_nonworse_aggregate_recall(): + rng = np.random.default_rng(123) + vectors = rng.normal(size=(300, 16)).astype(np.float32) + queries = rng.normal(size=(40, 16)).astype(np.float32) + index = HNSWIndex(16, M=8, ef_construction=40, seed=9) + normalized = vectors / np.linalg.norm(vectors, axis=1, keepdims=True) + for i, vector in enumerate(vectors): + index.add_item(str(i), vector) + small_recall = large_recall = 0.0 + for query in queries: + q = query / np.linalg.norm(query) + truth = {str(i) for i in np.argsort(-(normalized @ q))[:10]} + small = {item_id for item_id, _ in index.search(query, k=10, ef_search=10)} + large = {item_id for item_id, _ in index.search(query, k=10, ef_search=80)} + small_recall += len(truth & small) / 10 + large_recall += len(truth & large) / 10 + assert large_recall >= small_recall diff --git a/tests/unit/test_collections.py b/tests/unit/test_collections.py index 8808b7f..12f761a 100644 --- a/tests/unit/test_collections.py +++ b/tests/unit/test_collections.py @@ -1,13 +1,16 @@ """ Unit tests for the Collections and CollectionManager classes. """ + import os -import pytest -import tempfile import shutil -import numpy as np +import tempfile + +import pytest + from vecstream import Collection, CollectionManager + @pytest.fixture def temp_dir(): """Create a temporary directory for testing.""" @@ -15,75 +18,81 @@ def temp_dir(): yield temp_dir shutil.rmtree(temp_dir) + @pytest.fixture def collection_manager(temp_dir): """Fixture that provides a CollectionManager with a temporary directory.""" return CollectionManager(temp_dir, use_hnsw=True) + @pytest.fixture def sample_collection(collection_manager): """Fixture that provides a sample collection with vectors.""" collection = collection_manager.create_collection("test_collection") - + # Add sample vectors with metadata collection.add_vector("vec1", [1.0, 0.0, 0.0], {"tag": "red", "score": 10}) collection.add_vector("vec2", [0.0, 1.0, 0.0], {"tag": "green", "score": 20}) collection.add_vector("vec3", [0.0, 0.0, 1.0], {"tag": "blue", "score": 30}) collection.add_vector("vec4", [0.7, 0.7, 0.0], {"tag": "yellow", "score": 25}) - + return collection + def test_collection_initialization(temp_dir): """Test that Collection initializes correctly.""" collection = Collection("test", temp_dir) - + assert collection.name == "test" assert collection.collection_dir == os.path.join(temp_dir, "collections", "test") assert collection.use_hnsw is True assert collection.hnsw_index is None # Should be None until vectors are added assert os.path.exists(collection.collection_dir) + def test_collection_add_vector(sample_collection): """Test adding vectors to a collection.""" # Test existing vectors assert len(sample_collection.store.vectors) == 4 - + # Add a new vector sample_collection.add_vector("vec5", [0.5, 0.5, 0.5], {"tag": "gray", "score": 15}) - + # Verify it was added to the store assert "vec5" in sample_collection.store.vectors assert sample_collection.store.dimension == 3 - + # Get metadata and verify _, metadata = sample_collection.get_vector_with_metadata("vec5") assert metadata["tag"] == "gray" assert metadata["score"] == 15 - + # Verify it was added to the HNSW index if enabled if sample_collection.use_hnsw and sample_collection.hnsw_index: assert "vec5" in sample_collection.hnsw_index.nodes + def test_collection_remove_vector(sample_collection): """Test removing vectors from a collection.""" # Verify vec1 exists assert "vec1" in sample_collection.store.vectors - + # Remove it sample_collection.remove_vector("vec1") - + # Verify it's gone from the store assert "vec1" not in sample_collection.store.vectors - + # Verify it's gone from the HNSW index if enabled if sample_collection.use_hnsw and sample_collection.hnsw_index: assert "vec1" not in sample_collection.hnsw_index.nodes + def test_collection_search_similar(sample_collection): """Test searching for similar vectors in a collection.""" # Search for vectors similar to [1.0, 0.1, 0.0] results = sample_collection.search_similar([1.0, 0.1, 0.0], k=2) - + assert len(results) == 2 # Check that we get valid results with similarity scores result_ids = [id for id, _ in results] @@ -91,27 +100,29 @@ def test_collection_search_similar(sample_collection): # Check that similarity scores are positive for _, score in results: assert score >= 0.0 - + # Test with k parameter limited_results = sample_collection.search_similar([1.0, 0.1, 0.0], k=1) assert len(limited_results) == 1 + def test_collection_get_vector_count(sample_collection): """Test getting the vector count from a collection.""" assert sample_collection.get_vector_count() == 4 - + # Add another vector sample_collection.add_vector("vec5", [0.5, 0.5, 0.5], {"tag": "gray"}) assert sample_collection.get_vector_count() == 5 - + # Remove a vector sample_collection.remove_vector("vec1") assert sample_collection.get_vector_count() == 4 + def test_collection_manager_initialization(temp_dir): """Test that CollectionManager initializes correctly.""" manager = CollectionManager(temp_dir) - + assert manager.base_storage_dir == temp_dir assert manager.collections_dir == os.path.join(temp_dir, "collections") assert os.path.exists(manager.collections_dir) @@ -119,78 +130,81 @@ def test_collection_manager_initialization(temp_dir): assert manager.collections == {} assert manager.collections_metadata == {} + def test_collection_manager_create_collection(collection_manager): """Test creating collections with the CollectionManager.""" # Create a collection collection = collection_manager.create_collection("test1") - + # Verify it was created assert "test1" in collection_manager.collections assert "test1" in collection_manager.collections_metadata assert collection.name == "test1" assert collection.use_hnsw is True - + # Create another collection with custom HNSW settings collection2 = collection_manager.create_collection( - "test2", - use_hnsw=False, - hnsw_params={"M": 20, "ef_construction": 300} + "test2", use_hnsw=False, hnsw_params={"M": 20, "ef_construction": 300} ) - + assert "test2" in collection_manager.collections assert collection2.use_hnsw is False assert collection2.hnsw_params["M"] == 20 - + # Test that creating an existing collection raises ValueError with pytest.raises(ValueError): collection_manager.create_collection("test1") + def test_collection_manager_get_collection(collection_manager, sample_collection): """Test getting collections from the CollectionManager.""" # Check the existing collection is accessible retrieved = collection_manager.get_collection("test_collection") assert retrieved is sample_collection - + # Test getting a non-existent collection raises KeyError with pytest.raises(KeyError): collection_manager.get_collection("non_existent") + def test_collection_manager_list_collections(collection_manager): """Test listing collections from the CollectionManager.""" # Create a new collection for testing collection_manager.create_collection("test_list_collection") - + # Should have the new collection collections = collection_manager.list_collections() assert len(collections) >= 1 assert "test_list_collection" in collections + def test_collection_manager_delete_collection(collection_manager): """Test deleting collections with the CollectionManager.""" # Create a collection to delete collection_manager.create_collection("to_delete") assert "to_delete" in collection_manager.list_collections() - + # Delete it collection_manager.delete_collection("to_delete") - + # Verify it's gone assert "to_delete" not in collection_manager.list_collections() assert "to_delete" not in collection_manager.collections assert "to_delete" not in collection_manager.collections_metadata - + # Test that deleting a non-existent collection raises KeyError with pytest.raises(KeyError): collection_manager.delete_collection("non_existent") + def test_collection_manager_get_collection_stats(collection_manager, sample_collection): """Test getting collection statistics from the CollectionManager.""" stats = collection_manager.get_collection_stats("test_collection") - + assert stats["name"] == "test_collection" assert stats["vector_count"] == 4 assert stats["dimension"] == 3 assert stats["using_hnsw"] is True assert "vectors_size_bytes" in stats assert "metadata_size_bytes" in stats - assert "total_size_bytes" in stats \ No newline at end of file + assert "total_size_bytes" in stats diff --git a/tests/unit/test_hnsw_index.py b/tests/unit/test_hnsw_index.py index ee243e1..07ec668 100644 --- a/tests/unit/test_hnsw_index.py +++ b/tests/unit/test_hnsw_index.py @@ -1,15 +1,18 @@ """ Unit tests for the HNSW index implementation. """ -import pytest + import numpy as np +import pytest + from vecstream import HNSWIndex + @pytest.fixture def sample_hnsw_index(): """Fixture that provides a small HNSW index with sample vectors.""" index = HNSWIndex(dim=3, M=10, ef_construction=100) - + # Add some sample vectors vectors = [ ("vec1", np.array([1.0, 0.0, 0.0])), @@ -18,16 +21,17 @@ def sample_hnsw_index(): ("vec4", np.array([0.7, 0.7, 0.0])), ("vec5", np.array([0.3, 0.3, 0.3])), ] - + for id, vector in vectors: index.add_item(id, vector) - + return index + def test_hnsw_index_initialization(): """Test that HNSW index initializes correctly.""" index = HNSWIndex(dim=5, M=16, ef_construction=200, ml=4) - + assert index.dim == 5 assert index.M == 16 assert index.M_max0 == 32 # Should be 2*M @@ -37,95 +41,100 @@ def test_hnsw_index_initialization(): assert index.node_levels == {} assert len(index.graphs) == 5 # ml + 1 levels + def test_add_item(sample_hnsw_index): """Test adding items to the HNSW index.""" # Test that existing vectors are in the index assert "vec1" in sample_hnsw_index.nodes assert "vec2" in sample_hnsw_index.nodes assert "vec3" in sample_hnsw_index.nodes - + # Add a new vector new_vec = np.array([0.5, 0.5, 0.5]) sample_hnsw_index.add_item("new_vec", new_vec) - + # Verify it was added assert "new_vec" in sample_hnsw_index.nodes np.testing.assert_array_equal(sample_hnsw_index.nodes["new_vec"], new_vec) - + # Check that nodes dictionary has the item assert "new_vec" in sample_hnsw_index.nodes + def test_remove_item(sample_hnsw_index): """Test removing items from the HNSW index.""" # Verify vec1 exists before removal assert "vec1" in sample_hnsw_index.nodes - + # Remove vec1 sample_hnsw_index.remove_item("vec1") - + # Verify it was removed from nodes, levels, and graphs assert "vec1" not in sample_hnsw_index.nodes assert "vec1" not in sample_hnsw_index.node_levels - + # Check it's removed from all graph levels for level in sample_hnsw_index.graphs: assert "vec1" not in sample_hnsw_index.graphs[level] - + # Also verify no other nodes have connections to vec1 - for node_id, neighbors in sample_hnsw_index.graphs[level].items(): + for neighbors in sample_hnsw_index.graphs[level].values(): assert "vec1" not in neighbors - + # Test removing non-existent item raises KeyError with pytest.raises(KeyError): sample_hnsw_index.remove_item("non_existent_vec") + def test_search(sample_hnsw_index): """Test searching for nearest neighbors in the HNSW index.""" # Search for the nearest neighbor to [1.0, 0.1, 0.0] query = np.array([1.0, 0.1, 0.0]) results = sample_hnsw_index.search(query, k=1) - + # Just check that we get a result with valid ID assert len(results) == 1 assert isinstance(results[0][0], str) # ID should be a string - + # Search for nearest neighbors to [0.5, 0.5, 0.0] query2 = np.array([0.5, 0.5, 0.0]) results2 = sample_hnsw_index.search(query2, k=5) - + # Check that we get multiple results assert len(results2) > 0 - + # Test searching with custom ef_search parameter results3 = sample_hnsw_index.search(query2, k=3, ef_search=50) assert len(results3) <= 3 # May return fewer if not enough vectors + def test_search_empty_index(): """Test searching an empty index returns empty results.""" empty_index = HNSWIndex(dim=3) results = empty_index.search(np.array([1.0, 0.0, 0.0]), k=5) assert len(results) == 0 + def test_distance_calculation(sample_hnsw_index): """Test the distance calculation used by HNSW.""" vec1 = np.array([1.0, 0.0, 0.0]) vec2 = np.array([0.0, 1.0, 0.0]) - + # Cosine distance between orthogonal vectors should be 1.0 distance = sample_hnsw_index._distance(vec1, vec2) assert distance == pytest.approx(1.0) - + # Distance to self should be 0.0 self_distance = sample_hnsw_index._distance(vec1, vec1) assert self_distance == pytest.approx(0.0) - + # Test with non-normalized vectors vec3 = np.array([2.0, 0.0, 0.0]) vec4 = np.array([0.0, 2.0, 0.0]) distance2 = sample_hnsw_index._distance(vec3, vec4) assert distance2 == pytest.approx(1.0) - + # Test with zero vector zero_vec = np.array([0.0, 0.0, 0.0]) zero_distance = sample_hnsw_index._distance(vec1, zero_vec) - assert zero_distance == pytest.approx(1.0) # Should clamp to reasonable distance \ No newline at end of file + assert zero_distance == pytest.approx(1.0) # Should clamp to reasonable distance diff --git a/tests/unit/test_hnsw_internals.py b/tests/unit/test_hnsw_internals.py new file mode 100644 index 0000000..29ec5ba --- /dev/null +++ b/tests/unit/test_hnsw_internals.py @@ -0,0 +1,122 @@ +import numpy as np +import pytest + +from vecstream.errors import IndexInvariantError +from vecstream.hnsw_index import HNSWIndex +from vecstream.vector_store import normalize + + +def make_manual_index() -> HNSWIndex: + index = HNSWIndex(2, M=2, ef_construction=4, seed=1) + vectors = {"a": [1, 0], "b": [0.9, 0.1], "c": [0, 1], "d": [-1, 0]} + for item_id, vector in vectors.items(): + raw = np.asarray(vector, dtype=np.float32) + index.nodes[item_id] = raw + index._normalized_nodes[item_id] = normalize(raw) + index.node_levels[item_id] = 0 + index.graphs = {0: {"a": {"b"}, "b": {"a", "c"}, "c": {"b", "d"}, "d": {"c"}}} + index.ep = "a" + return index + + +def test_search_layer_orders_and_bounds_results(): + index = make_manual_index() + query = normalize(np.asarray([-1, 0], dtype=np.float32)) + one = index._search_layer(query, "a", ef=1, level=0) + many = index._search_layer(query, "a", ef=3, level=0) + assert one == [(0.0, "d")] + assert len(many) == 3 + assert many == sorted(many) + assert many[0] == (0.0, "d") + + +def test_search_layer_handles_sparse_or_missing_level(): + index = make_manual_index() + query = normalize(np.asarray([1, 0], dtype=np.float32)) + assert index._search_layer(query, "a", ef=4, level=7) == [(0.0, "a")] + assert index._search_layer(query, "missing", ef=4, level=0) == [] + + +@pytest.mark.parametrize("candidate_count", [1, 2, 5]) +def test_neighbor_selection_respects_capacity(candidate_count): + index = HNSWIndex(2, M=2, ef_construction=4, seed=1) + for i in range(candidate_count): + angle = i * 0.4 + raw = np.asarray([np.cos(angle), np.sin(angle)], dtype=np.float32) + index.nodes[str(i)] = raw + index._normalized_nodes[str(i)] = normalize(raw) + query = np.asarray([1, 0], dtype=np.float32) + candidates = [ + (index._normalized_distance(query, index._normalized_nodes[str(i)]), str(i)) + for i in range(candidate_count) + ] + selected = index._select_neighbors(query, candidates, M=2) + assert len(selected) == min(candidate_count, 2) + assert len(selected) == len(set(selected)) + + +def test_neighbor_selection_handles_duplicate_and_equidistant_vectors(): + index = HNSWIndex(2, M=2, ef_construction=4, seed=1) + for item_id, raw in {"a": [1, 1], "b": [1, 1], "c": [1, -1]}.items(): + index.nodes[item_id] = np.asarray(raw, dtype=np.float32) + index._normalized_nodes[item_id] = normalize(index.nodes[item_id]) + query = np.asarray([1, 0], dtype=np.float32) + candidates = [ + (index._normalized_distance(query, index._normalized_nodes[i]), i) for i in index.nodes + ] + assert index._select_neighbors(query, candidates, 2) == ["a", "c"] + + +def test_pruning_removes_reciprocal_edges(): + index = HNSWIndex(2, M=2, ef_construction=4, seed=1) + center = np.asarray([1, 0], dtype=np.float32) + for item_id, raw in {"x": center, "a": [1, 0.1], "b": [0, 1], "c": [-1, 0]}.items(): + index.nodes[item_id] = np.asarray(raw, dtype=np.float32) + index._normalized_nodes[item_id] = normalize(index.nodes[item_id]) + index.node_levels[item_id] = 1 + index.graphs = { + 0: {item_id: set() for item_id in index.nodes}, + 1: {item_id: set() for item_id in index.nodes}, + } + index.ep = "x" + index.graphs[1]["x"] = {"a", "b", "c"} + for neighbor in ("a", "b", "c"): + index.graphs[1][neighbor].add("x") + index._prune("x", 1) + removed = {"a", "b", "c"} - index.graphs[1]["x"] + assert len(index.graphs[1]["x"]) <= index.M + assert all("x" not in index.graphs[1][node] for node in removed) + index.validate() + + +def test_entry_point_lifecycle_and_update(): + index = HNSWIndex(2, M=2, ef_construction=4, seed=4) + index._get_random_level = iter([0, 2, 0, 1]).__next__ + index.add_item("first", [1, 0]) + assert index.ep == "first" + index.add_item("high", [0, 1]) + assert index.ep == "high" + index.add_item("ordinary", [-1, 0]) + assert index.ep == "high" + index.remove_item("ordinary") + assert index.ep == "high" + index.remove_item("high") + assert index.ep == "first" + index.add_item("first", [0, -1]) + assert index.nodes["first"].tolist() == [0.0, -1.0] + index.validate() + index.remove_item("first") + assert index.ep is None + index.validate() + + +def test_validate_rejects_self_edge_and_nonmax_entry_point(): + index = make_manual_index() + index.graphs[0]["a"].add("a") + with pytest.raises(IndexInvariantError, match="self-edge"): + index.validate() + index.graphs[0]["a"].remove("a") + index.node_levels["d"] = 1 + index.graphs[1] = {"d": set()} + with pytest.raises(IndexInvariantError, match="entry point.*maximum"): + index.validate() diff --git a/tests/unit/test_hnsw_rigour.py b/tests/unit/test_hnsw_rigour.py index 9984a1b..94fae83 100644 --- a/tests/unit/test_hnsw_rigour.py +++ b/tests/unit/test_hnsw_rigour.py @@ -36,7 +36,11 @@ def test_diversity_heuristic_does_not_only_choose_nearest(): "diverse": [0.940, -0.342], }.items(): index.nodes[item_id] = np.asarray(vector, dtype=np.float32) - index._normalized_nodes[item_id] = index.nodes[item_id] / np.linalg.norm(index.nodes[item_id]) + index._normalized_nodes[item_id] = index.nodes[item_id] / np.linalg.norm( + index.nodes[item_id] + ) query = np.array([1.0, 0.0], dtype=np.float32) - candidates = [(index._normalized_distance(query, index._normalized_nodes[i]), i) for i in index.nodes] + candidates = [ + (index._normalized_distance(query, index._normalized_nodes[i]), i) for i in index.nodes + ] assert index._select_neighbors(query, candidates, 2) == ["near", "diverse"] diff --git a/tests/unit/test_metadata_filtering.py b/tests/unit/test_metadata_filtering.py deleted file mode 100644 index 25c8e11..0000000 --- a/tests/unit/test_metadata_filtering.py +++ /dev/null @@ -1,258 +0,0 @@ -""" -Unit tests for metadata filtering functionality. -""" -import pytest -import tempfile -import shutil -import numpy as np -from vecstream import Collection, QueryEngine, IndexManager - -@pytest.fixture -def temp_dir(): - """Create a temporary directory for testing.""" - temp_dir = tempfile.mkdtemp() - yield temp_dir - shutil.rmtree(temp_dir) - -@pytest.fixture -def sample_collection(temp_dir): - """Fixture that provides a sample collection with various metadata.""" - collection = Collection("test", temp_dir) - - # Add vectors with rich metadata for testing filtering - collection.add_vector("doc1", [1.0, 0.0, 0.0], { - "category": "electronics", - "price": 100, - "in_stock": True, - "tags": ["laptop", "computer"], - "details": { - "brand": "Dell", - "color": "black", - "weight": 2.5 - } - }) - - collection.add_vector("doc2", [0.9, 0.1, 0.0], { - "category": "electronics", - "price": 200, - "in_stock": False, - "tags": ["phone", "mobile"], - "details": { - "brand": "Apple", - "color": "silver", - "weight": 0.2 - } - }) - - collection.add_vector("doc3", [0.0, 1.0, 0.0], { - "category": "clothing", - "price": 50, - "in_stock": True, - "tags": ["shirt", "apparel"], - "details": { - "brand": "Nike", - "color": "blue", - "size": "M" - } - }) - - collection.add_vector("doc4", [0.0, 0.0, 1.0], { - "category": "books", - "price": 15, - "in_stock": True, - "tags": ["fiction", "novel"], - "details": { - "author": "Jane Doe", - "pages": 300 - } - }) - - collection.add_vector("doc5", [0.7, 0.7, 0.0], { - "category": "electronics", - "price": 150, - "in_stock": True, - "tags": ["headphones", "audio"], - "details": { - "brand": "Sony", - "color": "black", - "wireless": True - } - }) - - return collection - -@pytest.fixture -def query_engine(sample_collection): - """Fixture that provides a QueryEngine with IndexManager for the sample collection.""" - index_manager = IndexManager(sample_collection.store) - return QueryEngine(index_manager) - -def test_matches_filter_simple(): - """Test the _matches_filter function with simple key-value filters.""" - from vecstream.query_engine import QueryEngine - query_engine = QueryEngine(None) # We only need the _matches_filter method - - metadata = { - "category": "electronics", - "price": 100, - "in_stock": True - } - - # Test match - assert query_engine._matches_filter(metadata, {"category": "electronics"}) is True - - # Test no match - assert query_engine._matches_filter(metadata, {"category": "books"}) is False - - # Test multiple conditions (AND) - assert query_engine._matches_filter(metadata, { - "category": "electronics", - "price": 100 - }) is True - - assert query_engine._matches_filter(metadata, { - "category": "electronics", - "price": 200 - }) is False - -def test_matches_filter_nested(): - """Test the _matches_filter function with nested dot notation filters.""" - from vecstream.query_engine import QueryEngine - query_engine = QueryEngine(None) # We only need the _matches_filter method - - metadata = { - "category": "electronics", - "details": { - "brand": "Sony", - "color": "black", - "specs": { - "weight": 1.5, - "dimensions": { - "height": 10, - "width": 20 - } - } - } - } - - # Test simple nested field - assert query_engine._matches_filter(metadata, {"details.brand": "Sony"}) is True - assert query_engine._matches_filter(metadata, {"details.brand": "Apple"}) is False - - # Test deeply nested field - assert query_engine._matches_filter(metadata, {"details.specs.weight": 1.5}) is True - assert query_engine._matches_filter(metadata, { - "details.specs.dimensions.height": 10 - }) is True - - # Test combined nested fields - assert query_engine._matches_filter(metadata, { - "category": "electronics", - "details.color": "black" - }) is True - - # Test nested field that doesn't exist - assert query_engine._matches_filter(metadata, {"details.nonexistent": "value"}) is False - assert query_engine._matches_filter(metadata, {"nonexistent.field": "value"}) is False - -def test_collection_search_with_filter(sample_collection): - """Test searching in a collection with metadata filtering.""" - query = [0.5, 0.5, 0.0] - - # Test simple filter - filter1 = {"category": "electronics"} - results1 = sample_collection.search_similar(query, k=5, filter_metadata=filter1) - - # Should return electronics items - assert len(results1) > 0 - result_ids = [id for id, _ in results1] - # Check that at least some of the expected electronics items are returned - electronics_items = set(result_ids).intersection({"doc1", "doc2", "doc5"}) - assert len(electronics_items) > 0 - - # Test combined filter if supported - try: - filter2 = {"category": "electronics", "in_stock": True} - results2 = sample_collection.search_similar(query, k=5, filter_metadata=filter2) - - # Should return in-stock electronics - result_ids = [id for id, _ in results2] - assert len(results2) > 0 - except Exception: - # If complex filtering not supported, this is okay - pass - - # Test nested filter if supported - try: - filter3 = {"details.brand": "Sony"} - results3 = sample_collection.search_similar(query, k=5, filter_metadata=filter3) - - # Should return Sony products - assert len(results3) > 0 - except Exception: - # If nested filtering not supported, this is okay - pass - -def test_query_engine_search_with_filter(query_engine, sample_collection): - """Test searching with the QueryEngine and metadata filtering.""" - query = [0.5, 0.5, 0.0] - - # First, update the index - query_engine.index_manager.update_index() - - # Test simple filter - filter1 = {"price": 100} - results1 = query_engine.search(query, k=5, filter_metadata=filter1) - - # Results may be empty depending on implementation - this is acceptable - if len(results1) > 0: - # Check that results have valid similarity scores - for _, score in results1: - assert score >= 0.0, "Expected positive similarity" - - # Test filter with threshold - filter2 = {"category": "electronics"} - threshold_value = 0.7 - results2 = query_engine.search(query, k=5, filter_metadata=filter2, threshold=threshold_value) - - # For this test we just verify the format of results is correct - for id, similarity in results2: - assert isinstance(id, str) - # Check that similarity is a numeric value (float or numpy.float32/64) - assert hasattr(similarity, "real"), "Similarity should be a numeric value" - assert float(similarity) >= 0.0, "Similarity should be non-negative" - -def test_filter_with_threshold(sample_collection): - """Test filtering with similarity threshold.""" - query = [0.9, 0.1, 0.0] # Very similar to doc1/doc2 - - # Apply filter with high threshold - filter1 = {"category": "electronics"} - results = sample_collection.search_similar( - query, k=5, threshold=0.95, filter_metadata=filter1 - ) - - # Should return only doc1 and doc2 (high similarity electronics) - assert len(results) == 2 - assert results[0][0] in ["doc1", "doc2"] - assert results[1][0] in ["doc1", "doc2"] - - # All results should have similarity >= threshold - for _, similarity in results: - assert similarity >= 0.95 - -def test_empty_filter(sample_collection): - """Test searching with an empty filter.""" - query = [0.5, 0.5, 0.0] - - # Search with empty filter (should match everything) - results = sample_collection.search_similar(query, k=5, filter_metadata={}) - - # Should return at least one item - assert len(results) > 0 - - # Check all results have valid format - for id, similarity in results: - assert isinstance(id, str) - assert isinstance(similarity, float) - assert similarity >= 0.0 \ No newline at end of file diff --git a/tests/unit/test_numerics.py b/tests/unit/test_numerics.py index 590f077..0320618 100644 --- a/tests/unit/test_numerics.py +++ b/tests/unit/test_numerics.py @@ -23,3 +23,24 @@ def test_zero_vector_similarity_is_zero(): store.add_vector("zero", [0, 0]) store.add_vector("unit", [1, 0]) assert dict(store.exact_search([0, 0], k=2, threshold=-1)) == {"zero": 0.0, "unit": 0.0} + + +@pytest.mark.parametrize( + "vector", + [ + [1e-40, -1e-40], + [np.finfo(np.float32).max, -np.finfo(np.float32).max], + [1.0, 1.0 + np.finfo(np.float32).eps], + [-1.0, 0.0], + ], +) +def test_extreme_finite_float32_vectors_remain_searchable(vector): + from vecstream import HNSWIndex + + index = HNSWIndex(2, M=2, ef_construction=4, seed=2) + index.add_item("value", vector) + index.add_item("duplicate", vector) + index.validate() + results = index.search(vector, k=2, ef_search=10) + assert {item_id for item_id, _ in results} == {"value", "duplicate"} + assert all(np.isfinite(score) for _, score in results) diff --git a/tests/unit/test_vector_store.py b/tests/unit/test_vector_store.py index 1ad3b9f..67fbb59 100644 --- a/tests/unit/test_vector_store.py +++ b/tests/unit/test_vector_store.py @@ -1,16 +1,19 @@ """ Unit tests for the VectorStore class. """ -import pytest + import numpy as np + from vecstream import VectorStore + def test_vector_store_initialization(): """Test that VectorStore initializes correctly.""" store = VectorStore() assert store.vectors == {} assert store.dimension is None + def test_add_vector(vector_store, sample_vectors): """Test adding vectors to the store.""" vec_id, vector = sample_vectors[0] @@ -19,25 +22,27 @@ def test_add_vector(vector_store, sample_vectors): np.testing.assert_array_equal(vector_store.vectors[vec_id], vector) assert vector_store.dimension == len(vector) + def test_add_multiple_vectors(vector_store, sample_vectors): """Test adding multiple vectors to the store.""" for vec_id, vector in sample_vectors: vector_store.add_vector(vec_id, vector) - + assert len(vector_store.vectors) == len(sample_vectors) for vec_id, vector in sample_vectors: assert vec_id in vector_store.vectors np.testing.assert_array_equal(vector_store.vectors[vec_id], vector) + def test_search_similar(vector_store, sample_vectors): """Test searching for similar vectors.""" # Add all sample vectors for vec_id, vector in sample_vectors: vector_store.add_vector(vec_id, vector) - + # Search for a vector similar to vec1 query = [1.0, 0.1, 0.1] results = vector_store.search_similar(query, k=2) - + assert len(results) == 2 - assert results[0][0] == "vec1" # Most similar should be vec1 \ No newline at end of file + assert results[0][0] == "vec1" # Most similar should be vec1 diff --git a/vecstream/__init__.py b/vecstream/__init__.py index 7272555..c4c87f4 100644 --- a/vecstream/__init__.py +++ b/vecstream/__init__.py @@ -1,17 +1,7 @@ -""" -VecStream - a compact vector-search engine. -""" +"""VecStream: compact exact and HNSW vector search.""" -__version__ = "0.4.0" - -from .vector_store import VectorStore from .binary_store import BinaryVectorStore -from .persistent_store import PersistentVectorStore -from .index_manager import IndexManager -from .query_engine import QueryEngine -from .hnsw_index import HNSWIndex from .collections import Collection, CollectionManager -from .client import ClientAPI from .errors import ( CorruptStoreError, DimensionMismatchError, @@ -20,21 +10,19 @@ InvalidVectorError, UnsupportedFormatVersionError, ) +from .hnsw_index import HNSWIndex +from .vector_store import VectorStore __all__ = [ - "VectorStore", "BinaryVectorStore", - "PersistentVectorStore", - "IndexManager", - "QueryEngine", - "HNSWIndex", "Collection", "CollectionManager", - "ClientAPI", "CorruptStoreError", "DimensionMismatchError", + "HNSWIndex", "IndexInvariantError", "InvalidManifestError", "InvalidVectorError", "UnsupportedFormatVersionError", + "VectorStore", ] diff --git a/vecstream/binary_store.py b/vecstream/binary_store.py index a0c4608..8d5b3b3 100644 --- a/vecstream/binary_store.py +++ b/vecstream/binary_store.py @@ -48,7 +48,9 @@ def _load_store(self) -> None: try: pointer = json.loads(self.current_file.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: - raise CorruptStoreError(f"cannot read checkpoint pointer {self.current_file}: {exc}") from exc + raise CorruptStoreError( + f"cannot read checkpoint pointer {self.current_file}: {exc}" + ) from exc candidates = [pointer.get("current"), pointer.get("previous")] failures = [] for generation in filter(None, candidates): @@ -70,7 +72,11 @@ def _load_generation(self, generation: str) -> None: f"format version {manifest.get('format_version')!r} is not supported" ) required = {"dimension", "dtype", "metric", "vector_count"} - if not required.issubset(manifest) or manifest["dtype"] != "float32" or manifest["metric"] != "cosine": + if ( + not required.issubset(manifest) + or manifest["dtype"] != "float32" + or manifest["metric"] != "cosine" + ): raise InvalidManifestError("manifest fields, dtype, or metric are invalid") try: ids = json.loads((directory / "ids.json").read_text(encoding="utf-8")) @@ -79,9 +85,14 @@ def _load_generation(self, generation: str) -> None: except (OSError, ValueError, json.JSONDecodeError) as exc: raise CorruptStoreError(f"checkpoint component is unreadable: {exc}") from exc expected_shape = (manifest["vector_count"], manifest["dimension"]) - if vectors.dtype != np.float32 or vectors.shape != expected_shape or len(ids) != manifest["vector_count"]: + if ( + vectors.dtype != np.float32 + or vectors.shape != expected_shape + or len(ids) != manifest["vector_count"] + ): raise CorruptStoreError( - f"vector checkpoint mismatch: expected {expected_shape} float32, got {vectors.shape} {vectors.dtype}" + "vector checkpoint mismatch: " + f"expected {expected_shape} float32, got {vectors.shape} {vectors.dtype}" ) if len(ids) != len(set(ids)) or not isinstance(metadata, dict): raise CorruptStoreError("IDs must be unique and metadata must be an object") @@ -98,8 +109,11 @@ def _save_store(self) -> None: directory.mkdir() ids = list(self.vectors) dimension = self.dimension or 0 - matrix = (np.stack([self.vectors[item_id] for item_id in ids]).astype(np.float32) - if ids else np.empty((0, dimension), dtype=np.float32)) + matrix = ( + np.stack([self.vectors[item_id] for item_id in ids]).astype(np.float32) + if ids + else np.empty((0, dimension), dtype=np.float32) + ) try: vectors_path = directory / "vectors.npy" with vectors_path.open("wb") as handle: @@ -108,13 +122,16 @@ def _save_store(self) -> None: os.fsync(handle.fileno()) _write_json(directory / "ids.json", ids) _write_json(directory / "metadata.json", self.metadata) - _write_json(directory / "manifest.json", { - "format_version": FORMAT_VERSION, - "dimension": dimension, - "dtype": "float32", - "metric": "cosine", - "vector_count": len(ids), - }) + _write_json( + directory / "manifest.json", + { + "format_version": FORMAT_VERSION, + "dimension": dimension, + "dtype": "float32", + "metric": "cosine", + "vector_count": len(ids), + }, + ) dir_fd = os.open(directory, os.O_RDONLY) try: os.fsync(dir_fd) @@ -128,8 +145,9 @@ def _save_store(self) -> None: shutil.rmtree(directory, ignore_errors=True) raise - def add_vector(self, id: str, vector: Sequence[float] | np.ndarray, - metadata: Optional[dict] = None) -> None: + def add_vector( + self, id: str, vector: Sequence[float] | np.ndarray, metadata: Optional[dict] = None + ) -> None: old_vector = self.vectors.get(id) old_metadata = self.metadata.get(id) super().add_vector(id, vector) @@ -175,5 +193,7 @@ def get_store_size(self) -> Tuple[int, int]: if self._generation is None: return 0, 0 directory = self.root / self._generation - return (directory.joinpath("vectors.npy").stat().st_size, - directory.joinpath("metadata.json").stat().st_size) + return ( + directory.joinpath("vectors.npy").stat().st_size, + directory.joinpath("metadata.json").stat().st_size, + ) diff --git a/vecstream/cli.py b/vecstream/cli.py deleted file mode 100644 index fd4ae6d..0000000 --- a/vecstream/cli.py +++ /dev/null @@ -1,541 +0,0 @@ -""" -VecStream CLI - A lightweight vector database with similarity search. -""" - -import os -import sys -import json -import click -import time -from pathlib import Path -from rich.console import Console -from rich import print as rprint -from rich.table import Table -from rich.progress import track -from rich.panel import Panel -from sentence_transformers import SentenceTransformer -import numpy as np - -from .vector_store import VectorStore -from .binary_store import BinaryVectorStore -from .collections import CollectionManager, Collection -from .index_manager import IndexManager -from .query_engine import QueryEngine -from .hnsw_index import HNSWIndex -from .persistent_store import PersistentVectorStore - - -def get_default_store_path(): - """Get the default storage path based on the OS.""" - if os.name == "nt": # Windows - base_dir = os.path.expandvars("%APPDATA%") - return os.path.join(base_dir, "VecStream", "store") - else: # Unix-like - home = str(Path.home()) - return os.path.join(home, ".vecstream", "store") - - -class VecStreamCLI: - def __init__(self, store_path=None, use_hnsw=True): - """Initialize CLI with binary storage. - - Args: - store_path: Optional path to store the vector database. - If not provided, uses the default OS-specific path. - use_hnsw: Whether to use HNSW indexing for better performance - """ - if store_path is None: - store_path = get_default_store_path() - - # Ensure storage directory exists - os.makedirs(store_path, exist_ok=True) - - # Initialize collection manager - self.collection_manager = CollectionManager(store_path, use_hnsw=use_hnsw) - - # Default collection - try: - self.default_collection = self.collection_manager.get_collection("default") - except KeyError: - # Create default collection if it doesn't exist - self.default_collection = self.collection_manager.create_collection( - "default" - ) - - # Store reference to the binary store for backward compatibility - self.store = self.default_collection.store - - # Initialize index and query engine on default collection - self.index_manager = IndexManager(self.store, use_hnsw=use_hnsw) - self.query_engine = QueryEngine(self.index_manager) - - self.model = None - self.store_path = store_path - self.use_hnsw = use_hnsw - - def load_model(self): - """Load the sentence transformer model if not already loaded.""" - if self.model is None: - rprint("[bold cyan]Loading embedding model...[/bold cyan]") - self.model = SentenceTransformer("all-MiniLM-L6-v2") - - def embed_text(self, text): - """Embed text using the sentence transformer model.""" - self.load_model() - return self.model.encode(text).tolist() - - def add_document(self, text, id, collection_name=None, metadata=None): - """Add a document with its embedding to the store.""" - # Get or create the collection - collection = self._get_collection(collection_name) - - # Generate embeddings - embedding = self.embed_text(text) - - # Create default metadata if none provided - if metadata is None: - metadata = {"text": text, "timestamp": time.time()} - elif isinstance(metadata, str): - # Try to parse as JSON - try: - metadata = json.loads(metadata) - except json.JSONDecodeError: - metadata = { - "text": text, - "metadata": metadata, - "timestamp": time.time(), - } - else: - # Ensure text is included in metadata - if "text" not in metadata: - metadata["text"] = text - - # Add to collection - collection.add_vector(id, embedding, metadata) - - # Update index if using default collection - if collection_name is None or collection_name == "default": - self.index_manager.update_index() - - return id - - def search( - self, query_text, collection_name=None, k=5, threshold=0.0, filter_metadata=None - ): - """Search for similar documents to the query text.""" - # Get the collection - collection = self._get_collection(collection_name) - - # Generate embedding for the query - query_embedding = self.embed_text(query_text) - - # Perform search with filtering if needed - if filter_metadata: - # Convert string filter to dict if provided as string - if isinstance(filter_metadata, str): - try: - filter_metadata = json.loads(filter_metadata) - except json.JSONDecodeError: - rprint(f"[bold red]Error: Invalid filter format.[/bold red]") - return [] - - # Use HNSW index directly if available - if collection.use_hnsw and collection.hnsw_index is not None: - results = collection.search_similar( - query_embedding, - k=k, - threshold=threshold, - filter_metadata=filter_metadata, - ) - # Use the query engine otherwise - elif collection_name is None or collection_name == "default": - results = self.query_engine.search( - query_embedding, - k=k, - threshold=threshold, - filter_metadata=filter_metadata, - ) - else: - # Fall back to standard search with manual filtering - all_results = collection.store.search_similar( - query_embedding, k=len(collection.store.vectors) - ) - results = [] - for id, score in all_results: - if score >= threshold: - _, meta = collection.store.get_vector_with_metadata(id) - if meta and self._matches_filter(meta, filter_metadata): - results.append((id, score)) - if len(results) >= k: - break - else: - # No filtering needed - if collection.use_hnsw and collection.hnsw_index is not None: - results = collection.search_similar( - query_embedding, k=k, threshold=threshold - ) - elif collection_name is None or collection_name == "default": - results = self.query_engine.search( - query_embedding, k=k, threshold=threshold - ) - else: - results = collection.store.search_similar( - query_embedding, k=k, threshold=threshold - ) - - return results - - def _matches_filter(self, metadata, filter_query): - """Check if metadata matches the filter query.""" - for key, value in filter_query.items(): - # Handle nested keys with dot notation - if "." in key: - parts = key.split(".") - current = metadata - match = True - - for part in parts[:-1]: - if part not in current or not isinstance(current[part], dict): - match = False - break - current = current[part] - - if not match or parts[-1] not in current: - return False - - # Check value comparison for the last part - if current[parts[-1]] != value: - return False - else: - # Simple key-value match - if key not in metadata or metadata[key] != value: - return False - - return True - - def get_document(self, id, collection_name=None): - """Get a document by ID.""" - collection = self._get_collection(collection_name) - try: - vector, metadata = collection.store.get_vector_with_metadata(id) - return vector, metadata - except KeyError: - return None, None - - def remove_document(self, id, collection_name=None): - """Remove a document by ID.""" - collection = self._get_collection(collection_name) - try: - collection.remove_vector(id) - # Update index if using default collection - if collection_name is None or collection_name == "default": - self.index_manager.update_index() - return True - except KeyError: - return False - - def _get_collection(self, collection_name=None): - """Get a collection by name, falling back to default.""" - if collection_name is None: - return self.default_collection - - try: - return self.collection_manager.get_collection(collection_name) - except KeyError: - # Create collection if it doesn't exist - return self.collection_manager.create_collection( - collection_name, use_hnsw=self.use_hnsw - ) - - def list_collections(self): - """List all available collections.""" - return self.collection_manager.list_collections() - - def get_collection_stats(self, collection_name): - """Get statistics for a collection.""" - try: - return self.collection_manager.get_collection_stats(collection_name) - except KeyError: - return None - - def delete_collection(self, collection_name): - """Delete a collection.""" - if collection_name == "default": - rprint("[bold red]Cannot delete default collection.[/bold red]") - return False - - try: - self.collection_manager.delete_collection(collection_name) - return True - except KeyError: - return False - - -@click.group() -@click.option("--db-path", help="Custom path for vector database storage.") -@click.option( - "--use-hnsw/--no-hnsw", - default=True, - help="Whether to use HNSW indexing for better performance.", -) -@click.pass_context -def cli(ctx, db_path, use_hnsw): - """VecStream - A lightweight vector database with similarity search.""" - ctx.ensure_object(dict) - ctx.obj["cli"] = VecStreamCLI(db_path, use_hnsw=use_hnsw) - - -@cli.command() -@click.argument("text") -@click.argument("id") -@click.option("--collection", "-c", help="Collection to add the document to.") -@click.option("--metadata", "-m", help="JSON metadata to associate with the document.") -@click.pass_context -def add(ctx, text, id, collection, metadata): - """Add a document with TEXT and ID to the database.""" - cli_instance = ctx.obj["cli"] - - try: - if metadata: - try: - metadata_dict = json.loads(metadata) - except json.JSONDecodeError: - rprint(f"[bold red]Error: Invalid JSON metadata.[/bold red]") - return - else: - metadata_dict = {"text": text, "timestamp": time.time()} - - cli_instance.add_document(text, id, collection, metadata_dict) - collection_info = f" in collection '{collection}'" if collection else "" - rprint( - f"[bold green]Added document with ID '{id}'{collection_info}[/bold green]" - ) - except Exception as e: - rprint(f"[bold red]Error adding document: {str(e)}[/bold red]") - - -@cli.command() -@click.argument("query") -@click.option("--k", default=5, help="Number of results to return.") -@click.option("--threshold", default=0.0, help="Minimum similarity threshold (0-1).") -@click.option("--collection", "-c", help="Collection to search in.") -@click.option("--filter", "-f", help="JSON metadata filter to apply to search results.") -@click.pass_context -def search(ctx, query, k, threshold, collection, filter): - """Search for documents similar to QUERY.""" - cli_instance = ctx.obj["cli"] - - try: - filter_dict = None - if filter: - try: - filter_dict = json.loads(filter) - except json.JSONDecodeError: - rprint(f"[bold red]Error: Invalid JSON filter.[/bold red]") - return - - results = cli_instance.search(query, collection, k, threshold, filter_dict) - - if not results: - rprint("[yellow]No matching documents found.[/yellow]") - return - - console = Console() - collection_info = f" in '{collection}'" if collection else "" - table = Table(title=f"Search Results{collection_info}", show_header=True) - table.add_column("ID", style="cyan") - table.add_column("Similarity", style="green") - table.add_column("Text", style="white") - - for id, similarity in results: - _, metadata = cli_instance.get_document(id, collection) - text = metadata.get("text", "N/A") if metadata else "N/A" - # Truncate text if too long - if len(text) > 60: - text = text[:57] + "..." - table.add_row(id, f"{similarity:.4f}", text) - - console.print(table) - except Exception as e: - rprint(f"[bold red]Error searching: {str(e)}[/bold red]") - - -@cli.command() -@click.argument("id") -@click.option("--collection", "-c", help="Collection to get the document from.") -@click.pass_context -def get(ctx, id, collection): - """Get a document by ID.""" - cli_instance = ctx.obj["cli"] - - try: - _, metadata = cli_instance.get_document(id, collection) - - if metadata is None: - collection_info = f" in collection '{collection}'" if collection else "" - rprint( - f"[yellow]Document with ID '{id}'{collection_info} not found.[/yellow]" - ) - return - - console = Console() - collection_info = f" from '{collection}'" if collection else "" - panel = Panel( - json.dumps(metadata, indent=2), - title=f"Document '{id}'{collection_info}", - border_style="green", - ) - console.print(panel) - except Exception as e: - rprint(f"[bold red]Error retrieving document: {str(e)}[/bold red]") - - -@cli.command() -@click.argument("id") -@click.option("--collection", "-c", help="Collection to remove the document from.") -@click.pass_context -def remove(ctx, id, collection): - """Remove a document by ID.""" - cli_instance = ctx.obj["cli"] - - try: - success = cli_instance.remove_document(id, collection) - - if success: - collection_info = f" from collection '{collection}'" if collection else "" - rprint( - f"[bold green]Removed document with ID '{id}'{collection_info}[/bold green]" - ) - else: - collection_info = f" in collection '{collection}'" if collection else "" - rprint( - f"[yellow]Document with ID '{id}'{collection_info} not found.[/yellow]" - ) - except Exception as e: - rprint(f"[bold red]Error removing document: {str(e)}[/bold red]") - - -@cli.command() -@click.option( - "--collection", - "-c", - help="Collection to display info for (omit for all collections).", -) -@click.pass_context -def info(ctx, collection): - """Display information about the vector database.""" - cli_instance = ctx.obj["cli"] - - try: - # If collection specified, show collection info - if collection: - stats = cli_instance.get_collection_stats(collection) - if not stats: - rprint(f"[yellow]Collection '{collection}' not found.[/yellow]") - return - - console = Console() - table = Table(title=f"Collection: {collection}", show_header=True) - table.add_column("Property", style="cyan") - table.add_column("Value", style="green") - - for k, v in stats.items(): - if k != "document_count": - table.add_row(f"{k.replace('_', ' ').title()}", str(v)) - - # Format sizes - vec_size_mb = stats.get("vectors_size_bytes", 0) / (1024 * 1024) - meta_size_mb = stats.get("metadata_size_bytes", 0) / (1024 * 1024) - total_size_mb = stats.get("total_size_bytes", 0) / (1024 * 1024) - - table.add_row("vectors_size", f"{vec_size_mb:.2f} MB") - table.add_row("metadata_size", f"{meta_size_mb:.2f} MB") - table.add_row("total_size", f"{total_size_mb:.2f} MB") - - console.print(table) - - # Show list of collections - else: - collections = cli_instance.list_collections() - - console = Console() - table = Table(title="VecStream Collections", show_header=True) - table.add_column("Collection", style="cyan") - table.add_column("Vectors", style="green") - table.add_column("Size", style="yellow") - table.add_column("HNSW Index", style="blue") - - for coll_name in collections: - stats = cli_instance.get_collection_stats(coll_name) - if stats: - vector_count = stats.get("vector_count", 0) - total_size_mb = stats.get("total_size_bytes", 0) / (1024 * 1024) - using_hnsw = "Yes" if stats.get("using_hnsw", False) else "No" - - table.add_row( - coll_name, - str(vector_count), - f"{total_size_mb:.2f} MB", - using_hnsw, - ) - - console.print(table) - - # Display collection list header - rprint("\n[bold]Available Collections:[/bold]") - - # Display storage path - rprint(f"\n[bold]Storage Location:[/bold] {cli_instance.store_path}") - except Exception as e: - rprint(f"[bold red]Error displaying information: {str(e)}[/bold red]") - - -@cli.command() -@click.argument("name") -@click.option( - "--use-hnsw/--no-hnsw", default=True, help="Whether to use HNSW indexing." -) -@click.pass_context -def create_collection(ctx, name, use_hnsw): - """Create a new collection.""" - cli_instance = ctx.obj["cli"] - - try: - collection = cli_instance.collection_manager.create_collection( - name, use_hnsw=use_hnsw - ) - rprint(f"[bold green]Created collection '{name}'[/bold green]") - except ValueError as e: - rprint(f"[bold red]{str(e)}[/bold red]") - except Exception as e: - rprint(f"[bold red]Error creating collection: {str(e)}[/bold red]") - - -@cli.command() -@click.argument("name") -@click.pass_context -def delete_collection(ctx, name): - """Delete a collection.""" - cli_instance = ctx.obj["cli"] - - try: - if name == "default": - rprint("[bold red]Cannot delete default collection.[/bold red]") - return - - success = cli_instance.delete_collection(name) - if success: - rprint(f"[bold green]Deleted collection '{name}'[/bold green]") - else: - rprint(f"[yellow]Collection '{name}' not found.[/yellow]") - except Exception as e: - rprint(f"[bold red]Error deleting collection: {str(e)}[/bold red]") - - -def main(): - """Main entry point for the CLI.""" - cli(obj={}) - - -if __name__ == "__main__": - main() diff --git a/vecstream/client.py b/vecstream/client.py deleted file mode 100644 index beceae0..0000000 --- a/vecstream/client.py +++ /dev/null @@ -1,86 +0,0 @@ -import socket -import json -import numpy as np -import struct - - -class ClientAPI: - def __init__(self, host="127.0.0.1", port=6379): - self.host = host - self.port = port - - def _send_request(self, request): - """Send a request to the server and get the response.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.connect((self.host, self.port)) - - # Convert request to JSON and encode - data = json.dumps(request).encode("utf-8") - - # Send message length first - sock.send(struct.pack(">I", len(data))) - - # Send the data - sock.send(data) - - # Receive response length - length_data = sock.recv(4) - if not length_data: - return {"status": "error", "message": "Connection closed"} - - msg_length = struct.unpack(">I", length_data)[0] - - # Receive the response in chunks - chunks = [] - bytes_received = 0 - while bytes_received < msg_length: - chunk = sock.recv(min(msg_length - bytes_received, 8192)) - if not chunk: - break - chunks.append(chunk) - bytes_received += len(chunk) - - response_data = b"".join(chunks) - return json.loads(response_data.decode("utf-8")) - - def add_vector(self, id, vector): - """Add a vector to the database.""" - request = { - "command": "add", - "id": id, - "vector": vector.tolist() if isinstance(vector, np.ndarray) else vector, - } - return self._send_request(request) - - def get_vector(self, id): - """Retrieve a vector by ID.""" - request = {"command": "get", "id": id} - response = self._send_request(request) - if response["status"] == "success": - return np.array(response["vector"]) - print(f"Error getting vector: {response['message']}") # Add error logging - return None - - def remove_vector(self, id): - """Remove a vector from the database.""" - request = {"command": "remove", "id": id} - return self._send_request(request) - - def search_similar(self, query_vector, k=10, metric="cosine"): - """Search for similar vectors.""" - request = { - "command": "search", - "query_vector": ( - query_vector.tolist() - if isinstance(query_vector, np.ndarray) - else query_vector - ), - "k": k, - "metric": metric, - } - return self._send_request(request) - - def clear_database(self): - """Clear all vectors from the database.""" - request = {"command": "clear"} - return self._send_request(request) diff --git a/vecstream/collections.py b/vecstream/collections.py index b391654..c6d1ffd 100644 --- a/vecstream/collections.py +++ b/vecstream/collections.py @@ -2,15 +2,15 @@ Collections module for managing multiple vector collections/namespaces. """ -import os import json -from typing import Dict, List, Tuple, Optional, Any, Set +import os import shutil import tempfile +from typing import Any, Dict, List, Optional, Tuple from .binary_store import BinaryVectorStore -from .hnsw_index import HNSWIndex from .errors import CorruptStoreError +from .hnsw_index import HNSWIndex class Collection: @@ -41,7 +41,7 @@ def __init__( # HNSW index settings self.use_hnsw = use_hnsw self.hnsw_params = hnsw_params or {} - self.hnsw_index = None + self.hnsw_index: Optional[HNSWIndex] = None # Initialize HNSW index if required if use_hnsw and self.store.dimension is not None: @@ -63,6 +63,7 @@ def _init_hnsw_index(self) -> None: # Add all existing vectors to the index for id, vector in self.store.vectors.items(): + assert self.hnsw_index is not None self.hnsw_index.add_item(id, vector) def add_vector( @@ -132,7 +133,7 @@ def search_similar( # When only using threshold filtering (no metadata), still use a larger k # to ensure we don't miss potential matches that meet the threshold search_k = max(k * 10, 100) - + results = self.hnsw_index.search(query, k=search_k, ef_search=ef_search) # Apply metadata filtering if needed @@ -170,9 +171,7 @@ def search_similar( # No filtering needed return self.store.search_similar(query, k=k, threshold=threshold) - def _matches_filter( - self, metadata: Dict[str, Any], filter_query: Dict[str, Any] - ) -> bool: + def _matches_filter(self, metadata: Dict[str, Any], filter_query: Dict[str, Any]) -> bool: """Check if metadata matches the filter query. Args: @@ -208,9 +207,7 @@ def _matches_filter( return True - def get_vector_with_metadata( - self, id: str - ) -> Tuple[List[float], Optional[Dict[str, Any]]]: + def get_vector_with_metadata(self, id: str) -> Tuple[List[float], Optional[Dict[str, Any]]]: """Get a vector and its metadata. Args: @@ -319,14 +316,10 @@ def create_collection( # Use defaults if not specified use_hnsw_flag = self.use_hnsw if use_hnsw is None else use_hnsw - hnsw_params_dict = ( - dict(self.default_hnsw_params) if hnsw_params is None else hnsw_params - ) + hnsw_params_dict = dict(self.default_hnsw_params) if hnsw_params is None else hnsw_params # Create and store collection - collection = Collection( - name, self.base_storage_dir, use_hnsw_flag, hnsw_params_dict - ) + collection = Collection(name, self.base_storage_dir, use_hnsw_flag, hnsw_params_dict) self.collections[name] = collection # Save metadata diff --git a/vecstream/hnsw_index.py b/vecstream/hnsw_index.py index e93ec98..2494c60 100644 --- a/vecstream/hnsw_index.py +++ b/vecstream/hnsw_index.py @@ -15,8 +15,14 @@ class HNSWIndex: """Approximate cosine search using the HNSW graph algorithm.""" - def __init__(self, dim: int, M: int = 16, ef_construction: int = 200, - ml: int | None = None, seed: int | None = None) -> None: + def __init__( + self, + dim: int, + M: int = 16, + ef_construction: int = 200, + ml: int | None = None, + seed: int | None = None, + ) -> None: if dim < 1 or M < 2 or ef_construction < M: raise ValueError("dim >= 1, M >= 2, and ef_construction >= M are required") self.dim, self.M, self.M_max0 = dim, M, 2 * M @@ -69,7 +75,9 @@ def _search_layer(self, q: np.ndarray, ep: str, ef: int, level: int) -> List[Tup heapq.heappop(results) return sorted([(-neg_distance, item_id) for neg_distance, item_id in results]) - def _select_neighbors(self, q: np.ndarray, candidates: List[Tuple[float, str]], M: int) -> List[str]: + def _select_neighbors( + self, q: np.ndarray, candidates: List[Tuple[float, str]], M: int + ) -> List[str]: """Select close but directionally diverse neighbors (Algorithm 4).""" by_id = {item_id: distance for distance, item_id in candidates} ordered = sorted(by_id.items(), key=lambda pair: (pair[1], pair[0])) @@ -77,7 +85,10 @@ def _select_neighbors(self, q: np.ndarray, candidates: List[Tuple[float, str]], rejected: List[str] = [] for candidate_id, query_distance in ordered: diverse = all( - self._normalized_distance(self._normalized_nodes[candidate_id], self._normalized_nodes[chosen]) > query_distance + self._normalized_distance( + self._normalized_nodes[candidate_id], self._normalized_nodes[chosen] + ) + > query_distance for chosen in selected ) (selected if diverse else rejected).append(candidate_id) @@ -92,7 +103,12 @@ def _prune(self, node_id: str, level: int) -> None: if len(graph[node_id]) <= limit: return candidates = [ - (self._normalized_distance(self._normalized_nodes[node_id], self._normalized_nodes[neighbor]), neighbor) + ( + self._normalized_distance( + self._normalized_nodes[node_id], self._normalized_nodes[neighbor] + ), + neighbor, + ) for neighbor in graph[node_id] ] keep = set(self._select_neighbors(self._normalized_nodes[node_id], candidates, limit)) @@ -146,10 +162,15 @@ def remove_item(self, id: str) -> None: del self.nodes[id] del self._normalized_nodes[id] del self.node_levels[id] - self.ep = max(self.node_levels, key=lambda item_id: self.node_levels[item_id]) if self.nodes else None + self.ep = ( + max(self.node_levels, key=lambda item_id: self.node_levels[item_id]) + if self.nodes + else None + ) - def search(self, query: Sequence[float] | np.ndarray, k: int = 10, - ef_search: int | None = None) -> List[Tuple[str, float]]: + def search( + self, query: Sequence[float] | np.ndarray, k: int = 10, ef_search: int | None = None + ) -> List[Tuple[str, float]]: if k < 1: raise ValueError("k must be at least 1") q = normalize(as_float32_vector(query, self.dim)) @@ -169,8 +190,19 @@ def validate(self) -> None: raise IndexInvariantError("entry point must exist iff graph is non-empty") if self.ep is not None and self.ep not in self.nodes: raise IndexInvariantError("entry point references a nonexistent node") - if set(self.nodes) != set(self.node_levels) or set(self.nodes) != set(self._normalized_nodes): - raise IndexInvariantError("nodes, normalized nodes, and node_levels contain different IDs") + if self.ep is not None: + maximum_level = max(self.node_levels.values()) + if self.node_levels[self.ep] != maximum_level: + raise IndexInvariantError( + f"entry point {self.ep!r} is at level {self.node_levels[self.ep]}, " + f"but maximum node level is {maximum_level}" + ) + if set(self.nodes) != set(self.node_levels) or set(self.nodes) != set( + self._normalized_nodes + ): + raise IndexInvariantError( + "nodes, normalized nodes, and node_levels contain different IDs" + ) for item_id, max_level in self.node_levels.items(): if max_level < 0: raise IndexInvariantError(f"node {item_id!r} has a negative level") @@ -183,11 +215,32 @@ def validate(self) -> None: if item_id not in self.nodes or level > self.node_levels[item_id]: raise IndexInvariantError(f"invalid node {item_id!r} at level {level}") if len(neighbors) > limit: - raise IndexInvariantError(f"degree {len(neighbors)} exceeds {limit} at level {level}") + raise IndexInvariantError( + f"degree {len(neighbors)} exceeds {limit} at level {level}" + ) for neighbor in neighbors: + if neighbor == item_id: + raise IndexInvariantError( + f"self-edge for node {item_id!r} at level {level}" + ) if neighbor not in self.nodes: raise IndexInvariantError(f"edge references missing node {neighbor!r}") if item_id not in graph.get(neighbor, set()): - raise IndexInvariantError(f"edge {item_id!r}-{neighbor!r} is not bidirectional") - if any(not np.all(np.isfinite(vector)) for vector in self.nodes.values()): - raise IndexInvariantError("stored vectors must be finite") + raise IndexInvariantError( + f"edge {item_id!r}-{neighbor!r} is not bidirectional" + ) + for item_id, vector in self.nodes.items(): + if vector.shape != (self.dim,) or not np.all(np.isfinite(vector)): + raise IndexInvariantError( + f"raw vector for node {item_id!r} must be finite with shape ({self.dim},)" + ) + normalized = self._normalized_nodes[item_id] + if normalized.shape != (self.dim,) or not np.all(np.isfinite(normalized)): + raise IndexInvariantError( + f"normalized vector for node {item_id!r} is malformed or non-finite" + ) + norm = float(np.linalg.norm(normalized.astype(np.float64))) + if not (np.isclose(norm, 0.0, atol=1e-7) or np.isclose(norm, 1.0, atol=1e-5)): + raise IndexInvariantError( + f"normalized vector for node {item_id!r} has norm {norm}, expected 0 or 1" + ) diff --git a/vecstream/index_manager.py b/vecstream/index_manager.py deleted file mode 100644 index fede0bf..0000000 --- a/vecstream/index_manager.py +++ /dev/null @@ -1,163 +0,0 @@ -""" -IndexManager class for managing vector indices for efficient similarity search. -""" - -import numpy as np -from typing import List, Tuple, Optional, Dict, Any -from .vector_store import VectorStore -from .hnsw_index import HNSWIndex - - -class IndexManager: - """A class for managing vector indices for efficient similarity search.""" - - def __init__( - self, - store: VectorStore, - use_hnsw: bool = True, - hnsw_params: Optional[Dict[str, Any]] = None, - ): - """Initialize the index manager. - - Args: - store: VectorStore instance to manage indices for - use_hnsw: Whether to use HNSW indexing (default: True) - hnsw_params: Optional parameters for HNSW index - """ - self.store = store - self.use_hnsw = use_hnsw - self.hnsw_params = hnsw_params or {"M": 16, "ef_construction": 200} - - # Initialize indices - self.index = None # Standard index (normalized vectors) - self.hnsw_index = None # HNSW index - self.indexed_ids = [] - - def create_index(self) -> None: - """Create or recreate the search index.""" - if not self.store.vectors: - return - - # Get all vectors and their IDs - vectors = [] - self.indexed_ids = [] - for id, vector in self.store.vectors.items(): - vectors.append(vector) - self.indexed_ids.append(id) - - # Convert to numpy array and normalize for standard index - vectors_array = np.array(vectors) - norms = np.linalg.norm(vectors_array, axis=1, keepdims=True) - norms[norms == 0] = 1 # Avoid division by zero - self.index = vectors_array / norms - - # Create HNSW index if enabled - if self.use_hnsw and self.store.dimension is not None: - self._create_hnsw_index() - - def _create_hnsw_index(self) -> None: - """Create HNSW index with current vectors.""" - if self.store.dimension is None: - return - - # Create new HNSW index with parameters - M = self.hnsw_params.get("M", 16) - ef_construction = self.hnsw_params.get("ef_construction", 200) - ml = self.hnsw_params.get("ml", None) - - self.hnsw_index = HNSWIndex( - dim=self.store.dimension, M=M, ef_construction=ef_construction, ml=ml - ) - - # Add all vectors to the index - for id, vector in self.store.vectors.items(): - # Convert to numpy array if needed - if not isinstance(vector, np.ndarray): - vector = np.array(vector, dtype=np.float32) - self.hnsw_index.add_item(id, vector) - - def update_index(self) -> None: - """Update the index with new vectors.""" - if self.index is None: - self.create_index() - return - - # Check for new vectors - current_ids = set(self.store.vectors.keys()) - indexed_ids = set(self.indexed_ids) - - if current_ids != indexed_ids: - # HNSW index can be updated incrementally, but for simplicity - # we'll recreate both indices for consistency when the vector set changes - self.create_index() - - def search( - self, - query: List[float], - k: int = 5, - threshold: float = 0.0, - ef_search: Optional[int] = None, - ) -> List[Tuple[str, float]]: - """Search the index for similar vectors. - - Args: - query: Query vector to find similar vectors for - k: Number of similar vectors to return - threshold: Minimum similarity score (0 to 1) for results - ef_search: HNSW search parameter (only used with HNSW) - - Returns: - List of (id, similarity) tuples, sorted by similarity (highest first) - """ - if (self.index is None and self.hnsw_index is None) or not self.indexed_ids: - return [] - - # Use HNSW index if available for better performance - if self.use_hnsw and self.hnsw_index is not None: - # Convert query to numpy array - query_array = np.array(query, dtype=np.float32) - - # Use HNSW index for search - results = self.hnsw_index.search(query_array, k=k, ef_search=ef_search) - - # Filter by threshold if needed - if threshold > 0: - results = [(id, sim) for id, sim in results if sim >= threshold] - - return results - else: - # Fall back to standard search - return self._search_standard(query, k, threshold) - - def _search_standard( - self, query: List[float], k: int = 5, threshold: float = 0.0 - ) -> List[Tuple[str, float]]: - """Search using standard normalized dot product. - - Args: - query: Query vector - k: Number of results to return - threshold: Minimum similarity threshold - - Returns: - List of (id, similarity) tuples - """ - # Normalize query vector - query_array = np.array(query, dtype=np.float32) - query_norm = np.linalg.norm(query_array) - if query_norm > 0: - query_array = query_array / query_norm - - # Compute similarities - similarities = np.dot(self.index, query_array) - - # Filter by threshold and get top k results - results = [] - for idx, sim in enumerate(similarities): - if sim >= threshold: - id = self.indexed_ids[idx] - results.append((id, float(sim))) - - # Sort by similarity (highest first) and return top k - results.sort(key=lambda x: x[1], reverse=True) - return results[:k] diff --git a/vecstream/persistent_store.py b/vecstream/persistent_store.py deleted file mode 100644 index 8a45fc4..0000000 --- a/vecstream/persistent_store.py +++ /dev/null @@ -1,79 +0,0 @@ -""" -PersistentVectorStore class for persistent vector storage. -""" - -import os -import json -import numpy as np -from typing import List, Tuple, Dict, Optional -from .vector_store import VectorStore - - -class PersistentVectorStore(VectorStore): - """A class for persistent vector storage that extends VectorStore.""" - - def __init__(self, filepath: str): - """Initialize a persistent vector store. - - Args: - filepath: Path to the file where vectors will be stored - """ - super().__init__() - self.filepath = filepath - if os.path.exists(filepath): - self.load() - - def add_vector(self, id: str, vector: List[float]) -> None: - """Add a vector to the store and save to disk. - - Args: - id: Unique identifier for the vector - vector: List of float values representing the vector - - Raises: - ValueError: If vector dimensions don't match existing vectors - """ - super().add_vector(id, vector) - self.save() - - def remove_vector(self, id: str) -> None: - """Remove a vector from the store and update the file. - - Args: - id: The vector's identifier - - Raises: - KeyError: If the ID is not found - """ - super().remove_vector(id) - self.save() - - def save(self) -> None: - """Save the current state to disk.""" - data = { - "dimension": self.dimension, - "vectors": {id: vector.tolist() for id, vector in self.vectors.items()}, - } - - # Create directory if it doesn't exist - parent = os.path.dirname(self.filepath) - if parent: - os.makedirs(parent, exist_ok=True) - - # Save to file - with open(self.filepath, "w") as f: - json.dump(data, f) - - def load(self) -> None: - """Load the state from disk.""" - if not os.path.exists(self.filepath): - return - - with open(self.filepath, "r") as f: - data = json.load(f) - - self.dimension = data.get("dimension") - self.vectors = { - id: np.array(vector, dtype=np.float32) - for id, vector in data.get("vectors", {}).items() - } diff --git a/vecstream/query_engine.py b/vecstream/query_engine.py deleted file mode 100644 index 6ac4f7e..0000000 --- a/vecstream/query_engine.py +++ /dev/null @@ -1,129 +0,0 @@ -""" -QueryEngine class for high-level vector query operations. -""" - -from typing import List, Tuple, Optional, Dict, Any -from .index_manager import IndexManager -import numpy as np - -from .vector_store import as_float32_vector, normalize - - -class QueryEngine: - """A class for high-level vector query operations.""" - - def __init__(self, index_manager: IndexManager): - """Initialize the query engine. - - Args: - index_manager: IndexManager instance to use for queries - """ - self.index_manager = index_manager - - def search( - self, - query: List[float], - k: int = 5, - threshold: float = 0.0, - filter_metadata: Optional[Dict[str, Any]] = None, - ) -> List[Tuple[str, float]]: - """Search for similar vectors with optional filtering. - - Args: - query: Query vector to find similar vectors for - k: Number of similar vectors to return - threshold: Minimum similarity score (0 to 1) for results - filter_metadata: Optional metadata filters to apply - - Returns: - List of (id, similarity) tuples, sorted by similarity (highest first) - """ - # Update index before searching - self.index_manager.update_index() - - # If we don't need to filter by metadata, use standard search - if filter_metadata is None: - return self.index_manager.search(query, k, threshold) - - # Get all results above threshold for filtering - # We need to retrieve more results than k because some might be filtered out - get_count = min(max(k * 4, 100), len(self.index_manager.store.vectors)) - results = self.index_manager.search(query, get_count, threshold) - - # Filter results by metadata - filtered_results = [] - for id, similarity in results: - _, metadata = self.index_manager.store.get_vector_with_metadata(id) - if metadata and self._matches_filter(metadata, filter_metadata): - filtered_results.append((id, similarity)) - # Stop once we have k results - if len(filtered_results) >= k: - break - - return filtered_results[:k] - - def _matches_filter( - self, metadata: Dict[str, Any], filter_query: Dict[str, Any] - ) -> bool: - """Check if metadata matches the filter query. - - Args: - metadata: Vector metadata - filter_query: Filter criteria - - Returns: - True if metadata matches filter, False otherwise - """ - for key, value in filter_query.items(): - # Handle nested keys with dot notation (e.g., "user.name") - if "." in key: - parts = key.split(".") - current = metadata - match = True - - for part in parts[:-1]: - if part not in current or not isinstance(current[part], dict): - match = False - break - current = current[part] - - if not match or parts[-1] not in current: - return False - - # Check value comparison for the last part - if current[parts[-1]] != value: - return False - else: - # Simple key-value match - if key not in metadata or metadata[key] != value: - return False - - return True - - def vector_similarity(self, vec1, vec2, metric="cosine"): - """ - Calculate the similarity between two vectors using the specified metric. - - Parameters: - ----------- - vec1 : array-like - First vector - vec2 : array-like - Second vector - metric : str - Similarity metric ('cosine', 'euclidean', etc.) - - Returns: - -------- - float - Similarity score (higher is more similar for cosine, lower for distance metrics) - """ - vec1 = as_float32_vector(vec1) - vec2 = as_float32_vector(vec2, len(vec1)) - - if metric == "cosine": - return float(np.dot(normalize(vec1), normalize(vec2))) - elif metric == "euclidean": - return -float(np.linalg.norm(vec1 - vec2)) - else: - raise ValueError(f"Unsupported metric: {metric}") diff --git a/vecstream/server.py b/vecstream/server.py deleted file mode 100644 index f42d714..0000000 --- a/vecstream/server.py +++ /dev/null @@ -1,127 +0,0 @@ -import socket -import json -import numpy as np -import struct -from threading import Thread -from .persistent_store import PersistentVectorStore -from .index_manager import IndexManager -from .query_engine import QueryEngine - - -class VectorDBServer: - def __init__(self, host="127.0.0.1", port=6379, storage_path="vector_db"): - self.host = host - self.port = port - self.store = PersistentVectorStore(storage_path) - self.index_manager = IndexManager(self.store) - self.query_engine = QueryEngine(self.index_manager) - self.running = False - - def start(self): - """Start the server.""" - self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.server_socket.bind((self.host, self.port)) - self.server_socket.listen(5) - self.running = True - - print(f"Server running on {self.host}:{self.port}") - - try: - while self.running: - client_socket, address = self.server_socket.accept() - print(f"Connection from {address}") - client_thread = Thread(target=self.handle_client, args=(client_socket,)) - client_thread.daemon = True - client_thread.start() - finally: - self.server_socket.close() - - def stop(self): - """Stop the server.""" - self.running = False - self.server_socket.close() - - def handle_client(self, client_socket): - """Handle client connections.""" - try: - while True: - # Receive message length first - length_data = client_socket.recv(4) - if not length_data: - break - - msg_length = struct.unpack(">I", length_data)[0] - - # Receive the full message in chunks - chunks = [] - bytes_received = 0 - while bytes_received < msg_length: - chunk = client_socket.recv(min(msg_length - bytes_received, 8192)) - if not chunk: - break - chunks.append(chunk) - bytes_received += len(chunk) - - if not chunks: - break - - data = b"".join(chunks) - - try: - request = json.loads(data.decode("utf-8")) - response = self.handle_request(request) - - # Send response with length prefix - response_data = json.dumps(response).encode("utf-8") - client_socket.send(struct.pack(">I", len(response_data))) - client_socket.sendall(response_data) - - except Exception as e: - error_response = {"status": "error", "message": str(e)} - response_data = json.dumps(error_response).encode("utf-8") - client_socket.send(struct.pack(">I", len(response_data))) - client_socket.sendall(response_data) - finally: - client_socket.close() - - def handle_request(self, request): - """Handle different types of requests.""" - command = request.get("command") - - if command == "add": - self.store.add_vector(request["id"], request["vector"]) - return {"status": "success", "message": "Vector added"} - - elif command == "get": - try: - return {"status": "success", "vector": self.store.get_vector(request["id"])} - except KeyError: - return {"status": "error", "message": "Vector not found"} - - elif command == "remove": - try: - self.store.remove_vector(request["id"]) - return {"status": "success", "message": "Vector removed"} - except KeyError: - return {"status": "error", "message": "Vector not found"} - - elif command == "search": - self.index_manager.update_index() - results = self.query_engine.search( - np.array(request["query_vector"]), - k=request.get("k", 10), - ) - return { - "status": "success", - "results": [(str(id), float(score)) for id, score in results], - } - - elif command == "clear": - self.store.vectors.clear() - self.store.dimension = None - self.store.save() - return {"status": "success", "message": "Database cleared"} - - else: - return {"status": "error", "message": f"Unknown command: {command}"} diff --git a/vecstream/vector_store.py b/vecstream/vector_store.py index ec5ae07..1899a83 100644 --- a/vecstream/vector_store.py +++ b/vecstream/vector_store.py @@ -9,7 +9,9 @@ from .errors import DimensionMismatchError, InvalidVectorError -def as_float32_vector(vector: Sequence[float] | np.ndarray, dimension: int | None = None) -> np.ndarray: +def as_float32_vector( + vector: Sequence[float] | np.ndarray, dimension: int | None = None +) -> np.ndarray: """Convert and validate a finite 1-D vector. Zero-vector cosine is 0.0.""" try: value: np.ndarray = np.asarray(vector, dtype=np.float32) @@ -28,10 +30,12 @@ def as_float32_vector(vector: Sequence[float] | np.ndarray, dimension: int | Non def normalize(vector: np.ndarray) -> np.ndarray: """Return a normalized float32 copy; a zero vector remains all-zero.""" - norm = float(np.linalg.norm(vector)) + # Accumulate in float64: squaring a finite float32 near its maximum can + # overflow in float32, while subnormal values can underflow to zero. + norm = float(np.linalg.norm(vector.astype(np.float64))) if norm == 0.0: return np.zeros_like(vector, dtype=np.float32) - return np.asarray(vector / norm, dtype=np.float32) + return np.asarray(vector.astype(np.float64) / norm, dtype=np.float32) class VectorStore: @@ -71,8 +75,13 @@ def search_similar( query_value = normalize(as_float32_vector(query, self.dimension)) ids = list(self.vectors) matrix = np.stack([self.vectors[item_id] for item_id in ids]) - norms = np.linalg.norm(matrix, axis=1, keepdims=True) - normalized = np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms != 0) + norms = np.linalg.norm(matrix.astype(np.float64), axis=1, keepdims=True) + normalized = np.divide( + matrix, + norms, + out=np.zeros_like(matrix), + where=norms != 0, + ) scores = normalized @ query_value eligible = np.flatnonzero(scores >= threshold) order = eligible[np.argsort(-scores[eligible], kind="stable")][:k]