UNPKG

@langchain/community

Version:
84 lines (83 loc) 3.65 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.GoogleVertexAIEmbeddings = void 0; const google_auth_library_1 = require("google-auth-library"); const embeddings_1 = require("@langchain/core/embeddings"); const chunk_array_1 = require("@langchain/core/utils/chunk_array"); const googlevertexai_connection_js_1 = require("../utils/googlevertexai-connection.cjs"); /** * Enables calls to the Google Cloud's Vertex AI API to access * the embeddings generated by Large Language Models. * * To use, you will need to have one of the following authentication * methods in place: * - You are logged into an account permitted to the Google Cloud project * using Vertex AI. * - You are running this on a machine using a service account permitted to * the Google Cloud project using Vertex AI. * - The `GOOGLE_APPLICATION_CREDENTIALS` environment variable is set to the * path of a credentials file for a service account permitted to the * Google Cloud project using Vertex AI. * @example * ```typescript * const model = new GoogleVertexAIEmbeddings(); * const res = await model.embedQuery( * "What would be a good company name for a company that makes colorful socks?" * ); * console.log({ res }); * ``` */ class GoogleVertexAIEmbeddings extends embeddings_1.Embeddings { constructor(fields) { super(fields ?? {}); Object.defineProperty(this, "model", { enumerable: true, configurable: true, writable: true, value: "textembedding-gecko" }); Object.defineProperty(this, "connection", { enumerable: true, configurable: true, writable: true, value: void 0 }); this.model = fields?.model ?? this.model; this.connection = new googlevertexai_connection_js_1.GoogleVertexAILLMConnection({ ...fields, ...this }, this.caller, new google_auth_library_1.GoogleAuth({ scopes: "https://www.googleapis.com/auth/cloud-platform", ...fields?.authOptions, })); } /** * Takes an array of documents as input and returns a promise that * resolves to a 2D array of embeddings for each document. It splits the * documents into chunks and makes requests to the Google Vertex AI API to * generate embeddings. * @param documents An array of documents to be embedded. * @returns A promise that resolves to a 2D array of embeddings for each document. */ async embedDocuments(documents) { const instanceChunks = (0, chunk_array_1.chunkArray)(documents.map((document) => ({ content: document, })), 5); // Vertex AI accepts max 5 instances per prediction const parameters = {}; const options = {}; const responses = await Promise.all(instanceChunks.map((instances) => this.connection.request(instances, parameters, options))); const result = responses ?.map((response) => response?.data?.predictions?.map((result) => result.embeddings.values) ?? []) .flat() ?? []; return result; } /** * Takes a document as input and returns a promise that resolves to an * embedding for the document. It calls the embedDocuments method with the * document as the input. * @param document A document to be embedded. * @returns A promise that resolves to an embedding for the document. */ async embedQuery(document) { const data = await this.embedDocuments([document]); return data[0]; } } exports.GoogleVertexAIEmbeddings = GoogleVertexAIEmbeddings;