bugfix(server): rename get_random_chunk_metadatas_query to get_random_chunk_qdrant_point_id_query and update its logic to reduce latency

Author: skeptrunedev

clients/search-component/example/src/routeTree.gen.ts

@@ -1,12 +1,12 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 
 // @ts-nocheck
 
 // noinspection JSUnusedGlobalSymbols
 
-// This file is auto-generated by TanStack Router
+// This file was automatically generated by TanStack Router.
+// You should NOT make any changes in this file as it will be overwritten.
+// Additionally, you should also exclude this file from your linter and/or formatter to prevent it from being checked or modified.
 
 import { createFileRoute } from "@tanstack/react-router";
 
@@ -24,21 +24,25 @@ const IndexLazyImport = createFileRoute("/")();
 // Create/Update Routes
 
 const SearchpageRoute = SearchpageImport.update({
+  id: "/searchpage",
   path: "/searchpage",
   getParentRoute: () => rootRoute,
 } as any);
 
 const RecsRoute = RecsImport.update({
+  id: "/recs",
   path: "/recs",
   getParentRoute: () => rootRoute,
 } as any);
 
 const EcommerceRoute = EcommerceImport.update({
+  id: "/ecommerce",
   path: "/ecommerce",
   getParentRoute: () => rootRoute,
 } as any);
 
 const IndexLazyRoute = IndexLazyImport.update({
+  id: "/",
   path: "/",
   getParentRoute: () => rootRoute,
 } as any).lazy(() => import("./routes/index.lazy").then((d) => d.Route));
@@ -129,8 +133,6 @@ export const routeTree = rootRoute
   ._addFileChildren(rootRouteChildren)
   ._addFileTypes<FileRouteTypes>();
 
-/* prettier-ignore-end */
-
 /* ROUTE_MANIFEST_START
 {
   "routes": {

server/src/operators/chunk_operator.rs

@@ -286,11 +286,10 @@ pub async fn get_content_chunk_from_point_ids_query(
     Ok(content_chunks)
 }
 
-pub async fn get_random_chunk_metadatas_query(
+pub async fn get_random_chunk_qdrant_point_id_query(
     dataset_id: uuid::Uuid,
-    limit: u64,
     pool: web::Data<Pool>,
-) -> Result<Vec<ChunkMetadata>, ServiceError> {
+) -> Result<uuid::Uuid, ServiceError> {
     use crate::data::schema::chunk_metadata::dsl as chunk_metadata_columns;
 
     let mut conn = pool
@@ -303,53 +302,39 @@ pub async fn get_random_chunk_metadatas_query(
         .select(chunk_metadata_columns::id)
         .order_by(chunk_metadata_columns::id.asc())
         .first::<uuid::Uuid>(&mut conn);
-    let get_highest_ids_future = chunk_metadata_columns::chunk_metadata
+    let get_highest_id_future = chunk_metadata_columns::chunk_metadata
         .filter(chunk_metadata_columns::dataset_id.eq(dataset_id))
         .select(chunk_metadata_columns::id)
         .order_by(chunk_metadata_columns::id.desc())
-        .limit(limit.try_into().unwrap_or(10))
-        .load::<uuid::Uuid>(&mut conn);
-    let (lowest_id, highest_ids) = futures::join!(get_lowest_id_future, get_highest_ids_future);
+        .first::<uuid::Uuid>(&mut conn);
+    let (lowest_id, highest_id) = futures::join!(get_lowest_id_future, get_highest_id_future);
     let lowest_id: uuid::Uuid = lowest_id.map_err(|err| {
         ServiceError::BadRequest(format!(
             "Failed to load chunk with the lowest id in the dataset for random range: {:?}",
             err
         ))
     })?;
-    let highest_ids: Vec<uuid::Uuid> = highest_ids.map_err(|err| {
+    let highest_id: uuid::Uuid = highest_id.map_err(|err| {
         ServiceError::BadRequest(format!(
-            "Failed to load chunks with the highest id in the dataset for random range: {:?}",
+            "Failed to load chunk with the lowest id in the dataset for random range: {:?}",
             err
         ))
     })?;
-    let highest_id = match highest_ids.get(0) {
-        Some(id) => *id,
-        None => {
-            return Err(ServiceError::NotFound(
-                "Chunk with the highest id in the dataset not found for random range".to_string(),
-            ))
-        }
-    };
     let random_uuid = uuid_between(lowest_id, highest_id);
 
-    let chunk_metadatas: Vec<ChunkMetadataTable> = chunk_metadata_columns::chunk_metadata
+    let qdrant_point_id: uuid::Uuid = chunk_metadata_columns::chunk_metadata
         .filter(chunk_metadata_columns::dataset_id.eq(dataset_id))
         .filter(chunk_metadata_columns::id.gt(random_uuid))
         .order_by(chunk_metadata_columns::id.asc())
-        .limit(limit.try_into().unwrap_or(10))
-        .select(ChunkMetadataTable::as_select())
-        .load::<ChunkMetadataTable>(&mut conn)
+        .select(chunk_metadata_columns::qdrant_point_id)
+        .first::<uuid::Uuid>(&mut conn)
         .await
-        .map_err(|_| ServiceError::BadRequest("Failed to load metadata".to_string()))?;
-
-    let chunk_metadatas = chunk_metadatas
-        .into_iter()
-        .map(|chunk_metadata_table| {
-            ChunkMetadata::from_table_and_tag_set(chunk_metadata_table, vec![])
-        })
-        .collect();
+        .map_err(|e| {
+            log::error!("Failed to load random qdrant_point_ids: {:?}", e);
+            ServiceError::BadRequest("Failed to load random qdrant_point_ids".to_string())
+        })?;
 
-    Ok(chunk_metadatas)
+    Ok(qdrant_point_id)
 }
 
 pub async fn get_metadata_from_id_query(

server/src/operators/message_operator.rs

@@ -1,4 +1,4 @@
-use crate::operators::chunk_operator::get_random_chunk_metadatas_query;
+use crate::operators::chunk_operator::get_random_chunk_qdrant_point_id_query;
 use crate::operators::message_operator::models::DatasetAndOrgWithSubAndPlan;
 use itertools::Itertools;
 use simple_server_timing_header::Timer;
@@ -1849,20 +1849,24 @@ pub async fn suggested_new_queries(
         }
         None => {
             let random_offset_id = match payload.filters {
-                Some(_) => Some(uuid::Uuid::nil()),
-                None => get_random_chunk_metadatas_query(dataset_id, 1, pool.clone())
-                    .await?
-                    .clone()
-                    .get(0)
-                    .cloned()
-                    .map(|chunk| chunk.qdrant_point_id),
+                Some(_) => uuid::Uuid::nil(),
+                None => get_random_chunk_qdrant_point_id_query(dataset_id, pool.clone()).await?,
             };
             let filter =
                 assemble_qdrant_filter(filters, None, None, dataset_id, pool.clone()).await?;
 
-            let (search_results, _) =
-                scroll_dataset_points(10, random_offset_id, None, dataset_config.clone(), filter)
-                    .await?;
+            let (search_results, _) = scroll_dataset_points(
+                payload
+                    .suggestions_to_create
+                    .unwrap_or(5)
+                    .try_into()
+                    .unwrap(),
+                Some(random_offset_id),
+                None,
+                dataset_config.clone(),
+                filter,
+            )
+            .await?;
             if qdrant_only {
                 search_results
                     .iter()
@@ -1890,11 +1894,9 @@ pub async fn suggested_new_queries(
 
     let rag_content = chunk_metadatas
         .iter()
-        .enumerate()
-        .map(|(idx, chunk)| {
+        .map(|chunk| {
             format!(
-                "Doc {}: {}",
-                idx + 1,
+                "- {}",
                 convert_html_to_text(&(chunk.chunk_html.clone().unwrap_or_default()))
             )
         })
@@ -1920,7 +1922,7 @@ pub async fn suggested_new_queries(
         None => "".to_string(),
     };
 
-    let number_of_suggestions_to_create = payload.suggestions_to_create.unwrap_or(10);
+    let number_of_suggestions_to_create = payload.suggestions_to_create.unwrap_or(5);
 
     let content = ChatMessageContent::Text(format!(
         "Here is some content which the user might be looking for: {rag_content}{context_sentence}. Generate {number_of_suggestions_to_create} varied followup {query_style} style queries based off the domain of this dataset. Your only response should be the {number_of_suggestions_to_create} followup {query_style} style queries which are separated by new lines and are just text and you do not add any other context or information about the followup {query_style} style queries. This should not be a list, so do not number each {query_style} style queries.",