Improvments to javascript and updates to the python sdk deploy script

postgresml · SilasMarvin · Aug 21, 2023 · Jul 14, 2023 · Jul 17, 2023 · Jul 19, 2023
commit 92c962315dbfe4545a7480c706cbcc000255506d
diff --git a/pgml-sdks/rust/pgml-macros/src/javascript.rs b/pgml-sdks/rust/pgml-macros/src/javascript.rs
@@ -123,7 +123,7 @@ pub fn generate_javascript_derive(parsed: DeriveInput) -> proc_macro::TokenStrea
         #[cfg(feature = "javascript")]
         impl FromJsType for & #wrapped_type_ident {
             type From = neon::types::JsObject;
-           fn from_js_type<'a, C: neon::context::Context<'a>>(cx: &mut C, arg: neon::handle::Handle<Self::From>) -> neon::result::NeonResult<Self> {
+            fn from_js_type<'a, C: neon::context::Context<'a>>(cx: &mut C, arg: neon::handle::Handle<Self::From>) -> neon::result::NeonResult<Self> {
                 use neon::prelude::*;
                 use core::ops::Deref;
                 let s: neon::handle::Handle<neon::types::JsBox<#name_ident>> = arg.get(cx, "s")?;
@@ -136,51 +136,6 @@ pub fn generate_javascript_derive(parsed: DeriveInput) -> proc_macro::TokenStrea
             }
         }
 
-        #[cfg(feature = "javascript")]
-        impl CustomInto<#wrapped_type_ident> for &#name_ident {
-            fn custom_into(self) -> #wrapped_type_ident {
-                *self.wrapped.clone()
-            }
-        }
-
-        #[cfg(feature = "javascript")]
-        impl CustomInto<&'static #wrapped_type_ident> for &#name_ident {
-            fn custom_into(self) -> &'static #wrapped_type_ident {
-                unsafe {
-                    let ptr = &*self.wrapped as *const #wrapped_type_ident;
-                    let ptr = ptr as *mut #wrapped_type_ident;
-                    let boxed = Box::from_raw(ptr);
-                    Box::leak(boxed)
-                }
-            }
-        }
-
-        #[cfg(feature = "javascript")]
-        impl CustomInto<&'static mut #wrapped_type_ident> for &#name_ident {
-            fn custom_into(self) -> &'static mut #wrapped_type_ident {
-                unsafe {
-                    let ptr = &*self.wrapped as *const #wrapped_type_ident;
-                    let ptr = ptr as *mut #wrapped_type_ident;
-                    let boxed = Box::from_raw(ptr);
-                    Box::leak(boxed)
-                }
-            }
-        }
-
-        // #[cfg(feature = "javascript")]
-        // impl FromJsType for #name_ident {
-        //     type From = neon::types::JsObject;
-        //     fn from_js_type<'a, C: neon::context::Context<'a>>(cx: &mut C, arg: neon::handle::Handle<Self::From>) -> neon::result::NeonResult<Self> {
-        //         use neon::prelude::*;
-        //         use core::ops::Deref;
-        //         let s: neon::handle::Handle<neon::types::JsBox<#name_ident>> = arg.get(cx, "s")?;
-        //         let wrapped = (*s).wrapped.clone();
-        //         Ok(Self {
-        //             wrapped
-        //         })
-        //     }
-        // }
-
         #[cfg(feature = "javascript")]
         impl IntoJsResult for #wrapped_type_ident {
             type Output = neon::types::JsObject;
@@ -309,9 +264,12 @@ pub fn generate_javascript_methods(
             } else {
                 quote! {
                     let this = this.into_inner(&mut cx);
-                    let s: neon::handle::Handle<neon::types::JsBox<#name_ident>> = this.get(&mut cx, "s")?;
-                    // let wrapped: &mut #wrapped_type_ident = <&mut #wrapped_type_ident>::from_js_type(&mut cx, s)?;
-                    let wrapped: &mut #wrapped_type_ident = s.custom_into();
+                    // let s: neon::handle::handle<neon::types::jsbox<#name_ident>> = this.get(&mut cx, "s")?;
+                    // let wrapped: &mut #wrapped_type_ident = s.custom_into();
+
+                    let s: neon::handle::Handle<neon::types::JsObject> = this.get(&mut cx, "s")?;
+                    let wrapped =  <&mut #wrapped_type_ident>::from_js_type(&mut cx, s)?;
+
                     #(#inner_prep_arguments)*
                 }
             }

diff --git a/pgml-sdks/rust/pgml/build.rs b/pgml-sdks/rust/pgml/build.rs
@@ -3,14 +3,14 @@ use std::fs::OpenOptions;
 use std::io::Write;
 
 const ADDITIONAL_DEFAULTS_FOR_PYTHON: &[u8] = br#"
-def py_init_logger(log_level: str) -> None
+def py_init_logger(level: Option[str], format: Option[str]) -> None
 
 Json = dict[str, Any]
 DateTime = int
 "#;
 
 const ADDITIONAL_DEFAULTS_FOR_JAVASCRIPT: &[u8] = br#"
-export function setupLogger(log_level: string): void;
+export function js_init_logger(level?: string, format?: string): void;
 
 export type Json = { [key: string]: any };
 export type DateTime = Date;
@@ -19,6 +19,7 @@ export function newCollection(name: string, database_url?: string): Collection;
 export function newModel(name?: string, task?: string, source?: string, parameters?: string, database_url?: string): Model;
 export function newSplitter(name?: string, parameters?: any, database_url?: string): Splitter;
 export function newBuiltins(database_url?: string): Builtins;
+export function newPipeline(name: string, model: Model, splitter: Splitter): Pipeline;
 "#;
 
 fn main() {

diff --git a/pgml-sdks/rust/pgml/javascript/tests/typescript-tests/test.ts b/pgml-sdks/rust/pgml/javascript/tests/typescript-tests/test.ts
@@ -2,14 +2,23 @@ import pgml from '../../index.js'
 
 ////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////
-// PLEASE BE AWARE THESE TESTS DO INVOLVE CHECKS ON LAZILY CREATD DATABASE ITEMS. //
+// PLEASE BE AWARE THESE TESTS DO INVOLVE CHECKS ON LAZILY CREATD DATABASE ITEMS  //
 // IF ANY OF THE COLLECTION NAMES ALREADY EXIST, SOME TESTS MAY FAIL              //
 // THIS DOES NOT MEAN THE SDK IS BROKEN. PLEASE CLEAR YOUR DATABASE INSTANCE      //
 // BEFORE RUNNING ANY TESTS                                                       //
 ////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////
 
-const generate_documents = (count: number) => {
+const DATABASE_URL = process.env.DATABASE_URL;
+if (!DATABASE_URL) {
+  console.log("No DATABASE_URL environment variable found. Please set one")
+  process.exit(1)
+}
+const LOG_LEVEL = process.env.LOG_LEVEL ? process.env.LOG_LEVEL : "ERROR";
+
+pgml.js_init_logger(DATABASE_URL, LOG_LEVEL);
+
+const generate_dummy_documents = (count: number) => {
   let docs = [];
   for (let i = 0; i < count; i++) {
     docs.push({
@@ -24,94 +33,49 @@ const generate_documents = (count: number) => {
   return docs;
 }
 
-it("can lazily create collection", async () => {
-  let collection_name = "j_ccc_test_2";
-  let collection = pgml.newCollection(collection_name);
-  let builtins = pgml.newBuiltins();
-  let does_collection_exist = await builtins.does_collection_exist(collection_name);
-  expect(does_collection_exist).toBe(false);
-  // Do something that requires the collection to be created
-  await collection.upsert_documents(generate_documents(1));
-  // Now the collection will exit because it had to be created to upsert documents
-  does_collection_exist = await builtins.does_collection_exist(collection_name);
-  await collection.archive();
-  expect(does_collection_exist).toBe(true);
+///////////////////////////////////////////////////
+// Test the API exposed is correct ////////////////
+///////////////////////////////////////////////////
+
+it("can create collection", () => {
+  let collection = pgml.newCollection("test_j_c_ccc_0");
+  expect(collection).toBeTruthy();
 });
 
-it("can lazily create model", async () => {
+it("can create model", () => {
   let model = pgml.newModel();
-  expect(model.get_verified_in_database()).toBe(false);
-  let id = await model.get_id();
-  expect(id).toBeDefined();
-  expect(model.get_verified_in_database()).toBe(true);
-})
+  expect(model).toBeTruthy();
+});
 
-it("can lazily create splitter", async () => {
+it("can create splitter", () => {
   let splitter = pgml.newSplitter();
-  expect(splitter.get_verified_in_database()).toBe(false);
-  let id = await splitter.get_id();
-  expect(id).toBeDefined();
-  expect(splitter.get_verified_in_database()).toBe(true);
-})
+  expect(splitter).toBeTruthy();
+});
 
-it("can vector search", async () => {
-  let collection_name = "j_cvs_test_0";
-  let collection = pgml.newCollection(collection_name);
+it("can create pipeline", () => {
   let model = pgml.newModel();
   let splitter = pgml.newSplitter();
-  await collection.upsert_documents(generate_documents(2));
-  // Splitter should not be verified in the database yet
-  expect(splitter.get_verified_in_database()).toBe(false);
-  await collection.generate_chunks(splitter);
-  // Now splitter should be verified in the database
-  expect(splitter.get_verified_in_database()).toBe(true);
-  // Model should not be verified in the database yet
-  expect(model.get_verified_in_database()).toBe(false);
-  await collection.generate_embeddings(model, splitter);
-  // Now model should be verified in the database
-  expect(model.get_verified_in_database()).toBe(true);
-  let results = await collection.vector_search("Here is some query", model, splitter);
-  await collection.archive();
-  expect(results.length).toBe(2);
-})
+  let pipeline = pgml.newPipeline("test_j_p_ccc_0", model, splitter);
+  expect(pipeline).toBeTruthy();
+});
 
-it("can vector search with remote embeddings", async () => {
-  let collection_name = "j_cvswre_test_0";
-  let collection = pgml.newCollection(collection_name);
-  let model = pgml.newModel("text-embedding-ada-002", "embeddings", "openai");
-  let splitter = pgml.newSplitter();
-  await collection.upsert_documents(generate_documents(2));
-  await collection.generate_chunks(splitter);
-  await collection.generate_embeddings(model, splitter);
-  let results = await collection.vector_search("Here is some query", model, splitter);
-  await collection.archive();
-  expect(results.length).toBe(2);
-})
+it("can create builtins", () => {
+  let builtins = pgml.newBuiltins();
+  expect(builtins).toBeTruthy();
+});
 
-it("can vector search with query builder", async () => {
-  let collection_name = "j_cvswqb_test_0";
-  let collection = pgml.newCollection(collection_name);
-  let model = pgml.newModel();
-  let splitter = pgml.newSplitter();
-  await collection.upsert_documents(generate_documents(2));
-  await collection.generate_chunks(splitter);
-  await collection.generate_embeddings(model, splitter);
-  await collection.generate_tsvectors();
-  let results = await collection.query().vector_recall("Here is some query", model, splitter).filter({
-        "metadata": {
-            "metadata": {
-                "$or": [
-                    {"uuid": {"$eq": 0 }},
-                    {"uuid": {"$eq": 10 }},
-                    {"category": {"$eq": [1, 2, 3]}}
-                ]
+///////////////////////////////////////////////////
+// Test various vector searches ///////////////////
+///////////////////////////////////////////////////
 
-            }
-        },
-        "full_text": {
-            "text": "Test document"
-        }
-    }).run()
-  await collection.archive();
-  expect(results.length).toBe(2);
-})
+// it("can vector search with local embeddings", async () => {
+//   let model = pgml.newModel();
+//   let splitter = pgml.newSplitter();
+//   let pipeline = pgml.newPipeline("test_j_p_cvswle_0", model, splitter);
+//   let collection = pgml.newCollection("test_j_c_cvswle_0");
+//   await collection.upsert_documents(generate_dummy_documents(3));
+//   await collection.add_pipeline(pipeline);
+//   let results = await collection.vector_search("Here is some query", pipeline);
+//   expect(results).toHaveLength(3);
+//   await collection.archive();
+// });
diff --git a/pgml-sdks/rust/pgml/src/builtins.rs b/pgml-sdks/rust/pgml/src/builtins.rs
@@ -10,7 +10,7 @@ pub struct Builtins {
 use crate::{get_or_initialize_pool, query_runner::QueryRunner, types::Json};
 
 #[cfg(feature = "javascript")]
-use crate::{languages::javascript::*, query_runner::QueryRunnerJavascript, languages::CustomInto};
+use crate::languages::javascript::*;
 
 #[cfg(feature = "python")]
 use crate::{languages::CustomInto, query_runner::QueryRunnerPython};

diff --git a/pgml-sdks/rust/pgml/src/collection.rs b/pgml-sdks/rust/pgml/src/collection.rs
@@ -23,10 +23,7 @@ use crate::types::Json;
 use crate::utils;
 
 #[cfg(feature = "javascript")]
-use crate::{
-    languages::javascript::*, languages::CustomInto, pipeline::PipelineJavascript,
-    query_builder::QueryBuilderJavascript,
-};
+use crate::languages::javascript::*;
 
 #[cfg(feature = "python")]
 use crate::{languages::CustomInto, pipeline::PipelinePython, query_builder::QueryBuilderPython};
@@ -529,10 +526,10 @@ impl Collection {
             Ok(Some((source_uuid, text, metadata)))
         }).collect();
 
-        // Yes we could continue chaining the above iterators but types become super annoying to
+        // We could continue chaining the above iterators but types become super annoying to
         // deal with, especially because we are dealing with async functions. This is much easier to read
         // Also, we may want to use a variant of chunks that is owned, I'm not 100% sure of what
-        // cloning happens when passing values into sqlx bind. itertools variant will not work as
+        // cloning happens when passing values into sqlx bind. itertools variants will not work as
         // it is not thread safe and pyo3 will get upset
         let mut document_ids = Vec::new();
         for chunk in documents?.chunks(10) {

diff --git a/pgml-sdks/rust/pgml/src/lib.rs b/pgml-sdks/rust/pgml/src/lib.rs
@@ -176,13 +176,26 @@ fn pgml(_py: pyo3::Python, m: &pyo3::types::PyModule) -> pyo3::PyResult<()> {
     Ok(())
 }
 
+#[cfg(feature = "javascript")]
+fn js_init_logger(mut cx: neon::context::FunctionContext) -> neon::result::JsResult<neon::types::JsUndefined> {
+    use crate::languages::javascript::*;
+    let level = cx.argument_opt(0);
+    let level = <Option<String>>::from_option_js_type(&mut cx, level)?;
+    let format = cx.argument_opt(1);
+    let format = <Option<String>>::from_option_js_type(&mut cx, format)?;
+    init_logger(level, format).ok();
+    ().into_js_result(&mut cx)
+}
+
 #[cfg(feature = "javascript")]
 #[neon::main]
 fn main(mut cx: neon::context::ModuleContext) -> neon::result::NeonResult<()> {
+    cx.export_function("js_init_logger", js_init_logger)?;
     cx.export_function("newCollection", collection::CollectionJavascript::new)?;
     cx.export_function("newModel", model::ModelJavascript::new)?;
     cx.export_function("newSplitter", splitter::SplitterJavascript::new)?;
     cx.export_function("newBuiltins", builtins::BuiltinsJavascript::new)?;
+    cx.export_function("newPipeline", pipeline::PipelineJavascript::new)?;
     Ok(())
 }
 

diff --git a/pgml-sdks/rust/pgml/src/model.rs b/pgml-sdks/rust/pgml/src/model.rs
@@ -10,7 +10,7 @@ use crate::{
 };
 
 #[cfg(feature = "javascript")]
-use crate::{languages::javascript::*, languages::CustomInto};
+use crate::languages::javascript::*;
 
 #[cfg(feature = "python")]
 use crate::languages::CustomInto;

diff --git a/pgml-sdks/rust/pgml/src/models.rs b/pgml-sdks/rust/pgml/src/models.rs
@@ -5,9 +5,6 @@ use serde::Serialize;
 
 use crate::types::{DateTime, Json};
 
-#[cfg(feature = "javascript")]
-use crate::languages::javascript::*;
-
 /// A pipeline
 #[enum_def]
 #[derive(FromRow)]

diff --git a/pgml-sdks/rust/pgml/src/pipeline.rs b/pgml-sdks/rust/pgml/src/pipeline.rs
@@ -22,8 +22,7 @@ use crate::{
 use crate::{languages::CustomInto, model::ModelPython, splitter::SplitterPython};
 
 #[cfg(feature = "javascript")]
-use crate::{languages::{javascript::*, CustomInto}, model::ModelJavascript, splitter::SplitterJavascript};
-
+use crate::languages::javascript::*;
 
 #[derive(Debug, Clone)]
 pub enum PipelineSyncStatus {

diff --git a/pgml-sdks/rust/pgml/src/query_builder.rs b/pgml-sdks/rust/pgml/src/query_builder.rs
@@ -14,10 +14,7 @@ use crate::{
 };
 
 #[cfg(feature = "javascript")]
-use crate::{
-    languages::{javascript::*, CustomInto},
-    pipeline::PipelineJavascript,
-};
+use crate::languages::javascript::*;
 
 #[cfg(feature = "python")]
 use crate::{languages::CustomInto, pipeline::PipelinePython};

diff --git a/pgml-sdks/rust/pgml/src/query_runner.rs b/pgml-sdks/rust/pgml/src/query_runner.rs
@@ -6,7 +6,7 @@ use sqlx::{Postgres, Row};
 use crate::{get_or_initialize_pool, types::Json};
 
 #[cfg(feature = "javascript")]
-use crate::{languages::javascript::*, languages::CustomInto};
+use crate::languages::javascript::*;
 
 #[cfg(feature = "python")]
 use crate::languages::CustomInto;

diff --git a/pgml-sdks/rust/pgml/src/splitter.rs b/pgml-sdks/rust/pgml/src/splitter.rs
@@ -10,7 +10,7 @@ use crate::{
 };
 
 #[cfg(feature = "javascript")]
-use crate::{languages::javascript::*, languages::CustomInto};
+use crate::languages::javascript::*;
 
 #[cfg(feature = "python")]
 use crate::languages::CustomInto;