community[patch]: Fix issue with inserting documents in Milvus with a…

…utoId enabled (#5829) * Fix issue with inserting documents in Milvus with autoId enabled * Lint and format --------- Co-authored-by: jacoblee93 <jacoblee93@gmail.com>
langchain-ai · Jun 20, 2024 · 5da535e · 5da535e
1 parent dd38360
commit 5da535e
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 7 deletions.
diff --git a/libs/langchain-community/src/vectorstores/milvus.ts b/libs/langchain-community/src/vectorstores/milvus.ts
@@ -276,9 +276,16 @@ export class Milvus extends VectorStore {
     if (this.partitionName !== undefined) {
       params.partition_name = this.partitionName;
     }
-    const insertResp = await this.client.upsert(params);
+    const insertResp = this.autoId
+      ? await this.client.insert(params)
+      : await this.client.upsert(params);
+
     if (insertResp.status.error_code !== ErrorCode.SUCCESS) {
-      throw new Error(`Error upserting data: ${JSON.stringify(insertResp)}`);
+      throw new Error(
+        `Error ${
+          this.autoId ? "inserting" : "upserting"
+        } data: ${JSON.stringify(insertResp)}`
+      );
     }
     await this.client.flushSync({ collection_names: [this.collectionName] });
   }

diff --git a/libs/langchain-community/src/vectorstores/tests/milvus.int.test.ts b/libs/langchain-community/src/vectorstores/tests/milvus.int.test.ts
@@ -1,6 +1,7 @@
 import { test, expect, afterAll, beforeAll } from "@jest/globals";
 import { ErrorCode, MilvusClient } from "@zilliz/milvus2-sdk-node";
 import { OpenAIEmbeddings } from "@langchain/openai";
+import { Document } from "@langchain/core/documents";
 import { Milvus } from "../milvus.js";
 
 let collectionName: string;
@@ -63,7 +64,7 @@ Harmonic Labyrinth of the dreaded Majotaur?`,
     { id: 5, other: objA },
   ]);
 
-  const resultThree = await milvus.similaritySearch(query, 1, "id == 1");
+  const resultThree = await milvus.similaritySearch(query, 1, 'id == "1"');
   const resultThreeMetadatas = resultThree.map(({ metadata }) => metadata);
   expect(resultThreeMetadatas).toEqual([{ id: 1, other: objB }]);
 });
@@ -92,6 +93,8 @@ Harmonic Labyrinth of the dreaded Majotaur?`,
   ];
   const milvus = await Milvus.fromTexts(texts, metadatas, embeddings, {
     collectionName,
+    autoId: false,
+    primaryField: "id",
     clientConfig: {
       address: MILVUS_ADDRESS,
       token: MILVUS_TOKEN,
@@ -111,14 +114,16 @@ Harmonic Labyrinth of the dreaded Majotaur?`,
     { id: 5, other: objA },
   ]);
 
-  const resultThree = await milvus.similaritySearch(query, 1, "id == 1");
+  const resultThree = await milvus.similaritySearch(query, 1, 'id == "1"');
   const resultThreeMetadatas = resultThree.map(({ metadata }) => metadata);
   expect(resultThreeMetadatas).toEqual([{ id: 1, other: objB }]);
 });
 
 test.skip("Test Milvus.fromExistingCollection", async () => {
   const milvus = await Milvus.fromExistingCollection(embeddings, {
     collectionName,
+    autoId: false,
+    primaryField: "id",
     clientConfig: {
       address: MILVUS_ADDRESS,
       token: MILVUS_TOKEN,
@@ -138,7 +143,7 @@ test.skip("Test Milvus.fromExistingCollection", async () => {
   expect(resultTwoMetadatas[1].id).toEqual(4);
   expect(resultTwoMetadatas[2].id).toEqual(5);
 
-  const resultThree = await milvus.similaritySearch(query, 1, "id == 1");
+  const resultThree = await milvus.similaritySearch(query, 1, 'id == "1"');
   const resultThreeMetadatas = resultThree.map(({ metadata }) => metadata);
   expect(resultThreeMetadatas.length).toBe(1);
   expect(resultThreeMetadatas[0].id).toEqual(1);
@@ -147,6 +152,8 @@ test.skip("Test Milvus.fromExistingCollection", async () => {
 test.skip("Test Milvus.deleteData with filter", async () => {
   const milvus = await Milvus.fromExistingCollection(embeddings, {
     collectionName,
+    autoId: false,
+    primaryField: "id",
     clientConfig: {
       address: MILVUS_ADDRESS,
       token: MILVUS_TOKEN,
@@ -160,16 +167,18 @@ test.skip("Test Milvus.deleteData with filter", async () => {
   expect(resultMetadatas.length).toBe(1);
   expect(resultMetadatas[0].id).toEqual(1);
 
-  await milvus.delete({ filter: `id in [${primaryId}]` });
+  await milvus.delete({ filter: `id in ["${primaryId}"]` });
 
   const resultTwo = await milvus.similaritySearch(query, 1);
   const resultTwoMetadatas = resultTwo.map(({ metadata }) => metadata);
-  expect(resultTwoMetadatas[0].id).not.toEqual(1);
+  expect(resultTwoMetadatas[0].id).not.toEqual(primaryId);
 });
 
 test.skip("Test Milvus.deleteData with ids", async () => {
   const milvus = await Milvus.fromExistingCollection(embeddings, {
     collectionName,
+    autoId: false,
+    primaryField: "id",
     clientConfig: {
       address: MILVUS_ADDRESS,
       token: MILVUS_TOKEN,
@@ -194,6 +203,29 @@ test.skip("Test Milvus.deleteData with ids", async () => {
   expect(resultTwoMetadatas[0].id).not.toEqual(5);
 });
 
+test.skip("Test Milvus.addDocuments with auto ID", async () => {
+  const vectorstore = new Milvus(embeddings, {
+    collectionName: `test_collection_${Math.random()
+      .toString(36)
+      .substring(7)}`,
+    clientConfig: {
+      address: MILVUS_ADDRESS,
+      token: MILVUS_TOKEN,
+    },
+  });
+
+  await vectorstore.addDocuments([
+    new Document({
+      pageContent: "test",
+      metadata: { test: "a" },
+    }),
+  ]);
+
+  const result = await vectorstore.similaritySearch("test", 1);
+  const resultMetadatas = result.map(({ metadata }) => metadata);
+  expect(resultMetadatas.length).toBe(1);
+});
+
 afterAll(async () => {
   // eslint-disable-next-line no-process-env
   if (!process.env.MILVUS_URL) return;