fix: fix python type annotations (#307)

python/zvec/__init__.py

@@ -112,7 +112,6 @@ __all__ = [
     "InvertIndexParam",
     "HnswIndexParam",
     "HnswRabitqIndexParam",
-    "HnswRabitqQueryParam",
     "FlatIndexParam",
     "IVFIndexParam",
     "CollectionOption",
@@ -121,6 +120,7 @@ __all__ = [
     "AddColumnOption",
     "AlterColumnOption",
     "HnswQueryParam",
+    "HnswRabitqQueryParam",
     "IVFQueryParam",
     # Extensions
     "DenseEmbeddingFunction",

python/zvec/__init__.pyi

@@ -19,6 +19,8 @@ from .model.param import (
     FlatIndexParam,
     HnswIndexParam,
     HnswQueryParam,
+    HnswRabitqIndexParam,
+    HnswRabitqQueryParam,
     IndexOption,
     InvertIndexParam,
     IVFIndexParam,
@@ -54,6 +56,8 @@ __all__: list = [
     "FlatIndexParam",
     "HnswIndexParam",
     "HnswQueryParam",
+    "HnswRabitqIndexParam",
+    "HnswRabitqQueryParam",
     "IVFIndexParam",
     "IVFQueryParam",
     "IndexOption",

python/zvec/model/collection.py

@@ -28,6 +28,7 @@ from .param import (
     CollectionOption,
     FlatIndexParam,
     HnswIndexParam,
+    HnswRabitqIndexParam,
     IndexOption,
     InvertIndexParam,
     IVFIndexParam,
@@ -38,7 +39,12 @@ from .schema import CollectionSchema, CollectionStats, FieldSchema
 
 __all__ = ["Collection"]
 
-_VECTOR_INDEX_TYPES = (HnswIndexParam, IVFIndexParam, FlatIndexParam)
+_VECTOR_INDEX_TYPES = (
+    HnswIndexParam,
+    HnswRabitqIndexParam,
+    IVFIndexParam,
+    FlatIndexParam,
+)
 
 
 class Collection:
@@ -107,7 +113,11 @@ class Collection:
         self,
         field_name: str,
         index_param: Union[
-            HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam
+            HnswIndexParam,
+            HnswRabitqIndexParam,
+            IVFIndexParam,
+            FlatIndexParam,
+            InvertIndexParam,
         ],
         option: IndexOption = IndexOption(),
     ) -> None:
@@ -118,7 +128,7 @@ class Collection:
 
         Args:
             field_name (str): Name of the field to index.
-            index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam]):
+            index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam]):
                 Index configuration.
             option (Optional[IndexOption], optional): Index creation options.
                 Defaults to ``IndexOption()``.

python/zvec/model/param/vector_query.py

@@ -17,7 +17,7 @@ from dataclasses import dataclass
 from typing import Optional, Union
 
 from ...common import VectorType
-from . import HnswQueryParam, IVFQueryParam
+from . import HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam
 
 __all__ = ["VectorQuery"]
 
@@ -55,7 +55,7 @@ class VectorQuery:
     field_name: str
     id: Optional[str] = None
     vector: VectorType = None
-    param: Optional[Union[HnswQueryParam, IVFQueryParam]] = None
+    param: Optional[Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam]] = None
 
     def has_id(self) -> bool:
         """Check if the query is based on a document ID.

python/zvec/model/schema/field_schema.py

@@ -21,6 +21,7 @@ from _zvec.schema import _FieldSchema
 from zvec.model.param import (
     FlatIndexParam,
     HnswIndexParam,
+    HnswRabitqIndexParam,
     InvertIndexParam,
     IVFIndexParam,
 )
@@ -209,7 +210,7 @@ class VectorSchema:
         data_type: DataType,
         dimension: Optional[int] = 0,
         index_param: Optional[
-            Union[HnswIndexParam, FlatIndexParam, IVFIndexParam]
+            Union[HnswIndexParam, HnswRabitqIndexParam, FlatIndexParam, IVFIndexParam]
         ] = None,
     ):
         if name is None or not isinstance(name, str):
@@ -263,8 +264,10 @@ class VectorSchema:
         return self._cpp_obj.dimension
 
     @property
-    def index_param(self) -> Union[HnswIndexParam, IVFIndexParam, FlatIndexParam]:
-        """Union[HnswIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector."""
+    def index_param(
+        self,
+    ) -> Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]:
+        """Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector."""
         return self._cpp_obj.index_param
 
     def __dict__(self) -> dict[str, Any]: