{"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"}