Enhance InferenceUsage handling and error processing

n0x29a · n0x29a · commit 0f109ffa239c · 2025-05-23T10:43:16.000+02:00
diff --git a/lib/api/src/rest/conversions.rs b/lib/api/src/rest/conversions.rs
@@ -15,7 +15,7 @@ use crate::rest::{DenseVector, NamedVectorStruct};
 
 impl From<InferenceUsage> for grpc::InferenceUsage {
     fn from(value: InferenceUsage) -> Self {
-        let mut grpc_usage_models = HashMap::default();
+        let mut grpc_usage_models = HashMap::with_capacity(value.models.len());
         for (model, usage) in value.models {
             grpc_usage_models.insert(
                 model,
diff --git a/lib/api/src/rest/models.rs b/lib/api/src/rest/models.rs
@@ -70,6 +70,18 @@ impl Usage {
     }
 }
 
+impl InferenceUsage {
+    pub fn merge(mut self, other: InferenceUsage) -> InferenceUsage {
+        for (model_name, other_usage) in other.models {
+            self.models
+                .entry(model_name)
+                .and_modify(|usage| usage.tokens += other_usage.tokens)
+                .or_insert(other_usage);
+        }
+        self
+    }
+}
+
 fn is_usage_none_or_empty(u: &Option<Usage>) -> bool {
     u.as_ref().is_none_or(|usage| usage.is_empty())
 }
diff --git a/lib/collection/src/operations/verification/query.rs b/lib/collection/src/operations/verification/query.rs
@@ -204,7 +204,12 @@ impl StrictModeVerification for CollectionQueryRequestWithUsage {
             // check query can perform fullscan when not rescoring
             if self.request.prefetch.is_empty() {
                 query
-                    .check_fullscan(&self.request.using, collection, strict_mode_config)
+                    .check_fullscan(
+                        &self.request.using,
+                        self.request.filter.as_ref(),
+                        collection,
+                        strict_mode_config,
+                    )
                     .await?;
             }
             // check for unindexed fields in formula
diff --git a/src/actix/api/snapshot_api.rs b/src/actix/api/snapshot_api.rs
@@ -616,7 +616,7 @@ async fn recover_partial_snapshot(
     let recovery_lock = match try_take_recovery_lock_future.await {
         Ok(recovery_lock) => recovery_lock,
         Err(err) => {
-            return helpers::process_response_error(err, tokio::time::Instant::now(), None);
+            return helpers::process_response_error(err, tokio::time::Instant::now(), None, None);
         }
     };
 
@@ -699,7 +699,7 @@ async fn recover_partial_snapshot_from(
     let recovery_lock = match try_take_recovery_lock_future.await {
         Ok(recovery_lock) => recovery_lock,
         Err(err) => {
-            return helpers::process_response_error(err, tokio::time::Instant::now(), None);
+            return helpers::process_response_error(err, tokio::time::Instant::now(), None, None);
         }
     };
 
diff --git a/src/common/inference/update_requests.rs b/src/common/inference/update_requests.rs
@@ -179,8 +179,11 @@ pub async fn convert_batch(
             for (name, vecs) in named {
                 let (converted_vectors, batch_usage) =
                     convert_vectors(vecs, InferenceType::Update, inference_token.clone()).await?;
-                if inference_usage.is_none() && batch_usage.is_some() {
-                    inference_usage = batch_usage;
+                if let Some(batch_usage) = batch_usage {
+                    inference_usage = match inference_usage.take() {
+                        Some(accum) => Some(accum.merge(batch_usage)),
+                        None => Some(batch_usage),
+                    };
                 }
                 named_vectors.insert(name, converted_vectors);
             }
diff --git a/src/common/update.rs b/src/common/update.rs
@@ -692,8 +692,11 @@ pub async fn do_batch_update_points(
         };
 
         results.push(current_update_result);
-        if inference_usage.is_none() && current_operation_usage_opt.is_some() {
-            inference_usage = current_operation_usage_opt;
+        if let Some(current_usage) = current_operation_usage_opt {
+            inference_usage = match inference_usage {
+                Some(accum) => Some(accum.merge(current_usage)),
+                None => Some(current_usage),
+            };
         }
     }
     Ok((results, inference_usage))

Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,18 @@ impl Usage {`
`70`	`70`	`}`
`71`	`71`	`}`
`72`	`72`
	`73`	`+impl InferenceUsage {`
	`74`	`+ pub fn merge(mut self, other: InferenceUsage) -> InferenceUsage {`
	`75`	`+ for (model_name, other_usage) in other.models {`
	`76`	`+ self.models`
	`77`	`+ .entry(model_name)`
	`78`	`+ .and_modify(\|usage\| usage.tokens += other_usage.tokens)`
	`79`	`+ .or_insert(other_usage);`
	`80`	`+ }`
	`81`	`+ self`
	`82`	`+ }`
	`83`	`+}`
	`84`	`+`
`73`	`85`	`fn is_usage_none_or_empty(u: &Option<Usage>) -> bool {`
`74`	`86`	`u.as_ref().is_none_or(\|usage\| usage.is_empty())`
`75`	`87`	`}`
Original file line number	Diff line number	Diff line change
`@@ -179,8 +179,11 @@ pub async fn convert_batch(`
`179`	`179`	`for (name, vecs) in named {`
`180`	`180`	`let (converted_vectors, batch_usage) =`
`181`	`181`	`convert_vectors(vecs, InferenceType::Update, inference_token.clone()).await?;`
`182`		`- if inference_usage.is_none() && batch_usage.is_some() {`
`183`		`- inference_usage = batch_usage;`
	`182`	`+ if let Some(batch_usage) = batch_usage {`
	`183`	`+ inference_usage = match inference_usage.take() {`
	`184`	`+ Some(accum) => Some(accum.merge(batch_usage)),`
	`185`	`+ None => Some(batch_usage),`
	`186`	`+ };`
`184`	`187`	`}`
`185`	`188`	`named_vectors.insert(name, converted_vectors);`
`186`	`189`	`}`
Original file line number	Diff line number	Diff line change
`@@ -692,8 +692,11 @@ pub async fn do_batch_update_points(`
`692`	`692`	`};`
`693`	`693`
`694`	`694`	`results.push(current_update_result);`
`695`		`- if inference_usage.is_none() && current_operation_usage_opt.is_some() {`
`696`		`- inference_usage = current_operation_usage_opt;`
	`695`	`+ if let Some(current_usage) = current_operation_usage_opt {`
	`696`	`+ inference_usage = match inference_usage {`
	`697`	`+ Some(accum) => Some(accum.merge(current_usage)),`
	`698`	`+ None => Some(current_usage),`
	`699`	`+ };`
`697`	`700`	`}`
`698`	`701`	`}`
`699`	`702`	`Ok((results, inference_usage))`