Generate gRPC types, update builders and conversions

qdrant · Jan 15, 2025 · a29cb7b · a29cb7b
1 parent 5924a9e
commit a29cb7b
Show file tree

Hide file tree

Showing 9 changed files with 420 additions and 33 deletions.
diff --git a/src/builders/max_optimization_threads_builder.rs b/src/builders/max_optimization_threads_builder.rs
@@ -0,0 +1,15 @@
+use crate::qdrant::*;
+
+impl From<u64> for MaxOptimizationThreads {
+    fn from(threads: u64) -> Self {
+        MaxOptimizationThreads {
+            variant: Some(max_optimization_threads::Variant::from(threads)),
+        }
+    }
+}
+
+impl From<u64> for max_optimization_threads::Variant {
+    fn from(threads: u64) -> Self {
+        Self::Value(threads)
+    }
+}
diff --git a/src/builders/mod.rs b/src/builders/mod.rs
@@ -70,6 +70,8 @@ pub use create_alias_builder::CreateAliasBuilder;
 mod rename_alias_builder;
 pub use rename_alias_builder::RenameAliasBuilder;
 
+mod max_optimization_threads_builder;
+
 mod move_shard_builder;
 pub use move_shard_builder::MoveShardBuilder;
 

diff --git a/src/builders/optimizers_config_diff_builder.rs b/src/builders/optimizers_config_diff_builder.rs
@@ -54,7 +54,7 @@ pub struct OptimizersConfigDiffBuilder {
     /// Note: each optimization job will also use `max_indexing_threads` threads by itself for index building.
     /// If null - have no limit and choose dynamically to saturate CPU.
     /// If 0 - no optimization threads, optimizations will be disabled.
-    pub(crate) max_optimization_threads: Option<Option<u64>>,
+    pub(crate) max_optimization_threads: Option<Option<MaxOptimizationThreads>>,
 }
 
 impl OptimizersConfigDiffBuilder {
@@ -147,9 +147,12 @@ impl OptimizersConfigDiffBuilder {
     /// If null - have no limit and choose dynamically to saturate CPU.
     /// If 0 - no optimization threads, optimizations will be disabled.
     #[allow(unused_mut)]
-    pub fn max_optimization_threads(self, value: u64) -> Self {
+    pub fn max_optimization_threads<VALUE: Into<MaxOptimizationThreads>>(
+        self,
+        value: VALUE,
+    ) -> Self {
         let mut new = self;
-        new.max_optimization_threads = Option::Some(Option::Some(value));
+        new.max_optimization_threads = Option::Some(Option::Some(value.into()));
         new
     }
 
@@ -163,6 +166,8 @@ impl OptimizersConfigDiffBuilder {
             indexing_threshold: self.indexing_threshold.unwrap_or_default(),
             flush_interval_sec: self.flush_interval_sec.unwrap_or_default(),
             max_optimization_threads: self.max_optimization_threads.unwrap_or_default(),
+            // Deprecated: replaced with max_optimization_threads
+            deprecated_max_optimization_threads: None,
         })
     }
     /// Create an empty builder, with all fields set to `None` or `PhantomData`.

diff --git a/src/builders/strict_mode_config_builder.rs b/src/builders/strict_mode_config_builder.rs
@@ -9,6 +9,13 @@ pub struct StrictModeConfigBuilder {
     pub(crate) search_max_hnsw_ef: Option<Option<u32>>,
     pub(crate) search_allow_exact: Option<Option<bool>>,
     pub(crate) search_max_oversampling: Option<Option<f32>>,
+    pub(crate) upsert_max_batchsize: Option<Option<u64>>,
+    pub(crate) max_collection_vector_size_bytes: Option<Option<u64>>,
+    pub(crate) read_rate_limit: Option<Option<u32>>,
+    pub(crate) write_rate_limit: Option<Option<u32>>,
+    pub(crate) max_collection_payload_size_bytes: Option<Option<u64>>,
+    pub(crate) filter_max_conditions: Option<Option<u64>>,
+    pub(crate) condition_max_size: Option<Option<u64>>,
 }
 
 impl StrictModeConfigBuilder {
@@ -60,6 +67,48 @@ impl StrictModeConfigBuilder {
         new.search_max_oversampling = Option::Some(Option::Some(value));
         new
     }
+    #[allow(unused_mut)]
+    pub fn upsert_max_batchsize(self, value: u64) -> Self {
+        let mut new = self;
+        new.upsert_max_batchsize = Option::Some(Option::Some(value));
+        new
+    }
+    #[allow(unused_mut)]
+    pub fn max_collection_vector_size_bytes(self, value: u64) -> Self {
+        let mut new = self;
+        new.max_collection_vector_size_bytes = Option::Some(Option::Some(value));
+        new
+    }
+    #[allow(unused_mut)]
+    pub fn read_rate_limit(self, value: u32) -> Self {
+        let mut new = self;
+        new.read_rate_limit = Option::Some(Option::Some(value));
+        new
+    }
+    #[allow(unused_mut)]
+    pub fn write_rate_limit(self, value: u32) -> Self {
+        let mut new = self;
+        new.write_rate_limit = Option::Some(Option::Some(value));
+        new
+    }
+    #[allow(unused_mut)]
+    pub fn max_collection_payload_size_bytes(self, value: u64) -> Self {
+        let mut new = self;
+        new.max_collection_payload_size_bytes = Option::Some(Option::Some(value));
+        new
+    }
+    #[allow(unused_mut)]
+    pub fn filter_max_conditions(self, value: u64) -> Self {
+        let mut new = self;
+        new.filter_max_conditions = Option::Some(Option::Some(value));
+        new
+    }
+    #[allow(unused_mut)]
+    pub fn condition_max_size(self, value: u64) -> Self {
+        let mut new = self;
+        new.condition_max_size = Option::Some(Option::Some(value));
+        new
+    }
 
     fn build_inner(self) -> Result<StrictModeConfig, std::convert::Infallible> {
         Ok(StrictModeConfig {
@@ -71,6 +120,17 @@ impl StrictModeConfigBuilder {
             search_max_hnsw_ef: self.search_max_hnsw_ef.unwrap_or_default(),
             search_allow_exact: self.search_allow_exact.unwrap_or_default(),
             search_max_oversampling: self.search_max_oversampling.unwrap_or_default(),
+            upsert_max_batchsize: self.upsert_max_batchsize.unwrap_or_default(),
+            max_collection_vector_size_bytes: self
+                .max_collection_vector_size_bytes
+                .unwrap_or_default(),
+            read_rate_limit: self.read_rate_limit.unwrap_or_default(),
+            write_rate_limit: self.write_rate_limit.unwrap_or_default(),
+            max_collection_payload_size_bytes: self
+                .max_collection_payload_size_bytes
+                .unwrap_or_default(),
+            filter_max_conditions: self.filter_max_conditions.unwrap_or_default(),
+            condition_max_size: self.condition_max_size.unwrap_or_default(),
         })
     }
     /// Create an empty builder, with all fields set to `None` or `PhantomData`.
@@ -84,6 +144,13 @@ impl StrictModeConfigBuilder {
             search_max_hnsw_ef: core::default::Default::default(),
             search_allow_exact: core::default::Default::default(),
             search_max_oversampling: core::default::Default::default(),
+            upsert_max_batchsize: core::default::Default::default(),
+            max_collection_vector_size_bytes: core::default::Default::default(),
+            read_rate_limit: core::default::Default::default(),
+            write_rate_limit: core::default::Default::default(),
+            max_collection_payload_size_bytes: core::default::Default::default(),
+            filter_max_conditions: core::default::Default::default(),
+            condition_max_size: core::default::Default::default(),
         }
     }
 }

diff --git a/src/builders/update_collection_builder.rs b/src/builders/update_collection_builder.rs
@@ -18,6 +18,8 @@ pub struct UpdateCollectionBuilder {
     quantization_config: Option<quantization_config_diff::Quantization>,
     /// New sparse vector parameters
     pub(crate) sparse_vectors_config: Option<Option<SparseVectorConfig>>,
+    /// New strict mode configuration
+    pub(crate) strict_mode_config: Option<Option<StrictModeConfig>>,
 }
 
 impl UpdateCollectionBuilder {
@@ -91,6 +93,16 @@ impl UpdateCollectionBuilder {
         new.sparse_vectors_config = Option::Some(Option::Some(value.into()));
         new
     }
+    /// New strict mode configuration
+    #[allow(unused_mut)]
+    pub fn strict_mode_config<VALUE: core::convert::Into<StrictModeConfig>>(
+        self,
+        value: VALUE,
+    ) -> Self {
+        let mut new = self;
+        new.strict_mode_config = Option::Some(Option::Some(value.into()));
+        new
+    }
 
     fn build_inner(self) -> Result<UpdateCollection, UpdateCollectionBuilderError> {
         Ok(UpdateCollection {
@@ -109,6 +121,7 @@ impl UpdateCollectionBuilder {
             vectors_config: self.vectors_config.unwrap_or_default(),
             quantization_config: { convert_option(&self.quantization_config) },
             sparse_vectors_config: self.sparse_vectors_config.unwrap_or_default(),
+            strict_mode_config: self.strict_mode_config.unwrap_or_default(),
         })
     }
     /// Create an empty builder, with all fields set to `None` or `PhantomData`.
@@ -122,6 +135,7 @@ impl UpdateCollectionBuilder {
             vectors_config: core::default::Default::default(),
             quantization_config: core::default::Default::default(),
             sparse_vectors_config: core::default::Default::default(),
+            strict_mode_config: core::default::Default::default(),
         }
     }
 }

diff --git a/src/client/collection.rs b/src/client/collection.rs
@@ -140,6 +140,7 @@ impl QdrantClient {
                         hnsw_config: hnsw_config.cloned(),
                         vectors_config: vectors_config.cloned(),
                         quantization_config: quantization_config.cloned(),
+                        strict_mode_config: None,
                     })
                     .await?;