Commita5ee9d7

authored

SDK - Added boosting for vector search (#1371)

1 parent9510b55 commita5ee9d7Copy full SHA for a5ee9d7

File tree

+10

-5

lines changed

+10

-5

lines changed

Lines changed: 5 additions & 3 deletions

Original file line number	Diff line number	Diff line change
`@@ -1052,7 +1052,7 @@ mod tests {`
`1052`	`1052`	`#[tokio::test]`
`1053`	`1053`	`asyncfncan_vector_search_with_local_embeddings() -> anyhow::Result<()>{`
`1054`	`1054`	`internal_init_logger(None,None).ok();`
`1055`		`-let collection_name ="testr_c_cvswle_9";`
	`1055`	`+let collection_name ="testr_c_cvswle_13";`
`1056`	`1056`	`letmut collection =Collection::new(collection_name,None)?;`
`1057`	`1057`	`let documents =generate_dummy_documents(10);`
`1058`	`1058`	`collection.upsert_documents(documents.clone(),None).await?;`
`@@ -1095,10 +1095,12 @@ mod tests {`
`1095`	`1095`	`"parameters":{`
`1096`	`1096`	`"instruction":"Represent the Wikipedia document for retrieval"`
`1097`	`1097`	`},`
`1098`		`-"full_text_filter":"test"`
	`1098`	`+"full_text_filter":"test",`
	`1099`	`+"boost":1.2`
`1099`	`1100`	`},`
`1100`	`1101`	`"body":{`
`1101`		`-"query":"Test document: 2"`
	`1102`	`+"query":"Test document: 2",`
	`1103`	`+"boost":1.0`
`1102`	`1104`	`},`
`1103`	`1105`	`},`
`1104`	`1106`	`"filter":{`

Lines changed: 5 additions & 2 deletions

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@ struct ValidField {`
`25`	`25`	`query:String,`
`26`	`26`	`parameters:Option<Json>,`
`27`	`27`	`full_text_filter:Option<String>,`
	`28`	`+boost:Option<f32>,`
`28`	`29`	`}`
`29`	`30`
`30`	`31`	`#[derive(Debug,Deserialize)]`
`@@ -98,6 +99,8 @@ pub async fn build_vector_search_query(`
`98`	`99`
`99`	`100`	`letmut query =Query::select();`
`100`	`101`
	`102`	`+let boost = vf.boost.unwrap_or(1.);`
	`103`	`+`
`101`	`104`	`match model_runtime{`
`102`	`105`	`ModelRuntime::Python =>{`
`103`	`106`	`// Build the embedding CTE`
`@@ -118,7 +121,7 @@ pub async fn build_vector_search_query(`
`118`	`121`
`119`	`122`	`query`
`120`	`123`	`.expr(Expr::cust(format!(`
`121`		`-r#"1 - (embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector) AS score"#`
	`124`	`+r#"(1 - (embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector)) * {boost} AS score"#`
`122`	`125`	`)))`
`123`	`126`	`.order_by_expr(Expr::cust(format!(`
`124`	`127`	`r#"embeddings.embedding <=> (SELECT embedding FROM "{key}_embedding")::vector"#`
`@@ -152,7 +155,7 @@ pub async fn build_vector_search_query(`
`152`	`155`	`// Build the score CTE`
`153`	`156`	`query`
`154`	`157`	`.expr(Expr::cust_with_values(`
`155`		`-r#"1 - (embeddings.embedding <=> $1::vector) AS score"#,`
	`158`	`+r#"(1 - (embeddings.embedding <=> $1::vector)) {boost} AS score"#,`
`156`	`159`	`[embedding.clone()],`
`157`	`160`	`))`
`158`	`161`	`.order_by_expr(`

Comments

(0)