pinecone[patch]: Add additional namespace customizing (#4254)

* pinecone[patch]: Add additional namespace customizing * chore: lint files * fix tests * skip * chore: lint files * chore: lint files
langchain-ai · Feb 5, 2024 · 2afaf88 · 2afaf88
1 parent 529bdec
commit 2afaf88
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 12 deletions.
diff --git a/libs/langchain-pinecone/src/tests/vectorstores.int.test.ts b/libs/langchain-pinecone/src/tests/vectorstores.int.test.ts
@@ -9,16 +9,24 @@ import { SyntheticEmbeddings } from "@langchain/core/utils/testing";
 import { Document } from "@langchain/core/documents";
 import { PineconeStoreParams, PineconeStore } from "../vectorstores.js";
 
-describe("PineconeStore", () => {
+function sleep(ms: number) {
+  // eslint-disable-next-line no-promise-executor-return
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+describe.skip("PineconeStore", () => {
   let pineconeStore: PineconeStore;
   const testIndexName = process.env.PINECONE_INDEX!;
+  let namespaces: string[] = [];
 
   beforeAll(async () => {
     const embeddings = new SyntheticEmbeddings({
       vectorSize: 1536,
     });
 
-    const pinecone = new Pinecone();
+    const pinecone = new Pinecone({
+      apiKey: process.env.PINECONE_API_KEY!,
+    });
 
     const pineconeIndex = pinecone.Index(testIndexName);
 
@@ -29,6 +37,18 @@ describe("PineconeStore", () => {
     pineconeStore = new PineconeStore(embeddings, pineconeArgs);
   });
 
+  afterEach(async () => {
+    if (namespaces.length) {
+      const delAllPromise = namespaces.map((namespace) =>
+        pineconeStore.delete({ deleteAll: true, namespace })
+      );
+      await Promise.all(delAllPromise);
+    } else {
+      await pineconeStore.delete({ deleteAll: true });
+    }
+    namespaces = [];
+  });
+
   test("user-provided ids", async () => {
     const documentId = uuid.v4();
     const pageContent = faker.lorem.sentence(5);
@@ -37,6 +57,7 @@ describe("PineconeStore", () => {
       [{ pageContent, metadata: {} }],
       [documentId]
     );
+    await sleep(35000);
 
     const results = await pineconeStore.similaritySearch(pageContent, 1);
 
@@ -46,6 +67,7 @@ describe("PineconeStore", () => {
       [{ pageContent: `${pageContent} upserted`, metadata: {} }],
       [documentId]
     );
+    await sleep(35000);
 
     const results2 = await pineconeStore.similaritySearch(pageContent, 1);
 
@@ -61,6 +83,7 @@ describe("PineconeStore", () => {
       { pageContent, metadata: { foo: "bar" } },
     ]);
 
+    await sleep(35000);
     const results = await pineconeStore.similaritySearch(pageContent, 1);
 
     expect(results).toEqual([
@@ -77,7 +100,7 @@ describe("PineconeStore", () => {
       { pageContent, metadata: { foo: id } },
       { pageContent, metadata: { foo: "qux" } },
     ]);
-
+    await sleep(35000);
     // If the filter wasn't working, we'd get all 3 documents back
     const results = await pineconeStore.similaritySearch(pageContent, 3, {
       foo: id,
@@ -97,7 +120,7 @@ describe("PineconeStore", () => {
       { pageContent, metadata: { foo: id } },
       { pageContent, metadata: { foo: id } },
     ]);
-
+    await sleep(35000);
     // If the filter wasn't working, we'd get all 3 documents back
     const results = await pineconeStore.maxMarginalRelevanceSearch(
       pageContent,
@@ -119,7 +142,7 @@ describe("PineconeStore", () => {
       { pageContent, metadata: { foo: id } },
       { pageContent, metadata: { foo: id } },
     ]);
-
+    await sleep(35000);
     const results = await pineconeStore.similaritySearch(pageContent, 2, {
       foo: id,
     });
@@ -145,7 +168,7 @@ describe("PineconeStore", () => {
       { pageContent, metadata: { foo: id } },
       { pageContent, metadata: { foo: id } },
     ]);
-
+    await sleep(35000);
     const results = await pineconeStore.similaritySearch(pageContent, 2, {
       foo: id,
     });
@@ -162,4 +185,29 @@ describe("PineconeStore", () => {
 
     expect(results2.length).toEqual(0);
   });
+
+  test("query based on passed namespace", async () => {
+    const pageContent = "Can we make namespaces work!";
+    const id1 = uuid.v4();
+    const id2 = uuid.v4();
+    namespaces = ["test-1", "test-2"];
+    await pineconeStore.addDocuments(
+      [{ pageContent, metadata: { foo: id1 } }],
+      {
+        namespace: namespaces[0],
+      }
+    );
+    await pineconeStore.addDocuments(
+      [{ pageContent, metadata: { foo: id2 } }],
+      {
+        namespace: namespaces[1],
+      }
+    );
+    await sleep(35000);
+    const results = await pineconeStore.similaritySearch(pageContent, 1, {
+      namespace: namespaces[0],
+    });
+    expect(results.length).toEqual(1);
+    expect(results[0].metadata.foo).toBe(id1);
+  });
 });
diff --git a/libs/langchain-pinecone/src/vectorstores.ts b/libs/langchain-pinecone/src/vectorstores.ts
@@ -86,7 +86,7 @@ export class PineconeStore extends VectorStore {
    */
   async addDocuments(
     documents: Document[],
-    options?: { ids?: string[] } | string[]
+    options?: { ids?: string[]; namespace?: string } | string[]
   ): Promise<string[]> {
     const texts = documents.map(({ pageContent }) => pageContent);
     return this.addVectors(
@@ -106,7 +106,7 @@ export class PineconeStore extends VectorStore {
   async addVectors(
     vectors: number[][],
     documents: Document[],
-    options?: { ids?: string[] } | string[]
+    options?: { ids?: string[]; namespace?: string } | string[]
   ) {
     const ids = Array.isArray(options) ? options : options?.ids;
     const documentIds = ids == null ? documents.map(() => uuid.v4()) : ids;
@@ -151,7 +151,11 @@ export class PineconeStore extends VectorStore {
       } as PineconeRecord<RecordMetadata>;
     });
 
-    const namespace = this.pineconeIndex.namespace(this.namespace ?? "");
+    const optionsNamespace =
+      !Array.isArray(options) && options?.namespace
+        ? options.namespace
+        : this.namespace;
+    const namespace = this.pineconeIndex.namespace(optionsNamespace ?? "");
     // Pinecone recommends a limit of 100 vectors per upsert request
     const chunkSize = 100;
     const chunkedVectors = chunkArray(pineconeVectors, chunkSize);
@@ -171,7 +175,8 @@ export class PineconeStore extends VectorStore {
    */
   async delete(params: PineconeDeleteParams): Promise<void> {
     const { deleteAll, ids, filter } = params;
-    const namespace = this.pineconeIndex.namespace(this.namespace ?? "");
+    const optionsNamespace = params.namespace ?? this.namespace;
+    const namespace = this.pineconeIndex.namespace(optionsNamespace ?? "");
 
     if (deleteAll) {
       await namespace.deleteAll();
@@ -198,7 +203,14 @@ export class PineconeStore extends VectorStore {
       throw new Error("cannot provide both `filter` and `this.filter`");
     }
     const _filter = filter ?? this.filter;
-    const namespace = this.pineconeIndex.namespace(this.namespace ?? "");
+
+    let optionsNamespace = this.namespace ?? "";
+    if (_filter && "namespace" in _filter) {
+      optionsNamespace = _filter.namespace;
+      delete _filter.namespace;
+    }
+
+    const namespace = this.pineconeIndex.namespace(optionsNamespace ?? "");
 
     const results = await namespace.query({
       includeMetadata: true,
@@ -207,7 +219,6 @@ export class PineconeStore extends VectorStore {
       filter: _filter,
       ...options,
     });
-
     return results;
   }