@langchain/community
Version:
Third-party integrations for LangChain.js
1 lines • 18.3 kB
Source Map (JSON)
{"version":3,"file":"rockset.cjs","names":["message: string","VectorStore","embeddings: EmbeddingsInterface","args: RocksetLibArgs","documents: Document[]","vectors: number[][]","ids: string[]","query: number[]","k: number","filter?: string","Document","rocksetDoc","texts: string[]","metadatas: object[] | object","dbConfig: RocksetLibArgs","docs: Document[]","Collection","waitUntilDeletion?: boolean","collectionOptions?: CreateCollectionRequest","metric: SimilarityMetric"],"sources":["../../src/vectorstores/rockset.ts"],"sourcesContent":["import { MainApi } from \"@rockset/client\";\nimport type { CreateCollectionRequest } from \"@rockset/client/dist/codegen/api.js\";\nimport { Collection } from \"@rockset/client/dist/codegen/api.js\";\n\nimport type { EmbeddingsInterface } from \"@langchain/core/embeddings\";\nimport { VectorStore } from \"@langchain/core/vectorstores\";\nimport { Document } from \"@langchain/core/documents\";\n/**\n * Generic Rockset vector storage error\n */\nexport class RocksetStoreError extends Error {\n /**\n * Constructs a RocksetStoreError\n * @param message The error message\n */\n constructor(message: string) {\n super(message);\n this.name = this.constructor.name;\n }\n}\n\n/**\n * Error that is thrown when a RocksetStore function is called\n * after `destroy()` is called (meaning the collection would be\n * deleted).\n */\nexport class RocksetStoreDestroyedError extends RocksetStoreError {\n constructor() {\n super(\"The Rockset store has been destroyed\");\n this.name = this.constructor.name;\n }\n}\n\n/**\n * Functions to measure vector distance/similarity by.\n * See https://rockset.com/docs/vector-functions/#vector-distance-functions\n * @enum SimilarityMetric\n */\nexport const SimilarityMetric = {\n CosineSimilarity: \"COSINE_SIM\",\n EuclideanDistance: \"EUCLIDEAN_DIST\",\n DotProduct: \"DOT_PRODUCT\",\n} as const;\n\nexport type SimilarityMetric =\n (typeof SimilarityMetric)[keyof typeof SimilarityMetric];\n\ninterface CollectionNotFoundError {\n message_key: string;\n}\n\n/**\n * Vector store arguments\n * @interface RocksetStore\n */\nexport interface RocksetLibArgs {\n /**\n * The rockset client object constructed with `rocksetConfigure`\n * @type {MainAPI}\n */\n client: MainApi;\n /**\n * The name of the Rockset collection to store vectors\n * @type {string}\n */\n collectionName: string;\n /**\n * The name of othe Rockset workspace that holds @member collectionName\n * @type {string}\n */\n workspaceName?: string;\n /**\n * The name of the collection column to contain page contnent of documents\n * @type {string}\n */\n textKey?: string;\n /**\n * The name of the collection column to contain vectors\n * @type {string}\n */\n embeddingKey?: string;\n /**\n * The SQL `WHERE` clause to filter by\n * @type {string}\n */\n filter?: string;\n /**\n * The metric used to measure vector relationship\n * @type {SimilarityMetric}\n */\n similarityMetric?: SimilarityMetric;\n}\n\n/**\n * Exposes Rockset's vector store/search functionality\n */\nexport class RocksetStore extends VectorStore {\n declare FilterType: string;\n\n client: MainApi;\n\n collectionName: string;\n\n workspaceName: string;\n\n textKey: string;\n\n embeddingKey: string;\n\n filter?: string;\n\n private _similarityMetric: SimilarityMetric;\n\n private similarityOrder: \"ASC\" | \"DESC\";\n\n private destroyed: boolean;\n\n /**\n * Gets a string representation of the type of this VectorStore\n * @returns {\"rockset\"}\n */\n _vectorstoreType(): \"rockset\" {\n return \"rockset\";\n }\n\n /**\n * Constructs a new RocksetStore\n * @param {Embeddings} embeddings Object used to embed queries and\n * page content\n * @param {RocksetLibArgs} args\n */\n constructor(embeddings: EmbeddingsInterface, args: RocksetLibArgs) {\n super(embeddings, args);\n\n this.embeddings = embeddings;\n this.client = args.client;\n this.collectionName = args.collectionName;\n this.workspaceName = args.workspaceName ?? \"commons\";\n this.textKey = args.textKey ?? \"text\";\n this.embeddingKey = args.embeddingKey ?? \"embedding\";\n this.filter = args.filter;\n this.similarityMetric =\n args.similarityMetric ?? SimilarityMetric.CosineSimilarity;\n this.setSimilarityOrder();\n }\n\n /**\n * Sets the object's similarity order based on what\n * SimilarityMetric is being used\n */\n private setSimilarityOrder() {\n this.checkIfDestroyed();\n this.similarityOrder =\n this.similarityMetric === SimilarityMetric.EuclideanDistance\n ? \"ASC\"\n : \"DESC\";\n }\n\n /**\n * Embeds and adds Documents to the store.\n * @param {Documents[]} documents The documents to store\n * @returns {Promise<string[]?>} The _id's of the documents added\n */\n async addDocuments(documents: Document[]): Promise<string[] | undefined> {\n const texts = documents.map(({ pageContent }) => pageContent);\n return await this.addVectors(\n await this.embeddings.embedDocuments(texts),\n documents\n );\n }\n\n /**\n * Adds vectors to the store given their corresponding Documents\n * @param {number[][]} vectors The vectors to store\n * @param {Document[]} documents The Documents they represent\n * @return {Promise<string[]?>} The _id's of the added documents\n */\n async addVectors(vectors: number[][], documents: Document[]) {\n this.checkIfDestroyed();\n const rocksetDocs = [];\n for (let i = 0; i < documents.length; i += 1) {\n const currDoc = documents[i];\n const currVector = vectors[i];\n rocksetDocs.push({\n [this.textKey]: currDoc.pageContent,\n [this.embeddingKey]: currVector,\n ...currDoc.metadata,\n });\n }\n\n return (\n await this.client.documents.addDocuments(\n this.workspaceName,\n this.collectionName,\n {\n data: rocksetDocs,\n }\n )\n ).data?.map((docStatus) => docStatus._id || \"\");\n }\n\n /**\n * Deletes Rockset documements given their _id's\n * @param {string[]} ids The IDS to remove documents with\n */\n async delete(ids: string[]): Promise<void> {\n this.checkIfDestroyed();\n await this.client.documents.deleteDocuments(\n this.workspaceName,\n this.collectionName,\n {\n data: ids.map((id) => ({ _id: id })),\n }\n );\n }\n\n /**\n * Gets the most relevant documents to a query along\n * with their similarity score. The returned documents\n * are ordered by similarity (most similar at the first\n * index)\n * @param {number[]} query The embedded query to search\n * the store by\n * @param {number} k The number of documents to retreive\n * @param {string?} filter The SQL `WHERE` clause to filter by\n */\n async similaritySearchVectorWithScore(\n query: number[],\n k: number,\n filter?: string\n ): Promise<[Document, number][]> {\n this.checkIfDestroyed();\n if (filter && this.filter) {\n throw new RocksetStoreError(\n \"cannot provide both `filter` and `this.filter`\"\n );\n }\n const similarityKey = \"similarity\";\n const _filter = filter ?? this.filter;\n return (\n (\n await this.client.queries.query({\n sql: {\n query: `\n SELECT\n * EXCEPT(\"${this.embeddingKey}\"),\n \"${this.textKey}\",\n ${this.similarityMetric}(:query, \"${\n this.embeddingKey\n }\") AS \"${similarityKey}\"\n FROM \n \"${this.workspaceName}\".\"${this.collectionName}\"\n ${_filter ? `WHERE ${_filter}` : \"\"}\n ORDER BY\n \"${similarityKey}\" ${this.similarityOrder}\n LIMIT\n ${k}\n `,\n parameters: [\n {\n name: \"query\",\n type: \"\",\n value: `[${query.toString()}]`,\n },\n ],\n },\n })\n ).results?.map((rocksetDoc) => [\n new Document<Record<string, object>>({\n pageContent: rocksetDoc[this.textKey],\n metadata: (({\n [this.textKey]: t,\n [similarityKey]: s,\n ...rocksetDoc\n }) => rocksetDoc)(rocksetDoc),\n }),\n rocksetDoc[similarityKey] as number,\n ]) ?? []\n );\n }\n\n /**\n * Constructs and returns a RocksetStore object given texts to store.\n * @param {string[]} texts The texts to store\n * @param {object[] | object} metadatas The metadatas that correspond\n * to @param texts\n * @param {Embeddings} embeddings The object used to embed queries\n * and page content\n * @param {RocksetLibArgs} dbConfig The options to be passed into the\n * RocksetStore constructor\n * @returns {RocksetStore}\n */\n static async fromTexts(\n texts: string[],\n metadatas: object[] | object,\n embeddings: EmbeddingsInterface,\n dbConfig: RocksetLibArgs\n ): Promise<RocksetStore> {\n const docs: Document[] = [];\n for (let i = 0; i < texts.length; i += 1) {\n const metadata = Array.isArray(metadatas) ? metadatas[i] : metadatas;\n const newDoc = new Document({\n pageContent: texts[i],\n metadata,\n });\n docs.push(newDoc);\n }\n\n return RocksetStore.fromDocuments(docs, embeddings, dbConfig);\n }\n\n /**\n * Constructs, adds docs to, and returns a RocksetStore object\n * @param {Document[]} docs The Documents to store\n * @param {Embeddings} embeddings The object used to embed queries\n * and page content\n * @param {RocksetLibArgs} dbConfig The options to be passed into the\n * RocksetStore constructor\n * @returns {RocksetStore}\n */\n static async fromDocuments(\n docs: Document[],\n embeddings: EmbeddingsInterface,\n dbConfig: RocksetLibArgs\n ): Promise<RocksetStore> {\n const args = { ...dbConfig, textKey: dbConfig.textKey ?? \"text\" };\n const instance = new this(embeddings, args);\n await instance.addDocuments(docs);\n return instance;\n }\n\n /**\n * Checks if a Rockset collection exists.\n * @param {RocksetLibArgs} dbConfig The object containing the collection\n * and workspace names\n * @return {boolean} whether the collection exists\n */\n private static async collectionExists(dbConfig: RocksetLibArgs) {\n try {\n await dbConfig.client.collections.getCollection(\n dbConfig.workspaceName ?? \"commons\",\n dbConfig.collectionName\n );\n } catch (err) {\n if (\n (err as CollectionNotFoundError).message_key ===\n \"COLLECTION_DOES_NOT_EXIST\"\n ) {\n return false;\n }\n throw err;\n }\n return true;\n }\n\n /**\n * Checks whether a Rockset collection is ready to be queried.\n * @param {RocksetLibArgs} dbConfig The object containing the collection\n * name and workspace\n * @return {boolean} whether the collection is ready\n */\n private static async collectionReady(dbConfig: RocksetLibArgs) {\n return (\n (\n await dbConfig.client.collections.getCollection(\n dbConfig.workspaceName ?? \"commons\",\n dbConfig.collectionName\n )\n ).data?.status === Collection.StatusEnum.READY\n );\n }\n\n /**\n * Deletes the collection this RocksetStore uses\n * @param {boolean?} waitUntilDeletion Whether to sleep until the\n * collection is ready to be\n * queried\n */\n async destroy(waitUntilDeletion?: boolean) {\n await this.client.collections.deleteCollection(\n this.workspaceName,\n this.collectionName\n );\n this.destroyed = true;\n if (waitUntilDeletion) {\n while (\n await RocksetStore.collectionExists({\n collectionName: this.collectionName,\n client: this.client,\n })\n );\n }\n }\n\n /**\n * Checks if this RocksetStore has been destroyed.\n * @throws {RocksetStoreDestroyederror} if it has.\n */\n private checkIfDestroyed() {\n if (this.destroyed) {\n throw new RocksetStoreDestroyedError();\n }\n }\n\n /**\n * Creates a new Rockset collection and returns a RocksetStore that\n * uses it\n * @param {Embeddings} embeddings Object used to embed queries and\n * page content\n * @param {RocksetLibArgs} dbConfig The options to be passed into the\n * RocksetStore constructor\n * @param {CreateCollectionRequest?} collectionOptions The arguments to sent with the\n * HTTP request when creating the\n * collection. Setting a field mapping\n * that `VECTOR_ENFORCE`s is recommended\n * when using this function. See\n * https://rockset.com/docs/vector-functions/#vector_enforce\n * @returns {RocsketStore}\n */\n static async withNewCollection(\n embeddings: EmbeddingsInterface,\n dbConfig: RocksetLibArgs,\n collectionOptions?: CreateCollectionRequest\n ): Promise<RocksetStore> {\n if (\n collectionOptions?.name &&\n dbConfig.collectionName !== collectionOptions?.name\n ) {\n throw new RocksetStoreError(\n \"`dbConfig.name` and `collectionOptions.name` do not match\"\n );\n }\n await dbConfig.client.collections.createCollection(\n dbConfig.workspaceName ?? \"commons\",\n collectionOptions || { name: dbConfig.collectionName }\n );\n while (\n !(await this.collectionExists(dbConfig)) ||\n !(await this.collectionReady(dbConfig))\n );\n return new this(embeddings, dbConfig);\n }\n\n public get similarityMetric() {\n return this._similarityMetric;\n }\n\n public set similarityMetric(metric: SimilarityMetric) {\n this._similarityMetric = metric;\n this.setSimilarityOrder();\n }\n}\n"],"mappings":";;;;;;;;;;;;;;;;AAUA,IAAa,oBAAb,cAAuC,MAAM;;;;;CAK3C,YAAYA,SAAiB;EAC3B,MAAM,QAAQ;EACd,KAAK,OAAO,KAAK,YAAY;CAC9B;AACF;;;;;;AAOD,IAAa,6BAAb,cAAgD,kBAAkB;CAChE,cAAc;EACZ,MAAM,uCAAuC;EAC7C,KAAK,OAAO,KAAK,YAAY;CAC9B;AACF;;;;;;AAOD,MAAa,mBAAmB;CAC9B,kBAAkB;CAClB,mBAAmB;CACnB,YAAY;AACb;;;;AAsDD,IAAa,eAAb,MAAa,qBAAqBC,0CAAY;CAG5C;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA,AAAQ;CAER,AAAQ;CAER,AAAQ;;;;;CAMR,mBAA8B;AAC5B,SAAO;CACR;;;;;;;CAQD,YAAYC,YAAiCC,MAAsB;EACjE,MAAM,YAAY,KAAK;EAEvB,KAAK,aAAa;EAClB,KAAK,SAAS,KAAK;EACnB,KAAK,iBAAiB,KAAK;EAC3B,KAAK,gBAAgB,KAAK,iBAAiB;EAC3C,KAAK,UAAU,KAAK,WAAW;EAC/B,KAAK,eAAe,KAAK,gBAAgB;EACzC,KAAK,SAAS,KAAK;EACnB,KAAK,mBACH,KAAK,oBAAoB,iBAAiB;EAC5C,KAAK,oBAAoB;CAC1B;;;;;CAMD,AAAQ,qBAAqB;EAC3B,KAAK,kBAAkB;EACvB,KAAK,kBACH,KAAK,qBAAqB,iBAAiB,oBACvC,QACA;CACP;;;;;;CAOD,MAAM,aAAaC,WAAsD;EACvE,MAAM,QAAQ,UAAU,IAAI,CAAC,EAAE,aAAa,KAAK,YAAY;AAC7D,SAAO,MAAM,KAAK,WAChB,MAAM,KAAK,WAAW,eAAe,MAAM,EAC3C,UACD;CACF;;;;;;;CAQD,MAAM,WAAWC,SAAqBD,WAAuB;EAC3D,KAAK,kBAAkB;EACvB,MAAM,cAAc,CAAE;AACtB,OAAK,IAAI,IAAI,GAAG,IAAI,UAAU,QAAQ,KAAK,GAAG;GAC5C,MAAM,UAAU,UAAU;GAC1B,MAAM,aAAa,QAAQ;GAC3B,YAAY,KAAK;KACd,KAAK,UAAU,QAAQ;KACvB,KAAK,eAAe;IACrB,GAAG,QAAQ;GACZ,EAAC;EACH;AAED,UACE,MAAM,KAAK,OAAO,UAAU,aAC1B,KAAK,eACL,KAAK,gBACL,EACE,MAAM,YACP,EACF,EACD,MAAM,IAAI,CAAC,cAAc,UAAU,OAAO,GAAG;CAChD;;;;;CAMD,MAAM,OAAOE,KAA8B;EACzC,KAAK,kBAAkB;EACvB,MAAM,KAAK,OAAO,UAAU,gBAC1B,KAAK,eACL,KAAK,gBACL,EACE,MAAM,IAAI,IAAI,CAAC,QAAQ,EAAE,KAAK,GAAI,GAAE,CACrC,EACF;CACF;;;;;;;;;;;CAYD,MAAM,gCACJC,OACAC,GACAC,QAC+B;EAC/B,KAAK,kBAAkB;AACvB,MAAI,UAAU,KAAK,OACjB,OAAM,IAAI,kBACR;EAGJ,MAAM,gBAAgB;EACtB,MAAM,UAAU,UAAU,KAAK;AAC/B,UAEI,MAAM,KAAK,OAAO,QAAQ,MAAM,EAC9B,KAAK;GACH,OAAO,CAAC;;sBAEE,EAAE,KAAK,aAAa;aAC7B,EAAE,KAAK,QAAQ;YAChB,EAAE,KAAK,iBAAiB,UAAU,EAChC,KAAK,aACN,OAAO,EAAE,cAAc;;aAEvB,EAAE,KAAK,cAAc,GAAG,EAAE,KAAK,eAAe;UACjD,EAAE,UAAU,CAAC,MAAM,EAAE,SAAS,GAAG,GAAG;;aAEjC,EAAE,cAAc,EAAE,EAAE,KAAK,gBAAgB;;YAE1C,EAAE,EAAE;QACR,CAAC;GACG,YAAY,CACV;IACE,MAAM;IACN,MAAM;IACN,OAAO,CAAC,CAAC,EAAE,MAAM,UAAU,CAAC,CAAC,CAAC;GAC/B,CACF;EACF,EACF,EAAC,EACF,SAAS,IAAI,CAAC,eAAe,CAC7B,IAAIC,oCAAiC;GACnC,aAAa,WAAW,KAAK;GAC7B,WAAW,CAAC,EACV,CAAC,KAAK,UAAU,GAChB,CAAC,gBAAgB,EACjB,GAAGC,cACJ,KAAKA,cAAY,WAAW;EAC9B,IACD,WAAW,cACZ,EAAC,IAAI,CAAE;CAEX;;;;;;;;;;;;CAaD,aAAa,UACXC,OACAC,WACAX,YACAY,UACuB;EACvB,MAAMC,OAAmB,CAAE;AAC3B,OAAK,IAAI,IAAI,GAAG,IAAI,MAAM,QAAQ,KAAK,GAAG;GACxC,MAAM,WAAW,MAAM,QAAQ,UAAU,GAAG,UAAU,KAAK;GAC3D,MAAM,SAAS,IAAIL,oCAAS;IAC1B,aAAa,MAAM;IACnB;GACD;GACD,KAAK,KAAK,OAAO;EAClB;AAED,SAAO,aAAa,cAAc,MAAM,YAAY,SAAS;CAC9D;;;;;;;;;;CAWD,aAAa,cACXK,MACAb,YACAY,UACuB;EACvB,MAAM,OAAO;GAAE,GAAG;GAAU,SAAS,SAAS,WAAW;EAAQ;EACjE,MAAM,WAAW,IAAI,KAAK,YAAY;EACtC,MAAM,SAAS,aAAa,KAAK;AACjC,SAAO;CACR;;;;;;;CAQD,aAAqB,iBAAiBA,UAA0B;AAC9D,MAAI;GACF,MAAM,SAAS,OAAO,YAAY,cAChC,SAAS,iBAAiB,WAC1B,SAAS,eACV;EACF,SAAQ,KAAK;AACZ,OACG,IAAgC,gBACjC,4BAEA,QAAO;AAET,SAAM;EACP;AACD,SAAO;CACR;;;;;;;CAQD,aAAqB,gBAAgBA,UAA0B;AAC7D,UAEI,MAAM,SAAS,OAAO,YAAY,cAChC,SAAS,iBAAiB,WAC1B,SAAS,eACV,EACD,MAAM,WAAWE,gDAAW,WAAW;CAE5C;;;;;;;CAQD,MAAM,QAAQC,mBAA6B;EACzC,MAAM,KAAK,OAAO,YAAY,iBAC5B,KAAK,eACL,KAAK,eACN;EACD,KAAK,YAAY;AACjB,MAAI,kBACF,QACE,MAAM,aAAa,iBAAiB;GAClC,gBAAgB,KAAK;GACrB,QAAQ,KAAK;EACd,EAAC;CAGP;;;;;CAMD,AAAQ,mBAAmB;AACzB,MAAI,KAAK,UACP,OAAM,IAAI;CAEb;;;;;;;;;;;;;;;;CAiBD,aAAa,kBACXf,YACAY,UACAI,mBACuB;AACvB,MACE,mBAAmB,QACnB,SAAS,mBAAmB,mBAAmB,KAE/C,OAAM,IAAI,kBACR;EAGJ,MAAM,SAAS,OAAO,YAAY,iBAChC,SAAS,iBAAiB,WAC1B,qBAAqB,EAAE,MAAM,SAAS,eAAgB,EACvD;AACD,SACE,CAAE,MAAM,KAAK,iBAAiB,SAAS,IACvC,CAAE,MAAM,KAAK,gBAAgB,SAAS;AAExC,SAAO,IAAI,KAAK,YAAY;CAC7B;CAED,IAAW,mBAAmB;AAC5B,SAAO,KAAK;CACb;CAED,IAAW,iBAAiBC,QAA0B;EACpD,KAAK,oBAAoB;EACzB,KAAK,oBAAoB;CAC1B;AACF"}