community[minor]: Add Upstash Embeddings Support (#5150)

* add: metadata filtering for UpstashVectorStore

* fmt

* skip tests

* add: upstash embeddings support

* add: tests

* fmt

* Update upstash.mdx

* Update upstash.mdx

* use fake embeddings

* fix: replace UpstashEmbeddings parameter with FakeEmbeddings class

* Naming

---------

Co-authored-by: Jacob Lee <[email protected]>

Author: fahreddinozcan

docs/core_docs/docs/integrations/vectorstores/upstash.mdx

@@ -2,6 +2,7 @@ import CodeBlock from "@theme/CodeBlock";
 import CreateClientExample from "@examples/indexes/vector_stores/upstash/create_client.ts";
 import IndexQueryExample from "@examples/indexes/vector_stores/upstash/index_and_query_docs.ts";
 import DeleteExample from "@examples/indexes/vector_stores/upstash/delete_docs.ts";
+import UpstashEmbeddingsExample from "@examples/indexes/vector_stores/upstash/upstash_embeddings.ts";
 import IntegrationInstallTooltip from "@mdx_components/integration_install_tooltip.mdx";
 
 # Upstash Vector
@@ -41,6 +42,12 @@ You can index the LangChain documents with any model of your choice, and perform
 
 <CodeBlock language="typescript">{IndexQueryExample}</CodeBlock>
 
+## Upstash embeddings
+
+It's possible to use the embeddings service of Upstash, which is based on the embedding model of choice when creating the vector database. You don't need to create the embeddings manually, as the Upstash Vector service will handle this for you.
+
+<CodeBlock language="typescript">{UpstashEmbeddingsExample}</CodeBlock>
+
 ## Delete Documents
 
 You can also delete the documents you've indexed previously.

examples/src/indexes/vector_stores/upstash/upstash_embeddings.ts

@@ -0,0 +1,66 @@
+import { Index } from "@upstash/vector";
+import { Document } from "@langchain/core/documents";
+import { UpstashVectorStore } from "@langchain/community/vectorstores/upstash";
+import { FakeEmbeddings } from "@langchain/core/utils/testing";
+
+const index = new Index({
+  url: process.env.UPSTASH_VECTOR_REST_URL as string,
+  token: process.env.UPSTASH_VECTOR_REST_TOKEN as string,
+});
+
+// Initializing the UpstashVectorStore with the Upstash Embeddings configuration.
+// Passing FakeEmbeddings here will enable the store to generate embeddings using Upstash Embeddings.
+const UpstashVector = new UpstashVectorStore(new FakeEmbeddings(), { index });
+
+// Creating the docs to be indexed.
+const id = new Date().getTime();
+const documents = [
+  new Document({
+    metadata: { name: id },
+    pageContent: "Hello there!",
+  }),
+  new Document({
+    metadata: { name: id },
+    pageContent: "What are you building?",
+  }),
+  new Document({
+    metadata: { time: id },
+    pageContent: "Upstash Vector is great for building AI applications.",
+  }),
+  new Document({
+    metadata: { time: id },
+    pageContent: "To be, or not to be, that is the question.",
+  }),
+];
+
+// Creating embeddings from the provided documents, and adding them to Upstash database.
+await UpstashVector.addDocuments(documents);
+
+// Waiting vectors to be indexed in the vector store.
+// eslint-disable-next-line no-promise-executor-return
+await new Promise((resolve) => setTimeout(resolve, 1000));
+
+const queryResult = await UpstashVector.similaritySearchWithScore(
+  "Vector database",
+  2
+);
+
+console.log(queryResult);
+/**
+[
+  [
+  Document {
+    pageContent: 'Upstash Vector is great for building AI applications.',
+    metadata: [Object]
+  },
+  0.9016147
+  ],
+  [
+  Document {
+    pageContent: 'What are you building?',
+    metadata: [Object]
+  },
+  0.8613077
+  ]
+]
+ */

libs/langchain-community/src/vectorstores/tests/upstash.int.test.ts

@@ -1,12 +1,15 @@
 /* eslint-disable no-process-env */
 import { Index } from "@upstash/vector";
 import { Document } from "@langchain/core/documents";
-import { SyntheticEmbeddings } from "@langchain/core/utils/testing";
+import {
+  SyntheticEmbeddings,
+  FakeEmbeddings,
+} from "@langchain/core/utils/testing";
 import { EmbeddingsInterface } from "@langchain/core/embeddings";
 import { UpstashVectorStore } from "../upstash.js";
 import { sleep } from "../../utils/time.js";
 
-describe.skip("UpstashVectorStore", () => {
+describe("UpstashVectorStore", () => {
   let store: UpstashVectorStore;
   let embeddings: EmbeddingsInterface;
   let index: Index;
@@ -17,8 +20,10 @@ describe.skip("UpstashVectorStore", () => {
       token: process.env.UPSTASH_VECTOR_REST_TOKEN,
     });
 
+    await index.reset();
+
     embeddings = new SyntheticEmbeddings({
-      vectorSize: 1536,
+      vectorSize: 384,
     });
 
     store = new UpstashVectorStore(embeddings, {
@@ -119,4 +124,52 @@ describe.skip("UpstashVectorStore", () => {
 
     expect(results2).toHaveLength(0);
   });
+
+  test("UpstashVectorStore with Upstash Embedding configuration, the embeddings will be created by Upstash's service", async () => {
+    const vectorStoreWithUpstashEmbeddings = new UpstashVectorStore(
+      new FakeEmbeddings(),
+      { index }
+    );
+
+    const createdAt = new Date().getTime();
+
+    const ids = await vectorStoreWithUpstashEmbeddings.addDocuments([
+      { pageContent: "hello", metadata: { a: createdAt + 1 } },
+      { pageContent: "car", metadata: { a: createdAt } },
+      { pageContent: "adjective", metadata: { a: createdAt } },
+      { pageContent: "hi", metadata: { a: createdAt } },
+    ]);
+
+    // Sleeping for a second to make sure that all the indexing operations are finished.
+    await sleep(1000);
+
+    const results1 =
+      await vectorStoreWithUpstashEmbeddings.similaritySearchVectorWithScore(
+        "hello!",
+        1
+      );
+    expect(results1).toHaveLength(1);
+
+    expect([results1[0][0]]).toEqual([
+      new Document({ metadata: { a: createdAt + 1 }, pageContent: "hello" }),
+    ]);
+
+    const results2 =
+      await vectorStoreWithUpstashEmbeddings.similaritySearchVectorWithScore(
+        "testing!",
+        6
+      );
+
+    expect(results2).toHaveLength(4);
+
+    await vectorStoreWithUpstashEmbeddings.delete({ ids: ids.slice(2) });
+
+    const results3 =
+      await vectorStoreWithUpstashEmbeddings.similaritySearchVectorWithScore(
+        "testing again!",
+        6
+      );
+
+    expect(results3).toHaveLength(2);
+  });
 });

libs/langchain-community/src/vectorstores/upstash.ts

@@ -1,9 +1,11 @@
 import * as uuid from "uuid";
 import { EmbeddingsInterface } from "@langchain/core/embeddings";
 import { VectorStore } from "@langchain/core/vectorstores";
-import { Index as UpstashIndex } from "@upstash/vector";
+import { Index as UpstashIndex, type QueryResult } from "@upstash/vector";
 import { Document, DocumentInterface } from "@langchain/core/documents";
 import { chunkArray } from "@langchain/core/utils/chunk_array";
+import { FakeEmbeddings } from "@langchain/core/utils/testing";
+
 import {
   AsyncCaller,
   AsyncCallerParams,
@@ -37,6 +39,7 @@ export type UpstashDeleteParams =
   | { deleteAll: boolean; ids?: never };
 
 const CONCURRENT_UPSERT_LIMIT = 1000;
+
 /**
  * The main class that extends the 'VectorStore' class. It provides
  * methods for interacting with Upstash index, such as adding documents,
@@ -49,7 +52,7 @@ export class UpstashVectorStore extends VectorStore {
 
   caller: AsyncCaller;
 
-  embeddings: EmbeddingsInterface;
+  useUpstashEmbeddings?: boolean;
 
   filter?: this["FilterType"];
 
@@ -59,8 +62,11 @@ export class UpstashVectorStore extends VectorStore {
 
   constructor(embeddings: EmbeddingsInterface, args: UpstashVectorLibArgs) {
     super(embeddings, args);
-
-    this.embeddings = embeddings;
+    // Special case where the embeddings instance is a FakeEmbeddings instance. In this case, we need to disable "instanceof" rule.
+    // eslint-disable-next-line no-instanceof/no-instanceof
+    if (embeddings instanceof FakeEmbeddings) {
+      this.useUpstashEmbeddings = true;
+    }
 
     const { index, ...asyncCallerArgs } = args;
 
@@ -78,10 +84,14 @@ export class UpstashVectorStore extends VectorStore {
    */
   async addDocuments(
     documents: DocumentInterface[],
-    options?: { ids?: string[] }
+    options?: { ids?: string[]; useUpstashEmbeddings?: boolean }
   ) {
     const texts = documents.map(({ pageContent }) => pageContent);
 
+    if (this.useUpstashEmbeddings || options?.useUpstashEmbeddings) {
+      return this._addData(documents, options);
+    }
+
     const embeddings = await this.embeddings.embedDocuments(texts);
 
     return this.addVectors(embeddings, documents, options);
@@ -128,6 +138,48 @@ export class UpstashVectorStore extends VectorStore {
     return documentIds;
   }
 
+  /**
+   * This method adds the provided documents to Upstash database. The pageContent of the documents will be embedded by Upstash Embeddings.
+   * @param documents Array of Document objects to be added to the Upstash database.
+   * @param options Optional object containing the array of ids for the documents.
+   * @returns Promise that resolves with the ids of the provided documents when the upsert operation is done.
+   */
+  protected async _addData(
+    documents: DocumentInterface[],
+    options?: { ids?: string[] }
+  ) {
+    const documentIds =
+      options?.ids ?? Array.from({ length: documents.length }, () => uuid.v4());
+
+    const upstashVectorsWithData = documents.map((document, index) => {
+      const metadata = {
+        _pageContentLC: documents[index].pageContent,
+        ...documents[index].metadata,
+      };
+
+      const id = documentIds[index];
+
+      return {
+        id,
+        data: document.pageContent,
+        metadata,
+      };
+    });
+
+    const vectorChunks = chunkArray(
+      upstashVectorsWithData,
+      CONCURRENT_UPSERT_LIMIT
+    );
+
+    const batchRequests = vectorChunks.map((chunk) =>
+      this.caller.call(async () => this.index.upsert(chunk))
+    );
+
+    await Promise.all(batchRequests);
+
+    return documentIds;
+  }
+
   /**
    * This method deletes documents from the Upstash database. You can either
    * provide the target ids, or delete all vectors in the database.
@@ -143,18 +195,30 @@ export class UpstashVectorStore extends VectorStore {
   }
 
   protected async _runUpstashQuery(
-    query: number[],
+    query: number[] | string,
     k: number,
     filter?: this["FilterType"],
     options?: { includeVectors: boolean }
   ) {
-    const queryResult = await this.index.query<UpstashQueryMetadata>({
-      vector: query,
-      topK: k,
-      includeMetadata: true,
-      filter,
-      ...options,
-    });
+    let queryResult: QueryResult<UpstashQueryMetadata>[] = [];
+
+    if (typeof query === "string") {
+      queryResult = await this.index.query<UpstashQueryMetadata>({
+        data: query,
+        topK: k,
+        includeMetadata: true,
+        filter,
+        ...options,
+      });
+    } else {
+      queryResult = await this.index.query<UpstashQueryMetadata>({
+        vector: query,
+        topK: k,
+        includeMetadata: true,
+        filter,
+        ...options,
+      });
+    }
 
     return queryResult;
   }
@@ -169,7 +233,7 @@ export class UpstashVectorStore extends VectorStore {
    *  maximum of 'k' and vectors in the index.
    */
   async similaritySearchVectorWithScore(
-    query: number[],
+    query: number[] | string,
     k: number,
     filter?: this["FilterType"]
   ): Promise<[DocumentInterface, number][]> {