Skip to main content

Rockset

Rockset is a real-time analyitics SQL database that runs in the cloud. Rockset provides vector search capabilities, in the form of SQL functions, to support AI applications that rely on text similarity.

Setup

Install the rockset client.

yarn add @rockset/client

Usage

npm install @langchain/openai @langchain/community

Below is an example showcasing how to use OpenAI and Rockset to answer questions about a text file:

import * as rockset from "@rockset/client";
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { RocksetStore } from "@langchain/community/vectorstores/rockset";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
import { readFileSync } from "fs";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { createStuffDocumentsChain } from "langchain/chains/combine_documents";
import { createRetrievalChain } from "langchain/chains/retrieval";

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({ model: "gpt-3.5-turbo-1106" });
const questionAnsweringPrompt = ChatPromptTemplate.fromMessages([
[
"system",
"Answer the user's questions based on the below context:\n\n{context}",
],
["human", "{input}"],
]);

const combineDocsChain = await createStuffDocumentsChain({
llm: model,
prompt: questionAnsweringPrompt,
});

const chain = await createRetrievalChain({
retriever: store.asRetriever(),
combineDocsChain,
});

const text = readFileSync("state_of_the_union.txt", "utf8");
const docs = await new RecursiveCharacterTextSplitter().createDocuments([text]);

await store.addDocuments(docs);
const response = await chain.invoke({
input: "When was America founded?",
});
console.log(response.answer);
await store.destroy();

API Reference:


Help us out by providing feedback on this documentation page: