Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitfd00bbf

Browse files
authored
Prepare pgml for publishing to crates.io (#1500)
1 parentb949d45 commitfd00bbf

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 aboutcustomizing 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-
pubstructBuiltins{
9-
database_url:Option<String>,
10-
}
11-
125
usecrate::{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
usecrate::{query_runner::QueryRunnerPython, types::JsonPython};
1612

1713
#[cfg(feature ="c")]
1814
usecrate::{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+
pubstructBuiltins{
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
implBuiltins{
2228
pubfnnew(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+
usecrate::languages::c::GeneralJsonAsyncIteratorC;
42+
3843
#[cfg(feature ="python")]
3944
usecrate::{
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
pubstructRAGStream{
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
implRAGStream{
6267
pubfnstream(&mutself) -> 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
pubstructCollection{
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
implCollection{
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
pubasyncfnrag(&self,query:Json,pipeline:&mutPipeline) -> 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
pubasyncfnrag_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
usecrate::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
pubstructModel{
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
implModel{
7477
/// Creates a new [Model]
7578
pubfnnew(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
usecrate::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
pubstructOpenSourceAI{
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
implOpenSourceAI{
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
usecrate::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
pubstructPipeline{
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
implPipeline{
210213
/// Creates a [Pipeline]
211214
///

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp