moves auto query to collections

oramasearch · Feb 20, 2025 · da880bb · da880bb
1 parent f78a790
commit da880bb
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 33 deletions.
diff --git a/src/collection_manager/sides/read/collection.rs b/src/collection_manager/sides/read/collection.rs
@@ -23,11 +23,11 @@ mod merge;
 mod uncommitted;
 
 use crate::{
-    ai::{AIService, OramaModel},
+    ai::{AIService, AutoQueryResponse, OramaModel},
     collection_manager::{
         dto::{
-            ApiKey, BM25Scorer, FacetDefinition, FacetResult, FieldId, Filter, Limit, NumberFilter,
-            Properties, SearchMode, SearchParams,
+            ApiKey, BM25Scorer, FacetDefinition, FacetResult, FieldId, Filter, FulltextMode,
+            HybridMode, Limit, NumberFilter, Properties, SearchMode, SearchParams, VectorMode,
         },
         sides::{CollectionWriteOperation, Offset, OramaModelSerializable, TypedFieldWrapper},
     },
@@ -715,9 +715,39 @@ impl CollectionReader {
             );
         }
 
+        // Manage the "auto" search mode. OramaCore will automatically determine
+        // wether to use full text search, vector search or hybrid search.
+        let search_mode: SearchMode = match mode.clone() {
+            SearchMode::Auto(mode_result) => {
+                let final_mode = self
+                    .ai_service
+                    .get_autoquery(mode_result.term.clone())
+                    .await
+                    .unwrap_or(AutoQueryResponse {
+                        // By default, we can use the hybrid mode when something goes wrong
+                        // with the AI service. This is the most versatile mode.
+                        mode: "hybrid".to_string(),
+                    });
+
+                match final_mode.mode.as_str() {
+                    "fulltext" => SearchMode::FullText(FulltextMode {
+                        term: mode_result.term,
+                    }),
+                    "hybrid" => SearchMode::Hybrid(HybridMode {
+                        term: mode_result.term,
+                    }),
+                    "vector" => SearchMode::Vector(VectorMode {
+                        term: mode_result.term,
+                    }),
+                    _ => anyhow::bail!("Invalid search mode"),
+                }
+            }
+            _ => mode.clone(),
+        };
+
         let boost = self.calculate_boost(boost);
 
-        let token_scores = match mode {
+        let token_scores = match search_mode {
             SearchMode::Default(search_params) | SearchMode::FullText(search_params) => {
                 let properties = self.calculate_string_properties(properties)?;
                 self.search_full_text(

diff --git a/src/web_server/api/collection/search.rs b/src/web_server/api/collection/search.rs
@@ -48,35 +48,7 @@ async fn search(
     let collection_id = CollectionId(id);
     let read_api_key = query.api_key;
 
-    let search_mode: SearchMode = match json.mode.clone() {
-        SearchMode::Auto(mode) => {
-            let final_mode = read_side
-                .get_search_mode(mode.term.clone())
-                .await
-                .context("Error getting search mode")
-                .unwrap_or(SearchMode::Hybrid(HybridMode { term: mode.term }));
-
-            match final_mode {
-                SearchMode::FullText(mode) => SearchMode::FullText(mode),
-                SearchMode::Hybrid(mode) => SearchMode::Hybrid(mode),
-                SearchMode::Vector(mode) => SearchMode::Vector(mode),
-                _ => Err((
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    Json(json!({ "error": "Invalid search mode" })),
-                ))?,
-            }
-        }
-        _ => json.mode.clone(),
-    };
-
-    let final_json = SearchParams {
-        mode: search_mode,
-        ..json
-    };
-
-    let output = read_side
-        .search(read_api_key, collection_id, final_json)
-        .await;
+    let output = read_side.search(read_api_key, collection_id, json).await;
 
     match output {
         Ok(data) => Ok((StatusCode::OK, Json(data))),