kitedb 0.2.5__cp313-cp313-win_amd64.whl

kitedb/vector_index.py

@@ -0,0 +1,472 @@
+"""High-level vector index API mirroring the TypeScript VectorIndex."""
+
+from __future__ import annotations
+
+import json
+import math
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple
+
+from kitedb._kitedb import IvfConfig, IvfIndex, SearchOptions, brute_force_search
+from kitedb.builders import NodeRef
+
+
+_METRIC_MAP = {
+    "cosine": "Cosine",
+    "euclidean": "Euclidean",
+    "dot": "DotProduct",
+    "dot_product": "DotProduct",
+    "dotproduct": "DotProduct",
+}
+
+
+def _validate_vector(vector: Sequence[float]) -> Optional[str]:
+    if len(vector) == 0:
+        return "Vector is empty"
+    for value in vector:
+        if not math.isfinite(value):
+            return "Vector contains NaN or infinity"
+    return None
+
+
+class _LRUCache:
+    def __init__(self, max_size: int = 10_000):
+        self._max_size = max_size
+        self._data: OrderedDict[int, NodeRef] = OrderedDict()
+
+    def get(self, key: int) -> Optional[NodeRef]:
+        value = self._data.get(key)
+        if value is not None:
+            self._data.move_to_end(key)
+        return value
+
+    def set(self, key: int, value: NodeRef) -> None:
+        if key in self._data:
+            self._data.move_to_end(key)
+        self._data[key] = value
+        if len(self._data) > self._max_size:
+            self._data.popitem(last=False)
+
+    def delete(self, key: int) -> None:
+        self._data.pop(key, None)
+
+    def clear(self) -> None:
+        self._data.clear()
+
+
+@dataclass
+class VectorIndexOptions:
+    dimensions: int
+    metric: str = "cosine"
+    row_group_size: int = 1024
+    fragment_target_size: int = 100_000
+    normalize: Optional[bool] = None
+    ivf: Optional[Dict[str, object]] = None
+    training_threshold: int = 1000
+    cache_max_size: int = 10_000
+
+
+@dataclass
+class SimilarOptions:
+    k: int
+    threshold: Optional[float] = None
+    n_probe: Optional[int] = None
+    filter: Optional[Callable[[int], bool]] = None
+
+
+@dataclass
+class VectorSearchHit:
+    node: NodeRef
+    distance: float
+    similarity: float
+
+
+class VectorIndex:
+    def __init__(self, options: VectorIndexOptions):
+        self._dimensions = options.dimensions
+        self._metric = options.metric.lower()
+        self._metric_enum = _METRIC_MAP.get(self._metric, "Cosine")
+        self._row_group_size = options.row_group_size
+        self._fragment_target_size = options.fragment_target_size
+        self._normalize = options.normalize
+        if self._normalize is None:
+            self._normalize = self._metric == "cosine"
+
+        self._ivf_config = options.ivf or {}
+        self._training_threshold = options.training_threshold
+        self._node_ref_cache = _LRUCache(options.cache_max_size)
+
+        self._vectors: Dict[int, List[float]] = {}
+        self._node_to_vector: Dict[int, int] = {}
+        self._vector_to_node: Dict[int, int] = {}
+        self._next_vector_id = 0
+
+        self._index: Optional[IvfIndex] = None
+        self._needs_training = True
+        self._is_building = False
+
+        self._manifest_cache: Optional[str] = None
+        self._manifest_dirty = True
+
+    def set(self, node_ref: NodeRef, vector: Sequence[float]) -> None:
+        if self._is_building:
+            raise ValueError("Cannot modify vectors while index is being built")
+
+        vec = self._coerce_vector(vector)
+        if len(vec) != self._dimensions:
+            raise ValueError(
+                f"Vector dimension mismatch: expected {self._dimensions}, got {len(vec)}"
+            )
+
+        validation_error = _validate_vector(vec)
+        if validation_error is not None:
+            raise ValueError(f"Invalid vector: {validation_error}")
+
+        if self._normalize:
+            vec = self._normalize_vector(vec)
+
+        node_id = node_ref.id
+
+        if node_id in self._node_to_vector:
+            old_vector_id = self._node_to_vector[node_id]
+            old_vector = self._vectors.get(node_id)
+            if old_vector is not None and self._index is not None and self._index.trained:
+                self._index.delete(old_vector_id, old_vector)
+            self._node_to_vector.pop(node_id, None)
+            self._vector_to_node.pop(old_vector_id, None)
+
+        vector_id = self._next_vector_id
+        self._next_vector_id += 1
+
+        self._vectors[node_id] = vec
+        self._node_to_vector[node_id] = vector_id
+        self._vector_to_node[vector_id] = node_id
+        self._manifest_dirty = True
+
+        self._node_ref_cache.set(node_id, node_ref)
+
+        if self._index is not None and self._index.trained:
+            self._index.insert(vector_id, vec)
+        else:
+            self._needs_training = True
+
+    def get(self, node_ref: NodeRef) -> Optional[List[float]]:
+        return self._vectors.get(node_ref.id)
+
+    def delete(self, node_ref: NodeRef) -> bool:
+        if self._is_building:
+            raise ValueError("Cannot modify vectors while index is being built")
+
+        node_id = node_ref.id
+        vector_id = self._node_to_vector.get(node_id)
+        if vector_id is None:
+            return False
+
+        vector = self._vectors.get(node_id)
+        if vector is None:
+            return False
+
+        if self._index is not None and self._index.trained:
+            self._index.delete(vector_id, vector)
+
+        self._vectors.pop(node_id, None)
+        self._node_to_vector.pop(node_id, None)
+        self._vector_to_node.pop(vector_id, None)
+        self._node_ref_cache.delete(node_id)
+        self._manifest_dirty = True
+        return True
+
+    def has(self, node_ref: NodeRef) -> bool:
+        return node_ref.id in self._node_to_vector
+
+    def build_index(self) -> None:
+        if self._is_building:
+            raise ValueError("Index build already in progress")
+
+        self._is_building = True
+        try:
+            live_vectors = len(self._vectors)
+            if live_vectors < self._training_threshold:
+                self._index = None
+                self._needs_training = False
+                return
+
+            n_clusters = self._ivf_config.get("n_clusters")
+            if n_clusters is None:
+                n_clusters = min(1024, max(16, int(math.sqrt(live_vectors))))
+            n_probe = self._ivf_config.get("n_probe")
+            ivf_config = IvfConfig(
+                n_clusters=n_clusters,
+                n_probe=int(n_probe) if n_probe is not None else None,
+                metric=self._metric,
+            )
+            index = IvfIndex(self._dimensions, ivf_config)
+
+            training_data, count = self._collect_training_vectors()
+            index.add_training_vectors(training_data, num_vectors=count)
+            index.train()
+
+            for node_id, vector in self._vectors.items():
+                vector_id = self._node_to_vector[node_id]
+                index.insert(vector_id, vector)
+
+            self._index = index
+            self._needs_training = False
+        finally:
+            self._is_building = False
+
+    def search(
+        self,
+        query: Sequence[float],
+        options: Optional[SimilarOptions] = None,
+        **kwargs: object,
+    ) -> List[VectorSearchHit]:
+        opts = self._coerce_similar_options(options, kwargs)
+
+        query_vec = self._coerce_vector(query)
+        if len(query_vec) != self._dimensions:
+            raise ValueError(
+                f"Query dimension mismatch: expected {self._dimensions}, got {len(query_vec)}"
+            )
+
+        validation_error = _validate_vector(query_vec)
+        if validation_error is not None:
+            raise ValueError(f"Invalid query vector: {validation_error}")
+
+        if self._normalize:
+            query_vec = self._normalize_vector(query_vec)
+
+        if self._needs_training:
+            self.build_index()
+
+        if self._index is not None and self._index.trained:
+            results = self._search_ivf(query_vec, opts)
+        else:
+            results = self._search_brute_force(query_vec, opts)
+
+        hits: List[VectorSearchHit] = []
+        for node_id, distance, similarity in results:
+            node_ref = self._node_ref_cache.get(node_id)
+            if node_ref is None:
+                continue
+            hits.append(VectorSearchHit(node=node_ref, distance=distance, similarity=similarity))
+            if len(hits) >= opts.k:
+                break
+
+        return hits
+
+    def stats(self) -> Dict[str, object]:
+        total_vectors = len(self._node_to_vector)
+        live_vectors = len(self._vectors)
+        return {
+            "totalVectors": total_vectors,
+            "liveVectors": live_vectors,
+            "dimensions": self._dimensions,
+            "metric": self._metric,
+            "indexTrained": self._index.trained if self._index is not None else False,
+            "indexClusters": self._index.config.n_clusters if self._index is not None else None,
+        }
+
+    def clear(self) -> None:
+        self._vectors.clear()
+        self._node_to_vector.clear()
+        self._vector_to_node.clear()
+        self._node_ref_cache.clear()
+        self._next_vector_id = 0
+        self._index = None
+        self._needs_training = True
+        self._manifest_cache = None
+        self._manifest_dirty = True
+
+    def buildIndex(self) -> None:  # noqa: N802
+        self.build_index()
+
+    def _coerce_vector(self, vector: Sequence[float]) -> List[float]:
+        return [float(v) for v in vector]
+
+    def _normalize_vector(self, vector: List[float]) -> List[float]:
+        norm = math.sqrt(sum(v * v for v in vector))
+        if norm <= 1e-12:
+            return vector
+        return [v / norm for v in vector]
+
+    def _collect_training_vectors(self) -> Tuple[List[float], int]:
+        flat: List[float] = []
+        count = 0
+        for vec in self._vectors.values():
+            flat.extend(vec)
+            count += 1
+        return flat, count
+
+    def _search_brute_force(
+        self,
+        query_vec: List[float],
+        options: SimilarOptions,
+    ) -> List[Tuple[int, float, float]]:
+        vectors: List[List[float]] = []
+        node_ids: List[int] = []
+
+        for node_id, vec in self._vectors.items():
+            if options.filter is not None:
+                try:
+                    if not options.filter(node_id):
+                        continue
+                except Exception:
+                    continue
+            node_ids.append(node_id)
+            vectors.append(vec)
+
+        if not vectors:
+            return []
+
+        results = brute_force_search(
+            vectors=vectors,
+            node_ids=node_ids,
+            query=query_vec,
+            k=max(options.k * 2, options.k),
+            metric=self._metric,
+        )
+
+        hits: List[Tuple[int, float, float]] = []
+        for result in results:
+            if options.threshold is not None and result.similarity < options.threshold:
+                continue
+            hits.append((int(result.node_id), float(result.distance), float(result.similarity)))
+        return hits
+
+    def _search_ivf(
+        self,
+        query_vec: List[float],
+        options: SimilarOptions,
+    ) -> List[Tuple[int, float, float]]:
+        if self._index is None:
+            return []
+
+        search_options = None
+        if options.n_probe is not None or options.threshold is not None:
+            search_options = SearchOptions(
+                n_probe=options.n_probe,
+                threshold=options.threshold,
+            )
+
+        manifest_json = self._build_manifest_json()
+        results = self._index.search(
+            manifest_json=manifest_json,
+            query=query_vec,
+            k=max(options.k * 2, options.k),
+            options=search_options,
+        )
+
+        hits: List[Tuple[int, float, float]] = []
+        for result in results:
+            node_id = int(result.node_id)
+            if options.filter is not None:
+                try:
+                    if not options.filter(node_id):
+                        continue
+                except Exception:
+                    continue
+            if options.threshold is not None and result.similarity < options.threshold:
+                continue
+            hits.append((node_id, float(result.distance), float(result.similarity)))
+        return hits
+
+    def _build_manifest_json(self) -> str:
+        if not self._manifest_dirty and self._manifest_cache is not None:
+            return self._manifest_cache
+
+        vector_ids = sorted(self._vector_to_node.keys())
+        row_groups: List[Dict[str, object]] = []
+        vector_locations: Dict[str, Dict[str, int]] = {}
+
+        current_data: List[float] = []
+        current_count = 0
+        row_group_id = 0
+        local_index = 0
+
+        for vector_id in vector_ids:
+            node_id = self._vector_to_node[vector_id]
+            vector = self._vectors.get(node_id)
+            if vector is None:
+                continue
+
+            if current_count >= self._row_group_size:
+                row_groups.append(
+                    {
+                        "id": row_group_id,
+                        "count": current_count,
+                        "data": current_data,
+                    }
+                )
+                row_group_id += 1
+                current_data = []
+                current_count = 0
+
+            current_data.extend(vector)
+            current_count += 1
+            vector_locations[str(vector_id)] = {
+                "fragment_id": 0,
+                "local_index": local_index,
+            }
+            local_index += 1
+
+        if current_count > 0 or not row_groups:
+            row_groups.append(
+                {
+                    "id": row_group_id,
+                    "count": current_count,
+                    "data": current_data,
+                }
+            )
+
+        fragment = {
+            "id": 0,
+            "state": "Active",
+            "row_groups": row_groups,
+            "total_vectors": local_index,
+            "deletion_bitmap": [],
+            "deleted_count": 0,
+        }
+
+        manifest = {
+            "config": {
+                "dimensions": self._dimensions,
+                "metric": self._metric_enum,
+                "row_group_size": self._row_group_size,
+                "fragment_target_size": self._fragment_target_size,
+                "normalize_on_insert": self._normalize,
+            },
+            "fragments": [fragment],
+            "active_fragment_id": 0,
+            "total_vectors": local_index,
+            "total_deleted": 0,
+            "next_vector_id": self._next_vector_id,
+            "node_to_vector": {str(k): v for k, v in self._node_to_vector.items()},
+            "vector_to_node": {str(k): v for k, v in self._vector_to_node.items()},
+            "vector_locations": vector_locations,
+        }
+
+        self._manifest_cache = json.dumps(manifest)
+        self._manifest_dirty = False
+        return self._manifest_cache
+
+    def _coerce_similar_options(
+        self,
+        options: Optional[SimilarOptions],
+        kwargs: Dict[str, object],
+    ) -> SimilarOptions:
+        if options is None:
+            if "k" not in kwargs:
+                raise ValueError("search requires k or SimilarOptions")
+            return SimilarOptions(
+                k=int(kwargs["k"]),
+                threshold=kwargs.get("threshold"),
+                n_probe=kwargs.get("n_probe"),
+                filter=kwargs.get("filter"),
+            )
+        return options
+
+
+def create_vector_index(options: VectorIndexOptions) -> VectorIndex:
+    return VectorIndex(options)

kitedb-0.2.5.dist-info/METADATA

@@ -0,0 +1,217 @@
+Metadata-Version: 2.4
+Name: kitedb
+Version: 0.2.5
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Rust
+Classifier: Topic :: Database
+Classifier: Topic :: Database :: Database Engines/Servers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: pytest>=7.0 ; extra == 'dev'
+Requires-Dist: pytest-benchmark>=4.0 ; extra == 'dev'
+Requires-Dist: numpy>=1.20 ; extra == 'dev'
+Requires-Dist: numpy>=1.20 ; extra == 'numpy'
+Provides-Extra: dev
+Provides-Extra: numpy
+License-File: LICENSE
+Summary: High-performance embedded graph database with vector search
+Keywords: graph,database,embedded,vector,search,pathfinding
+Home-Page: https://github.com/mask-software/kitedb
+Author: KiteDB Contributors
+License: MIT
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Documentation, https://github.com/mask-software/kitedb#readme
+Project-URL: Homepage, https://github.com/mask-software/kitedb
+Project-URL: Repository, https://github.com/mask-software/kitedb
+
+# KiteDB for Python
+
+KiteDB is a high-performance embedded graph database with built-in vector search.
+This package provides the Python bindings to the Rust core.
+
+## Features
+
+- ACID transactions with commit/rollback
+- Node and edge CRUD operations with properties
+- Labels, edge types, and property keys
+- Fluent traversal and pathfinding (BFS, Dijkstra, A\*)
+- Vector embeddings with IVF and IVF-PQ indexes
+- Single-file storage format
+
+## Install
+
+### From PyPI
+
+```bash
+pip install kitedb
+```
+
+### From source
+
+```bash
+# Install maturin (Rust extension build tool)
+python -m pip install -U maturin
+
+# Build and install in development mode
+cd ray-rs
+maturin develop --features python
+
+# Or build a wheel
+maturin build --features python --release
+pip install target/wheels/kitedb-*.whl
+```
+
+## Quick start (fluent API)
+
+The fluent API provides a high-level, type-safe interface:
+
+```python
+from kitedb import ray, node, edge, prop, optional
+
+# Define your schema
+User = node("user",
+    key=lambda id: f"user:{id}",
+    props={
+        "name": prop.string("name"),
+        "email": prop.string("email"),
+        "age": optional(prop.int("age")),
+    }
+)
+
+Knows = edge("knows", {
+    "since": prop.int("since"),
+})
+
+# Open database
+with ray("./social.kitedb", nodes=[User], edges=[Knows]) as db:
+    # Insert nodes
+    alice = db.insert(User).values(key="alice", name="Alice", email="alice@example.com").returning()
+    bob = db.insert(User).values(key="bob", name="Bob", email="bob@example.com").returning()
+
+    # Create edges
+    db.link(alice, Knows, bob, since=2024)
+
+    # Traverse
+    friends = db.from_(alice).out(Knows).nodes().to_list()
+
+    # Pathfinding
+    path = db.shortest_path(alice).via(Knows).to(bob).dijkstra()
+```
+
+## Quick start (low-level API)
+
+For direct control, use the low-level `Database` class:
+
+```python
+from kitedb import Database, PropValue
+
+with Database("my_graph.kitedb") as db:
+    db.begin()
+
+    alice = db.create_node("user:alice")
+    bob = db.create_node("user:bob")
+
+    name_key = db.get_or_create_propkey("name")
+    db.set_node_prop(alice, name_key, PropValue.string("Alice"))
+    db.set_node_prop(bob, name_key, PropValue.string("Bob"))
+
+    knows = db.get_or_create_etype("knows")
+    db.add_edge(alice, knows, bob)
+
+    db.commit()
+
+    print("nodes:", db.count_nodes())
+    print("edges:", db.count_edges())
+```
+
+## Fluent traversal
+
+```python
+from kitedb import TraverseOptions
+
+friends = db.from_(alice).out(knows).to_list()
+
+results = db.from_(alice).traverse(
+    knows,
+    TraverseOptions(max_depth=3, min_depth=1, direction="out", unique=True),
+).to_list()
+```
+
+## Concurrent Access
+
+KiteDB supports concurrent read operations from multiple threads. Read operations don't block each other:
+
+```python
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+# Multiple threads can read concurrently
+def read_user(key):
+    return db.get_node_by_key(key)
+
+with ThreadPoolExecutor(max_workers=4) as executor:
+    futures = [executor.submit(read_user, f"user:{i}") for i in range(100)]
+    results = [f.result() for f in futures]
+
+# Or with asyncio (reads run concurrently)
+import asyncio
+
+async def read_users():
+    loop = asyncio.get_event_loop()
+    tasks = [
+        loop.run_in_executor(None, db.get_node_by_key, f"user:{i}")
+        for i in range(100)
+    ]
+    return await asyncio.gather(*tasks)
+```
+
+**Concurrency model:**
+
+- **Reads are concurrent**: Multiple `get_node_by_key()`, `get_neighbors()`, traversals, etc. can run in parallel
+- **Writes are exclusive**: Write operations (`create_node()`, `add_edge()`, etc.) require exclusive access
+- **Thread safety**: The `Database` object is safe to share across threads
+
+Note: Python's GIL is released during Rust operations, allowing true parallelism for I/O-bound database access.
+
+## Vector search
+
+```python
+from kitedb import IvfIndex, IvfConfig, SearchOptions
+
+index = IvfIndex(dimensions=128, config=IvfConfig(n_clusters=100))
+
+training_data = [0.1] * (128 * 1000)
+index.add_training_vectors(training_data, num_vectors=1000)
+index.train()
+
+index.insert(vector_id=1, vector=[0.1] * 128)
+
+results = index.search(
+    manifest_json='{"vectors": {...}}',
+    query=[0.1] * 128,
+    k=10,
+    options=SearchOptions(n_probe=20),
+)
+
+for result in results:
+    print(result.node_id, result.distance)
+```
+
+## Documentation
+
+```text
+https://ray-kwaf.vercel.app/docs
+```
+
+## License
+
+MIT License - see the main project LICENSE file for details.
+

kitedb-0.2.5.dist-info/RECORD

@@ -0,0 +1,12 @@
+kitedb\__init__.py,sha256=Bv89fIrMT104bfqygV_wtHiUp9IACTJLGGyRtxAhprc,6191
+kitedb\_kitedb.cp313-win_amd64.pyd,sha256=BkCqrfpZLDbT3cNIPTEx8d8a1gq_UgHfMb3YyfG7CNA,3515392
+kitedb\_kitedb.pyi,sha256=Jyw5ujVUc1jBh5Kkgl75vO5m2xiod9TaUXaSMusE534,23624
+kitedb\builders.py,sha256=Pg7e4sClo7km50bitPaeuuyGDkNgmKRkbtjNb1IOru4,28552
+kitedb\fluent.py,sha256=vMcmbnve57gmg_hj99z6hWoQrLya1tOq6GdA9IcWuxM,27545
+kitedb\schema.py,sha256=h_H0S_Ti-PdI_34Cv-Y3NL2t9CtZqT3-h6D9sZb-SMY,8974
+kitedb\traversal.py,sha256=OzenuP2KFQWjeq_M7pXHzNkSz8yB-1-Bxvr_XXzvApc,54624
+kitedb\vector_index.py,sha256=IF411qHuPp926FfoPglUF1DbydhtbKJtFVoB4zEryjY,15995
+kitedb-0.2.5.dist-info\METADATA,sha256=BdxvCKkt6peNr954jBpKOoBeIp8fxrLFSV2iapZAeq4,6110
+kitedb-0.2.5.dist-info\WHEEL,sha256=n_BmF69IyGtioVWE9c3M_zsEfe6-xMZy1v5HCL_6qE0,97
+kitedb-0.2.5.dist-info\licenses\LICENSE,sha256=hOXoX-b3Si1uhntZcUxKoeyHDL7TIec1lCdmv1ib24c,1092
+kitedb-0.2.5.dist-info\RECORD,,

kitedb-0.2.5.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: maturin (1.11.5)
+Root-Is-Purelib: false
+Tag: cp313-cp313-win_amd64