Integrate Rockset as a vector store (langchain-ai#3231)

* Integrate Rockset as a vector store * address comments * Mark Rockset as requiring optional dep * Fix lint * Fix build * Format --------- Co-authored-by: jacoblee93 <[email protected]>
malandis · Nov 17, 2023 · d625ddd · d625ddd
1 parent 02f6662
commit d625ddd
Show file tree

Hide file tree

Showing 14 changed files with 675 additions and 2 deletions.
diff --git a/docs/api_refs/typedoc.json b/docs/api_refs/typedoc.json
@@ -120,6 +120,7 @@
     "./langchain/src/vectorstores/typeorm.ts",
     "./langchain/src/vectorstores/myscale.ts",
     "./langchain/src/vectorstores/redis.ts",
+    "./langchain/src/vectorstores/rockset.ts",
     "./langchain/src/vectorstores/typesense.ts",
     "./langchain/src/vectorstores/singlestore.ts",
     "./langchain/src/vectorstores/tigris.ts",

diff --git a/docs/core_docs/docs/integrations/vectorstores/rockset.mdx b/docs/core_docs/docs/integrations/vectorstores/rockset.mdx
@@ -0,0 +1,26 @@
+---
+sidebar_class_name: node-only
+---
+
+import CodeBlock from "@theme/CodeBlock";
+
+# Rockset
+
+[Rockset](https://rockset.com) is a real-time analyitics SQL database that runs in the cloud.
+Rockset provides vector search capabilities, in the form of [SQL functions](https://rockset.com/docs/vector-functions/#vector-distance-functions), to support AI applications that rely on text similarity.
+
+## Setup
+
+Install the rockset client.
+
+```bash
+yarn add @rockset/client
+```
+
+### Usage
+
+import UsageExample from "@examples/indexes/vector_stores/rockset.ts";
+
+Below is an example showcasing how to use OpenAI and Rockset to answer questions about a text file:
+
+<CodeBlock language="typescript">{UsageExample}</CodeBlock>
diff --git a/examples/package.json b/examples/package.json
@@ -32,6 +32,7 @@
     "@planetscale/database": "^1.8.0",
     "@prisma/client": "^4.11.0",
     "@raycast/api": "^1.55.2",
+    "@rockset/client": "^0.9.1",
     "@supabase/supabase-js": "^2.10.0",
     "@tensorflow/tfjs-backend-cpu": "^4.4.0",
     "@upstash/redis": "^1.20.6",

diff --git a/examples/src/indexes/vector_stores/rockset.ts b/examples/src/indexes/vector_stores/rockset.ts
@@ -0,0 +1,31 @@
+import * as rockset from "@rockset/client";
+import { ChatOpenAI } from "langchain/chat_models/openai";
+import { RetrievalQAChain } from "langchain/chains";
+import { OpenAIEmbeddings } from "langchain/embeddings/openai";
+import { RocksetStore } from "langchain/vectorstores/rockset";
+import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
+import { readFileSync } from "fs";
+
+export const run = async () => {
+  const store = await RocksetStore.withNewCollection(new OpenAIEmbeddings(), {
+    client: rockset.default.default(
+      process.env.ROCKSET_API_KEY ?? "",
+      `https://api.${process.env.ROCKSET_API_REGION ?? "usw2a1"}.rockset.com`
+    ),
+    collectionName: "langchain_demo",
+  });
+
+  const model = new ChatOpenAI({ modelName: "gpt-3.5-turbo" });
+  const chain = RetrievalQAChain.fromLLM(model, store.asRetriever());
+  const text = readFileSync("state_of_the_union.txt", "utf8");
+  const docs = await new RecursiveCharacterTextSplitter().createDocuments([
+    text,
+  ]);
+
+  await store.addDocuments(docs);
+  const response = await chain.call({
+    query: "What is America's role in Ukraine?",
+  });
+  console.log(response.text);
+  await store.destroy();
+};
diff --git a/langchain/.env.example b/langchain/.env.example
@@ -40,6 +40,9 @@ MYSCALE_USERNAME=ADD_YOURS_HERE
 MYSCALE_PASSWORD=ADD_YOURS_HERE
 FIGMA_ACCESS_TOKEN=ADD_YOURS_HERE
 REDIS_URL=ADD_YOURS_HERE
+ROCKSET_API_KEY=ADD_YOURS_HERE
+# defaults to "usw2a1" (oregon)
+ROCKSET_REGION=ADD_YOURS_HERE
 SINGLESTORE_HOST=ADD_YOURS_HERE
 SINGLESTORE_PORT=ADD_YOURS_HERE
 SINGLESTORE_USERNAME=ADD_YOURS_HERE

diff --git a/langchain/.gitignore b/langchain/.gitignore
@@ -304,6 +304,9 @@ vectorstores/myscale.d.ts
 vectorstores/redis.cjs
 vectorstores/redis.js
 vectorstores/redis.d.ts
+vectorstores/rockset.cjs
+vectorstores/rockset.js
+vectorstores/rockset.d.ts
 vectorstores/typesense.cjs
 vectorstores/typesense.js
 vectorstores/typesense.d.ts

diff --git a/langchain/package.json b/langchain/package.json
@@ -316,6 +316,9 @@
     "vectorstores/redis.cjs",
     "vectorstores/redis.js",
     "vectorstores/redis.d.ts",
+    "vectorstores/rockset.cjs",
+    "vectorstores/rockset.js",
+    "vectorstores/rockset.d.ts",
     "vectorstores/typesense.cjs",
     "vectorstores/typesense.js",
     "vectorstores/typesense.d.ts",
@@ -857,6 +860,7 @@
     "@planetscale/database": "^1.8.0",
     "@qdrant/js-client-rest": "^1.2.0",
     "@raycast/api": "^1.55.2",
+    "@rockset/client": "^0.9.1",
     "@smithy/eventstream-codec": "^2.0.5",
     "@smithy/protocol-http": "^3.0.6",
     "@smithy/signature-v4": "^2.0.10",
@@ -993,6 +997,7 @@
     "@planetscale/database": "^1.8.0",
     "@qdrant/js-client-rest": "^1.2.0",
     "@raycast/api": "^1.55.2",
+    "@rockset/client": "^0.9.1",
     "@smithy/eventstream-codec": "^2.0.5",
     "@smithy/protocol-http": "^3.0.6",
     "@smithy/signature-v4": "^2.0.10",
@@ -1148,6 +1153,9 @@
     "@raycast/api": {
       "optional": true
     },
+    "@rockset/client": {
+      "optional": true
+    },
     "@smithy/eventstream-codec": {
       "optional": true
     },
@@ -1907,6 +1915,11 @@
       "import": "./vectorstores/redis.js",
       "require": "./vectorstores/redis.cjs"
     },
+    "./vectorstores/rockset": {
+      "types": "./vectorstores/rockset.d.ts",
+      "import": "./vectorstores/rockset.js",
+      "require": "./vectorstores/rockset.cjs"
+    },
     "./vectorstores/typesense": {
       "types": "./vectorstores/typesense.d.ts",
       "import": "./vectorstores/typesense.js",

diff --git a/langchain/scripts/check-tree-shaking.js b/langchain/scripts/check-tree-shaking.js
@@ -31,6 +31,7 @@ export function listExternals() {
     "convex",
     "convex/server",
     "convex/values",
+    "@rockset/client/dist/codegen/api.js",
     "mysql2/promise",
     "pdf-parse/lib/pdf.js/v1.10.100/build/pdf.js",
     "@zilliz/milvus2-sdk-node/dist/milvus/const/Milvus.js",

diff --git a/langchain/scripts/create-entrypoints.js b/langchain/scripts/create-entrypoints.js
@@ -119,6 +119,7 @@ const entrypoints = {
   "vectorstores/typeorm": "vectorstores/typeorm",
   "vectorstores/myscale": "vectorstores/myscale",
   "vectorstores/redis": "vectorstores/redis",
+  "vectorstores/rockset": "vectorstores/rockset",
   "vectorstores/typesense": "vectorstores/typesense",
   "vectorstores/singlestore": "vectorstores/singlestore",
   "vectorstores/tigris": "vectorstores/tigris",
@@ -393,6 +394,7 @@ const requiresOptionalDependency = [
   "vectorstores/myscale",
   "vectorstores/neo4j_vector",
   "vectorstores/redis",
+  "vectorstores/rockset",
   "vectorstores/singlestore",
   "vectorstores/typesense",
   "vectorstores/tigris",

diff --git a/langchain/src/load/import_constants.ts b/langchain/src/load/import_constants.ts
@@ -65,6 +65,7 @@ export const optionalImportEntrypoints = [
   "langchain/vectorstores/typeorm",
   "langchain/vectorstores/myscale",
   "langchain/vectorstores/redis",
+  "langchain/vectorstores/rockset",
   "langchain/vectorstores/typesense",
   "langchain/vectorstores/singlestore",
   "langchain/vectorstores/tigris",

diff --git a/langchain/src/load/import_type.d.ts b/langchain/src/load/import_type.d.ts
@@ -193,6 +193,9 @@ export interface OptionalImportMap {
   "langchain/vectorstores/redis"?:
     | typeof import("../vectorstores/redis.js")
     | Promise<typeof import("../vectorstores/redis.js")>;
+  "langchain/vectorstores/rockset"?:
+    | typeof import("../vectorstores/rockset.js")
+    | Promise<typeof import("../vectorstores/rockset.js")>;
   "langchain/vectorstores/typesense"?:
     | typeof import("../vectorstores/typesense.js")
     | Promise<typeof import("../vectorstores/typesense.js")>;