chroma_types/
collection_configuration.rs

1use crate::{
2    collection_schema::is_embedding_function_default, default_batch_size, default_construction_ef,
3    default_construction_ef_spann, default_initial_lambda, default_m, default_m_spann,
4    default_merge_threshold, default_nreplica_count, default_num_centers_to_merge_to,
5    default_num_samples_kmeans, default_num_threads, default_reassign_neighbor_count,
6    default_resize_factor, default_search_ef, default_search_ef_spann, default_search_nprobe,
7    default_search_rng_epsilon, default_search_rng_factor, default_space, default_split_threshold,
8    default_sync_threshold, default_write_nprobe, default_write_rng_epsilon,
9    default_write_rng_factor,
10};
11use crate::{
12    HnswConfiguration, HnswParametersFromSegmentError, InternalHnswConfiguration,
13    InternalSpannConfiguration, Metadata, Segment, SpannConfiguration, UpdateHnswConfiguration,
14    UpdateSpannConfiguration,
15};
16use chroma_error::{ChromaError, ErrorCodes};
17use serde::{Deserialize, Serialize};
18use thiserror::Error;
19
20#[derive(Deserialize, Serialize, Clone, Debug, Copy)]
21pub enum KnnIndex {
22    #[serde(alias = "hnsw")]
23    Hnsw,
24    #[serde(alias = "spann")]
25    Spann,
26}
27
28pub fn default_default_knn_index() -> KnnIndex {
29    KnnIndex::Hnsw
30}
31
32#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
33#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
34#[serde(tag = "type")]
35pub enum EmbeddingFunctionConfiguration {
36    #[serde(rename = "legacy")]
37    Legacy,
38    #[serde(rename = "known")]
39    Known(EmbeddingFunctionNewConfiguration),
40    #[serde(rename = "unknown")]
41    Unknown,
42}
43
44impl EmbeddingFunctionConfiguration {
45    pub fn is_default(&self) -> bool {
46        match self {
47            EmbeddingFunctionConfiguration::Legacy => false,
48            EmbeddingFunctionConfiguration::Unknown => false,
49            EmbeddingFunctionConfiguration::Known(config) => config.name == "default",
50        }
51    }
52}
53
54#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
55#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
56pub struct EmbeddingFunctionNewConfiguration {
57    pub name: String,
58    pub config: serde_json::Value,
59}
60
61#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
62#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
63#[serde(rename_all = "snake_case")]
64pub enum VectorIndexConfiguration {
65    Hnsw(InternalHnswConfiguration),
66    Spann(InternalSpannConfiguration),
67}
68
69impl VectorIndexConfiguration {
70    pub fn update(&mut self, vector_index: &VectorIndexConfiguration) {
71        match (self, vector_index) {
72            (VectorIndexConfiguration::Hnsw(hnsw), VectorIndexConfiguration::Hnsw(hnsw_new)) => {
73                *hnsw = hnsw_new.clone();
74            }
75            (
76                VectorIndexConfiguration::Spann(spann),
77                VectorIndexConfiguration::Spann(spann_new),
78            ) => {
79                *spann = spann_new.clone();
80            }
81            (VectorIndexConfiguration::Hnsw(_), VectorIndexConfiguration::Spann(_)) => {
82                // For now, we don't support converting between different index types
83                // This could be implemented in the future if needed
84            }
85            (VectorIndexConfiguration::Spann(_), VectorIndexConfiguration::Hnsw(_)) => {
86                // For now, we don't support converting between different index types
87                // This could be implemented in the future if needed
88            }
89        }
90    }
91}
92impl From<InternalHnswConfiguration> for VectorIndexConfiguration {
93    fn from(config: InternalHnswConfiguration) -> Self {
94        VectorIndexConfiguration::Hnsw(config)
95    }
96}
97
98impl From<InternalSpannConfiguration> for VectorIndexConfiguration {
99    fn from(config: InternalSpannConfiguration) -> Self {
100        VectorIndexConfiguration::Spann(config)
101    }
102}
103
104fn default_vector_index_config() -> VectorIndexConfiguration {
105    VectorIndexConfiguration::Hnsw(InternalHnswConfiguration::default())
106}
107
108#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
109#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
110pub struct InternalCollectionConfiguration {
111    #[serde(default = "default_vector_index_config")]
112    pub vector_index: VectorIndexConfiguration,
113    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
114}
115
116impl InternalCollectionConfiguration {
117    pub fn from_legacy_metadata(
118        metadata: Metadata,
119    ) -> Result<Self, HnswParametersFromSegmentError> {
120        let hnsw = InternalHnswConfiguration::from_legacy_segment_metadata(&Some(metadata))?;
121        Ok(Self {
122            vector_index: VectorIndexConfiguration::Hnsw(hnsw),
123            embedding_function: None,
124        })
125    }
126
127    pub fn default_hnsw() -> Self {
128        Self {
129            vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration::default()),
130            embedding_function: None,
131        }
132    }
133
134    pub fn default_spann() -> Self {
135        Self {
136            vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration::default()),
137            embedding_function: None,
138        }
139    }
140
141    /// Check if this collection configuration is default
142    pub fn is_default(&self) -> bool {
143        if !is_embedding_function_default(&self.embedding_function) {
144            return false;
145        }
146
147        // Check vector index configuration
148        match &self.vector_index {
149            VectorIndexConfiguration::Hnsw(hnsw_config) => {
150                hnsw_config.ef_construction == default_construction_ef()
151                    && hnsw_config.ef_search == default_search_ef()
152                    && hnsw_config.max_neighbors == default_m()
153                    && hnsw_config.num_threads == default_num_threads()
154                    && hnsw_config.batch_size == default_batch_size()
155                    && hnsw_config.sync_threshold == default_sync_threshold()
156                    && hnsw_config.resize_factor == default_resize_factor()
157                    && hnsw_config.space == default_space()
158            }
159            VectorIndexConfiguration::Spann(spann_config) => {
160                spann_config.search_nprobe == default_search_nprobe()
161                    && spann_config.search_rng_factor == default_search_rng_factor()
162                    && spann_config.search_rng_epsilon == default_search_rng_epsilon()
163                    && spann_config.write_nprobe == default_write_nprobe()
164                    && spann_config.nreplica_count == default_nreplica_count()
165                    && spann_config.write_rng_factor == default_write_rng_factor()
166                    && spann_config.write_rng_epsilon == default_write_rng_epsilon()
167                    && spann_config.split_threshold == default_split_threshold()
168                    && spann_config.num_samples_kmeans == default_num_samples_kmeans()
169                    && spann_config.initial_lambda == default_initial_lambda()
170                    && spann_config.reassign_neighbor_count == default_reassign_neighbor_count()
171                    && spann_config.merge_threshold == default_merge_threshold()
172                    && spann_config.num_centers_to_merge_to == default_num_centers_to_merge_to()
173                    && spann_config.ef_construction == default_construction_ef_spann()
174                    && spann_config.ef_search == default_search_ef_spann()
175                    && spann_config.max_neighbors == default_m_spann()
176                    && spann_config.space == default_space()
177            }
178        }
179    }
180
181    pub fn get_hnsw_config_with_legacy_fallback(
182        &self,
183        segment: &Segment,
184    ) -> Result<Option<InternalHnswConfiguration>, HnswParametersFromSegmentError> {
185        self.get_hnsw_config_from_legacy_metadata(&segment.metadata)
186    }
187
188    pub fn get_hnsw_config_from_legacy_metadata(
189        &self,
190        metadata: &Option<Metadata>,
191    ) -> Result<Option<InternalHnswConfiguration>, HnswParametersFromSegmentError> {
192        if let Some(config) = self.get_hnsw_config() {
193            let config_from_metadata =
194                InternalHnswConfiguration::from_legacy_segment_metadata(metadata)?;
195
196            if config == InternalHnswConfiguration::default() && config != config_from_metadata {
197                return Ok(Some(config_from_metadata));
198            }
199
200            return Ok(Some(config));
201        }
202
203        Ok(None)
204    }
205
206    pub fn get_spann_config(&self) -> Option<InternalSpannConfiguration> {
207        match &self.vector_index {
208            VectorIndexConfiguration::Spann(config) => Some(config.clone()),
209            _ => None,
210        }
211    }
212
213    fn get_hnsw_config(&self) -> Option<InternalHnswConfiguration> {
214        match &self.vector_index {
215            VectorIndexConfiguration::Hnsw(config) => Some(config.clone()),
216            _ => None,
217        }
218    }
219
220    pub fn update(&mut self, configuration: &InternalUpdateCollectionConfiguration) {
221        // Update vector_index if it exists in the update configuration
222
223        if let Some(vector_index) = &configuration.vector_index {
224            match vector_index {
225                UpdateVectorIndexConfiguration::Hnsw(hnsw_config) => {
226                    if let VectorIndexConfiguration::Hnsw(current_config) = &mut self.vector_index {
227                        if let Some(update_config) = hnsw_config {
228                            if let Some(ef_search) = update_config.ef_search {
229                                current_config.ef_search = ef_search;
230                            }
231                            if let Some(max_neighbors) = update_config.max_neighbors {
232                                current_config.max_neighbors = max_neighbors;
233                            }
234                            if let Some(num_threads) = update_config.num_threads {
235                                current_config.num_threads = num_threads;
236                            }
237                            if let Some(resize_factor) = update_config.resize_factor {
238                                current_config.resize_factor = resize_factor;
239                            }
240                            if let Some(sync_threshold) = update_config.sync_threshold {
241                                current_config.sync_threshold = sync_threshold;
242                            }
243                            if let Some(batch_size) = update_config.batch_size {
244                                current_config.batch_size = batch_size;
245                            }
246                        }
247                    }
248                }
249                UpdateVectorIndexConfiguration::Spann(spann_config) => {
250                    if let VectorIndexConfiguration::Spann(current_config) = &mut self.vector_index
251                    {
252                        if let Some(update_config) = spann_config {
253                            if let Some(search_nprobe) = update_config.search_nprobe {
254                                current_config.search_nprobe = search_nprobe;
255                            }
256                            if let Some(ef_search) = update_config.ef_search {
257                                current_config.ef_search = ef_search;
258                            }
259                        }
260                    }
261                }
262            }
263        }
264        // Update embedding_function if it exists in the update configuration
265        if let Some(embedding_function) = &configuration.embedding_function {
266            self.embedding_function = Some(embedding_function.clone());
267        }
268    }
269
270    pub fn try_from_config(
271        value: CollectionConfiguration,
272        default_knn_index: KnnIndex,
273        metadata: Option<Metadata>,
274    ) -> Result<Self, CollectionConfigurationToInternalConfigurationError> {
275        let mut hnsw: Option<HnswConfiguration> = value.hnsw;
276        let spann: Option<SpannConfiguration> = value.spann;
277
278        // if neither hnsw nor spann is provided, use the collection metadata to build an hnsw configuration
279        // the match then handles cases where hnsw is provided, and correctly routes to either spann or hnsw configuration
280        // based on the default_knn_index
281        if hnsw.is_none() && spann.is_none() {
282            let hnsw_config_from_metadata =
283            InternalHnswConfiguration::from_legacy_segment_metadata(&metadata).map_err(|e| {
284                CollectionConfigurationToInternalConfigurationError::HnswParametersFromSegmentError(
285                    e,
286                )
287            })?;
288            hnsw = Some(hnsw_config_from_metadata.into());
289        }
290
291        match (hnsw, spann) {
292            (Some(_), Some(_)) => Err(CollectionConfigurationToInternalConfigurationError::MultipleVectorIndexConfigurations),
293            (Some(hnsw), None) => {
294                match default_knn_index {
295                    // Create a spann index. Only inherit the space if it exists in the hnsw config.
296                    // This is for backwards compatibility so that users who migrate to distributed
297                    // from local don't break their code.
298                    KnnIndex::Spann => {
299                        let internal_config = if let Some(space) = hnsw.space {
300                            InternalSpannConfiguration {
301                                space,
302                                ..Default::default()
303                            }
304                        } else {
305                            InternalSpannConfiguration::default()
306                        };
307
308                        Ok(InternalCollectionConfiguration {
309                            vector_index: VectorIndexConfiguration::Spann(internal_config),
310                            embedding_function: value.embedding_function,
311                        })
312                    },
313                    KnnIndex::Hnsw => {
314                        let hnsw: InternalHnswConfiguration = hnsw.into();
315                        Ok(InternalCollectionConfiguration {
316                            vector_index: hnsw.into(),
317                            embedding_function: value.embedding_function,
318                        })
319                    }
320                }
321            }
322            (None, Some(spann)) => {
323                match default_knn_index {
324                    // Create a hnsw index. Only inherit the space if it exists in the spann config.
325                    // This is for backwards compatibility so that users who migrate to local
326                    // from distributed don't break their code.
327                    KnnIndex::Hnsw => {
328                        let internal_config = if let Some(space) = spann.space {
329                            InternalHnswConfiguration {
330                                space,
331                                ..Default::default()
332                            }
333                        } else {
334                            InternalHnswConfiguration::default()
335                        };
336                        Ok(InternalCollectionConfiguration {
337                            vector_index: VectorIndexConfiguration::Hnsw(internal_config),
338                            embedding_function: value.embedding_function,
339                        })
340                    }
341                    KnnIndex::Spann => {
342                        let spann: InternalSpannConfiguration = spann.into();
343                        Ok(InternalCollectionConfiguration {
344                            vector_index: spann.into(),
345                            embedding_function: value.embedding_function,
346                        })
347                    }
348                }
349            }
350            (None, None) => {
351                let vector_index = match default_knn_index {
352                    KnnIndex::Hnsw => InternalHnswConfiguration::default().into(),
353                    KnnIndex::Spann => InternalSpannConfiguration::default().into(),
354                };
355                Ok(InternalCollectionConfiguration {
356                    vector_index,
357                    embedding_function: value.embedding_function,
358                })
359            }
360        }
361    }
362}
363
364impl TryFrom<CollectionConfiguration> for InternalCollectionConfiguration {
365    type Error = CollectionConfigurationToInternalConfigurationError;
366
367    fn try_from(value: CollectionConfiguration) -> Result<Self, Self::Error> {
368        match (value.hnsw, value.spann) {
369            (Some(_), Some(_)) => Err(Self::Error::MultipleVectorIndexConfigurations),
370            (Some(hnsw), None) => {
371                let hnsw: InternalHnswConfiguration = hnsw.into();
372                Ok(InternalCollectionConfiguration {
373                    vector_index: hnsw.into(),
374                    embedding_function: value.embedding_function,
375                })
376            }
377            (None, Some(spann)) => {
378                let spann: InternalSpannConfiguration = spann.into();
379                Ok(InternalCollectionConfiguration {
380                    vector_index: spann.into(),
381                    embedding_function: value.embedding_function,
382                })
383            }
384            (None, None) => Ok(InternalCollectionConfiguration {
385                vector_index: InternalHnswConfiguration::default().into(),
386                embedding_function: value.embedding_function,
387            }),
388        }
389    }
390}
391
392#[derive(Debug, Error)]
393pub enum CollectionConfigurationToInternalConfigurationError {
394    #[error("Multiple vector index configurations provided")]
395    MultipleVectorIndexConfigurations,
396    #[error("Failed to parse hnsw parameters from segment metadata")]
397    HnswParametersFromSegmentError(#[from] HnswParametersFromSegmentError),
398}
399
400impl ChromaError for CollectionConfigurationToInternalConfigurationError {
401    fn code(&self) -> ErrorCodes {
402        match self {
403            Self::MultipleVectorIndexConfigurations => ErrorCodes::InvalidArgument,
404            Self::HnswParametersFromSegmentError(_) => ErrorCodes::InvalidArgument,
405        }
406    }
407}
408
409#[derive(Default, Deserialize, Serialize, Debug, Clone)]
410#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
411#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
412pub struct CollectionConfiguration {
413    pub hnsw: Option<HnswConfiguration>,
414    pub spann: Option<SpannConfiguration>,
415    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
416}
417
418impl From<InternalCollectionConfiguration> for CollectionConfiguration {
419    fn from(value: InternalCollectionConfiguration) -> Self {
420        Self {
421            hnsw: match value.vector_index.clone() {
422                VectorIndexConfiguration::Hnsw(config) => Some(config.into()),
423                _ => None,
424            },
425            spann: match value.vector_index {
426                VectorIndexConfiguration::Spann(config) => Some(config.into()),
427                _ => None,
428            },
429            embedding_function: value.embedding_function,
430        }
431    }
432}
433
434#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
435#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
436#[serde(rename_all = "snake_case")]
437pub enum UpdateVectorIndexConfiguration {
438    Hnsw(Option<UpdateHnswConfiguration>),
439    Spann(Option<UpdateSpannConfiguration>),
440}
441
442impl From<UpdateHnswConfiguration> for UpdateVectorIndexConfiguration {
443    fn from(config: UpdateHnswConfiguration) -> Self {
444        UpdateVectorIndexConfiguration::Hnsw(Some(config))
445    }
446}
447
448impl From<UpdateSpannConfiguration> for UpdateVectorIndexConfiguration {
449    fn from(config: UpdateSpannConfiguration) -> Self {
450        UpdateVectorIndexConfiguration::Spann(Some(config))
451    }
452}
453
454#[derive(Debug, Error)]
455pub enum UpdateCollectionConfigurationToInternalConfigurationError {
456    #[error("Multiple vector index configurations provided")]
457    MultipleVectorIndexConfigurations,
458}
459
460impl ChromaError for UpdateCollectionConfigurationToInternalConfigurationError {
461    fn code(&self) -> ErrorCodes {
462        match self {
463            Self::MultipleVectorIndexConfigurations => ErrorCodes::InvalidArgument,
464        }
465    }
466}
467
468#[derive(Deserialize, Serialize, Debug, Clone)]
469#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
470#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
471pub struct UpdateCollectionConfiguration {
472    pub hnsw: Option<UpdateHnswConfiguration>,
473    pub spann: Option<UpdateSpannConfiguration>,
474    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
475}
476
477#[derive(Deserialize, Serialize, Debug, Clone)]
478#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
479pub struct InternalUpdateCollectionConfiguration {
480    pub vector_index: Option<UpdateVectorIndexConfiguration>,
481    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
482}
483
484#[derive(Debug, Error)]
485pub enum UpdateCollectionConfigurationToInternalUpdateConfigurationError {
486    #[error("Multiple vector index configurations provided")]
487    MultipleVectorIndexConfigurations,
488}
489
490impl ChromaError for UpdateCollectionConfigurationToInternalUpdateConfigurationError {
491    fn code(&self) -> ErrorCodes {
492        match self {
493            Self::MultipleVectorIndexConfigurations => ErrorCodes::InvalidArgument,
494        }
495    }
496}
497
498impl TryFrom<UpdateCollectionConfiguration> for InternalUpdateCollectionConfiguration {
499    type Error = UpdateCollectionConfigurationToInternalUpdateConfigurationError;
500
501    fn try_from(value: UpdateCollectionConfiguration) -> Result<Self, Self::Error> {
502        match (value.hnsw, value.spann) {
503            (Some(_), Some(_)) => Err(Self::Error::MultipleVectorIndexConfigurations),
504            (Some(hnsw), None) => Ok(InternalUpdateCollectionConfiguration {
505                vector_index: Some(UpdateVectorIndexConfiguration::Hnsw(Some(hnsw))),
506                embedding_function: value.embedding_function,
507            }),
508            (None, Some(spann)) => Ok(InternalUpdateCollectionConfiguration {
509                vector_index: Some(UpdateVectorIndexConfiguration::Spann(Some(spann))),
510                embedding_function: value.embedding_function,
511            }),
512            (None, None) => Ok(InternalUpdateCollectionConfiguration {
513                vector_index: None,
514                embedding_function: value.embedding_function,
515            }),
516        }
517    }
518}
519
520#[cfg(test)]
521mod tests {
522
523    use crate::hnsw_configuration::HnswConfiguration;
524    use crate::hnsw_configuration::Space;
525    use crate::spann_configuration::SpannConfiguration;
526    use crate::{test_segment, CollectionUuid, Metadata};
527
528    use super::*;
529
530    #[test]
531    fn metadata_overrides_parameter() {
532        let mut metadata = Metadata::new();
533        metadata.insert(
534            "hnsw:construction_ef".to_string(),
535            crate::MetadataValue::Int(1),
536        );
537
538        let mut segment = test_segment(CollectionUuid::new(), crate::SegmentScope::VECTOR);
539        segment.metadata = Some(metadata);
540
541        let config = InternalCollectionConfiguration::default_hnsw();
542        let overridden_config = config
543            .get_hnsw_config_with_legacy_fallback(&segment)
544            .unwrap()
545            .unwrap();
546
547        assert_eq!(overridden_config.ef_construction, 1);
548    }
549
550    #[test]
551    fn metadata_ignored_when_config_is_not_default() {
552        let mut metadata = Metadata::new();
553        metadata.insert(
554            "hnsw:construction_ef".to_string(),
555            crate::MetadataValue::Int(1),
556        );
557
558        let mut segment = test_segment(CollectionUuid::new(), crate::SegmentScope::VECTOR);
559        segment.metadata = Some(metadata);
560
561        let config = InternalCollectionConfiguration {
562            vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
563                ef_construction: 2,
564                ..Default::default()
565            }),
566            embedding_function: None,
567        };
568
569        let overridden_config = config
570            .get_hnsw_config_with_legacy_fallback(&segment)
571            .unwrap()
572            .unwrap();
573
574        // Setting from metadata is ignored since the config is not default
575        assert_eq!(overridden_config.ef_construction, 2);
576    }
577
578    #[test]
579    fn test_hnsw_config_with_hnsw_default() {
580        let hnsw_config = HnswConfiguration {
581            max_neighbors: Some(16),
582            ef_construction: Some(100),
583            ef_search: Some(10),
584            batch_size: Some(100),
585            num_threads: Some(4),
586            sync_threshold: Some(500),
587            resize_factor: Some(1.2),
588            space: Some(Space::Cosine),
589        };
590
591        let collection_config = CollectionConfiguration {
592            hnsw: Some(hnsw_config.clone()),
593            spann: None,
594            embedding_function: None,
595        };
596
597        let internal_config_result = InternalCollectionConfiguration::try_from_config(
598            collection_config,
599            KnnIndex::Hnsw,
600            None,
601        );
602
603        assert!(internal_config_result.is_ok());
604        let internal_config = internal_config_result.unwrap();
605
606        let expected_vector_index = VectorIndexConfiguration::Hnsw(hnsw_config.into());
607        assert_eq!(internal_config.vector_index, expected_vector_index);
608    }
609
610    #[test]
611    fn test_hnsw_config_with_spann_default() {
612        let hnsw_config = HnswConfiguration {
613            max_neighbors: Some(16),
614            ef_construction: Some(100),
615            ef_search: Some(10),
616            batch_size: Some(100),
617            num_threads: Some(4),
618            sync_threshold: Some(500),
619            resize_factor: Some(1.2),
620            space: Some(Space::Cosine),
621        };
622
623        let collection_config = CollectionConfiguration {
624            hnsw: Some(hnsw_config.clone()),
625            spann: None,
626            embedding_function: None,
627        };
628
629        let internal_config_result = InternalCollectionConfiguration::try_from_config(
630            collection_config,
631            KnnIndex::Spann,
632            None,
633        );
634
635        assert!(internal_config_result.is_ok());
636        let internal_config = internal_config_result.unwrap();
637
638        let expected_vector_index = VectorIndexConfiguration::Spann(InternalSpannConfiguration {
639            space: hnsw_config.space.unwrap_or(Space::L2),
640            ..Default::default()
641        });
642        assert_eq!(internal_config.vector_index, expected_vector_index);
643    }
644
645    #[test]
646    fn test_spann_config_with_spann_default() {
647        let spann_config = SpannConfiguration {
648            ef_construction: Some(100),
649            ef_search: Some(10),
650            max_neighbors: Some(16),
651            search_nprobe: Some(1),
652            write_nprobe: Some(1),
653            space: Some(Space::Cosine),
654            reassign_neighbor_count: Some(64),
655            split_threshold: Some(200),
656            merge_threshold: Some(100),
657        };
658
659        let collection_config = CollectionConfiguration {
660            hnsw: None,
661            spann: Some(spann_config.clone()),
662            embedding_function: None,
663        };
664
665        let internal_config_result = InternalCollectionConfiguration::try_from_config(
666            collection_config,
667            KnnIndex::Spann,
668            None,
669        );
670
671        assert!(internal_config_result.is_ok());
672        let internal_config = internal_config_result.unwrap();
673
674        let expected_vector_index = VectorIndexConfiguration::Spann(spann_config.into());
675        assert_eq!(internal_config.vector_index, expected_vector_index);
676    }
677
678    #[test]
679    fn test_spann_config_with_hnsw_default() {
680        let spann_config = SpannConfiguration {
681            ef_construction: Some(100),
682            ef_search: Some(10),
683            max_neighbors: Some(16),
684            search_nprobe: Some(1),
685            write_nprobe: Some(1),
686            space: Some(Space::Cosine),
687            reassign_neighbor_count: Some(64),
688            split_threshold: Some(200),
689            merge_threshold: Some(100),
690        };
691
692        let collection_config = CollectionConfiguration {
693            hnsw: None,
694            spann: Some(spann_config.clone()),
695            embedding_function: None,
696        };
697
698        let internal_config_result = InternalCollectionConfiguration::try_from_config(
699            collection_config,
700            KnnIndex::Hnsw,
701            None,
702        );
703
704        let expected_vector_index = VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
705            space: spann_config.space.unwrap_or(Space::L2),
706            ..Default::default()
707        });
708        assert_eq!(
709            internal_config_result.unwrap().vector_index,
710            expected_vector_index
711        );
712    }
713
714    #[test]
715    fn test_no_config_with_metadata_default_hnsw() {
716        let metadata = Metadata::new();
717        let collection_config = CollectionConfiguration {
718            hnsw: None,
719            spann: None,
720            embedding_function: None,
721        };
722
723        let internal_config_result = InternalCollectionConfiguration::try_from_config(
724            collection_config,
725            KnnIndex::Hnsw,
726            Some(metadata),
727        );
728
729        assert!(internal_config_result.is_ok());
730        let internal_config = internal_config_result.unwrap();
731
732        assert_eq!(
733            internal_config.vector_index,
734            VectorIndexConfiguration::Hnsw(InternalHnswConfiguration::default())
735        );
736    }
737
738    #[test]
739    fn test_no_config_with_metadata_default_spann() {
740        let metadata = Metadata::new();
741        let collection_config = CollectionConfiguration {
742            hnsw: None,
743            spann: None,
744            embedding_function: None,
745        };
746
747        let internal_config_result = InternalCollectionConfiguration::try_from_config(
748            collection_config,
749            KnnIndex::Spann,
750            Some(metadata),
751        );
752
753        assert!(internal_config_result.is_ok());
754        let internal_config = internal_config_result.unwrap();
755
756        assert_eq!(
757            internal_config.vector_index,
758            VectorIndexConfiguration::Spann(InternalSpannConfiguration::default())
759        );
760    }
761
762    #[test]
763    fn test_legacy_metadata_with_hnsw_config() {
764        let mut metadata = Metadata::new();
765        metadata.insert(
766            "hnsw:space".to_string(),
767            crate::MetadataValue::Str("cosine".to_string()),
768        );
769        metadata.insert(
770            "hnsw:construction_ef".to_string(),
771            crate::MetadataValue::Int(1),
772        );
773
774        let collection_config = CollectionConfiguration {
775            hnsw: None,
776            spann: None,
777            embedding_function: None,
778        };
779
780        let internal_config_result = InternalCollectionConfiguration::try_from_config(
781            collection_config,
782            KnnIndex::Hnsw,
783            Some(metadata),
784        );
785
786        assert!(internal_config_result.is_ok());
787        let internal_config = internal_config_result.unwrap();
788
789        assert_eq!(
790            internal_config.vector_index,
791            VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
792                space: Space::Cosine,
793                ef_construction: 1,
794                ..Default::default()
795            })
796        );
797    }
798
799    #[test]
800    fn test_legacy_metadata_with_spann_config() {
801        let mut metadata = Metadata::new();
802        metadata.insert(
803            "hnsw:space".to_string(),
804            crate::MetadataValue::Str("cosine".to_string()),
805        );
806        metadata.insert(
807            "hnsw:construction_ef".to_string(),
808            crate::MetadataValue::Int(1),
809        );
810
811        let collection_config = CollectionConfiguration {
812            hnsw: None,
813            spann: None,
814            embedding_function: None,
815        };
816
817        let internal_config_result = InternalCollectionConfiguration::try_from_config(
818            collection_config,
819            KnnIndex::Spann,
820            Some(metadata),
821        );
822
823        assert!(internal_config_result.is_ok());
824
825        let internal_config = internal_config_result.unwrap();
826
827        assert_eq!(
828            internal_config.vector_index,
829            VectorIndexConfiguration::Spann(InternalSpannConfiguration {
830                space: Space::Cosine,
831                ..Default::default()
832            })
833        );
834    }
835
836    #[test]
837    fn test_update_collection_configuration_with_hnsw() {
838        let mut config = InternalCollectionConfiguration {
839            vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
840                space: Space::Cosine,
841                ..Default::default()
842            }),
843            embedding_function: Some(EmbeddingFunctionConfiguration::Known(
844                EmbeddingFunctionNewConfiguration {
845                    name: "test".to_string(),
846                    config: serde_json::Value::Null,
847                },
848            )),
849        };
850        let update_config = UpdateCollectionConfiguration {
851            hnsw: Some(UpdateHnswConfiguration {
852                ef_search: Some(1),
853                ..Default::default()
854            }),
855            spann: None,
856            embedding_function: None,
857        };
858        config.update(&update_config.try_into().unwrap());
859        assert_eq!(
860            config.vector_index,
861            VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
862                space: Space::Cosine,
863                ef_search: 1,
864                ..Default::default()
865            })
866        );
867
868        assert_eq!(
869            config.embedding_function,
870            Some(EmbeddingFunctionConfiguration::Known(
871                EmbeddingFunctionNewConfiguration {
872                    name: "test".to_string(),
873                    config: serde_json::Value::Null,
874                },
875            ))
876        );
877    }
878
879    #[test]
880    fn test_update_collection_configuration_with_spann() {
881        let mut config = InternalCollectionConfiguration {
882            vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
883                space: Space::Cosine,
884                ..Default::default()
885            }),
886            embedding_function: Some(EmbeddingFunctionConfiguration::Known(
887                EmbeddingFunctionNewConfiguration {
888                    name: "test".to_string(),
889                    config: serde_json::Value::Null,
890                },
891            )),
892        };
893        let update_config = UpdateCollectionConfiguration {
894            hnsw: None,
895            spann: Some(UpdateSpannConfiguration {
896                ef_search: Some(1),
897                ..Default::default()
898            }),
899            embedding_function: None,
900        };
901        config.update(&update_config.try_into().unwrap());
902        assert_eq!(
903            config.vector_index,
904            VectorIndexConfiguration::Spann(InternalSpannConfiguration {
905                space: Space::Cosine,
906                ef_search: 1,
907                ..Default::default()
908            })
909        );
910
911        assert_eq!(
912            config.embedding_function,
913            Some(EmbeddingFunctionConfiguration::Known(
914                EmbeddingFunctionNewConfiguration {
915                    name: "test".to_string(),
916                    config: serde_json::Value::Null,
917                },
918            ))
919        );
920    }
921
922    #[test]
923    fn test_update_collection_configuration_with_embedding_function() {
924        let mut config = InternalCollectionConfiguration {
925            vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration::default()),
926            embedding_function: Some(EmbeddingFunctionConfiguration::Known(
927                EmbeddingFunctionNewConfiguration {
928                    name: "test".to_string(),
929                    config: serde_json::Value::Null,
930                },
931            )),
932        };
933        let emb_fn_config = EmbeddingFunctionNewConfiguration {
934            name: "test2".to_string(),
935            config: serde_json::Value::Object(serde_json::Map::from_iter([(
936                "test".to_string(),
937                serde_json::Value::String("test".to_string()),
938            )])),
939        };
940        let update_config = UpdateCollectionConfiguration {
941            hnsw: None,
942            spann: None,
943            embedding_function: Some(EmbeddingFunctionConfiguration::Known(emb_fn_config)),
944        };
945        config.update(&update_config.try_into().unwrap());
946        assert_eq!(
947            config.embedding_function,
948            Some(EmbeddingFunctionConfiguration::Known(
949                EmbeddingFunctionNewConfiguration {
950                    name: "test2".to_string(),
951                    config: serde_json::Value::Object(serde_json::Map::from_iter([(
952                        "test".to_string(),
953                        serde_json::Value::String("test".to_string()),
954                    )])),
955                },
956            ))
957        );
958    }
959}
chroma_types/collection_configuration.rs

chroma_types/
collection_configuration.rs
