fix(community): Add id in chroma search result (#7120)

Co-authored-by: jacoblee93 <[email protected]>
langchain-ai · Nov 5, 2024 · fcef784 · fcef784
1 parent 2e19277
commit fcef784
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/libs/langchain-community/src/vectorstores/chroma.ts b/libs/langchain-community/src/vectorstores/chroma.ts
@@ -398,6 +398,7 @@ export class Chroma extends VectorStore {
         new Document({
           pageContent: firstDocuments?.[i] ?? "",
           metadata,
+          id: firstIds[i],
         }),
         firstDistances[i],
       ]);

diff --git a/libs/langchain-community/src/vectorstores/tests/chroma.test.ts b/libs/langchain-community/src/vectorstores/tests/chroma.test.ts
@@ -132,4 +132,23 @@ describe("Chroma", () => {
     });
     expect(results).toHaveLength(5);
   });
+
+  test("should return id properly when adding documents", async () => {
+    const document1 = {
+      pageContent: "Document 1",
+      metadata: { source: "https://example.com" },
+    };
+
+    const documents = [document1];
+
+    const chroma = new Chroma(new FakeEmbeddings(), {
+      collectionName: "new-test-collection",
+      index: mockClient,
+    });
+
+    await chroma.addDocuments(documents, { ids: ["0"] });
+    const result = await chroma.similaritySearch(document1.pageContent, 1);
+
+    expect(result[0]).toHaveProperty("id", "0");
+  });
 });