Skip to content

Commit fd00bbf

Browse files
SilasMarvinSilasMarvin
authored
Prepare pgml for publishing to crates.io (#1500)
1 parent b949d45 commit fd00bbf

File tree

12 files changed

+192
-76
lines changed

12 files changed

+192
-76
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: 6 additions & 4 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,6 +51,7 @@ serde_with = "3.8.1"
5051

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

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,29 @@
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

1713
#[cfg(feature = "c")]
1814
use crate::{languages::c::JsonC, query_runner::QueryRunnerC};
1915

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

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

Lines changed: 96 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,12 @@ use crate::{
3534
utils,
3635
};
3736

37+
#[cfg(feature = "rust_bridge")]
38+
use rust_bridge::{alias, alias_methods};
39+
40+
#[cfg(feature = "c")]
41+
use crate::languages::c::GeneralJsonAsyncIteratorC;
42+
3843
#[cfg(feature = "python")]
3944
use crate::{
4045
pipeline::PipelinePython,
@@ -43,7 +48,7 @@ use crate::{
4348
};
4449

4550
/// A RAGStream Struct
46-
#[derive(alias)]
51+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
4752
#[allow(dead_code)]
4853
pub struct RAGStream {
4954
general_json_async_iterator: Option<GeneralJsonAsyncIterator>,
@@ -57,7 +62,7 @@ impl Clone for RAGStream {
5762
}
5863
}
5964

60-
#[alias_methods(stream, sources)]
65+
#[cfg_attr(feature = "rust_bridge", alias_methods(stream, sources))]
6166
impl RAGStream {
6267
pub fn stream(&mut self) -> anyhow::Result<GeneralJsonAsyncIterator> {
6368
self.general_json_async_iterator
@@ -140,7 +145,8 @@ pub(crate) struct CollectionDatabaseData {
140145
}
141146

142147
/// A collection of documents
143-
#[derive(alias, Debug, Clone)]
148+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
149+
#[derive(Debug, Clone)]
144150
pub struct Collection {
145151
pub(crate) name: String,
146152
pub(crate) database_url: Option<String>,
@@ -149,29 +155,32 @@ pub struct Collection {
149155
pub(crate) database_data: Option<CollectionDatabaseData>,
150156
}
151157

152-
#[alias_methods(
153-
new,
154-
upsert_documents,
155-
get_documents,
156-
delete_documents,
157-
get_pipelines,
158-
get_pipeline,
159-
add_pipeline,
160-
remove_pipeline,
161-
enable_pipeline,
162-
disable_pipeline,
163-
search,
164-
add_search_event,
165-
vector_search,
166-
query,
167-
rag,
168-
rag_stream,
169-
exists,
170-
archive,
171-
upsert_directory,
172-
upsert_file,
173-
generate_er_diagram,
174-
get_pipeline_status
158+
#[cfg_attr(
159+
feature = "rust_bridge",
160+
alias_methods(
161+
new,
162+
upsert_documents,
163+
get_documents,
164+
delete_documents,
165+
get_pipelines,
166+
get_pipeline,
167+
add_pipeline,
168+
remove_pipeline,
169+
enable_pipeline,
170+
disable_pipeline,
171+
search,
172+
add_search_event,
173+
vector_search,
174+
query,
175+
rag,
176+
rag_stream,
177+
exists,
178+
archive,
179+
upsert_directory,
180+
upsert_file,
181+
generate_er_diagram,
182+
get_pipeline_status
183+
)
175184
)]
176185
impl Collection {
177186
/// Creates a new [Collection]
@@ -1128,6 +1137,65 @@ impl Collection {
11281137
.collect())
11291138
}
11301139

1140+
/// Performs rag on the [Collection]
1141+
///
1142+
/// # Arguments
1143+
/// * `query` - The query to search for
1144+
/// * `pipeline` - The [Pipeline] to use for the search
1145+
///
1146+
/// # Example
1147+
/// ```
1148+
/// use pgml::Collection;
1149+
/// use pgml::Pipeline;
1150+
/// use serde_json::json;
1151+
/// use anyhow::Result;
1152+
/// async fn run() -> anyhow::Result<()> {
1153+
/// let mut collection = Collection::new("my_collection", None)?;
1154+
/// let mut pipeline = Pipeline::new("my_pipeline", None)?;
1155+
/// let results = collection.rag(json!({
1156+
/// "CONTEXT": {
1157+
/// "vector_search": {
1158+
/// "query": {
1159+
/// "fields": {
1160+
/// "body": {
1161+
/// "query": "Test document: 2",
1162+
/// "parameters": {
1163+
/// "prompt": "query: "
1164+
/// }
1165+
/// },
1166+
/// },
1167+
/// },
1168+
/// "document": {
1169+
/// "keys": [
1170+
/// "id"
1171+
/// ]
1172+
/// },
1173+
/// "limit": 2
1174+
/// },
1175+
/// "aggregate": {
1176+
/// "join": "\n"
1177+
/// }
1178+
/// },
1179+
/// "CUSTOM": {
1180+
/// "sql": "SELECT 'test'"
1181+
/// },
1182+
/// "chat": {
1183+
/// "model": "meta-llama/Meta-Llama-3-8B-Instruct",
1184+
/// "messages": [
1185+
/// {
1186+
/// "role": "system",
1187+
/// "content": "You are a friendly and helpful chatbot"
1188+
/// },
1189+
/// {
1190+
/// "role": "user",
1191+
/// "content": "Some text with {CONTEXT} - {CUSTOM}",
1192+
/// }
1193+
/// ],
1194+
/// "max_tokens": 10
1195+
/// }
1196+
/// }).into(), &mut pipeline).await?;
1197+
/// Ok(())
1198+
/// }
11311199
#[instrument(skip(self))]
11321200
pub async fn rag(&self, query: Json, pipeline: &mut Pipeline) -> anyhow::Result<Json> {
11331201
let pool = get_or_initialize_pool(&self.database_url).await?;
@@ -1138,6 +1206,7 @@ impl Collection {
11381206
Ok(std::mem::take(&mut results[0].0))
11391207
}
11401208

1209+
/// Same as rag buit returns a stream of results
11411210
#[instrument(skip(self))]
11421211
pub async fn rag_stream(
11431212
&self,

‎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

@@ -14,6 +13,9 @@ use crate::types::JsonPython;
1413
#[cfg(feature = "c")]
1514
use crate::languages::c::JsonC;
1615

16+
#[cfg(feature = "rust_bridge")]
17+
use rust_bridge::{alias, alias_methods};
18+
1719
/// A few notes on the following enums:
1820
/// - Sqlx does provide type derivation for enums, but it's not very good
1921
/// - Queries using these enums require a number of additional queries to get their oids and
@@ -55,7 +57,8 @@ pub(crate) struct ModelDatabaseData {
5557
}
5658

5759
/// A model used for embedding, inference, etc...
58-
#[derive(alias, Debug, Clone)]
60+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
61+
#[derive(Debug, Clone)]
5962
pub struct Model {
6063
pub(crate) name: String,
6164
pub(crate) runtime: ModelRuntime,
@@ -69,7 +72,7 @@ impl Default for Model {
6972
}
7073
}
7174

72-
#[alias_methods(new, transform)]
75+
#[cfg_attr(feature = "rust_bridge", alias_methods(new, transform))]
7376
impl Model {
7477
/// Creates a new [Model]
7578
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,6 +9,9 @@ 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

@@ -20,7 +22,8 @@ use crate::{
2022
};
2123

2224
/// A drop in replacement for OpenAI
23-
#[derive(alias, Debug, Clone)]
25+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
26+
#[derive(Debug, Clone)]
2427
pub struct OpenSourceAI {
2528
database_url: Option<String>,
2629
}
@@ -166,12 +169,15 @@ impl Iterator for AsyncToSyncJsonIterator {
166169
}
167170
}
168171

169-
#[alias_methods(
170-
new,
171-
chat_completions_create,
172-
chat_completions_create_async,
173-
chat_completions_create_stream,
174-
chat_completions_create_stream_async
172+
#[cfg_attr(
173+
feature = "rust_bridge",
174+
alias_methods(
175+
new,
176+
chat_completions_create,
177+
chat_completions_create_async,
178+
chat_completions_create_stream,
179+
chat_completions_create_stream_async
180+
)
175181
)]
176182
impl OpenSourceAI {
177183
/// 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

@@ -179,7 +181,8 @@ pub struct PipelineDatabaseData {
179181
}
180182

181183
/// A pipeline that describes transformations to documents
182-
#[derive(alias, Debug, Clone)]
184+
#[cfg_attr(feature = "rust_bridge", derive(alias))]
185+
#[derive(Debug, Clone)]
183186
pub struct Pipeline {
184187
pub(crate) name: String,
185188
pub(crate) schema: Option<Json>,
@@ -205,7 +208,7 @@ fn json_to_schema(schema: &Json) -> anyhow::Result<ParsedSchema> {
205208
})
206209
}
207210

208-
#[alias_methods(new)]
211+
#[cfg_attr(feature = "rust_bridge", alias_methods(new))]
209212
impl Pipeline {
210213
/// Creates a [Pipeline]
211214
///

0 commit comments

Comments
 (0)