Skip to content

Commit fe593a9

Browse files
committed
Gated rust_bridge behind a feature flag
1 parent 5661eec commit fe593a9

File tree

12 files changed

+128
-75
lines changed

12 files changed

+128
-75
lines changed

pgml-sdks/pgml/Cargo.lock

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pgml-sdks/pgml/Cargo.toml

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@ name = "pgml"
1414
crate-type = ["lib", "cdylib"]
1515

1616
[dependencies]
17-
rust_bridge = {path = "../rust-bridge/rust-bridge", version = "0.1.0"}
17+
# rust_bridge = {path = "../rust-bridge/rust-bridge", version = "0.1.0", optional = true }
18+
rust_bridge = {git = "https://github.com/postgresml/postgresml", version = "0.1.0", optional = true }
1819
sqlx = { version = "0.7.3", features = [ "runtime-tokio-rustls", "postgres", "json", "time", "uuid"] }
1920
serde_json = "1.0.9"
2021
anyhow = "1.0.9"
@@ -50,5 +51,6 @@ serde_with = "3.8.1"
5051

5152
[features]
5253
default = []
53-
python = ["dep:pyo3", "dep:pyo3-asyncio"]
54-
javascript = ["dep:neon"]
54+
rust_bridge = ["dep:rust_bridge"]
55+
python = ["rust_bridge", "dep:pyo3", "dep:pyo3-asyncio"]
56+
javascript = ["rust_bridge", "dep:neon"]

pgml-sdks/pgml/src/builtins.rs

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,26 @@
11
use anyhow::Context;
2-
use rust_bridge::{alias, alias_methods};
32
use sqlx::Row;
43
use tracing::instrument;
54

6-
/// Provides access to builtin database methods
7-
#[derive(alias, Debug, Clone)]
8-
pub struct Builtins {
9-
database_url: Option<String>,
10-
}
11-
125
use crate::{get_or_initialize_pool, query_runner::QueryRunner, types::Json};
136

7+
#[cfg(feature = "rust_bridge")]
8+
use rust_bridge::{alias, alias_methods};
9+
1410
#[cfg(feature = "python")]
1511
use crate::{query_runner::QueryRunnerPython, types::JsonPython};
1612

17-
#[alias_methods(new, query, transform, embed, embed_batch)]
13+
/// Provides access to builtin database methods
14+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
15+
#[derive(Debug, Clone)]
16+
pub struct Builtins {
17+
database_url: Option<String>,
18+
}
19+
20+
#[cfg_attr(
21+
feature = "rust_bridge",
22+
alias_methods(new, query, transform, embed, embed_batch)
23+
)]
1824
impl Builtins {
1925
pub fn new(database_url: Option<String>) -> Self {
2026
Self { database_url }

pgml-sdks/pgml/src/collection.rs

Lines changed: 33 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ use anyhow::Context;
22
use indicatif::MultiProgress;
33
use itertools::Itertools;
44
use regex::Regex;
5-
use rust_bridge::{alias, alias_methods};
65
use sea_query::Alias;
76
use sea_query::{Expr, NullOrdering, Order, PostgresQueryBuilder, Query};
87
use sea_query_binder::SqlxBinder;
@@ -35,6 +34,9 @@ use crate::{
3534
utils,
3635
};
3736

37+
#[cfg(feature = "rust_bridge")]
38+
use rust_bridge::{alias, alias_methods};
39+
3840
#[cfg(feature = "python")]
3941
use crate::{
4042
pipeline::PipelinePython,
@@ -43,7 +45,7 @@ use crate::{
4345
};
4446

4547
/// A RAGStream Struct
46-
#[derive(alias)]
48+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
4749
#[allow(dead_code)]
4850
pub struct RAGStream {
4951
general_json_async_iterator: Option<GeneralJsonAsyncIterator>,
@@ -57,7 +59,7 @@ impl Clone for RAGStream {
5759
}
5860
}
5961

60-
#[alias_methods(stream, sources)]
62+
#[cfg_attr(feature = "rust_bridge", alias_methods(stream, sources))]
6163
impl RAGStream {
6264
pub fn stream(&mut self) -> anyhow::Result<GeneralJsonAsyncIterator> {
6365
self.general_json_async_iterator
@@ -137,7 +139,8 @@ pub(crate) struct CollectionDatabaseData {
137139
}
138140

139141
/// A collection of documents
140-
#[derive(alias, Debug, Clone)]
142+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
143+
#[derive(Debug, Clone)]
141144
pub struct Collection {
142145
pub(crate) name: String,
143146
pub(crate) database_url: Option<String>,
@@ -146,29 +149,32 @@ pub struct Collection {
146149
pub(crate) database_data: Option<CollectionDatabaseData>,
147150
}
148151

149-
#[alias_methods(
150-
new,
151-
upsert_documents,
152-
get_documents,
153-
delete_documents,
154-
get_pipelines,
155-
get_pipeline,
156-
add_pipeline,
157-
remove_pipeline,
158-
enable_pipeline,
159-
disable_pipeline,
160-
search,
161-
add_search_event,
162-
vector_search,
163-
query,
164-
rag,
165-
rag_stream,
166-
exists,
167-
archive,
168-
upsert_directory,
169-
upsert_file,
170-
generate_er_diagram,
171-
get_pipeline_status
152+
#[cfg_attr(
153+
feature = "rust_bridge",
154+
alias_methods(
155+
new,
156+
upsert_documents,
157+
get_documents,
158+
delete_documents,
159+
get_pipelines,
160+
get_pipeline,
161+
add_pipeline,
162+
remove_pipeline,
163+
enable_pipeline,
164+
disable_pipeline,
165+
search,
166+
add_search_event,
167+
vector_search,
168+
query,
169+
rag,
170+
rag_stream,
171+
exists,
172+
archive,
173+
upsert_directory,
174+
upsert_file,
175+
generate_er_diagram,
176+
get_pipeline_status
177+
)
172178
)]
173179
impl Collection {
174180
/// Creates a new [Collection]

pgml-sdks/pgml/src/model.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use rust_bridge::{alias, alias_methods};
21
use sqlx::{Pool, Postgres};
32
use tracing::instrument;
43

@@ -11,6 +10,9 @@ use crate::{
1110
#[cfg(feature = "python")]
1211
use crate::types::JsonPython;
1312

13+
#[cfg(feature = "rust_bridge")]
14+
use rust_bridge::{alias, alias_methods};
15+
1416
/// A few notes on the following enums:
1517
/// - Sqlx does provide type derivation for enums, but it's not very good
1618
/// - Queries using these enums require a number of additional queries to get their oids and
@@ -52,7 +54,8 @@ pub(crate) struct ModelDatabaseData {
5254
}
5355

5456
/// A model used for embedding, inference, etc...
55-
#[derive(alias, Debug, Clone)]
57+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
58+
#[derive(Debug, Clone)]
5659
pub struct Model {
5760
pub(crate) name: String,
5861
pub(crate) runtime: ModelRuntime,
@@ -66,7 +69,7 @@ impl Default for Model {
6669
}
6770
}
6871

69-
#[alias_methods(new, transform)]
72+
#[cfg_attr(feature = "rust_bridge", alias_methods(new, transform))]
7073
impl Model {
7174
/// Creates a new [Model]
7275
pub fn new(name: Option<String>, source: Option<String>, parameters: Option<Json>) -> Self {

pgml-sdks/pgml/src/open_source_ai.rs

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
use anyhow::Context;
22
use futures::{Stream, StreamExt};
3-
use rust_bridge::{alias, alias_methods};
43
use std::time::{SystemTime, UNIX_EPOCH};
54
use uuid::Uuid;
65

@@ -10,11 +9,15 @@ use crate::{
109
TransformerPipeline,
1110
};
1211

12+
#[cfg(feature = "rust_bridge")]
13+
use rust_bridge::{alias, alias_methods};
14+
1315
#[cfg(feature = "python")]
1416
use crate::types::{GeneralJsonAsyncIteratorPython, GeneralJsonIteratorPython, JsonPython};
1517

1618
/// A drop in replacement for OpenAI
17-
#[derive(alias, Debug, Clone)]
19+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
20+
#[derive(Debug, Clone)]
1821
pub struct OpenSourceAI {
1922
database_url: Option<String>,
2023
}
@@ -160,12 +163,15 @@ impl Iterator for AsyncToSyncJsonIterator {
160163
}
161164
}
162165

163-
#[alias_methods(
164-
new,
165-
chat_completions_create,
166-
chat_completions_create_async,
167-
chat_completions_create_stream,
168-
chat_completions_create_stream_async
166+
#[cfg_attr(
167+
feature = "rust_bridge",
168+
alias_methods(
169+
new,
170+
chat_completions_create,
171+
chat_completions_create_async,
172+
chat_completions_create_stream,
173+
chat_completions_create_stream_async
174+
)
169175
)]
170176
impl OpenSourceAI {
171177
/// Creates a new [OpenSourceAI]

pgml-sdks/pgml/src/pipeline.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use anyhow::Context;
2-
use rust_bridge::{alias, alias_methods};
32
use serde::Deserialize;
43
use serde_json::json;
54
use sqlx::{Executor, PgConnection, Pool, Postgres, Transaction};
@@ -16,6 +15,9 @@ use crate::{
1615
types::{DateTime, Json, TryToNumeric},
1716
};
1817

18+
#[cfg(feature = "rust_bridge")]
19+
use rust_bridge::{alias, alias_methods};
20+
1921
#[cfg(feature = "python")]
2022
use crate::types::JsonPython;
2123

@@ -176,7 +178,8 @@ pub struct PipelineDatabaseData {
176178
}
177179

178180
/// A pipeline that describes transformations to documents
179-
#[derive(alias, Debug, Clone)]
181+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
182+
#[derive(Debug, Clone)]
180183
pub struct Pipeline {
181184
pub(crate) name: String,
182185
pub(crate) schema: Option<Json>,
@@ -202,7 +205,7 @@ fn json_to_schema(schema: &Json) -> anyhow::Result<ParsedSchema> {
202205
})
203206
}
204207

205-
#[alias_methods(new)]
208+
#[cfg_attr(feature = "rust_bridge", alias_methods(new))]
206209
impl Pipeline {
207210
/// Creates a [Pipeline]
208211
///

pgml-sdks/pgml/src/query_builder.rs

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,29 @@
33
// No new things should be added here, instead add new items to collection.vector_search
44

55
use anyhow::Context;
6-
use rust_bridge::{alias, alias_methods};
76
use serde_json::json;
87
use tracing::instrument;
98

109
use crate::{pipeline::Pipeline, types::Json, Collection};
1110

11+
#[cfg(feature = "rust_bridge")]
12+
use rust_bridge::{alias, alias_methods};
13+
1214
#[cfg(feature = "python")]
1315
use crate::{pipeline::PipelinePython, types::JsonPython};
1416

15-
#[derive(alias, Clone, Debug)]
17+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
18+
#[derive(Clone, Debug)]
1619
pub struct QueryBuilder {
1720
collection: Collection,
1821
query: Json,
1922
pipeline: Option<Pipeline>,
2023
}
2124

22-
#[alias_methods(limit, filter, vector_recall, to_full_string, fetch_all)]
25+
#[cfg_attr(
26+
feature = "rust_bridge",
27+
alias_methods(limit, filter, vector_recall, to_full_string, fetch_all(skip = "C"))
28+
)]
2329
impl QueryBuilder {
2430
pub fn new(collection: Collection) -> Self {
2531
let query = json!({

pgml-sdks/pgml/src/query_runner.rs

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use rust_bridge::{alias, alias_methods};
21
use sqlx::postgres::PgArguments;
32
use sqlx::query::Query;
43
use sqlx::{Postgres, Row};
@@ -8,6 +7,9 @@ use crate::{get_or_initialize_pool, types::Json};
87
#[cfg(feature = "python")]
98
use crate::types::JsonPython;
109

10+
#[cfg(feature = "rust_bridge")]
11+
use rust_bridge::{alias, alias_methods};
12+
1113
#[derive(Clone, Debug)]
1214
enum BindValue {
1315
String(String),
@@ -17,21 +19,25 @@ enum BindValue {
1719
Json(Json),
1820
}
1921

20-
#[derive(alias, Clone, Debug)]
22+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
23+
#[derive(Clone, Debug)]
2124
pub struct QueryRunner {
2225
query: String,
2326
bind_values: Vec<BindValue>,
2427
database_url: Option<String>,
2528
}
2629

27-
#[alias_methods(
28-
fetch_all,
29-
execute,
30-
bind_string,
31-
bind_int,
32-
bind_float,
33-
bind_bool,
34-
bind_json
30+
#[cfg_attr(
31+
feature = "rust_bridge",
32+
alias_methods(
33+
fetch_all,
34+
execute,
35+
bind_string,
36+
bind_int,
37+
bind_float,
38+
bind_bool,
39+
bind_json
40+
)
3541
)]
3642
impl QueryRunner {
3743
pub fn new(query: &str, database_url: Option<String>) -> Self {

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy