postgresml · SilasMarvin · Aug 25, 2023 · Aug 25, 2023
diff --git a/pgml-sdks/rust/pgml/javascript/README.md b/pgml-sdks/rust/pgml/javascript/README.md
@@ -375,7 +375,7 @@ const model = pgml.newModel(
 
 **Use an OpenAI model**
 ```javascript
-const model = pgml.newModel(name="text-embedding-ada-002", source="openai")
+const model = pgml.newModel("text-embedding-ada-002", "openai")
 ```
 
 ### Splitters

diff --git a/pgml-sdks/rust/pgml/javascript/tests/typescript-tests/test.ts b/pgml-sdks/rust/pgml/javascript/tests/typescript-tests/test.ts
@@ -1,4 +1,4 @@
-import pgml from '../../index.js'
+import pgml from "../../index.js";
 
 ////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////
@@ -9,29 +9,21 @@ import pgml from '../../index.js'
 ////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////
 
-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(LOG_LEVEL);
 
 const generate_dummy_documents = (count: number) => {
   let docs = [];
   for (let i = 0; i < count; i++) {
     docs.push({
-      "id": i,
-      "text": `This is a test document: ${i}`,
-      "metadata": {
-        "uuid": i * 10,
-        "name": `Test Document ${i}`
-      }
+      id: i,
+      text: `This is a test document: ${i}`,
+      uuid: i * 10,
+      name: `Test Document ${i}`,
     });
   }
   return docs;
-}
+};
 
 ///////////////////////////////////////////////////
 // Test the API exposed is correct ////////////////
@@ -118,6 +110,25 @@ it("can vector search with query builder with remote embeddings", async() => {
   await collection.archive();
 });
 
+it("can vector search with query builder and metadata filtering", async () => {
+  let model = pgml.newModel();
+  let splitter = pgml.newSplitter();
+  let pipeline = pgml.newPipeline("test_j_p_cvswqbamf_0", model, splitter);
+  let collection = pgml.newCollection("test_j_c_cvswqbamf_4");
+  await collection.upsert_documents(generate_dummy_documents(3));
+  await collection.add_pipeline(pipeline);
+  let results = await collection
+    .query()
+    .vector_recall("Here is some query", pipeline)
+    .filter({
+      metadata: {
+        $or: [{ uuid: { $eq: 0 } }, { uuid: { $eq: 20 } }],
+      },
+    })
+    .limit(10).fetch_all();
+  expect(results).toHaveLength(2);
+  await collection.archive();
+});
 
 ///////////////////////////////////////////////////
 // Test user output facing functions //////////////

diff --git a/pgml-sdks/rust/pgml/python/README.md b/pgml-sdks/rust/pgml/python/README.md
@@ -315,13 +315,13 @@ results = (
                 {
                     "$and": [
                         {
-                            "$eq": {
-                                "uuid": 1
+                            "uuid": {
+                                "$eq": 1
                             }    
                         },
                         {
-                            "$lt": {
-                                "index": 100 
+                            "index": {
+                                "$lt": 100 
                             }
                         }
                     ] 

diff --git a/pgml-sdks/rust/pgml/python/tests/test.py b/pgml-sdks/rust/pgml/python/tests/test.py
@@ -30,7 +30,8 @@ def generate_dummy_documents(count: int) -> List[Dict[str, Any]]:
                 "id": i,
                 "text": "This is a test document: {}".format(i),
                 "some_random_thing": "This will be metadata on it",
-                "metadata": {"uuid": i * 10, "name": "Test Document {}".format(i)},
+                "uuid": i * 10,
+                "name": "Test Document {}".format(i),
             }
         )
     return dummy_documents
@@ -135,6 +136,31 @@ async def test_can_vector_search_with_query_builder_with_remote_embeddings():
     await collection.archive()
 
 
+@pytest.mark.asyncio
+async def test_can_vector_search_with_query_builder_and_metadata_filtering():
+    model = pgml.Model()
+    splitter = pgml.Splitter()
+    pipeline = pgml.Pipeline("test_p_p_tcvswqbamf_1", model, splitter)
+    collection = pgml.Collection(name="test_p_c_tcvswqbamf_2")
+    await collection.upsert_documents(generate_dummy_documents(3))
+    await collection.add_pipeline(pipeline)
+    results = (
+        await collection.query()
+        .vector_recall("Here is some query", pipeline)
+        .filter({
+            "metadata": {
+                "uuid": {
+                    "$eq": 0
+                }
+            }
+        })
+        .limit(10)
+        .fetch_all()
+    )
+    assert len(results) == 1
+    await collection.archive()
+
+
 ###################################################
 ## Test user output facing functions ##############
 ###################################################
@@ -158,17 +184,17 @@ async def test_pipeline_to_dict():
 ###################################################
 
 
-def vector_search(collection_name, pipeline_name):
-    collection = pgml.Collection(collection_name)
-    pipeline = pgml.Pipeline(pipeline_name)
-    result = asyncio.run(
-        collection.query()
-        .vector_recall("Here is some query", pipeline)
-        .limit(10)
-        .fetch_all()
-    )
-    print(result)
-    return [0, 1, 2]
+# def vector_search(collection_name, pipeline_name):
+#     collection = pgml.Collection(collection_name)
+#     pipeline = pgml.Pipeline(pipeline_name)
+#     result = asyncio.run(
+#         collection.query()
+#         .vector_recall("Here is some query", pipeline)
+#         .limit(10)
+#         .fetch_all()
+#     )
+#     print(result)
+#     return [0, 1, 2]
 
 
 # @pytest.mark.asyncio
@@ -200,23 +226,23 @@ def vector_search(collection_name, pipeline_name):
 ###################################################
 
 
-async def silas_test_add_pipeline():
-    model = pgml.Model()
-    splitter = pgml.Splitter()
-    pipeline = pgml.Pipeline("silas_test_p_1", model, splitter)
-    collection = pgml.Collection(name="silas_test_c_10")
-    await collection.add_pipeline(pipeline)
-
-async def silas_test_upsert_documents():
-    collection = pgml.Collection(name="silas_test_c_9")
-    await collection.upsert_documents(generate_dummy_documents(10))
-
-async def silas_test_vector_search():
-    pipeline = pgml.Pipeline("silas_test_p_1")
-    collection = pgml.Collection(name="silas_test_c_9")
-    results = await collection.vector_search("Here is some query", pipeline)
-    print(results)
-
-# asyncio.run(silas_test_add_pipeline())
-# asyncio.run(silas_test_upsert_documents())
-# asyncio.run(silas_test_vector_search())
+# async def test_add_pipeline():
+#     model = pgml.Model()
+#     splitter = pgml.Splitter()
+#     pipeline = pgml.Pipeline("silas_test_p_1", model, splitter)
+#     collection = pgml.Collection(name="silas_test_c_10")
+#     await collection.add_pipeline(pipeline)
+#
+# async def test_upsert_documents():
+#     collection = pgml.Collection(name="silas_test_c_9")
+#     await collection.upsert_documents(generate_dummy_documents(10))
+#
+# async def test_vector_search():
+#     pipeline = pgml.Pipeline("silas_test_p_1")
+#     collection = pgml.Collection(name="silas_test_c_9")
+#     results = await collection.vector_search("Here is some query", pipeline)
+#     print(results)
+
+# asyncio.run(test_add_pipeline())
+# asyncio.run(test_upsert_documents())
+# asyncio.run(test_vector_search())