Skip to content

Commit 027080f

Browse files
SilasMarvinSilasMarvin
committed
collection query_builder now a wrapper around collection.vector_search
1 parent ec351ff commit 027080f

14 files changed

+208
-479
lines changed

‎pgml-sdks/pgml/src/collection.rs

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,11 @@ use indicatif::MultiProgress;
33
use itertools::Itertools;
44
use regex::Regex;
55
use rust_bridge::{alias, alias_methods};
6-
use sea_query::{Alias, Expr, JoinType, NullOrdering, Order, PostgresQueryBuilder, Query};
6+
use sea_query::{Expr, NullOrdering, Order, PostgresQueryBuilder, Query};
77
use sea_query_binder::SqlxBinder;
88
use serde_json::json;
9-
use sqlx::postgres::PgPool;
109
use sqlx::Executor;
1110
use sqlx::PgConnection;
12-
use sqlx::Transaction;
1311
use std::borrow::Cow;
1412
use std::path::Path;
1513
use std::sync::Arc;
@@ -22,22 +20,20 @@ use crate::filter_builder::FilterBuilder;
2220
use crate::search_query_builder::build_search_query;
2321
use crate::vector_search_query_builder::build_vector_search_query;
2422
use crate::{
25-
filter_builder, get_or_initialize_pool,
26-
model::ModelRuntime,
27-
models,
23+
get_or_initialize_pool, models,
2824
multi_field_pipeline::MultiFieldPipeline,
29-
order_by_builder,
30-
pipeline::Pipeline,
31-
queries, query_builder,
25+
order_by_builder, queries, query_builder,
3226
query_builder::QueryBuilder,
33-
remote_embeddings::build_remote_embeddings,
3427
splitter::Splitter,
3528
types::{DateTime, IntoTableNameAndSchema, Json, SIden, TryToNumeric},
3629
utils,
3730
};
3831

3932
#[cfg(feature = "python")]
40-
use crate::{pipeline::PipelinePython, query_builder::QueryBuilderPython, types::JsonPython};
33+
use crate::{
34+
multi_field_pipeline::MultiFieldPipelinePython, query_builder::QueryBuilderPython,
35+
types::JsonPython,
36+
};
4137

4238
/// Our project tasks
4339
#[derive(Debug, Clone)]
@@ -738,7 +734,6 @@ impl Collection {
738734
/// * `query` - The query to search for
739735
/// * `pipeline` - The [Pipeline] used for the search
740736
/// * `query_paramaters` - The query parameters passed to the model for search
741-
/// * `top_k` - How many results to limit on.
742737
///
743738
/// # Example
744739
///
@@ -758,7 +753,6 @@ impl Collection {
758753
&mut self,
759754
query: Json,
760755
pipeline: &mut MultiFieldPipeline,
761-
top_k: Option<i64>,
762756
) -> anyhow::Result<Vec<Json>> {
763757
let pool = get_or_initialize_pool(&self.database_url).await?;
764758

@@ -1113,7 +1107,7 @@ documents ||..|{{ {nice_name_key}_chunks
11131107
);
11141108
uml_relations.push_str(&relations);
11151109

1116-
if let Some(_embed_action) = &field_action.embed {
1110+
if let Some(_embed_action) = &field_action.semantic_search {
11171111
let entites = format!(
11181112
r#"
11191113
entity "{schema}.{key}_chunks" as {nice_name_key}_chunks {{

‎pgml-sdks/pgml/src/languages/javascript.rs

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,7 @@ use rust_bridge::javascript::{FromJsType, IntoJsResult};
44
use std::cell::RefCell;
55
use std::sync::Arc;
66

7-
use crate::{
8-
pipeline::PipelineSyncData,
9-
types::{DateTime, GeneralJsonAsyncIterator, GeneralJsonIterator, Json},
10-
};
7+
use crate::types::{DateTime, GeneralJsonAsyncIterator, GeneralJsonIterator, Json};
118

129
////////////////////////////////////////////////////////////////////////////////
1310
// Rust to JS //////////////////////////////////////////////////////////////////
@@ -63,16 +60,6 @@ impl IntoJsResult for Json {
6360
}
6461
}
6562

66-
impl IntoJsResult for PipelineSyncData {
67-
type Output = JsValue;
68-
fn into_js_result<'a, 'b, 'c: 'b, C: Context<'c>>(
69-
self,
70-
cx: &mut C,
71-
) -> JsResult<'b, Self::Output> {
72-
Json::from(self).into_js_result(cx)
73-
}
74-
}
75-
7663
#[derive(Clone)]
7764
struct GeneralJsonAsyncIteratorJavaScript(Arc<tokio::sync::Mutex<GeneralJsonAsyncIterator>>);
7865

‎pgml-sdks/pgml/src/languages/python.rs

Lines changed: 1 addition & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,7 @@ use std::sync::Arc;
66

77
use rust_bridge::python::CustomInto;
88

9-
use crate::{
10-
pipeline::PipelineSyncData,
11-
types::{GeneralJsonAsyncIterator, GeneralJsonIterator, Json},
12-
};
9+
use crate::types::{GeneralJsonAsyncIterator, GeneralJsonIterator, Json};
1310

1411
////////////////////////////////////////////////////////////////////////////////
1512
// Rust to PY //////////////////////////////////////////////////////////////////
@@ -50,12 +47,6 @@ impl IntoPy<PyObject> for Json {
5047
}
5148
}
5249

53-
impl IntoPy<PyObject> for PipelineSyncData {
54-
fn into_py(self, py: Python) -> PyObject {
55-
Json::from(self).into_py(py)
56-
}
57-
}
58-
5950
#[pyclass]
6051
#[derive(Clone)]
6152
struct GeneralJsonAsyncIteratorPython {
@@ -177,13 +168,6 @@ impl FromPyObject<'_> for Json {
177168
}
178169
}
179170

180-
impl FromPyObject<'_> for PipelineSyncData {
181-
fn extract(ob: &PyAny) -> PyResult<Self> {
182-
let json = Json::extract(ob)?;
183-
Ok(json.into())
184-
}
185-
}
186-
187171
impl FromPyObject<'_> for GeneralJsonAsyncIterator {
188172
fn extract(_ob: &PyAny) -> PyResult<Self> {
189173
panic!("We must implement this, but this is impossible to be reached")
@@ -199,9 +183,3 @@ impl FromPyObject<'_> for GeneralJsonIterator {
199183
////////////////////////////////////////////////////////////////////////////////
200184
// Rust to Rust //////////////////////////////////////////////////////////////////
201185
////////////////////////////////////////////////////////////////////////////////
202-
203-
impl CustomInto<Json> for PipelineSyncData {
204-
fn custom_into(self) -> Json {
205-
Json::from(self)
206-
}
207-
}

0 commit comments

Comments
 (0)