llama-index-vector-stores-rabitqlib 0.1.0__tar.gz

llama_index_vector_stores_rabitqlib-0.1.0/.gitignore

@@ -0,0 +1,8 @@
+.venv/
+__pycache__/
+.coverage
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+dist/
+*.egg-info/

llama_index_vector_stores_rabitqlib-0.1.0/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

llama_index_vector_stores_rabitqlib-0.1.0/PKG-INFO

@@ -0,0 +1,108 @@
+Metadata-Version: 2.5
+Name: llama-index-vector-stores-rabitqlib
+Version: 0.1.0
+Summary: LlamaIndex vector store backed by RaBitQ-quantized HNSW (rabitqlib)
+License-Expression: Apache-2.0
+License-File: LICENSE
+Requires-Python: <4.0,>=3.10
+Requires-Dist: llama-index-core<0.15,>=0.13.0
+Requires-Dist: numpy>=1.24
+Requires-Dist: rabitqlib<0.9,>=0.8.3
+Description-Content-Type: text/markdown
+
+# llama-index-vector-stores-rabitqlib
+
+A [LlamaIndex](https://github.com/run-llama/llama_index) vector store backed by [RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library)-quantized HNSW (`rabitqlib.HnswIndex`). It's a local, in-process index, a drop-in for `SimpleVectorStore` or `FaissVectorStore`, and stores compressed vectors (default 7 bits per dimension).
+
+Status: v0.1, not yet published to PyPI.
+
+## Install
+
+```bash
+pip install -e .   # from this folder; pulls llama-index-core and rabitqlib>=0.8.3
+```
+
+## Quickstart
+
+```python
+from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage
+from llama_index.vector_stores.rabitqlib import RabitqHnswVectorStore
+
+vector_store = RabitqHnswVectorStore(dim=1536)  # your embedding model's dimension
+storage_context = StorageContext.from_defaults(vector_store=vector_store)
+index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
+
+nodes = index.as_retriever(similarity_top_k=5).retrieve("my question")
+index.delete_ref_doc("some-doc-id", delete_from_docstore=True)
+
+# Persist and reload. The vector store has to be loaded explicitly, because
+# StorageContext.from_defaults(persist_dir=...) alone would load SimpleVectorStore.
+index.storage_context.persist(persist_dir="./storage")
+vector_store = RabitqHnswVectorStore.from_persist_dir("./storage")
+storage_context = StorageContext.from_defaults(
+    vector_store=vector_store, persist_dir="./storage"
+)
+index = load_index_from_storage(storage_context)
+```
+
+`examples/quickstart.py` is a runnable offline version that uses a toy embedding and no LLM. It passes `min_train_size=1` so its three documents go into the HNSW index. With the default of 1024, a store that small is answered by an exact scan of the buffer and never reaches RaBitQ (see "Training is deferred" below).
+
+## Options
+
+| Option | Default | Meaning |
+|---|---|---|
+| `dim` | required | Embedding dimension, 64 to 65536. |
+| `metric` | `"cosine"` | `"cosine"` (normalizes vectors, then uses inner product), `"ip"`, or `"l2"`. |
+| `nbits` | 7 | RaBitQ bits per dimension, 1 to 9. Fewer bits means less memory and lower recall. |
+| `M`, `ef_construction` | 16, 200 | HNSW graph parameters. |
+| `num_clusters` | 16 | RaBitQ centroids, trained once. |
+| `seed` | 100 | Graph, rotation and clustering seed. |
+| `ef` | `max(k, 64)` | Search breadth; always raised to at least `similarity_top_k`. Can be changed at any time (`store.ef = 128`). |
+| `num_threads` | library default | Threads for add, train and search; `0` means all CPUs. With more than one thread the graph is built in a nondeterministic order, so results can differ slightly between runs. Use `1` for a reproducible graph. |
+| `min_train_size` | 1024 | Embeddings buffered before the index is trained (see below). |
+| `train_vectors` / `centroids` | none | Train up front instead of buffering. Centroids are used as given. |
+
+## How it works
+
+- **Text lives in LlamaIndex's docstore** (`stores_text = False`), as with FAISS. The store keeps only embeddings plus a node-ID ↔ index-ID map and a ref-doc ↔ node map.
+- **Training is deferred.** RaBitQ trains its centroids once, on the first vectors it gets, and can't retrain afterwards. Training on a handful of vectors would hurt every later query. So until `min_train_size` embeddings have arrived they're kept in a float32 buffer and searched exactly. When the threshold is reached, the whole buffer is added to a new `HnswIndex` in one call, which trains on it. Pass `train_vectors=` or `centroids=` to skip the buffer.
+- **Adds are batched.** Each `add(nodes)` makes one `HnswIndex.add` call. Re-adding an existing node ID replaces its embedding.
+- **Scores are higher-is-better**, computed from RaBitQ distance *estimates*. They're close to exact at `nbits=7`, but not exact:
+  - `cosine`: cosine similarity, clipped to [-1, 1];
+  - `ip`: inner product (unbounded);
+  - `l2`: `1 / (1 + squared_l2)`, in (0, 1].
+- **Persistence.** `persist(path)` writes JSON metadata to `path` (for example `default__vector_store.json`). Alongside it go `path.hnsw` (the index), `path.buffer.npy` (vectors still buffered) or `path.train.npy` (training data not yet used). The JSON has a `format_version`. Only the local filesystem is supported.
+- **Thread safety.** Queries and `persist` run concurrently with each other, though two `persist` calls run one at a time. Add, delete and clear are exclusive. Waiting writers don't block new queries, so a continuous stream of overlapping queries can hold an add back indefinitely. That's unlikely with RAG-style traffic.
+
+## Limitations (v0.1)
+
+- **No metadata filters.** A query with `filters`, or `delete_nodes(filters=...)`, raises `ValueError`, and so does a query restricted to `doc_ids`. Only the default query mode is supported, so no MMR, hybrid or sparse search.
+- **`node_ids` on a query is ignored**, as in `FaissVectorStore`, so `index.as_retriever(node_ids=subset)` searches every node. `as_retriever()` always passes the node IDs it knew about when it was created. Honouring that list would break long-lived retrievers after an insert, and would cost O(n) per query.
+- **Deletes don't reclaim memory.** Deleted nodes are excluded from results but stay in storage. To reclaim it, rebuild from your documents.
+- **Accuracy is bounded by quantization.** Recall can plateau below 1.0 however large `ef` is (see below). Raise `nbits` for higher accuracy.
+- **No compaction or retraining.** Centroids are fixed after the first training. If your data drifts a lot, rebuild.
+
+## Measured recall
+
+Run with `examples/recall.py`. Setup: 10,000 base and 200 query vectors, dim 128, drawn from 50 Gaussian clusters (seed 0); default `M=16`, `ef_construction=200`, `num_clusters=16`; graph built with `num_threads=1`, so runs are reproducible; top-10 compared against an exact NumPy scan. Machine: 96-core Linux, rabitqlib 0.8.3, llama-index-core 0.14.25. With multi-threaded builds, expect differences of about ±0.002.
+
+| metric | nbits | ef=64 | ef=128 | ef=256 |
+|---|---|---|---|---|
+| cosine | 4 | 0.7335 | 0.7335 | 0.7335 |
+| cosine | 7 | 0.9510 | 0.9510 | 0.9510 |
+| l2 | 4 | 0.7685 | 0.7680 | 0.7680 |
+| l2 | 7 | 0.9680 | 0.9685 | 0.9685 |
+
+On this dataset, recall is limited by the quantized distance estimates, not by graph search: calling `rabitqlib.HnswIndex` directly gives the same numbers, and they level off from about `ef=64`. `HnswIndex` has no exact re-ranking step, so raising `ef` can't push recall past this ceiling; only a higher `nbits` can. This is one synthetic dataset, so measure on your own embeddings.
+
+## Development
+
+```bash
+python -m pytest --cov=llama_index.vector_stores.rabitqlib
+ruff check . && ruff format --check .
+mypy llama_index tests
+```
+
+## License
+
+Apache-2.0, the same as RaBitQ-Library.

llama_index_vector_stores_rabitqlib-0.1.0/README.md

@@ -0,0 +1,96 @@
+# llama-index-vector-stores-rabitqlib
+
+A [LlamaIndex](https://github.com/run-llama/llama_index) vector store backed by [RaBitQ](https://github.com/VectorDB-NTU/RaBitQ-Library)-quantized HNSW (`rabitqlib.HnswIndex`). It's a local, in-process index, a drop-in for `SimpleVectorStore` or `FaissVectorStore`, and stores compressed vectors (default 7 bits per dimension).
+
+Status: v0.1, not yet published to PyPI.
+
+## Install
+
+```bash
+pip install -e .   # from this folder; pulls llama-index-core and rabitqlib>=0.8.3
+```
+
+## Quickstart
+
+```python
+from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage
+from llama_index.vector_stores.rabitqlib import RabitqHnswVectorStore
+
+vector_store = RabitqHnswVectorStore(dim=1536)  # your embedding model's dimension
+storage_context = StorageContext.from_defaults(vector_store=vector_store)
+index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
+
+nodes = index.as_retriever(similarity_top_k=5).retrieve("my question")
+index.delete_ref_doc("some-doc-id", delete_from_docstore=True)
+
+# Persist and reload. The vector store has to be loaded explicitly, because
+# StorageContext.from_defaults(persist_dir=...) alone would load SimpleVectorStore.
+index.storage_context.persist(persist_dir="./storage")
+vector_store = RabitqHnswVectorStore.from_persist_dir("./storage")
+storage_context = StorageContext.from_defaults(
+    vector_store=vector_store, persist_dir="./storage"
+)
+index = load_index_from_storage(storage_context)
+```
+
+`examples/quickstart.py` is a runnable offline version that uses a toy embedding and no LLM. It passes `min_train_size=1` so its three documents go into the HNSW index. With the default of 1024, a store that small is answered by an exact scan of the buffer and never reaches RaBitQ (see "Training is deferred" below).
+
+## Options
+
+| Option | Default | Meaning |
+|---|---|---|
+| `dim` | required | Embedding dimension, 64 to 65536. |
+| `metric` | `"cosine"` | `"cosine"` (normalizes vectors, then uses inner product), `"ip"`, or `"l2"`. |
+| `nbits` | 7 | RaBitQ bits per dimension, 1 to 9. Fewer bits means less memory and lower recall. |
+| `M`, `ef_construction` | 16, 200 | HNSW graph parameters. |
+| `num_clusters` | 16 | RaBitQ centroids, trained once. |
+| `seed` | 100 | Graph, rotation and clustering seed. |
+| `ef` | `max(k, 64)` | Search breadth; always raised to at least `similarity_top_k`. Can be changed at any time (`store.ef = 128`). |
+| `num_threads` | library default | Threads for add, train and search; `0` means all CPUs. With more than one thread the graph is built in a nondeterministic order, so results can differ slightly between runs. Use `1` for a reproducible graph. |
+| `min_train_size` | 1024 | Embeddings buffered before the index is trained (see below). |
+| `train_vectors` / `centroids` | none | Train up front instead of buffering. Centroids are used as given. |
+
+## How it works
+
+- **Text lives in LlamaIndex's docstore** (`stores_text = False`), as with FAISS. The store keeps only embeddings plus a node-ID ↔ index-ID map and a ref-doc ↔ node map.
+- **Training is deferred.** RaBitQ trains its centroids once, on the first vectors it gets, and can't retrain afterwards. Training on a handful of vectors would hurt every later query. So until `min_train_size` embeddings have arrived they're kept in a float32 buffer and searched exactly. When the threshold is reached, the whole buffer is added to a new `HnswIndex` in one call, which trains on it. Pass `train_vectors=` or `centroids=` to skip the buffer.
+- **Adds are batched.** Each `add(nodes)` makes one `HnswIndex.add` call. Re-adding an existing node ID replaces its embedding.
+- **Scores are higher-is-better**, computed from RaBitQ distance *estimates*. They're close to exact at `nbits=7`, but not exact:
+  - `cosine`: cosine similarity, clipped to [-1, 1];
+  - `ip`: inner product (unbounded);
+  - `l2`: `1 / (1 + squared_l2)`, in (0, 1].
+- **Persistence.** `persist(path)` writes JSON metadata to `path` (for example `default__vector_store.json`). Alongside it go `path.hnsw` (the index), `path.buffer.npy` (vectors still buffered) or `path.train.npy` (training data not yet used). The JSON has a `format_version`. Only the local filesystem is supported.
+- **Thread safety.** Queries and `persist` run concurrently with each other, though two `persist` calls run one at a time. Add, delete and clear are exclusive. Waiting writers don't block new queries, so a continuous stream of overlapping queries can hold an add back indefinitely. That's unlikely with RAG-style traffic.
+
+## Limitations (v0.1)
+
+- **No metadata filters.** A query with `filters`, or `delete_nodes(filters=...)`, raises `ValueError`, and so does a query restricted to `doc_ids`. Only the default query mode is supported, so no MMR, hybrid or sparse search.
+- **`node_ids` on a query is ignored**, as in `FaissVectorStore`, so `index.as_retriever(node_ids=subset)` searches every node. `as_retriever()` always passes the node IDs it knew about when it was created. Honouring that list would break long-lived retrievers after an insert, and would cost O(n) per query.
+- **Deletes don't reclaim memory.** Deleted nodes are excluded from results but stay in storage. To reclaim it, rebuild from your documents.
+- **Accuracy is bounded by quantization.** Recall can plateau below 1.0 however large `ef` is (see below). Raise `nbits` for higher accuracy.
+- **No compaction or retraining.** Centroids are fixed after the first training. If your data drifts a lot, rebuild.
+
+## Measured recall
+
+Run with `examples/recall.py`. Setup: 10,000 base and 200 query vectors, dim 128, drawn from 50 Gaussian clusters (seed 0); default `M=16`, `ef_construction=200`, `num_clusters=16`; graph built with `num_threads=1`, so runs are reproducible; top-10 compared against an exact NumPy scan. Machine: 96-core Linux, rabitqlib 0.8.3, llama-index-core 0.14.25. With multi-threaded builds, expect differences of about ±0.002.
+
+| metric | nbits | ef=64 | ef=128 | ef=256 |
+|---|---|---|---|---|
+| cosine | 4 | 0.7335 | 0.7335 | 0.7335 |
+| cosine | 7 | 0.9510 | 0.9510 | 0.9510 |
+| l2 | 4 | 0.7685 | 0.7680 | 0.7680 |
+| l2 | 7 | 0.9680 | 0.9685 | 0.9685 |
+
+On this dataset, recall is limited by the quantized distance estimates, not by graph search: calling `rabitqlib.HnswIndex` directly gives the same numbers, and they level off from about `ef=64`. `HnswIndex` has no exact re-ranking step, so raising `ef` can't push recall past this ceiling; only a higher `nbits` can. This is one synthetic dataset, so measure on your own embeddings.
+
+## Development
+
+```bash
+python -m pytest --cov=llama_index.vector_stores.rabitqlib
+ruff check . && ruff format --check .
+mypy llama_index tests
+```
+
+## License
+
+Apache-2.0, the same as RaBitQ-Library.

llama_index_vector_stores_rabitqlib-0.1.0/llama_index/vector_stores/rabitqlib/__init__.py

@@ -0,0 +1,3 @@
+from llama_index.vector_stores.rabitqlib.base import RabitqHnswVectorStore
+
+__all__ = ["RabitqHnswVectorStore"]

llama_index_vector_stores_rabitqlib-0.1.0/llama_index/vector_stores/rabitqlib/base.py

@@ -0,0 +1,625 @@
+"""LlamaIndex vector store backed by ``rabitqlib.HnswIndex``."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import threading
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from typing import Any, Literal
+
+import fsspec
+import numpy as np
+import numpy.typing as npt
+from fsspec.implementations.local import LocalFileSystem
+from llama_index.core.bridge.pydantic import Field, PrivateAttr
+from llama_index.core.schema import BaseNode
+from llama_index.core.vector_stores.simple import DEFAULT_VECTOR_STORE, NAMESPACE_SEP
+from llama_index.core.vector_stores.types import (
+    DEFAULT_PERSIST_DIR,
+    DEFAULT_PERSIST_FNAME,
+    BasePydanticVectorStore,
+    MetadataFilters,
+    VectorStoreQuery,
+    VectorStoreQueryMode,
+    VectorStoreQueryResult,
+)
+from rabitqlib import HnswIndex
+
+logger = logging.getLogger(__name__)
+
+FORMAT_VERSION = 1
+DEFAULT_PERSIST_PATH = os.path.join(
+    DEFAULT_PERSIST_DIR, f"{DEFAULT_VECTOR_STORE}{NAMESPACE_SEP}{DEFAULT_PERSIST_FNAME}"
+)
+_MISSING_ID = 0xFFFFFFFF
+_HNSW_SUFFIX = ".hnsw"
+_BUFFER_SUFFIX = ".buffer.npy"
+_TRAIN_SUFFIX = ".train.npy"
+
+Metric = Literal["cosine", "l2", "ip"]
+FloatArray = npt.NDArray[np.float32]
+
+
+class _RWLock:
+    """
+    Many concurrent readers or one writer (not reentrant).
+
+    Readers are not blocked by a waiting writer, so a continuous stream of
+    overlapping queries can delay an add or delete indefinitely. That is
+    unlikely with RAG-style traffic, where queries have gaps between them.
+    """
+
+    def __init__(self) -> None:
+        self._cond = threading.Condition()
+        self._readers = 0
+        self._writer = False
+
+    @contextmanager
+    def read(self) -> Iterator[None]:
+        with self._cond:
+            while self._writer:
+                self._cond.wait()
+            self._readers += 1
+        try:
+            yield
+        finally:
+            with self._cond:
+                self._readers -= 1
+                if self._readers == 0:
+                    self._cond.notify_all()
+
+    @contextmanager
+    def write(self) -> Iterator[None]:
+        with self._cond:
+            while self._writer or self._readers:
+                self._cond.wait()
+            self._writer = True
+        try:
+            yield
+        finally:
+            with self._cond:
+                self._writer = False
+                self._cond.notify_all()
+
+
+def _check_local_fs(fs: fsspec.AbstractFileSystem | None) -> None:
+    if fs is not None and not isinstance(fs, LocalFileSystem):
+        raise NotImplementedError(
+            "RabitqHnswVectorStore only supports the local filesystem."
+        )
+
+
+class RabitqHnswVectorStore(BasePydanticVectorStore):
+    """
+    Vector store backed by a RaBitQ-quantized HNSW graph (``rabitqlib.HnswIndex``).
+
+    Only embeddings are stored (``stores_text = False``); LlamaIndex keeps node
+    text and metadata in its docstore. Metadata filters are not supported, and
+    ``VectorStoreQuery.node_ids`` is ignored (as in ``FaissVectorStore``), so
+    ``as_retriever(node_ids=subset)`` does not restrict results.
+
+    Until ``min_train_size`` embeddings have been added, they are buffered and
+    queried by exact search. Once the threshold is reached the whole buffer is
+    added to a new ``HnswIndex`` in one call, which trains its centroids on it.
+    Pass ``train_vectors`` or ``centroids`` to skip buffering.
+
+    Similarity scores are computed from RaBitQ distance estimates (close to
+    exact at the default ``nbits=7``, but not exact) and are higher-is-better:
+
+    - ``cosine``: cosine similarity in ``[-1, 1]``.
+    - ``ip``: inner product (unbounded).
+    - ``l2``: ``1 / (1 + squared_l2)`` in ``(0, 1]``.
+
+    Deleting nodes hides them from results but does not reclaim storage.
+
+    Examples:
+        ```python
+        from llama_index.core import StorageContext, VectorStoreIndex
+        from llama_index.vector_stores.rabitqlib import RabitqHnswVectorStore
+
+        vector_store = RabitqHnswVectorStore(dim=1536)
+        storage_context = StorageContext.from_defaults(vector_store=vector_store)
+        index = VectorStoreIndex.from_documents(
+            documents, storage_context=storage_context
+        )
+        ```
+
+    """
+
+    stores_text: bool = False
+
+    dim: int = Field(ge=64, le=65536, description="Embedding dimension.")
+    metric: Metric = Field(default="cosine", description="cosine, l2 or ip.")
+    nbits: int = Field(default=7, ge=1, le=9, description="RaBitQ bits per dim.")
+    M: int = Field(default=16, ge=2, description="HNSW graph connectivity.")
+    ef_construction: int = Field(default=200, ge=1)
+    num_clusters: int = Field(default=16, ge=1)
+    seed: int = Field(default=100, description="Graph and rotation seed.")
+    ef: int | None = Field(
+        default=None, ge=1, description="Search breadth; default max(k, 64)."
+    )
+    num_threads: int | None = Field(
+        default=None, ge=0, description="Threads for add and search; 0 = all CPUs."
+    )
+    min_train_size: int = Field(
+        default=1024, ge=1, description="Embeddings buffered before training."
+    )
+
+    _index: HnswIndex | None = PrivateAttr(default=None)
+    _node_to_label: dict[str, int] = PrivateAttr(default_factory=dict)
+    _label_to_node: dict[int, str] = PrivateAttr(default_factory=dict)
+    _node_to_ref: dict[str, str] = PrivateAttr(default_factory=dict)
+    _ref_to_nodes: dict[str, set[str]] = PrivateAttr(default_factory=dict)
+    _buf_vecs: Any = PrivateAttr(default=None)
+    _buf_ids: list[str] = PrivateAttr(default_factory=list)
+    _pending_train: tuple[str, Any] | None = PrivateAttr(default=None)
+    _lock: _RWLock = PrivateAttr(default_factory=_RWLock)
+    # Serializes persist() calls; they share temp-file names.
+    _persist_lock: Any = PrivateAttr(default_factory=threading.Lock)
+
+    def __init__(
+        self,
+        dim: int,
+        *,
+        metric: Metric = "cosine",
+        nbits: int = 7,
+        M: int = 16,
+        ef_construction: int = 200,
+        num_clusters: int = 16,
+        seed: int = 100,
+        ef: int | None = None,
+        num_threads: int | None = None,
+        min_train_size: int = 1024,
+        train_vectors: npt.ArrayLike | None = None,
+        centroids: npt.ArrayLike | None = None,
+    ) -> None:
+        """
+        Create an empty store.
+
+        Args:
+            dim: Embedding dimension, 64 through 65536.
+            metric: ``"cosine"`` (default), ``"l2"`` or ``"ip"``.
+            nbits: RaBitQ quantization bits, 1 through 9.
+            M: HNSW connectivity.
+            ef_construction: HNSW construction breadth.
+            num_clusters: Number of RaBitQ centroids to train.
+            seed: Graph, rotation and clustering seed.
+            ef: Search breadth. Defaults to ``max(k, 64)`` and is raised to ``k``.
+            num_threads: Threads for add, train and search (library default if None).
+            min_train_size: Embeddings to buffer before training the index.
+            train_vectors: Optional sample to train centroids on; disables buffering.
+            centroids: Optional ``(num_clusters, dim)`` centroids, used as given;
+                disables buffering. Overrides ``num_clusters``.
+
+        """
+        pending: tuple[str, FloatArray] | None = None
+        if train_vectors is not None and centroids is not None:
+            raise ValueError("Pass at most one of train_vectors and centroids.")
+        if centroids is not None:
+            c = self._as_matrix(centroids, dim, "centroids")
+            num_clusters = c.shape[0]
+            pending = ("centroids", c)
+        elif train_vectors is not None:
+            t = self._as_matrix(train_vectors, dim, "train_vectors")
+            if metric == "cosine":
+                t = self._normalize(t)
+            pending = ("train", t)
+
+        super().__init__(
+            stores_text=False,
+            dim=dim,
+            metric=metric,
+            nbits=nbits,
+            M=M,
+            ef_construction=ef_construction,
+            num_clusters=num_clusters,
+            seed=seed,
+            ef=ef,
+            num_threads=num_threads,
+            min_train_size=min_train_size,
+        )
+        self._buf_vecs = np.empty((0, dim), dtype=np.float32)
+        self._pending_train = pending
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "RabitqHnswVectorStore"
+
+    @property
+    def client(self) -> HnswIndex | None:
+        """The underlying ``HnswIndex``, or ``None`` while the store is buffering."""
+        return self._index
+
+    # No __len__: StorageContext.from_defaults tests `if vector_store:`, and an
+    # empty store must not be falsy or it is silently replaced.
+    @property
+    def num_nodes(self) -> int:
+        """Number of live (searchable) nodes."""
+        with self._lock.read():
+            return len(self._buf_ids) + len(self._node_to_label)
+
+    # ------------------------------------------------------------------ helpers
+
+    @staticmethod
+    def _as_matrix(data: npt.ArrayLike, dim: int, what: str) -> FloatArray:
+        arr = np.ascontiguousarray(data, dtype=np.float32)
+        if arr.ndim != 2 or arr.shape[1] != dim or arr.shape[0] == 0:
+            raise ValueError(f"{what} must have shape (n, {dim}) with n >= 1.")
+        if not np.all(np.isfinite(arr)):
+            raise ValueError(f"{what} must contain only finite values.")
+        return arr
+
+    @staticmethod
+    def _normalize(x: FloatArray) -> FloatArray:
+        norms = np.linalg.norm(x, axis=1, keepdims=True)
+        if np.any(norms == 0):
+            raise ValueError("Cannot use zero vectors with the cosine metric.")
+        return np.ascontiguousarray(x / norms, dtype=np.float32)
+
+    def _prepare(self, x: npt.ArrayLike, what: str) -> FloatArray:
+        arr = self._as_matrix(x, self.dim, what)
+        return self._normalize(arr) if self.metric == "cosine" else arr
+
+    def _hnsw_metric(self) -> Literal["l2", "ip"]:
+        return "l2" if self.metric == "l2" else "ip"
+
+    def _new_index(self) -> HnswIndex:
+        index = HnswIndex(
+            self.dim,
+            nbits=self.nbits,
+            num_clusters=self.num_clusters,
+            M=self.M,
+            ef_construction=self.ef_construction,
+            metric=self._hnsw_metric(),
+            seed=self.seed,
+        )
+        if self._pending_train is not None:
+            kind, data = self._pending_train
+            if kind == "centroids":
+                index.set_centroids(data)
+            else:
+                index.train(data, num_threads=self.num_threads)
+        return index
+
+    def _to_similarity(self, d: FloatArray) -> list[float]:
+        d = d.astype(np.float64)
+        if self.metric == "cosine":
+            sims = np.clip(1.0 - d, -1.0, 1.0)
+        elif self.metric == "ip":
+            sims = 1.0 - d
+        else:
+            sims = 1.0 / (1.0 + np.maximum(d, 0.0))
+        return [float(s) for s in sims]
+
+    def _exact_distances(self, x: FloatArray, q: FloatArray) -> FloatArray:
+        if self.metric == "l2":
+            diff = x - q
+            return np.einsum("ij,ij->i", diff, diff).astype(np.float32)
+        return (1.0 - x @ q).astype(np.float32)
+
+    def _unlink_ref(self, node_id: str) -> None:
+        ref = self._node_to_ref.pop(node_id, None)
+        if ref is not None:
+            members = self._ref_to_nodes.get(ref)
+            if members is not None:
+                members.discard(node_id)
+                if not members:
+                    del self._ref_to_nodes[ref]
+
+    def _remove_unlocked(self, node_ids: Sequence[str]) -> None:
+        targets = set(node_ids)
+        if not targets:
+            return
+        if self._buf_ids:
+            keep = [i for i, nid in enumerate(self._buf_ids) if nid not in targets]
+            if len(keep) != len(self._buf_ids):
+                self._buf_vecs = self._buf_vecs[keep]
+                self._buf_ids = [self._buf_ids[i] for i in keep]
+        labels = [self._node_to_label[n] for n in targets if n in self._node_to_label]
+        if labels and self._index is not None:
+            self._index.remove(np.asarray(labels, dtype=np.int64))
+            for label in labels:
+                del self._node_to_label[self._label_to_node.pop(label)]
+        for nid in targets:
+            self._unlink_ref(nid)
+
+    # --------------------------------------------------------------- interface
+
+    def add(self, nodes: Sequence[BaseNode], **add_kwargs: Any) -> list[str]:
+        """
+        Add nodes with embeddings; returns their node IDs.
+
+        Re-adding an existing node ID replaces its embedding.
+        """
+        if not nodes:
+            return []
+        last: dict[str, int] = {}
+        for i, node in enumerate(nodes):
+            last[node.node_id] = i
+        unique = [nodes[i] for i in sorted(last.values())]
+        rows = []
+        for node in unique:
+            if node.embedding is None:
+                raise ValueError(f"Node {node.node_id} has no embedding.")
+            if len(node.embedding) != self.dim:
+                raise ValueError(
+                    f"Node {node.node_id} embedding has dimension "
+                    f"{len(node.embedding)}, expected {self.dim}."
+                )
+            rows.append(node.embedding)
+        vecs = self._prepare(rows, "embeddings")
+        ids = [node.node_id for node in unique]
+        new_set = set(ids)
+
+        with self._lock.write():
+            if self._index is None:
+                keep = [i for i, nid in enumerate(self._buf_ids) if nid not in new_set]
+                buf_vecs = self._buf_vecs[keep]
+                buf_ids = [self._buf_ids[i] for i in keep]
+                all_vecs = np.ascontiguousarray(np.vstack([buf_vecs, vecs]))
+                all_ids = buf_ids + ids
+                if self._pending_train is None and len(all_ids) < self.min_train_size:
+                    self._buf_vecs, self._buf_ids = all_vecs, all_ids
+                else:
+                    index = self._new_index()
+                    labels = index.add(all_vecs, num_threads=self.num_threads)
+                    self._index = index
+                    self._pending_train = None
+                    self._buf_vecs = np.empty((0, self.dim), dtype=np.float32)
+                    self._buf_ids = []
+                    for nid, label in zip(all_ids, labels.tolist(), strict=True):
+                        self._node_to_label[nid] = label
+                        self._label_to_node[label] = nid
+            else:
+                labels = self._index.add(vecs, num_threads=self.num_threads)
+                old = [self._node_to_label[n] for n in ids if n in self._node_to_label]
+                if old:
+                    self._index.remove(np.asarray(old, dtype=np.int64))
+                    for label in old:
+                        del self._label_to_node[label]
+                for nid, label in zip(ids, labels.tolist(), strict=True):
+                    self._node_to_label[nid] = label
+                    self._label_to_node[label] = nid
+            for node in unique:
+                self._unlink_ref(node.node_id)
+                ref = node.ref_doc_id
+                if ref is not None:
+                    self._node_to_ref[node.node_id] = ref
+                    self._ref_to_nodes.setdefault(ref, set()).add(node.node_id)
+        return [node.node_id for node in nodes]
+
+    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
+        """Delete every node that belongs to ``ref_doc_id``."""
+        with self._lock.write():
+            node_ids = list(self._ref_to_nodes.get(ref_doc_id, ()))
+            self._remove_unlocked(node_ids)
+
+    def delete_nodes(
+        self,
+        node_ids: list[str] | None = None,
+        filters: MetadataFilters | None = None,
+        **delete_kwargs: Any,
+    ) -> None:
+        """Delete nodes by ID. Metadata filters are not supported."""
+        if filters is not None:
+            raise ValueError("Metadata filters are not supported.")
+        if not node_ids:
+            return
+        with self._lock.write():
+            self._remove_unlocked(node_ids)
+
+    def clear(self) -> None:
+        """
+        Remove everything and return to the empty, buffering state.
+
+        Any ``train_vectors`` or ``centroids`` given at construction are not kept.
+        """
+        with self._lock.write():
+            self._index = None
+            self._pending_train = None
+            self._node_to_label = {}
+            self._label_to_node = {}
+            self._node_to_ref = {}
+            self._ref_to_nodes = {}
+            self._buf_vecs = np.empty((0, self.dim), dtype=np.float32)
+            self._buf_ids = []
+
+    def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
+        """Return the ``similarity_top_k`` nearest node IDs and similarities."""
+        if query.filters is not None:
+            raise ValueError("Metadata filters are not supported.")
+        if query.doc_ids:
+            raise ValueError("Restricting queries by doc_ids is not supported.")
+        if query.mode != VectorStoreQueryMode.DEFAULT:
+            raise ValueError(f"Query mode {query.mode!r} is not supported.")
+        if query.query_embedding is None:
+            raise ValueError("query_embedding is required.")
+        if len(query.query_embedding) != self.dim:
+            raise ValueError(
+                f"Query embedding has dimension {len(query.query_embedding)}, "
+                f"expected {self.dim}."
+            )
+        q = self._prepare([query.query_embedding], "query_embedding")
+        k = query.similarity_top_k
+
+        ids: list[str] = []
+        dists = np.empty(0, dtype=np.float32)
+        # query.node_ids is deliberately ignored: VectorIndexRetriever always
+        # passes every node ID it knew of when it was created, which goes stale
+        # after an insert, and checking it would cost O(n) per query.
+        with self._lock.read():
+            if k < 1:
+                pass
+            elif self._index is None:
+                n = len(self._buf_ids)
+                if n:
+                    d = self._exact_distances(self._buf_vecs, q[0])
+                    order = np.argsort(d, kind="stable")[: min(k, n)]
+                    ids = [self._buf_ids[i] for i in order.tolist()]
+                    dists = d[order]
+            elif len(self._node_to_label):
+                kk = min(k, len(self._node_to_label))
+                ef = max(self.ef if self.ef is not None else max(kk, 64), kk)
+                labels, d = self._index.search(
+                    q, kk, ef=ef, num_threads=self.num_threads
+                )
+                keep = [
+                    j
+                    for j, (label, dist) in enumerate(
+                        zip(labels[0].tolist(), d[0], strict=True)
+                    )
+                    if label != _MISSING_ID
+                    and np.isfinite(dist)
+                    and label in self._label_to_node
+                ]
+                ids = [self._label_to_node[int(labels[0][j])] for j in keep]
+                dists = d[0][keep]
+        return VectorStoreQueryResult(
+            nodes=None, similarities=self._to_similarity(dists), ids=ids
+        )
+
+    # ------------------------------------------------------------- persistence
+
+    def persist(
+        self,
+        persist_path: str = DEFAULT_PERSIST_PATH,
+        fs: fsspec.AbstractFileSystem | None = None,
+    ) -> None:
+        """
+        Save the store.
+
+        Writes JSON metadata to ``persist_path`` and, as needed, the index to
+        ``<persist_path>.hnsw``, the buffer to ``<persist_path>.buffer.npy`` and
+        pending training data to ``<persist_path>.train.npy``.
+
+        Queries keep running while saving; adds and deletes wait for it.
+        """
+        _check_local_fs(fs)
+        dirpath = os.path.dirname(persist_path)
+        if dirpath:
+            os.makedirs(dirpath, exist_ok=True)
+
+        with self._persist_lock, self._lock.read():
+            files: dict[str, str] = {}
+            if self._index is not None:
+                files["hnsw"] = persist_path + _HNSW_SUFFIX
+                tmp = files["hnsw"] + ".tmp"
+                self._index.save(tmp)
+                os.replace(tmp, files["hnsw"])
+            if self._buf_ids:
+                files["buffer"] = persist_path + _BUFFER_SUFFIX
+                self._save_npy(files["buffer"], self._buf_vecs)
+            pending_kind = None
+            if self._pending_train is not None:
+                pending_kind, data = self._pending_train
+                files["train"] = persist_path + _TRAIN_SUFFIX
+                self._save_npy(files["train"], data)
+
+            meta = {
+                "format_version": FORMAT_VERSION,
+                "config": self._config(),
+                "pending_train": pending_kind,
+                "files": {k: os.path.basename(v) for k, v in files.items()},
+                "labels": [[n, lbl] for n, lbl in self._node_to_label.items()],
+                "buffer_ids": list(self._buf_ids),
+                "refs": dict(self._node_to_ref),
+            }
+            tmp = persist_path + ".tmp"
+            with open(tmp, "w", encoding="utf-8") as f:
+                json.dump(meta, f)
+            os.replace(tmp, persist_path)
+
+            for suffix in (_HNSW_SUFFIX, _BUFFER_SUFFIX, _TRAIN_SUFFIX):
+                stale = persist_path + suffix
+                if stale not in files.values() and os.path.exists(stale):
+                    os.remove(stale)
+
+    def _config(self) -> dict[str, Any]:
+        return {
+            "dim": self.dim,
+            "metric": self.metric,
+            "nbits": self.nbits,
+            "M": self.M,
+            "ef_construction": self.ef_construction,
+            "num_clusters": self.num_clusters,
+            "seed": self.seed,
+            "ef": self.ef,
+            "num_threads": self.num_threads,
+            "min_train_size": self.min_train_size,
+        }
+
+    @staticmethod
+    def _save_npy(path: str, arr: Any) -> None:
+        tmp = path + ".tmp"
+        with open(tmp, "wb") as f:
+            np.save(f, np.asarray(arr, dtype=np.float32), allow_pickle=False)
+        os.replace(tmp, path)
+
+    @classmethod
+    def from_persist_path(
+        cls,
+        persist_path: str,
+        fs: fsspec.AbstractFileSystem | None = None,
+        **overrides: Any,
+    ) -> RabitqHnswVectorStore:
+        """
+        Load a store saved by :meth:`persist`.
+
+        ``overrides`` may set the search-time options ``ef`` and ``num_threads``.
+        """
+        _check_local_fs(fs)
+        if not os.path.exists(persist_path):
+            raise ValueError(f"No RabitqHnswVectorStore found at {persist_path}.")
+        with open(persist_path, encoding="utf-8") as f:
+            meta = json.load(f)
+        version = meta.get("format_version")
+        if version != FORMAT_VERSION:
+            raise ValueError(
+                f"Unsupported format_version {version!r}; expected {FORMAT_VERSION}."
+            )
+        unknown = set(overrides) - {"ef", "num_threads"}
+        if unknown:
+            raise TypeError(f"Unexpected overrides: {sorted(unknown)}")
+        config = {**meta["config"], **overrides}
+        store = cls(**config)
+        base = os.path.dirname(persist_path)
+        files = meta["files"]
+
+        if "hnsw" in files:
+            index = HnswIndex.load(os.path.join(base, files["hnsw"]))
+            if index.dim != store.dim or index.metric != store._hnsw_metric():
+                raise ValueError("Saved index does not match the saved config.")
+            store._index = index
+            for nid, label in meta["labels"]:
+                store._node_to_label[nid] = int(label)
+                store._label_to_node[int(label)] = nid
+        if "buffer" in files:
+            buf = np.load(os.path.join(base, files["buffer"]), allow_pickle=False)
+            store._buf_vecs = np.ascontiguousarray(buf, dtype=np.float32)
+            store._buf_ids = list(meta["buffer_ids"])
+            if store._buf_vecs.shape != (len(store._buf_ids), store.dim):
+                raise ValueError("Saved buffer does not match its IDs.")
+        if "train" in files:
+            data = np.load(os.path.join(base, files["train"]), allow_pickle=False)
+            store._pending_train = (meta["pending_train"], data)
+        for nid, ref in meta["refs"].items():
+            store._node_to_ref[nid] = ref
+            store._ref_to_nodes.setdefault(ref, set()).add(nid)
+        return store
+
+    @classmethod
+    def from_persist_dir(
+        cls,
+        persist_dir: str = DEFAULT_PERSIST_DIR,
+        fs: fsspec.AbstractFileSystem | None = None,
+        **overrides: Any,
+    ) -> RabitqHnswVectorStore:
+        """Load the default-namespace store saved by ``StorageContext.persist``."""
+        _check_local_fs(fs)
+        persist_path = os.path.join(
+            persist_dir, f"{DEFAULT_VECTOR_STORE}{NAMESPACE_SEP}{DEFAULT_PERSIST_FNAME}"
+        )
+        return cls.from_persist_path(persist_path, **overrides)

llama_index_vector_stores_rabitqlib-0.1.0/pyproject.toml

@@ -0,0 +1,55 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "llama-index-vector-stores-rabitqlib"
+version = "0.1.0"
+description = "LlamaIndex vector store backed by RaBitQ-quantized HNSW (rabitqlib)"
+requires-python = ">=3.10,<4.0"
+readme = "README.md"
+license = "Apache-2.0"
+dependencies = [
+    "llama-index-core>=0.13.0,<0.15",
+    "rabitqlib>=0.8.3,<0.9",
+    "numpy>=1.24",
+]
+
+[dependency-groups]
+dev = ["pytest", "pytest-cov", "ruff", "mypy"]
+bench = ["faiss-cpu", "llama-index-vector-stores-faiss"]
+
+[tool.hatch.build.targets.sdist]
+include = ["llama_index/"]
+
+[tool.hatch.build.targets.wheel]
+include = ["llama_index/"]
+
+[tool.llamahub]
+contains_example = false
+import_path = "llama_index.vector_stores.rabitqlib"
+
+[tool.llamahub.class_authors]
+RabitqHnswVectorStore = "rabitq"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "B", "UP"]
+
+[tool.ruff.lint.isort]
+known-third-party = ["rabitqlib"]
+
+[tool.mypy]
+disallow_untyped_defs = true
+ignore_missing_imports = true
+python_version = "3.10"
+exclude = ["examples", ".venv"]
+explicit_package_bases = true
+namespace_packages = true
+plugins = ["pydantic.mypy"]