Skip to content

Commit 8957d61

Browse files
robbespo00rlmanrique
robbespo00
authored and
rlmanrique
committed
add spfresh to python client
1 parent e51f0f3 commit 8957d61

File tree

9 files changed

+120
-4
lines changed

9 files changed

+120
-4
lines changed

‎weaviate/collections/classes/config.py‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
_VectorIndexConfigDynamicUpdate,
4646
_VectorIndexConfigFlatUpdate,
4747
_VectorIndexConfigHNSWUpdate,
48+
_VectorIndexConfigSPFreshUpdate,
4849
_VectorIndexConfigUpdate,
4950
)
5051
from weaviate.collections.classes.config_vector_index import (
@@ -1593,6 +1594,21 @@ def vector_index_type() -> str:
15931594

15941595
VectorIndexConfigHNSW = _VectorIndexConfigHNSW
15951596

1597+
@dataclass
1598+
class _VectorIndexConfigSPFresh(_VectorIndexConfig):
1599+
distance_metric: VectorDistances
1600+
max_posting_size: int
1601+
min_posting_size: int
1602+
replicas: int
1603+
rng_factor: int
1604+
search_probe: int
1605+
centroids_index_type: str
1606+
1607+
@staticmethod
1608+
def vector_index_type() -> str:
1609+
return VectorIndexType.SPFRESH.value
1610+
1611+
VectorIndexConfigSPFresh = _VectorIndexConfigSPFresh
15961612

15971613
@dataclass
15981614
class _VectorIndexConfigFlat(_VectorIndexConfig):
@@ -1667,7 +1683,7 @@ def to_dict(self) -> Dict[str, Any]:
16671683
class _NamedVectorConfig(_ConfigBase):
16681684
vectorizer: _NamedVectorizerConfig
16691685
vector_index_config: Union[
1670-
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic
1686+
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic, VectorIndexConfigSPFresh
16711687
]
16721688

16731689
def to_dict(self) -> Dict:
@@ -1692,7 +1708,7 @@ class _CollectionConfig(_ConfigBase):
16921708
reranker_config: Optional[RerankerConfig]
16931709
sharding_config: Optional[ShardingConfig]
16941710
vector_index_config: Union[
1695-
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic, None
1711+
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic, VectorIndexConfigSPFresh, None
16961712
]
16971713
vector_index_type: Optional[VectorIndexType]
16981714
vectorizer_config: Optional[VectorizerConfig]

‎weaviate/collections/classes/config_methods.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
_VectorIndexConfigDynamic,
4040
_VectorIndexConfigFlat,
4141
_VectorIndexConfigHNSW,
42+
_VectorIndexConfigSPFresh,
4243
_VectorizerConfig,
4344
)
4445

@@ -210,6 +211,18 @@ def __get_hnsw_config(config: Dict[str, Any]) -> _VectorIndexConfigHNSW:
210211
multi_vector=__get_multivector(config),
211212
)
212213

214+
def __get_spfresh_config(config: Dict[str, Any]) -> _VectorIndexConfigSPFresh:
215+
quantizer = __get_quantizer_config(config)
216+
return _VectorIndexConfigSPFresh(
217+
distance_metric=VectorDistances(config.get("distance")),
218+
max_posting_size=config["maxPostingSize"],
219+
min_posting_size=config["minPostingSize"],
220+
replicas=config["replicas"],
221+
rng_factor=config["rngFactor"],
222+
search_probe=config["searchProbe"],
223+
centroids_index_type=config["centroidsIndexType"],
224+
quantizer=quantizer,
225+
)
213226

214227
def __get_flat_config(config: Dict[str, Any]) -> _VectorIndexConfigFlat:
215228
quantizer = __get_quantizer_config(config)
@@ -237,6 +250,8 @@ def __get_vector_index_config(
237250
hnsw=__get_hnsw_config(schema["vectorIndexConfig"]["hnsw"]),
238251
flat=__get_flat_config(schema["vectorIndexConfig"]["flat"]),
239252
)
253+
elif schema["vectorIndexType"] == "spfresh":
254+
return __get_spfresh_config(schema["vectorIndexConfig"])
240255
else:
241256
return None
242257

‎weaviate/collections/classes/config_named_vectors.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
_VectorIndexConfigDynamicUpdate,
1616
_VectorIndexConfigFlatUpdate,
1717
_VectorIndexConfigHNSWUpdate,
18+
_VectorIndexConfigSPFreshUpdate,
1819
_VectorIndexConfigUpdate,
1920
)
2021
from weaviate.collections.classes.config_vectorizers import (
@@ -1330,6 +1331,7 @@ def update(
13301331
*,
13311332
vector_index_config: Union[
13321333
_VectorIndexConfigHNSWUpdate,
1334+
_VectorIndexConfigSPFreshUpdate,
13331335
_VectorIndexConfigFlatUpdate,
13341336
_VectorIndexConfigDynamicUpdate,
13351337
],

‎weaviate/collections/classes/config_vector_index.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,14 @@ class VectorIndexType(str, Enum):
3434
Attributes:
3535
HNSW: Hierarchical Navigable Small World (HNSW) index.
3636
FLAT: Flat index.
37+
DYNAMIC: Dynamic index.
38+
SPFRESH: SPFRESH index.
3739
"""
3840

3941
HNSW = "hnsw"
4042
FLAT = "flat"
4143
DYNAMIC = "dynamic"
44+
SPFRESH = "spfresh"
4245

4346

4447
class _MultiVectorConfigCreateBase(_ConfigCreateModel):
@@ -127,6 +130,19 @@ def vector_index_type() -> VectorIndexType:
127130
return VectorIndexType.HNSW
128131

129132

133+
class _VectorIndexConfigSPFreshCreate(_VectorIndexConfigCreate):
134+
maxPostingSize: Optional[int]
135+
minPostingSize: Optional[int]
136+
replicas: Optional[int]
137+
rngFactor: Optional[int]
138+
searchProbe: Optional[int]
139+
centroidsIndexType: Optional[str]
140+
141+
@staticmethod
142+
def vector_index_type() -> VectorIndexType:
143+
return VectorIndexType.SPFRESH
144+
145+
130146
class _VectorIndexConfigFlatCreate(_VectorIndexConfigCreate):
131147
vectorCacheMaxObjects: Optional[int]
132148

@@ -149,6 +165,17 @@ def vector_index_type() -> VectorIndexType:
149165
return VectorIndexType.HNSW
150166

151167

168+
class _VectorIndexConfigSPFreshUpdate(_VectorIndexConfigUpdate):
169+
maxPostingSize: Optional[int]
170+
minPostingSize: Optional[int]
171+
rngFactor: Optional[int]
172+
searchProbe: Optional[int]
173+
174+
@staticmethod
175+
def vector_index_type() -> VectorIndexType:
176+
return VectorIndexType.SPFRESH
177+
178+
152179
class _VectorIndexConfigFlatUpdate(_VectorIndexConfigUpdate):
153180
vectorCacheMaxObjects: Optional[int]
154181

@@ -564,6 +591,36 @@ def hnsw(
564591
multivector=multi_vector,
565592
)
566593

594+
@staticmethod
595+
def spfresh(
596+
distance_metric: Optional[VectorDistances] = None,
597+
max_posting_size: Optional[int] = None,
598+
min_posting_size: Optional[int] = None,
599+
replicas: Optional[int] = None,
600+
rng_factor: Optional[int] = None,
601+
search_probe: Optional[int] = None,
602+
centroids_index_type: Optional[str] = None,
603+
quantizer: Optional[_QuantizerConfigCreate] = None,
604+
) -> _VectorIndexConfigSPFreshCreate:
605+
"""Create a `_VectorIndexConfigSPFreshCreate` object to be used when defining the SPFresh vector index configuration of Weaviate.
606+
607+
Use this method when defining the `vector_index_config` argument in `collections.create()`.
608+
609+
Args:
610+
See [the docs](https://weaviate.io/developers/weaviate/configuration/indexes#how-to-configure-spfresh) for a more detailed view!
611+
"""
612+
return _VectorIndexConfigSPFreshCreate(
613+
distance=distance_metric,
614+
maxPostingSize=max_posting_size,
615+
minPostingSize=min_posting_size,
616+
replicas=replicas,
617+
rngFactor=rng_factor,
618+
searchProbe=search_probe,
619+
centroidsIndexType=centroids_index_type,
620+
quantizer=quantizer,
621+
multivector=None,
622+
)
623+
567624
@staticmethod
568625
def flat(
569626
distance_metric: Optional[VectorDistances] = None,

‎weaviate/collections/classes/config_vectors.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
_VectorIndexConfigFlatUpdate,
2121
_VectorIndexConfigHNSWCreate,
2222
_VectorIndexConfigHNSWUpdate,
23+
_VectorIndexConfigSPFreshCreate,
24+
_VectorIndexConfigSPFreshUpdate,
2325
_VectorIndexConfigUpdate,
2426
)
2527
from weaviate.collections.classes.config_vectorizers import (
@@ -126,6 +128,19 @@ def __hnsw(
126128
multivector=multivector,
127129
)
128130

131+
@staticmethod
132+
def __spfresh(*, quantizer: Optional[_QuantizerConfigCreate]) -> _VectorIndexConfigSPFreshCreate:
133+
return _VectorIndexConfigSPFreshCreate(
134+
distance_metric=None,
135+
maxPostingSize=None,
136+
minPostingSize=None,
137+
replicas=None,
138+
rngFactor=None,
139+
searchProbe=None,
140+
centroidsIndexType=None,
141+
quantizer=quantizer,
142+
)
143+
129144
@staticmethod
130145
def __flat(*, quantizer: Optional[_QuantizerConfigCreate]) -> _VectorIndexConfigFlatCreate:
131146
return _VectorIndexConfigFlatCreate(
@@ -1504,6 +1519,7 @@ def update(
15041519
name: Optional[str] = None,
15051520
vector_index_config: Union[
15061521
_VectorIndexConfigHNSWUpdate,
1522+
_VectorIndexConfigSPFreshUpdate,
15071523
_VectorIndexConfigFlatUpdate,
15081524
_VectorIndexConfigDynamicUpdate,
15091525
],

‎weaviate/collections/config/async_.pyi‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ from weaviate.collections.classes.config import (
2121
_VectorConfigUpdate,
2222
_VectorIndexConfigFlatUpdate,
2323
_VectorIndexConfigHNSWUpdate,
24+
_VectorIndexConfigSPFreshUpdate,
2425
)
2526
from weaviate.collections.classes.config_vector_index import _VectorIndexConfigDynamicUpdate
2627
from weaviate.connect.v4 import ConnectionAsync
@@ -45,13 +46,14 @@ class _ConfigCollectionAsync(_ConfigCollectionExecutor[ConnectionAsync]):
4546
multi_tenancy_config: Optional[_MultiTenancyConfigUpdate] = None,
4647
replication_config: Optional[_ReplicationConfigUpdate] = None,
4748
vector_index_config: Optional[
48-
Union[_VectorIndexConfigHNSWUpdate, _VectorIndexConfigFlatUpdate]
49+
Union[_VectorIndexConfigHNSWUpdate, _VectorIndexConfigFlatUpdate, _VectorIndexConfigSPFreshUpdate]
4950
] = None,
5051
vectorizer_config: Optional[
5152
Union[
5253
_VectorIndexConfigHNSWUpdate,
5354
_VectorIndexConfigFlatUpdate,
5455
_VectorIndexConfigDynamicUpdate,
56+
_VectorIndexConfigSPFreshUpdate,
5557
List[_NamedVectorConfigUpdate],
5658
]
5759
] = None,

‎weaviate/collections/config/executor.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
_VectorConfigUpdate,
3939
_VectorIndexConfigFlatUpdate,
4040
_VectorIndexConfigHNSWUpdate,
41+
_VectorIndexConfigSPFreshUpdate,
4142
)
4243
from weaviate.collections.classes.config_methods import (
4344
_collection_config_from_json,
@@ -134,13 +135,15 @@ def update(
134135
Union[
135136
_VectorIndexConfigHNSWUpdate,
136137
_VectorIndexConfigFlatUpdate,
138+
_VectorIndexConfigSPFreshUpdate,
137139
]
138140
] = None,
139141
vectorizer_config: Optional[
140142
Union[
141143
_VectorIndexConfigHNSWUpdate,
142144
_VectorIndexConfigFlatUpdate,
143145
_VectorIndexConfigDynamicUpdate,
146+
_VectorIndexConfigSPFreshUpdate,
144147
List[_NamedVectorConfigUpdate],
145148
]
146149
] = None,
@@ -184,6 +187,7 @@ def update(
184187
_VectorIndexConfigHNSWUpdate,
185188
_VectorIndexConfigFlatUpdate,
186189
_VectorIndexConfigDynamicUpdate,
190+
_VectorIndexConfigSPFreshUpdate,
187191
),
188192
):
189193
_Warnings.vectorizer_config_in_config_update()

‎weaviate/collections/config/sync.pyi‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ from weaviate.collections.classes.config import (
2121
_VectorConfigUpdate,
2222
_VectorIndexConfigFlatUpdate,
2323
_VectorIndexConfigHNSWUpdate,
24+
_VectorIndexConfigSPFreshUpdate,
2425
)
2526
from weaviate.collections.classes.config_vector_index import _VectorIndexConfigDynamicUpdate
2627
from weaviate.connect.v4 import ConnectionSync
@@ -43,13 +44,14 @@ class _ConfigCollection(_ConfigCollectionExecutor[ConnectionSync]):
4344
multi_tenancy_config: Optional[_MultiTenancyConfigUpdate] = None,
4445
replication_config: Optional[_ReplicationConfigUpdate] = None,
4546
vector_index_config: Optional[
46-
Union[_VectorIndexConfigHNSWUpdate, _VectorIndexConfigFlatUpdate]
47+
Union[_VectorIndexConfigHNSWUpdate, _VectorIndexConfigFlatUpdate, _VectorIndexConfigSPFreshUpdate]
4748
] = None,
4849
vectorizer_config: Optional[
4950
Union[
5051
_VectorIndexConfigHNSWUpdate,
5152
_VectorIndexConfigFlatUpdate,
5253
_VectorIndexConfigDynamicUpdate,
54+
_VectorIndexConfigSPFreshUpdate,
5355
List[_NamedVectorConfigUpdate],
5456
]
5557
] = None,

‎weaviate/outputs/config.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
VectorDistances,
2424
VectorIndexConfigFlat,
2525
VectorIndexConfigHNSW,
26+
VectorIndexConfigSPFresh,
2627
VectorIndexType,
2728
VectorizerConfig,
2829
Vectorizers,
@@ -52,6 +53,7 @@
5253
"ShardTypes",
5354
"VectorDistances",
5455
"VectorIndexConfigHNSW",
56+
"VectorIndexConfigSPFresh",
5557
"VectorIndexConfigFlat",
5658
"VectorIndexType",
5759
"Vectorizers",

0 commit comments

Comments
 (0)