diff --git a/pgml-sdks/pgml/javascript/README.md b/pgml-sdks/pgml/javascript/README.md index 63b6c41a8..bbf43be7f 100644 --- a/pgml-sdks/pgml/javascript/README.md +++ b/pgml-sdks/pgml/javascript/README.md @@ -4,15 +4,20 @@ - [Overview](#overview) - [Quickstart](#quickstart) -- [Usage](#usage) -- [Examples](./examples/README.md) +- [Upgrading](#upgrading) - [Developer setup](#developer-setup) - [Roadmap](#roadmap) +- [Documentation](https://postgresml.org/docs/guides/sdks/overview) +- [Examples](./examples/README.md) # Overview JavaScript SDK is designed to facilitate the development of scalable vector search applications on PostgreSQL databases. With this SDK, you can seamlessly manage various database tables related to documents, text chunks, text splitters, LLM (Language Model) models, and embeddings. By leveraging the SDK's capabilities, you can efficiently index LLM embeddings using PgVector for fast and accurate queries. +Documentation: [PostgresML SDK Docs](https://postgresml.org/docs/guides/sdks/overview) + +Examples Folder: [Examples](./examples/README.md) + ## Key Features - **Automated Database Management**: With the SDK, you can easily handle the management of database tables related to documents, text chunks, text splitters, LLM models, and embeddings. This automated management system simplifies the process of setting up and maintaining your vector search application's data structure. @@ -181,413 +186,15 @@ You should see the search results printed in the terminal. As you can see, our v ] ``` - -# Usage - -## High-level Description - -The JavaScript SDK provides a set of functionalities to build scalable vector search applications on PostgresQL databases. It enables users to create a collection, which represents a schema in the database, to store tables for documents, chunks, models, splitters, and embeddings. The Collection class in the SDK handles all operations related to these tables, allowing users to interact with the collection and perform various tasks. - -## Collections - -Collections are the organizational building blocks of the SDK. They manage all documents and related chunks, embeddings, tsvectors, and pipelines. - -### Creating Collections - -By default, collections will read and write to the database specified by `DATABASE_URL`. - -**Create a Collection that uses the default `DATABASE_URL` environment variable.** -```javascript -const collection = pgml.newCollection("test_collection") -``` - -**Create a Collection that reads from a different database than that set by the environment variable `DATABASE_URL`.** -```javascript -const collection = pgml.newCollection("test_collection", CUSTOM_DATABASE_URL) -``` - -### Upserting Documents - -The `upsert_documents` method can be used to insert new documents and update existing documents. - -New documents are dictionaries with two required keys: `id` and `text`. All other keys/value pairs are stored as metadata for the document. - -**Upsert new documents with metadata** -```javascript -const documents = [ - { - id: "Document 1", - text: "Here are the contents of Document 1", - random_key: "this will be metadata for the document" - }, - { - id: "Document 2", - text: "Here are the contents of Document 2", - random_key: "this will be metadata for the document" - } -] -const collection = pgml.newCollection("test_collection") -await collection.upsert_documents(documents) -``` - -Document metadata can be updated by upserting the document without the `text` key. - -**Update document metadata** -```javascript -documents = [ - { - id: "Document 1", - random_key: "this will be NEW metadata for the document" - }, - { - id: "Document 2", - random_key: "this will be NEW metadata for the document" - } -] -collection = pgml.newCollection("test_collection") -await collection.upsert_documents(documents) -``` - -### Getting Documents - -Documents can be retrieved using the `get_documents` method on the collection object - -**Get the first 100 documents** -```javascript -collection = pgml.newCollection("test_collection") -documents = await collection.get_documents({ limit: 100 }) -``` - -#### Pagination - -The JavaScript SDK supports limit-offset pagination and keyset pagination - -**Limit-Offset pagination** -```javascript -collection = pgml.newCollection("test_collection") -documents = await collection.get_documents({ limit: 100, offset: 10 }) -``` - -**Keyset pagination** -```javascript -collection = pgml.newCollection("test_collection") -documents = await collection.get_documents({ limit: 100, last_row_id: 10 }) -``` - -The `last_row_id` can be taken from the `row_id` field in the returned document's dictionary. - -#### Filtering - -Metadata and full text filtering are supported just like they are in vector recall. - -**Metadata and full text filtering** -```javascript -collection = pgml.newCollection("test_collection") -documents = await collection.get_documents({ - limit: 100, - offset: 10, - filter: { - metadata: { - id: { - $eq: 1 - } - }, - full_text_search: { - configuration: "english", - text: "Some full text query" - } - } -}) - -``` - -### Deleting Documents - -Documents can be deleted with the `delete_documents` method on the collection object. - -Metadata and full text filtering are supported just like they are in vector recall. - -```javascript -collection = pgml.newCollection("test_collection") -documents = await collection.delete_documents({ - metadata: { - id: { - $eq: 1 - } - }, - full_text_search: { - configuration: "english", - text: "Some full text query" - } -}) -``` - -### Searching Collections - -The JavaScript SDK is specifically designed to provide powerful, flexible vector search. - -Pipelines are required to perform search. See the [Pipelines Section](#pipelines) for more information about using Pipelines. - -**Basic vector search** -```javascript -const collection = pgml.newCollection("test_collection") -const pipeline = pgml.newPipeline("test_pipeline") -const results = await collection.query().vector_recall("Why is PostgresML the best?", pipeline).fetch_all() -``` - -**Vector search with custom limit** -```javascript -const collection = pgml.newCollection("test_collection") -const pipeline = pgml.newPipeline("test_pipeline") -const results = await collection.query().vector_recall("Why is PostgresML the best?", pipeline).limit(10).fetch_all() -``` - -#### Metadata Filtering - -We provide powerful and flexible arbitrarly nested metadata filtering based off of [MongoDB Comparison Operators](https://www.mongodb.com/docs/manual/reference/operator/query-comparison/). We support each operator mentioned except the `$nin`. - -**Vector search with $eq metadata filtering** -```javascript -const collection = pgml.newCollection("test_collection") -const pipeline = pgml.newPipeline("test_pipeline") -const results = await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - metadata: { - uuid: { - $eq: 1 - } - } - }) - .fetch_all() -``` - -The above query would filter out all documents that do not contain a key `uuid` equal to `1`. - -**Vector search with $gte metadata filtering** -```javascript -const collection = pgml.newCollection("test_collection") -const pipeline = pgml.newPipeline("test_pipeline") -const results = await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - metadata: { - index: { - $gte: 3 - } - } - }) - .fetch_all() -) -``` - -The above query would filter out all documents that do not contain a key `index` with a value greater than `3`. - -**Vector search with $or and $and metadata filtering** -```javascript -const collection = pgml.newCollection("test_collection") -const pipeline = pgml.newPipeline("test_pipeline") -const results = await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - metadata: { - $or: [ - { - $and: [ - { - uuid: { - $eq: 1 - } - }, - { - index: { - $lt: 100 - } - } - ] - }, - { - special: { - $ne: true - } - } - ] - } - }) - .fetch_all() -``` - -The above query would filter out all documents that do not have a key `special` with a value `true` or (have a key `uuid` equal to 1 and a key `index` less than 100). - -#### Full Text Filtering - -If full text search is enabled for the associated Pipeline, documents can be first filtered by full text search and then recalled by embedding similarity. - -```javascript -const collection = pgml.newCollection("test_collection") -const pipeline = pgml.newPipeline("test_pipeline") -const results = await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - full_text: { - configuration: "english", - text: "Match Me" - } - }) - .fetch_all() -``` - -The above query would first filter out all documents that do not match the full text search criteria, and then perform vector recall on the remaining documents. - -## Pipelines - -Collections can have any number of Pipelines. Each Pipeline is ran everytime documents are upserted. - -Pipelines are composed of a Model, Splitter, and additional optional arguments. - -### Models - -Models are used for embedding chuncked documents. We support most every open source model on [Hugging Face](https://huggingface.co/), and also OpenAI's embedding models. - -**Create a default Model "intfloat/e5-small" with default parameters: {}** -```javascript -const model = pgml.newModel() -``` - -**Create a Model with custom parameters** -```javascript -const model = pgml.newModel( - "hkunlp/instructor-base", - {instruction: "Represent the Wikipedia document for retrieval: "} -) -``` - -**Use an OpenAI model** -```javascript -const model = pgml.newModel("text-embedding-ada-002", "openai") -``` - -### Splitters - -Splitters are used to split documents into chunks before embedding them. We support splitters found in [LangChain](https://www.langchain.com/). - -**Create a default Splitter "recursive_character" with default parameters: {}** -```javascript -const splitter = pgml.newSplitter() -``` - -**Create a Splitter with custom parameters** -```javascript -const splitter = pgml.newSplitter( - "recursive_character", - {chunk_size: 1500, chunk_overlap: 40} -) -``` - -### Adding Pipelines to a Collection - -When adding a Pipeline to a collection it is required that Pipeline has a Model and Splitter. - -The first time a Pipeline is added to a Collection it will automatically chunk and embed any documents already in that Collection. - -```javascript -const model = pgml.newModel() -const splitter = pgml.newSplitter() -const pipeline = pgml.newPipeline("test_pipeline", model, splitter) -await collection.add_pipeline(pipeline) -``` - -### Enabling full text search - -Pipelines can take additional arguments enabling full text search. When full text search is enabled, in addition to automatically chunking and embedding, the Pipeline will create the necessary tsvectors to perform full text search. - -For more information on full text search please see: [Postgres Full Text Search](https://www.postgresql.org/docs/15/textsearch.html). - -```javascript -const model = pgml.newModel() -const splitter = pgml.newSplitter() -const pipeline = pgml.newPipeline("test_pipeline", model, splitter, { - "full_text_search": { - active: true, - configuration: "english" - } -}) -await collection.add_pipeline(pipeline) -``` - -### Configuring HNSW Indexing Parameters - -Our SDK utilizes [pgvector](https://github.com/pgvector/pgvector) for storing vectors and performing recall. We use HNSW indexing as it is the most performant mix of performance and recall. - -Our SDK allows for configuration of `m` (the maximum number of connections per layer (16 by default)) and `ef_construction` (the size of the dynamic candidate list when constructing the graph (64 by default)) per pipeline. - -```javascript -const model = pgml.newModel() -const splitter = pgml.newSplitter() -const pipeline = pgml.newPipeline("test_pipeline", model, splitter, { - hnsw: { - m: 100, - ef_construction: 200 - } -}) -await collection.add_pipeline(pipeline) -``` - -### Searching with Pipelines - -Pipelines are a required argument when performing vector search. After a Pipeline has been added to a Collection, the Model and Splitter can be omitted when instantiating it. - -```javascript -const pipeline = pgml.newPipeline("test_pipeline") -const collection = pgml.newCollection("test_collection") -const results = await collection.query().vector_recall("Why is PostgresML the best?", pipeline).fetch_all() -``` - -### Enabling, Disabling, and Removing Pipelines - -Pipelines can be disabled or removed to prevent them from running automatically when documents are upserted. - -**Disable a Pipeline** -```javascript -const pipeline = pgml.newPipeline("test_pipeline") -const collection = pgml.newCollection("test_collection") -await collection.disable_pipeline(pipeline) -``` - -Disabling a Pipeline prevents it from running automatically, but leaves all chunks and embeddings already created by that Pipeline in the database. - -**Enable a Pipeline** -```javascript -const pipeline = pgml.newPipeline("test_pipeline") -const collection = pgml.newCollection("test_collection") -await collection.enable_pipeline(pipeline) -``` - -Enabling a Pipeline will cause it to automatically run and chunk and embed all documents it may have missed while disabled. - -**Remove a Pipeline** -```javascript -const pipeline = pgml.newPipeline("test_pipeline") -const collection = pgml.newCollection("test_collection") -await collection.remove_pipeline(pipeline) -``` - -Removing a Pipeline deletes it and all associated data from the database. Removed Pipelines cannot be re-enabled but can be recreated. - ## Upgrading Changes between SDK versions are not necessarily backwards compatible. We provide a migrate function to help transition smoothly. ```javascript +const pgml = require("pgml"); await pgml.migrate() ``` -This will migrate all collections to be compatible with the latest SDK version. - ## Developer Setup This javascript library is generated from our core rust-sdk. Please check [rust-sdk documentation](../README.md) for developer setup. diff --git a/pgml-sdks/pgml/javascript/package.json b/pgml-sdks/pgml/javascript/package.json index 3aeeebfcc..dd3e59426 100644 --- a/pgml-sdks/pgml/javascript/package.json +++ b/pgml-sdks/pgml/javascript/package.json @@ -1,6 +1,6 @@ { "name": "pgml", - "version": "0.9.3", + "version": "0.9.4", "description": "Open Source Alternative for Building End-to-End Vector Search Applications without OpenAI & Pinecone", "keywords": [ "postgres", diff --git a/pgml-sdks/pgml/python/README.md b/pgml-sdks/pgml/python/README.md index c2b3f0b0c..0d1aad825 100644 --- a/pgml-sdks/pgml/python/README.md +++ b/pgml-sdks/pgml/python/README.md @@ -4,15 +4,20 @@ - [Overview](#overview) - [Quickstart](#quickstart) -- [Usage](#usage) -- [Examples](./examples/README.md) +- [Upgrading](#upgrading) - [Developer setup](#developer-setup) - [Roadmap](#roadmap) +- [Documentation](https://postgresml.org/docs/guides/sdks/overview) +- [Examples](./examples/README.md) # Overview Python SDK is designed to facilitate the development of scalable vector search applications on PostgreSQL databases. With this SDK, you can seamlessly manage various database tables related to documents, text chunks, text splitters, LLM (Language Model) models, and embeddings. By leveraging the SDK's capabilities, you can efficiently index LLM embeddings using PgVector for fast and accurate queries. +Documentation: [PostgresML SDK Docs](https://postgresml.org/docs/guides/sdks/overview) + +Examples Folder: [Examples](./examples/README.md) + ## Key Features - **Automated Database Management**: With the SDK, you can easily handle the management of database tables related to documents, text chunks, text splitters, LLM models, and embeddings. This automated management system simplifies the process of setting up and maintaining your vector search application's data structure. @@ -187,408 +192,6 @@ powerful female musician of 2015.', ] ``` -# Usage - -## High-level Description - -The Python SDK provides a set of functionalities to build scalable vector search applications on PostgresQL databases. It enables users to create a collection, which represents a schema in the database, to store tables for documents, chunks, models, splitters, and embeddings. The Collection class in the SDK handles all operations related to these tables, allowing users to interact with the collection and perform various tasks. - -## Collections - -Collections are the organizational building blocks of the SDK. They manage all documents and related chunks, embeddings, tsvectors, and pipelines. - -### Creating Collections - -By default, collections will read and write to the database specified by `DATABASE_URL`. - -**Create a Collection that uses the default `DATABASE_URL` environment variable.** -```python -collection = Collection("test_collection") -``` - -**Create a Collection that reads from a different database than that set by the environment variable `DATABASE_URL`.** -```python -collection = Collection("test_collection", CUSTOM_DATABASE_URL) -``` - -### Upserting Documents - -The `upsert_documents` method can be used to insert new documents and update existing documents. - -New documents are dictionaries with two required keys: `id` and `text`. All other keys/value pairs are stored as metadata for the document. - -**Upsert new documents with metadata** -```python -documents = [ - { - "id": "Document 1", - "text": "Here are the contents of Document 1", - "random_key": "this will be metadata for the document" - }, - { - "id": "Document 2", - "text": "Here are the contents of Document 2", - "random_key": "this will be metadata for the document" - } -] -collection = Collection("test_collection") -await collection.upsert_documents(documents) -``` - -Document metadata can be updated by upserting the document without the `text` key. - -**Update document metadata** -```python -documents = [ - { - "id": "Document 1", - "random_key": "this will be NEW metadata for the document" - }, - { - "id": "Document 2", - "random_key": "this will be NEW metadata for the document" - } -] -collection = Collection("test_collection") -await collection.upsert_documents(documents) -``` - -### Getting Documents - -Documents can be retrieved using the `get_documents` method on the collection object - -**Get the first 100 documents** -```python -collection = Collection("test_collection") -documents = await collection.get_documents({ "limit": 100 }) -``` - -#### Pagination - -The Python SDK supports limit-offset pagination and keyset pagination - -**Limit-Offset pagination** -```python -collection = Collection("test_collection") -documents = await collection.get_documents({ "limit": 100, "offset": 10 }) -``` - -**Keyset pagination** -```python -collection = Collection("test_collection") -documents = await collection.get_documents({ "limit": 100, "last_row_id": 10 }) -``` - -The `last_row_id` can be taken from the `row_id` field in the returned document's dictionary. - -#### Filtering - -Metadata and full text filtering are supported just like they are in vector recall. - -**Metadata and full text filtering** -```python -collection = Collection("test_collection") -documents = await collection.get_documents({ - "limit": 100, - "offset": 10, - "filter": { - "metadata": { - "id": { - "$eq": 1 - } - }, - "full_text_search": { - "configuration": "english", - "text": "Some full text query" - } - } -}) - -``` - -### Deleting Documents - -Documents can be deleted with the `delete_documents` method on the collection object. - -Metadata and full text filtering are supported just like they are in vector recall. - -```python -documents = await collection.delete_documents({ - "metadata": { - "id": { - "$eq": 1 - } - }, - "full_text_search": { - "configuration": "english", - "text": "Some full text query" - } -}) -``` - -### Searching Collections - -The Python SDK is specifically designed to provide powerful, flexible vector search. - -Pipelines are required to perform search. See the [Pipelines Section](#pipelines) for more information about using Pipelines. - -**Basic vector search** -```python -collection = Collection("test_collection") -pipeline = Pipeline("test_pipeline") -results = await collection.query().vector_recall("Why is PostgresML the best?", pipeline).fetch_all() -``` - -**Vector search with custom limit** -```python -collection = Collection("test_collection") -pipeline = Pipeline("test_pipeline") -results = await collection.query().vector_recall("Why is PostgresML the best?", pipeline).limit(10).fetch_all() -``` - -#### Metadata Filtering - -We provide powerful and flexible arbitrarly nested metadata filtering based off of [MongoDB Comparison Operators](https://www.mongodb.com/docs/manual/reference/operator/query-comparison/). We support each operator mentioned except the `$nin`. - -**Vector search with $eq metadata filtering** -```python -collection = Collection("test_collection") -pipeline = Pipeline("test_pipeline") -results = ( - await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - "metadata": { - "uuid": { - "$eq": 1 - } - } - }) - .fetch_all() -) -``` - -The above query would filter out all documents that do not contain a key `uuid` equal to `1`. - -**Vector search with $gte metadata filtering** -```python -collection = Collection("test_collection") -pipeline = Pipeline("test_pipeline") -results = ( - await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - "metadata": { - "index": { - "$gte": 3 - } - } - }) - .fetch_all() -) -``` - -The above query would filter out all documents that do not contain a key `index` with a value greater than `3`. - -**Vector search with $or and $and metadata filtering** -```python -collection = Collection("test_collection") -pipeline = Pipeline("test_pipeline") -results = ( - await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - "metadata": { - "$or": [ - { - "$and": [ - { - "uuid": { - "$eq": 1 - } - }, - { - "index": { - "$lt": 100 - } - } - ] - }, - { - "special": { - "$ne": True - } - } - ] - } - }) - .fetch_all() -) -``` - -The above query would filter out all documents that do not have a key `special` with a value `True` or (have a key `uuid` equal to 1 and a key `index` less than 100). - -#### Full Text Filtering - -If full text search is enabled for the associated Pipeline, documents can be first filtered by full text search and then recalled by embedding similarity. - -```python -collection = Collection("test_collection") -pipeline = Pipeline("test_pipeline") -results = ( - await collection.query() - .vector_recall("Here is some query", pipeline) - .limit(10) - .filter({ - "full_text_search": { - "configuration": "english", - "text": "Match Me" - } - }) - .fetch_all() -) -``` - -The above query would first filter out all documents that do not match the full text search criteria, and then perform vector recall on the remaining documents. - -## Pipelines - -Collections can have any number of Pipelines. Each Pipeline is ran everytime documents are upserted. - -Pipelines are composed of a Model, Splitter, and additional optional arguments. - -### Models - -Models are used for embedding chuncked documents. We support most every open source model on [Hugging Face](https://huggingface.co/), and also OpenAI's embedding models. - -**Create a default Model "intfloat/e5-small" with default parameters: {}** -```python -model = Model() -``` - -**Create a Model with custom parameters** -```python -model = Model( - name="hkunlp/instructor-base", - parameters={"instruction": "Represent the Wikipedia document for retrieval: "} -) -``` - -**Use an OpenAI model** -```python -model = Model(name="text-embedding-ada-002", source="openai") -``` - -### Splitters - -Splitters are used to split documents into chunks before embedding them. We support splitters found in [LangChain](https://www.langchain.com/). - -**Create a default Splitter "recursive_character" with default parameters: {}** -```python -splitter = Splitter() -``` - -**Create a Splitter with custom parameters** -```python -splitter = Splitter( - name="recursive_character", - parameters={"chunk_size": 1500, "chunk_overlap": 40} -) -``` - -### Adding Pipelines to a Collection - -When adding a Pipeline to a collection it is required that Pipeline has a Model and Splitter. - -The first time a Pipeline is added to a Collection it will automatically chunk and embed any documents already in that Collection. - -```python -model = Model() -splitter = Splitter() -pipeline = Pipeline("test_pipeline", model, splitter) -await collection.add_pipeline(pipeline) -``` - -### Enabling full text search - -Pipelines can take additional arguments enabling full text search. When full text search is enabled, in addition to automatically chunking and embedding, the Pipeline will create the necessary tsvectors to perform full text search. - -For more information on full text search please see: [Postgres Full Text Search](https://www.postgresql.org/docs/15/textsearch.html). - -```python -model = Model() -splitter = Splitter() -pipeline = Pipeline("test_pipeline", model, splitter, { - "full_text_search": { - "active": True, - "configuration": "english" - } -}) -await collection.add_pipeline(pipeline) -``` - -### Configuring HNSW Indexing Parameters - -Our SDK utilizes [pgvector](https://github.com/pgvector/pgvector) for storing vectors and performing recall. We use HNSW indexing as it is the most performant mix of performance and recall. - -Our SDK allows for configuration of `m` (the maximum number of connections per layer (16 by default)) and `ef_construction` (the size of the dynamic candidate list when constructing the graph (64 by default)) per pipeline. - -```python -model = Model() -splitter = Splitter() -pipeline = Pipeline("test_pipeline", model, splitter, { - "hnsw": { - "m": 100, - "ef_construction": 200 - } -}) -await collection.add_pipeline(pipeline) -``` - -### Searching with Pipelines - -Pipelines are a required argument when performing vector search. After a Pipeline has been added to a Collection, the Model and Splitter can be omitted when instantiating it. - -```python -pipeline = Pipeline("test_pipeline") -collection = Collection("test_collection") -results = await collection.query().vector_recall("Why is PostgresML the best?", pipeline).fetch_all() -``` - -### Enabling, Disabling, and Removing Pipelines - -Pipelines can be disabled or removed to prevent them from running automatically when documents are upserted. - -**Disable a Pipeline** -```python -pipeline = Pipeline("test_pipeline") -collection = Collection("test_collection") -await collection.disable_pipeline(pipeline) -``` - -Disabling a Pipeline prevents it from running automatically, but leaves all chunks and embeddings already created by that Pipeline in the database. - -**Enable a Pipeline** -```python -pipeline = Pipeline("test_pipeline") -collection = Collection("test_collection") -await collection.enable_pipeline(pipeline) -``` - -Enabling a Pipeline will cause it to automatically run and chunk and embed all documents it may have missed while disabled. - -**Remove a Pipeline** -```python -pipeline = Pipeline("test_pipeline") -collection = Collection("test_collection") -await collection.remove_pipeline(pipeline) -``` - -Removing a Pipeline deletes it and all associated data from the database. Removed Pipelines cannot be re-enabled but can be recreated. - ## Upgrading Changes between SDK versions are not necessarily backwards compatible. We provide a migrate function to help transition smoothly. diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs index 0e1ca7243..d146a38fd 100644 --- a/pgml-sdks/pgml/src/lib.rs +++ b/pgml-sdks/pgml/src/lib.rs @@ -8,7 +8,7 @@ use sqlx::PgPool; use std::collections::HashMap; use std::env; use std::sync::RwLock; -use tokio::runtime::{Builder, Runtime}; +use tokio::runtime::Runtime; use tracing::Level; use tracing_subscriber::FmtSubscriber; @@ -116,11 +116,7 @@ fn get_or_set_runtime<'a>() -> &'a Runtime { if let Some(r) = &RUNTIME { r } else { - let runtime = Builder::new_current_thread() - .worker_threads(1) - .enable_all() - .build() - .unwrap(); + let runtime = Runtime::new().unwrap(); RUNTIME = Some(runtime); get_or_set_runtime() } diff --git a/pgml-sdks/pgml/src/model.rs b/pgml-sdks/pgml/src/model.rs index 8a663d120..0a4f62d24 100644 --- a/pgml-sdks/pgml/src/model.rs +++ b/pgml-sdks/pgml/src/model.rs @@ -72,7 +72,7 @@ impl Default for Model { } } -#[alias_methods(new)] +#[alias_methods(new, transform)] impl Model { /// Creates a new [Model] /// @@ -202,6 +202,7 @@ impl Model { let task = json!({ "task": task, "model": self.name, + "trust_remote_code": true }); let args = args.unwrap_or_default(); let query = sqlx::query("SELECT pgml.transform(task => $1, inputs => $2, args => $3)"); @@ -256,7 +257,12 @@ mod tests { #[sqlx::test] async fn model_can_transform() -> anyhow::Result<()> { internal_init_logger(None, None).ok(); - let model = Model::new(Some("Helsinki-NLP/opus-mt-en-fr".to_string()), Some("pgml".to_string()), None, None); + let model = Model::new( + Some("Helsinki-NLP/opus-mt-en-fr".to_string()), + Some("pgml".to_string()), + None, + None, + ); let results = model .transform( "translation", diff --git a/pgml-sdks/rust-bridge/rust-bridge-macros/src/javascript.rs b/pgml-sdks/rust-bridge/rust-bridge-macros/src/javascript.rs index 004a78d4e..5d7f76b02 100644 --- a/pgml-sdks/rust-bridge/rust-bridge-macros/src/javascript.rs +++ b/pgml-sdks/rust-bridge/rust-bridge-macros/src/javascript.rs @@ -235,8 +235,9 @@ pub fn generate_javascript_methods( let middle = if method.is_async { quote! { - let runtime = crate::get_or_set_runtime(); - let x = runtime.block_on(#wrapped_call); + // let runtime = crate::get_or_set_runtime(); + // let x = runtime.block_on(#wrapped_call); + let x = #wrapped_call.await; } } else { @@ -265,14 +266,19 @@ pub fn generate_javascript_methods( #signature { use neon::prelude::*; use rust_bridge::javascript::{IntoJsResult, FromJsType}; + #outer_prepared + #inner_prepared + let channel = cx.channel(); let (deferred, promise) = cx.promise(); - deferred.try_settle_with(&channel, move |mut cx| { - #inner_prepared + crate::get_or_set_runtime().spawn(async move { #middle - x.into_js_result(&mut cx) - }).expect("Error sending js"); + deferred.try_settle_with(&channel, move |mut cx| { + x.into_js_result(&mut cx) + }).expect("Error sending js"); + }); + Ok(promise) } }