add enum for centroids index type in spfresh

Author: robbespo00

weaviate/classes/config.py

@@ -17,6 +17,7 @@
 from weaviate.collections.classes.config_vector_index import (
     MultiVectorAggregation,
     VectorFilterStrategy,
+    VectorCentroidsIndexType
 )
 from weaviate.collections.classes.config_vectorizers import Multi2VecField, Vectorizers
 from weaviate.connect.integrations import Integrations
@@ -41,4 +42,5 @@
     "Vectorizers",
     "VectorDistances",
     "VectorFilterStrategy",
+    "VectorCentroidsIndexType",
 ]

weaviate/collections/classes/config.py

@@ -35,6 +35,7 @@
     PQEncoderDistribution,
     PQEncoderType,
     VectorFilterStrategy,
+    VectorCentroidsIndexType,
     _BQConfigUpdate,
     _PQConfigUpdate,
     _PQEncoderConfigUpdate,
@@ -1599,7 +1600,7 @@ class _VectorIndexConfigSPFresh(_VectorIndexConfig):
     replicas: int
     rng_factor: int
     search_probe: int
-    centroids_index_type: str
+    centroids_index_type: VectorCentroidsIndexType
 
     @staticmethod
     def vector_index_type() -> str:

weaviate/collections/classes/config_methods.py

@@ -12,6 +12,7 @@
     VectorDistances,
     VectorFilterStrategy,
     VectorIndexType,
+    VectorCentroidsIndexType,
     Vectorizers,
     _BM25Config,
     _BQConfig,
@@ -219,7 +220,11 @@ def __get_spfresh_config(config: Dict[str, Any]) -> _VectorIndexConfigSPFresh:
         replicas=config["replicas"],
         rng_factor=config["rngFactor"],
         search_probe=config["searchProbe"],
-        centroids_index_type=config["centroidsIndexType"],
+        centroids_index_type=(
+            VectorCentroidsIndexType(config["centroidsIndexType"])
+            if "centroidsIndexType" in config
+            else VectorCentroidsIndexType.HNSW
+        ),
         quantizer=quantizer,
     )
 

weaviate/collections/classes/config_vector_index.py

@@ -44,6 +44,18 @@ class VectorIndexType(str, Enum):
     SPFRESH = "spfresh"
 
 
+class VectorCentroidsIndexType(str, Enum):
+    """The available index type that can be used for centroids in SPFresh.
+    
+    Attributes:
+        HNSW: Hierarchical Navigable Small World (HNSW) index.
+        FLAT: flat index. 
+    """
+
+    HNSW = "hnsw"
+    FLAT = "bruteforce"
+
+
 class _MultiVectorConfigCreateBase(_ConfigCreateModel):
     enabled: bool = Field(default=True)
 
@@ -136,7 +148,7 @@ class _VectorIndexConfigSPFreshCreate(_VectorIndexConfigCreate):
     replicas: Optional[int]
     rngFactor: Optional[int]
     searchProbe: Optional[int]
-    centroidsIndexType: Optional[str]
+    centroidsIndexType: Optional[VectorCentroidsIndexType]
 
     @staticmethod
     def vector_index_type() -> VectorIndexType:
@@ -596,7 +608,7 @@ def spfresh(
         replicas: Optional[int] = None,
         rng_factor: Optional[int] = None,
         search_probe: Optional[int] = None,
-        centroids_index_type: Optional[str] = None,
+        centroids_index_type: Optional[VectorCentroidsIndexType] = None,
         quantizer: Optional[_QuantizerConfigCreate] = None,
     ) -> _VectorIndexConfigSPFreshCreate:
         """Create a `_VectorIndexConfigSPFreshCreate` object to be used when defining the SPFresh vector index configuration of Weaviate.