Add inference usage in errors as well

n0x29a · n0x29a · commit d043cd1e0e09 · 2025-05-22T17:20:45.000+02:00
diff --git a/src/actix/api/query_api.rs b/src/actix/api/query_api.rs
@@ -3,7 +3,9 @@ use actix_web_validator::{Json, Path, Query};
 use api::rest::models::{InferenceUsage, ModelUsage};
 use api::rest::{QueryGroupsRequest, QueryRequest, QueryRequestBatch, QueryResponse};
 use collection::operations::shard_selector_internal::ShardSelectorInternal;
-use collection::operations::universal_query::collection_query::CollectionQueryGroupsRequestWithUsage;
+use collection::operations::universal_query::collection_query::{
+    CollectionQueryGroupsRequestWithUsage, CollectionQueryRequestWithUsage,
+};
 use itertools::Itertools;
 use storage::content_manager::collection_verification::{
     check_strict_mode, check_strict_mode_batch,
@@ -159,7 +161,7 @@ async fn query_points_batch(
                 shard_key,
             } = request_item;
 
-            let collection::operations::universal_query::collection_query::CollectionQueryRequestWithUsage { request, usage } =
+            let CollectionQueryRequestWithUsage { request, usage } =
                 convert_query_request_from_rest(internal, &inference_token).await?;
 
             inference_usages.push(usage);
@@ -176,17 +178,17 @@ async fn query_points_batch(
         let mut total_usage = InferenceUsage::default();
         for inference_usage in inference_usages.iter_mut().flatten() {
             let usage = inference_usage;
-                for (model, usage) in usage.models.iter() {
-                    total_usage
-                        .models
-                        .entry(model.clone())
-                        .and_modify(|e| {
-                            e.tokens += usage.tokens;
-                        })
-                        .or_insert_with(|| ModelUsage {
-                            tokens: usage.tokens,
-                        });
-                }
+            for (model, usage) in usage.models.iter() {
+                total_usage
+                    .models
+                    .entry(model.clone())
+                    .and_modify(|e| {
+                        e.tokens += usage.tokens;
+                    })
+                    .or_insert_with(|| ModelUsage {
+                        tokens: usage.tokens,
+                    });
+            }
         }
 
         let inference_usage: Option<InferenceUsage> = {
@@ -204,7 +206,11 @@ async fn query_points_batch(
             &dispatcher,
             &access,
         )
-        .await?;
+        .await
+        .map_err(|err| StorageError::InferenceError {
+            description: err.to_string(),
+            usage: inference_usage.clone(),
+        })?;
 
         let res = dispatcher
             .toc(&access, &pass)
@@ -216,7 +222,11 @@ async fn query_points_batch(
                 params.timeout(),
                 hw_measurement_acc,
             )
-            .await?
+            .await
+            .map_err(|err| StorageError::InferenceError {
+                description: err.to_string(),
+                usage: inference_usage.clone(),
+            })?
             .into_iter()
             .map(|response| QueryResponse {
                 points: response
@@ -287,7 +297,11 @@ async fn query_points_groups(
             &dispatcher,
             &access,
         )
-        .await?;
+        .await
+        .map_err(|err| StorageError::InferenceError {
+            description: err.to_string(),
+            usage: usage.clone(),
+        })?;
 
         let query_result = do_query_point_groups(
             dispatcher.toc(&access, &pass),
@@ -299,7 +313,11 @@ async fn query_points_groups(
             params.timeout(),
             hw_measurement_acc,
         )
-        .await?;
+        .await
+        .map_err(|err| StorageError::InferenceError {
+            description: err.to_string(),
+            usage: usage.clone(),
+        })?;
         Ok((query_result, usage))
     }
     .await;
