diff --git a/pgml-sdks/pgml/src/collection.rs b/pgml-sdks/pgml/src/collection.rs index 69dd8574a..f8107d050 100644 --- a/pgml-sdks/pgml/src/collection.rs +++ b/pgml-sdks/pgml/src/collection.rs @@ -935,7 +935,6 @@ impl Collection { /// Ok(()) /// } #[instrument(skip(self))] - #[allow(clippy::type_complexity)] pub async fn vector_search( &mut self, query: Json, @@ -995,6 +994,32 @@ impl Collection { } } + /// Same as vector_search but assumes embeddings are done locally + #[instrument(skip(self))] + pub async fn vector_search_local( + &self, + query: Json, + pipeline: &Pipeline, + ) -> anyhow::Result> { + let pool = get_or_initialize_pool(&self.database_url).await?; + let (built_query, values) = + build_vector_search_query(query.clone(), self, pipeline).await?; + let results: Vec<(Json, String, f64)> = sqlx::query_as_with(&built_query, values) + .fetch_all(&pool) + .await?; + Ok(results + .into_iter() + .map(|v| { + serde_json::json!({ + "document": v.0, + "chunk": v.1, + "score": v.2 + }) + .into() + }) + .collect()) + } + /// Archives a [Collection] /// This will free up the name to be reused. It does not delete it. ///