Working site search with doc type filtering

SilasMarvin · SilasMarvin · commit ec351ffce454 · 2024-02-28T08:57:57.000-08:00
diff --git a/pgml-sdks/pgml/src/collection.rs b/pgml-sdks/pgml/src/collection.rs
@@ -718,6 +718,19 @@ impl Collection {
         }
     }
 
+    #[instrument(skip(self))]
+    pub async fn search_local(
+        &self,
+        query: Json,
+        pipeline: &MultiFieldPipeline,
+    ) -> anyhow::Result<Vec<Json>> {
+        let pool = get_or_initialize_pool(&self.database_url).await?;
+        let (built_query, values) = build_search_query(self, query.clone(), pipeline).await?;
+        let results: Vec<(Json,)> = sqlx::query_as_with(&built_query, values)
+            .fetch_all(&pool)
+            .await?;
+        Ok(results.into_iter().map(|v| v.0).collect())
+    }
     /// Performs vector search on the [Collection]
     ///
     /// # Arguments
diff --git a/pgml-sdks/pgml/src/search_query_builder.rs b/pgml-sdks/pgml/src/search_query_builder.rs
@@ -307,7 +307,7 @@ pub async fn build_search_query(
             .expr(Expr::cust("*"))
             .from_subquery(query, Alias::new("q1"))
             .order_by(SIden::Str("score"), Order::Desc)
-            .limit(5);
+            .limit(limit);
 
         let mut combined_query = Query::select();
         combined_query
