feature: add timings to all search and recommendation routes for both…

… chunk's and chunk_group's

server/src/bin/ingestion-microservice.rs

@@ -15,7 +15,7 @@ use trieve_server::operators::chunk_operator::{
     update_chunk_metadata_query,
 };
 use trieve_server::operators::event_operator::create_event_query;
-use trieve_server::operators::model_operator::{create_embeddings, get_splade_embedding};
+use trieve_server::operators::model_operator::{create_embeddings, get_sparse_vector};
 use trieve_server::operators::parse_operator::{
     average_embeddings, coarse_doc_chunker, convert_html_to_text,
 };
@@ -420,7 +420,7 @@ async fn upload_chunk(
     };
 
     let splade_vector = if dataset_config.FULLTEXT_ENABLED {
-        match get_splade_embedding(&content.clone(), "doc").await {
+        match get_sparse_vector(&content.clone(), "doc").await {
             Ok(v) => v,
             Err(_) => vec![(0, 0.0)],
         }
@@ -606,7 +606,7 @@ async fn update_chunk(
         .map_err(|_| ServiceError::BadRequest("chunk not found".into()))?;
 
     let splade_vector = if server_dataset_config.FULLTEXT_ENABLED {
-        match get_splade_embedding(&content, "doc").await {
+        match get_sparse_vector(&content, "doc").await {
             Ok(v) => v,
             Err(_) => vec![(0, 0.0)],
         }

server/src/handlers/chunk_handler.rs

@@ -1286,6 +1286,7 @@ pub async fn search_chunk(
                 pool,
                 dataset_org_plan_sub.dataset,
                 server_dataset_config,
+                &mut timer,
             )
             .await?
         }
@@ -1297,6 +1298,7 @@ pub async fn search_chunk(
                 pool,
                 dataset_org_plan_sub.dataset,
                 server_dataset_config,
+                &mut timer,
             )
             .await?
         }
@@ -1307,8 +1309,8 @@ pub async fn search_chunk(
                 page,
                 pool,
                 dataset_org_plan_sub.dataset,
-                &mut timer,
                 server_dataset_config,
+                &mut timer,
             )
             .await?
         }
@@ -1507,6 +1509,10 @@ pub async fn get_recommended_chunks(
 
     let mut positive_qdrant_ids = vec![];
 
+    let mut timer = Timer::new();
+
+    timer.add("start extending tracking_ids and chunk_ids to qdrant_point_ids");
+
     if let Some(positive_chunk_ids) = positive_chunk_ids {
         positive_qdrant_ids.extend(
             get_point_ids_from_unified_chunk_ids(
@@ -1589,6 +1595,8 @@ pub async fn get_recommended_chunks(
         )
     }
 
+    timer.add("finish extending tracking_ids and chunk_ids to qdrant_point_ids; start recommend_qdrant_query");
+
     let recommended_qdrant_point_ids = recommend_qdrant_query(
         positive_qdrant_ids,
         negative_qdrant_ids,
@@ -1603,6 +1611,8 @@ pub async fn get_recommended_chunks(
         ServiceError::BadRequest(format!("Could not get recommended chunks: {}", err))
     })?;
 
+    timer.add("finish recommend_qdrant_query; start get_metadata_from_point_ids");
+
     let recommended_chunk_metadatas =
         get_metadata_from_point_ids(recommended_qdrant_point_ids, pool)
             .await
@@ -1613,6 +1623,8 @@ pub async fn get_recommended_chunks(
                 ))
             })?;
 
+    timer.add("finish get_metadata_from_point_ids and return results");
+
     if data.slim_chunks.unwrap_or(false) {
         let res = recommended_chunk_metadatas
             .into_iter()
@@ -1622,7 +1634,9 @@ pub async fn get_recommended_chunks(
         return Ok(HttpResponse::Ok().json(res));
     }
 
-    Ok(HttpResponse::Ok().json(recommended_chunk_metadatas))
+    Ok(HttpResponse::Ok()
+        .insert_header((Timer::header_key(), timer.header_value()))
+        .json(recommended_chunk_metadatas))
 }
 
 #[derive(Debug, Serialize, Deserialize, ToSchema)]

server/src/handlers/group_handler.rs

@@ -26,6 +26,7 @@ use crate::{
 };
 use actix_web::{web, HttpResponse};
 use serde::{Deserialize, Serialize};
+use simple_server_timing_header::Timer;
 use utoipa::{IntoParams, ToSchema};
 
 #[tracing::instrument(skip(pool))]
@@ -955,6 +956,10 @@ pub async fn get_recommended_groups(
         ServerDatasetConfiguration::from_json(dataset_org_plan_sub.dataset.server_configuration);
     let dataset_id = dataset_org_plan_sub.dataset.id;
 
+    let mut timer = Timer::new();
+
+    timer.add("start to extend qdrant_point_ids for group_tracking_ids and group_ids");
+
     let mut positive_qdrant_ids = vec![];
 
     if let Some(positive_group_ids) = positive_group_ids {
@@ -1039,6 +1044,8 @@ pub async fn get_recommended_groups(
         );
     }
 
+    timer.add("finish to extend qdrant_point_ids for group_tracking_ids and group_ids; start to recommend_qdrant_groups_query from qdrant");
+
     let recommended_qdrant_point_ids = recommend_qdrant_groups_query(
         positive_qdrant_ids,
         negative_qdrant_ids,
@@ -1059,9 +1066,13 @@ pub async fn get_recommended_groups(
         total_chunk_pages: (recommended_qdrant_point_ids.len() as f64 / 10.0).ceil() as i64,
     };
 
+    timer.add("finish to recommend_qdrant_groups_query from qdrant; start to get_metadata_from_groups from postgres");
+
     let recommended_chunk_metadatas =
         get_metadata_from_groups(group_query_result, Some(false), pool).await?;
 
+    timer.add("finish to get_metadata_from_groups from postgres and return results");
+
     if data.slim_chunks.unwrap_or(false) {
         let res = recommended_chunk_metadatas
             .into_iter()
@@ -1075,10 +1086,14 @@ pub async fn get_recommended_groups(
             })
             .collect::<Vec<GroupSlimChunksDTO>>();
 
-        return Ok(HttpResponse::Ok().json(res));
+        return Ok(HttpResponse::Ok()
+            .insert_header((Timer::header_key(), timer.header_value()))
+            .json(res));
     }
 
-    Ok(HttpResponse::Ok().json(recommended_chunk_metadatas))
+    Ok(HttpResponse::Ok()
+        .insert_header((Timer::header_key(), timer.header_value()))
+        .json(recommended_chunk_metadatas))
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, ToSchema, IntoParams)]
@@ -1332,6 +1347,8 @@ pub async fn search_over_groups(
 
     let parsed_query = parse_query(data.query.clone());
 
+    let mut timer = Timer::new();
+
     let result_chunks = match data.search_type.as_str() {
         "fulltext" => {
             if !server_dataset_config.FULLTEXT_ENABLED {
@@ -1348,6 +1365,7 @@ pub async fn search_over_groups(
                 pool,
                 dataset_org_plan_sub.dataset,
                 server_dataset_config,
+                &mut timer,
             )
             .await?
         }
@@ -1359,6 +1377,7 @@ pub async fn search_over_groups(
                 pool,
                 dataset_org_plan_sub.dataset,
                 server_dataset_config,
+                &mut timer,
             )
             .await?
         }
@@ -1370,6 +1389,7 @@ pub async fn search_over_groups(
                 pool,
                 dataset_org_plan_sub.dataset,
                 server_dataset_config,
+                &mut timer,
             )
             .await?
         }
@@ -1394,8 +1414,12 @@ pub async fn search_over_groups(
             total_chunk_pages: result_chunks.total_chunk_pages,
         };
 
-        return Ok(HttpResponse::Ok().json(res));
+        return Ok(HttpResponse::Ok()
+            .insert_header((Timer::header_key(), timer.header_value()))
+            .json(res));
     }
 
-    Ok(HttpResponse::Ok().json(result_chunks))
+    Ok(HttpResponse::Ok()
+        .insert_header((Timer::header_key(), timer.header_value()))
+        .json(result_chunks))
 }

server/src/operators/model_operator.rs

@@ -164,7 +164,7 @@ pub struct CustomSparseEmbedData {
 }
 
 #[tracing::instrument]
-pub async fn get_splade_embedding(
+pub async fn get_sparse_vector(
     message: &str,
     embed_type: &str,
 ) -> Result<Vec<(u32, f32)>, ServiceError> {