Skip to content

Commit 6cb1d30

Browse files
committed
add spfresh to python client
1 parent 5b75e5d commit 6cb1d30

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 (
@@ -1590,6 +1591,21 @@ def vector_index_type() -> str:
15901591

15911592
VectorIndexConfigHNSW = _VectorIndexConfigHNSW
15921593

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

15941610
@dataclass
15951611
class _VectorIndexConfigFlat(_VectorIndexConfig):
@@ -1664,7 +1680,7 @@ def to_dict(self) -> Dict[str, Any]:
16641680
class _NamedVectorConfig(_ConfigBase):
16651681
vectorizer: _NamedVectorizerConfig
16661682
vector_index_config: Union[
1667-
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic
1683+
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic, VectorIndexConfigSPFresh
16681684
]
16691685

16701686
def to_dict(self) -> Dict:
@@ -1689,7 +1705,7 @@ class _CollectionConfig(_ConfigBase):
16891705
reranker_config: Optional[RerankerConfig]
16901706
sharding_config: Optional[ShardingConfig]
16911707
vector_index_config: Union[
1692-
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic, None
1708+
VectorIndexConfigHNSW, VectorIndexConfigFlat, VectorIndexConfigDynamic, VectorIndexConfigSPFresh, None
16931709
]
16941710
vector_index_type: Optional[VectorIndexType]
16951711
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

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

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

213226
def __get_flat_config(config: Dict[str, Any]) -> _VectorIndexConfigFlat:
214227
quantizer = __get_quantizer_config(config)
@@ -236,6 +249,8 @@ def __get_vector_index_config(
236249
hnsw=__get_hnsw_config(schema["vectorIndexConfig"]["hnsw"]),
237250
flat=__get_flat_config(schema["vectorIndexConfig"]["flat"]),
238251
)
252+
elif schema["vectorIndexType"] == "spfresh":
253+
return __get_spfresh_config(schema["vectorIndexConfig"])
239254
else:
240255
return None
241256

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

@@ -561,6 +588,36 @@ def hnsw(
561588
multivector=multi_vector,
562589
)
563590

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