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)
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