{"version":3,"file":"embeddings.cjs","names":["Embeddings","GoogleGenerativeAI"],"sources":["../src/embeddings.ts"],"sourcesContent":["import { GoogleGenerativeAI, GenerativeModel } from \"@google/generative-ai\";\nimport type { TaskType, EmbedContentRequest } from \"@google/generative-ai\";\nimport { getEnvironmentVariable } from \"@langchain/core/utils/env\";\nimport { Embeddings, EmbeddingsParams } from \"@langchain/core/embeddings\";\nimport { chunkArray } from \"@langchain/core/utils/chunk_array\";\n\n/**\n * Interface that extends EmbeddingsParams and defines additional\n * parameters specific to the GoogleGenerativeAIEmbeddings class.\n */\nexport interface GoogleGenerativeAIEmbeddingsParams extends EmbeddingsParams {\n  /**\n   * Model Name to use\n   *\n   * Alias for `model`\n   *\n   * Note: The format must follow the pattern - `{model}`\n   */\n  modelName?: string;\n  /**\n   * Model Name to use\n   *\n   * Note: The format must follow the pattern - `{model}`\n   */\n  model?: string;\n\n  /**\n   * Type of task for which the embedding will be used\n   *\n   * Note: currently only supported by `embedding-001` model\n   */\n  taskType?: TaskType;\n\n  /**\n   * An optional title for the text. Only applicable when TaskType is\n   * `RETRIEVAL_DOCUMENT`\n   *\n   * Note: currently only supported by `embedding-001` model\n   */\n  title?: string;\n\n  /**\n   * Whether to strip new lines from the input text. Default to true\n   */\n  stripNewLines?: boolean;\n\n  /**\n   * Google API key to use\n   */\n  apiKey?: string;\n\n  /**\n   * Google API base URL to use\n   */\n  baseUrl?: string;\n}\n\n/**\n * Class that extends the Embeddings class and provides methods for\n * generating embeddings using the Google Palm API.\n * @example\n * ```typescript\n * const model = new GoogleGenerativeAIEmbeddings({\n *   apiKey: \"<YOUR API KEY>\",\n *   modelName: \"embedding-001\",\n * });\n *\n * // Embed a single query\n * const res = await model.embedQuery(\n *   \"What would be a good company name for a company that makes colorful socks?\"\n * );\n * console.log({ res });\n *\n * // Embed multiple documents\n * const documentRes = await model.embedDocuments([\"Hello world\", \"Bye bye\"]);\n * console.log({ documentRes });\n * ```\n */\nexport class GoogleGenerativeAIEmbeddings\n  extends Embeddings\n  implements GoogleGenerativeAIEmbeddingsParams\n{\n  apiKey?: string;\n\n  modelName = \"embedding-001\";\n\n  model = \"embedding-001\";\n\n  taskType?: TaskType;\n\n  title?: string;\n\n  stripNewLines = true;\n\n  maxBatchSize = 100; // Max batch size for embedDocuments set by GenerativeModel client's batchEmbedContents call\n\n  private client: GenerativeModel;\n\n  constructor(fields?: GoogleGenerativeAIEmbeddingsParams) {\n    super(fields ?? {});\n\n    this.modelName =\n      fields?.model?.replace(/^models\\//, \"\") ??\n      fields?.modelName?.replace(/^models\\//, \"\") ??\n      this.modelName;\n    this.model = this.modelName;\n\n    this.taskType = fields?.taskType ?? this.taskType;\n\n    this.title = fields?.title ?? this.title;\n\n    if (this.title && this.taskType !== \"RETRIEVAL_DOCUMENT\") {\n      throw new Error(\n        \"title can only be sepcified with TaskType.RETRIEVAL_DOCUMENT\"\n      );\n    }\n\n    this.apiKey = fields?.apiKey ?? getEnvironmentVariable(\"GOOGLE_API_KEY\");\n    if (!this.apiKey) {\n      throw new Error(\n        \"Please set an API key for Google GenerativeAI \" +\n          \"in the environmentb variable GOOGLE_API_KEY \" +\n          \"or in the `apiKey` field of the \" +\n          \"GoogleGenerativeAIEmbeddings constructor\"\n      );\n    }\n\n    this.client = new GoogleGenerativeAI(this.apiKey).getGenerativeModel(\n      {\n        model: this.model,\n      },\n      {\n        baseUrl: fields?.baseUrl,\n      }\n    );\n  }\n\n  private _convertToContent(text: string): EmbedContentRequest {\n    const cleanedText = this.stripNewLines ? text.replace(/\\n/g, \" \") : text;\n    return {\n      content: { role: \"user\", parts: [{ text: cleanedText }] },\n      taskType: this.taskType,\n      title: this.title,\n    };\n  }\n\n  protected async _embedQueryContent(text: string): Promise<number[]> {\n    const req = this._convertToContent(text);\n    const res = await this.client.embedContent(req);\n    return res.embedding.values ?? [];\n  }\n\n  protected async _embedDocumentsContent(\n    documents: string[]\n  ): Promise<number[][]> {\n    const batchEmbedChunks: string[][] = chunkArray<string>(\n      documents,\n      this.maxBatchSize\n    );\n\n    const batchEmbedRequests = batchEmbedChunks.map((chunk) => ({\n      requests: chunk.map((doc) => this._convertToContent(doc)),\n    }));\n\n    const responses = await Promise.allSettled(\n      batchEmbedRequests.map((req) => this.client.batchEmbedContents(req))\n    );\n\n    const embeddings = responses.flatMap((res, idx) => {\n      if (res.status === \"fulfilled\") {\n        return res.value.embeddings.map((e) => e.values || []);\n      } else {\n        return Array(batchEmbedChunks[idx].length).fill([]);\n      }\n    });\n\n    return embeddings;\n  }\n\n  /**\n   * Method that takes a document as input and returns a promise that\n   * resolves to an embedding for the document. It calls the _embedText\n   * method with the document as the input.\n   * @param document Document for which to generate an embedding.\n   * @returns Promise that resolves to an embedding for the input document.\n   */\n  embedQuery(document: string): Promise<number[]> {\n    return this.caller.call(this._embedQueryContent.bind(this), document);\n  }\n\n  /**\n   * Method that takes an array of documents as input and returns a promise\n   * that resolves to a 2D array of embeddings for each document. It calls\n   * the _embedText method for each document in the array.\n   * @param documents Array of documents for which to generate embeddings.\n   * @returns Promise that resolves to a 2D array of embeddings for each input document.\n   */\n  embedDocuments(documents: string[]): Promise<number[][]> {\n    return this.caller.call(this._embedDocumentsContent.bind(this), documents);\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;AA8EA,IAAa,+BAAb,cACUA,2BAAAA,WAEV;CACE;CAEA,YAAY;CAEZ,QAAQ;CAER;CAEA;CAEA,gBAAgB;CAEhB,eAAe;CAEf;CAEA,YAAY,QAA6C;AACvD,QAAM,UAAU,EAAE,CAAC;AAEnB,OAAK,YACH,QAAQ,OAAO,QAAQ,aAAa,GAAG,IACvC,QAAQ,WAAW,QAAQ,aAAa,GAAG,IAC3C,KAAK;AACP,OAAK,QAAQ,KAAK;AAElB,OAAK,WAAW,QAAQ,YAAY,KAAK;AAEzC,OAAK,QAAQ,QAAQ,SAAS,KAAK;AAEnC,MAAI,KAAK,SAAS,KAAK,aAAa,qBAClC,OAAM,IAAI,MACR,+DACD;AAGH,OAAK,SAAS,QAAQ,WAAA,GAAA,0BAAA,wBAAiC,iBAAiB;AACxE,MAAI,CAAC,KAAK,OACR,OAAM,IAAI,MACR,qKAID;AAGH,OAAK,SAAS,IAAIC,sBAAAA,mBAAmB,KAAK,OAAO,CAAC,mBAChD,EACE,OAAO,KAAK,OACb,EACD,EACE,SAAS,QAAQ,SAClB,CACF;;CAGH,kBAA0B,MAAmC;AAE3D,SAAO;GACL,SAAS;IAAE,MAAM;IAAQ,OAAO,CAAC,EAAE,MAFjB,KAAK,gBAAgB,KAAK,QAAQ,OAAO,IAAI,GAAG,MAEZ,CAAC;IAAE;GACzD,UAAU,KAAK;GACf,OAAO,KAAK;GACb;;CAGH,MAAgB,mBAAmB,MAAiC;EAClE,MAAM,MAAM,KAAK,kBAAkB,KAAK;AAExC,UADY,MAAM,KAAK,OAAO,aAAa,IAAI,EACpC,UAAU,UAAU,EAAE;;CAGnC,MAAgB,uBACd,WACqB;EACrB,MAAM,oBAAA,GAAA,kCAAA,YACJ,WACA,KAAK,aACN;EAED,MAAM,qBAAqB,iBAAiB,KAAK,WAAW,EAC1D,UAAU,MAAM,KAAK,QAAQ,KAAK,kBAAkB,IAAI,CAAC,EAC1D,EAAE;AAcH,UAZkB,MAAM,QAAQ,WAC9B,mBAAmB,KAAK,QAAQ,KAAK,OAAO,mBAAmB,IAAI,CAAC,CACrE,EAE4B,SAAS,KAAK,QAAQ;AACjD,OAAI,IAAI,WAAW,YACjB,QAAO,IAAI,MAAM,WAAW,KAAK,MAAM,EAAE,UAAU,EAAE,CAAC;OAEtD,QAAO,MAAM,iBAAiB,KAAK,OAAO,CAAC,KAAK,EAAE,CAAC;IAErD;;;;;;;;;CAYJ,WAAW,UAAqC;AAC9C,SAAO,KAAK,OAAO,KAAK,KAAK,mBAAmB,KAAK,KAAK,EAAE,SAAS;;;;;;;;;CAUvE,eAAe,WAA0C;AACvD,SAAO,KAAK,OAAO,KAAK,KAAK,uBAAuB,KAAK,KAAK,EAAE,UAAU"}