SDK - Added boosting for vector search (#1371)

SilasMarvin · web-flow · commit a5ee9d7d6e28 · 2024-03-14T10:37:44.000-07:00
diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs
@@ -1052,7 +1052,7 @@ mod tests {
     #[tokio::test]
     async fn can_vector_search_with_local_embeddings() -> anyhow::Result<()> {
         internal_init_logger(None, None).ok();
-        let collection_name = "test r_c_cvswle_9";
+        let collection_name = "test r_c_cvswle_13";
         let mut collection = Collection::new(collection_name, None)?;
         let documents = generate_dummy_documents(10);
         collection.upsert_documents(documents.clone(), None).await?;
@@ -1095,10 +1095,12 @@ mod tests {
                                 "parameters": {
                                     "instruction": "Represent the Wikipedia document for retrieval"
                                 },
-                                "full_text_filter": "test"
+                                "full_text_filter": "test",
+                                "boost": 1.2
                             },
                             "body": {
-                                "query": "Test document: 2"
+                                "query": "Test document: 2",
+                                "boost": 1.0
                             },
                         },
                         "filter": {
diff --git a/pgml-sdks/pgml/src/vector_search_query_builder.rs b/pgml-sdks/pgml/src/vector_search_query_builder.rs
@@ -25,6 +25,7 @@ struct ValidField {
     query: String,
     parameters: Option<Json>,
     full_text_filter: Option<String>,
+    boost: Option<f32>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -98,6 +99,8 @@ pub async fn build_vector_search_query(
 
         let mut query = Query::select();
 
+        let boost = vf.boost.unwrap_or(1.);
+
         match model_runtime {
             ModelRuntime::Python => {
                 // Build the embedding CTE
@@ -118,7 +121,7 @@ pub async fn build_vector_search_query(
 
                 query
                     .expr(Expr::cust(format!(
-                        r#"1 - (embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector) AS score"#
+                        r#"(1 - (embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector)) * {boost} AS score"#
                     )))
                     .order_by_expr(Expr::cust(format!(
                         r#"embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector"#
@@ -152,7 +155,7 @@ pub async fn build_vector_search_query(
                 // Build the score CTE
                 query
                     .expr(Expr::cust_with_values(
-                        r#"1 - (embeddings.embedding <=> $1::vector) AS score"#,
+                        r#"(1 - (embeddings.embedding <=> $1::vector)) {boost} AS score"#,
                         [embedding.clone()],
                     ))
                     .order_by_expr(
