postgresml · SilasMarvin · Jun 29, 2023 · Jun 29, 2023 · Jun 29, 2023
diff --git a/.github/workflows/javascript-sdk.yml b/.github/workflows/javascript-sdk.yml
@@ -41,38 +41,38 @@ jobs:
           name: node-artifacts 
           path: pgml-sdks/rust/pgml/javascript/${{ matrix.neon-out-name }}
           retention-days: 1
-  publish-javascript-sdk:
-    needs: build-javascript-sdk
-    runs-on: "ubuntu-22.04"
-    defaults:
-      run:
-        working-directory: pgml-sdks/rust/pgml/javascript
-    steps:
-      - uses: actions/checkout@v3
-      - uses: actions-rs/toolchain@v1
-        with:
-          toolchain: stable
-      - name: Validate cargo is working
-        uses: postgresml/gh-actions-cargo@master
-        with:
-          command: version
-      - name: Create artifact directory
-        run: mkdir dist
-      - name: Download artifacts
-        uses: actions/download-artifact@v3
-        with:
-          name: node-artifacts
-          path: pgml-sdks/rust/pgml/javascript/dist
-      - uses: actions/setup-node@v3
-        with:
-          node-version: '20.x'
-          registry-url: 'https://registry.npmjs.org'
-      - name: Generate types declaration
-        run: |
-          npm i
-          npm run build 
-          rm index.node
-      - run: npm ci
-      - run: npm publish
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
+  # publish-javascript-sdk:
+  #   needs: build-javascript-sdk
+  #   runs-on: "ubuntu-22.04"
+  #   defaults:
+  #     run:
+  #       working-directory: pgml-sdks/rust/pgml/javascript
+  #   steps:
+  #     - uses: actions/checkout@v3
+  #     - uses: actions-rs/toolchain@v1
+  #       with:
+  #         toolchain: stable
+  #     - name: Validate cargo is working
+  #       uses: postgresml/gh-actions-cargo@master
+  #       with:
+  #         command: version
+  #     - name: Create artifact directory
+  #       run: mkdir dist
+  #     - name: Download artifacts
+  #       uses: actions/download-artifact@v3
+  #       with:
+  #         name: node-artifacts
+  #         path: pgml-sdks/rust/pgml/javascript/dist
+  #     - uses: actions/setup-node@v3
+  #       with:
+  #         node-version: '20.x'
+  #         registry-url: 'https://registry.npmjs.org'
+  #     - name: Generate types declaration
+  #       run: |
+  #         npm i
+  #         npm run build 
+  #         rm index.node
+  #     - run: npm ci
+  #     - run: npm publish
+  #       env:
+  #         NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
diff --git a/pgml-sdks/rust/pgml/javascript/index.js b/pgml-sdks/rust/pgml/javascript/index.js
@@ -10,6 +10,9 @@ try {
   if (type == "Darwin" && arch == "x64") {
     const pgml = require("./dist/x86_64-apple-darwin-index.node")
     module.exports = pgml
+  } else if (type == "Darwin" && arch == "arm64") {
+    const pgml = require("./dist/aarch64-apple-darwin-index.node")
+    module.exports = pgml
   } else if (type == "Windows" && arch == "x64") {
     const pgml = require("./dist/x86_64-pc-windows-gnu-index.node")
     module.exports = pgml

diff --git a/pgml-sdks/rust/pgml/javascript/package.json b/pgml-sdks/rust/pgml/javascript/package.json
@@ -1,6 +1,6 @@
 {
   "name": "pgml",
-  "version": "0.1.5",
+  "version": "0.1.6",
   "description": "Open Source Alternative for Building End-to-End Vector Search Applications without OpenAI & Pinecone",
   "keywords": ["postgres", "machine learning", "vector databases", "embeddings"],
   "main": "index.js",

diff --git a/pgml-sdks/rust/pgml/python/tests/test.py b/pgml-sdks/rust/pgml/python/tests/test.py
@@ -14,12 +14,12 @@ async def main():
     await collection.register_text_splitter("recursive_character", {"chunk_size": 1500, "chunk_overlap": 40})
     splitters = await collection.get_text_splitters()
     print(splitters)
-    await collection.generate_chunks(2)
+    await collection.generate_chunks()
     await collection.register_model("embedding", "intfloat/e5-small")
     models = await collection.get_models()
     print(models)
     await collection.generate_embeddings()
-    results = await collection.vector_search("small", {}, 2);
+    results = await collection.vector_search("small")
     print(results)
     await db.archive_collection(collection_name)
 

diff --git a/pgml-sdks/rust/pgml/src/collection.rs b/pgml-sdks/rust/pgml/src/collection.rs
@@ -731,16 +731,17 @@ impl Collection {
             }
         };
 
+        let model_name = self.get_model_name(model_id).await?.expect("Model with id: {} does not exist");
+
         let results: Vec<(f64, String, sqlx::types::Json<HashMap<String, String>>)> =
             sqlx::query_as(&query_builder!(
                 queries::VECTOR_SEARCH,
-                self.models_table_name,
                 embeddings_table_name,
                 embeddings_table_name,
                 self.chunks_table_name,
                 self.documents_table_name
             ))
-            .bind(model_id)
+            .bind(model_name)
             .bind(query)
             .bind(query_params)
             .bind(top_k)
@@ -751,6 +752,20 @@ impl Collection {
         Ok(results)
     }
 
+    async fn get_model_name(&self, model_id: i64) -> anyhow::Result<Option<String>> {
+        let model: Option<models::Model> = sqlx::query_as(&query_builder!(
+            "SELECT * from %s WHERE id = $1",
+            self.models_table_name 
+        ))
+        .bind(model_id)
+        .fetch_optional(self.pool.borrow())
+        .await?;
+        match model {
+            Some(m) => Ok(Some(m.name)),
+            None => Ok(None),
+        }
+    }
+
     fn generate_table_names(name: &str) -> (String, String, String, String, String) {
         [
             ".documents",

diff --git a/pgml-sdks/rust/pgml/src/queries.rs b/pgml-sdks/rust/pgml/src/queries.rs
@@ -122,22 +122,11 @@ WHERE
 "#;
 
 pub const VECTOR_SEARCH: &str = r#"
-WITH model AS (
-  SELECT 
-    name 
-  FROM 
-    %s 
-  WHERE 
-    id = $1 
-), 
-query_cte AS (
+WITH query_cte AS (
   SELECT 
     pgml.embed(
       transformer => (
-        SELECT 
-          name 
-        from 
-          model
+        $1
       ), 
       text => $2, 
       kwargs => $3