Add vectorization via Vertex AI (#4311) * Add vectorization via Vertex AI * Enable batching for Google vectors * Add batch embeddings for Vertex * Split embed methods for Vertex/AI Studio

82d9fa79e830eaac87f73e80fe499564557905ac

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
Parents: c094e5d8efaa3d4595e80a1c9bb360e6b6f6f80c
5 files changed, +127 -66Showing whitespace changes
public/scripts/extensions/vectors/index.js+12 -2
@@ -36,6 +36,7 @@ import { slashCommandReturnHelper } from '../../slash-commands/SlashCommandRetur
3636import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js';
3737import { WebLlmVectorProvider } from './webllm.js';
3838import { removeReasoningFromString } from '../../reasoning.js';
39+import { oai_settings } from '../../openai.js';
3940
4041/**
4142 * @typedef {object} HashedMessage
@@ -50,7 +51,7 @@ export const EXTENSION_PROMPT_TAG = '3_vectors';
5051export const EXTENSION_PROMPT_TAG_DB = '4_vectors_data_bank';
5152
5253// Force solo chunks for sources that don't support batching.
5354const getBatchSize = () => ['transformers', 'palm', 'ollama'].includes(settings.source) ? 1 : 5;
5455
5556const settings = {
5657 // For both
@@ -796,6 +797,14 @@ function getVectorsRequestBody(args = {}) {
796797 break;
797798 case 'palm':
798799 body.model = extension_settings.vectors.google_model;
800+ body.api = 'makersuite';
801+ break;
802+ case 'vertexai':
803+ body.model = extension_settings.vectors.google_model;
804+ body.api = 'vertexai';
805+ body.vertexai_auth_mode = oai_settings.vertexai_auth_mode;
806+ body.vertexai_region = oai_settings.vertexai_region;
807+ body.vertexai_express_project_id = oai_settings.vertexai_express_project_id;
799808 break;
800809 default:
801810 break;
@@ -881,6 +890,7 @@ async function insertVectorItems(collectionId, items) {
881890function throwIfSourceInvalid() {
882891 if (settings.source === 'openai' && !secret_state[SECRET_KEYS.OPENAI] ||
883892 settings.source === 'palm' && !secret_state[SECRET_KEYS.MAKERSUITE] ||
893+ settings.source === 'vertexai' && !secret_state[SECRET_KEYS.VERTEXAI] && !secret_state[SECRET_KEYS.VERTEXAI_SERVICE_ACCOUNT] ||
884894 settings.source === 'mistral' && !secret_state[SECRET_KEYS.MISTRALAI] ||
885895 settings.source === 'togetherai' && !secret_state[SECRET_KEYS.TOGETHERAI] ||
886896 settings.source === 'nomicai' && !secret_state[SECRET_KEYS.NOMICAI] ||
@@ -1098,7 +1108,7 @@ function toggleSettings() {
10981108 $('#nomicai_apiKey').toggle(settings.source === 'nomicai');
10991109 $('#webllm_vectorsModel').toggle(settings.source === 'webllm');
11001110 $('#koboldcpp_vectorsModel').toggle(settings.source === 'koboldcpp');
11011111 $('#google_vectorsModel').toggle(settings.source === 'palm' || settings.source === 'vertexai');
11021112 $('#vector_altEndpointUrl').toggle(vectorApiRequiresUrl.includes(settings.source));
11031113 if (settings.source === 'webllm') {
11041114 loadWebLlmModels();
public/scripts/extensions/vectors/settings.html+2 -0
@@ -13,6 +13,7 @@
1313 <option value="cohere">Cohere</option>
1414 <option value="extras">Extras (deprecated)</option>
1515 <option value="palm">Google AI Studio</option>
16+ <option value="vertexai">Google Vertex AI</option>
1617 <option value="koboldcpp">KoboldCpp</option>
1718 <option value="llamacpp">llama.cpp</option>
1819 <option value="transformers" data-i18n="Local (Transformers)">Local (Transformers)</option>
@@ -136,6 +137,7 @@
136137 Vectorization Model
137138 </label>
138139 <select id="vectors_google_model" class="text_pole">
140+ <option value="gemini-embedding-001">gemini-embedding-001</option>
139141 <option value="gemini-embedding-exp-03-07">gemini-embedding-exp-03-07</option>
140142 <option value="text-embedding-004">text-embedding-004</option>
141143 <option value="embedding-001">embedding-001</option>
src/endpoints/vectors.js+12 -3
@@ -11,7 +11,8 @@ import { getNomicAIBatchVector, getNomicAIVector } from '../vectors/nomicai-vect
1111import { getOpenAIVector, getOpenAIBatchVector } from '../vectors/openai-vectors.js';
1212import { getTransformersVector, getTransformersBatchVector } from '../vectors/embedding.js';
1313import { getExtrasVector, getExtrasBatchVector } from '../vectors/extras-vectors.js';
1414import { getMakerSuiteVector, getMakerSuiteBatchVector } from '../vectors/makersuitegoogle-vectors.js';
15+import { getVertexVector, getVertexBatchVector } from '../vectors/google-vectors.js';
1516import { getCohereVector, getCohereBatchVector } from '../vectors/cohere-vectors.js';
1617import { getLlamaCppVector, getLlamaCppBatchVector } from '../vectors/llamacpp-vectors.js';
1718import { getVllmVector, getVllmBatchVector } from '../vectors/vllm-vectors.js';
@@ -32,6 +33,7 @@ const SOURCES = [
3233 'vllm',
3334 'webllm',
3435 'koboldcpp',
36+ 'vertexai',
3537];
3638
3739/**
@@ -56,7 +58,9 @@ async function getVector(source, sourceSettings, text, isQuery, directories) {
5658 case 'extras':
5759 return getExtrasVector(text, sourceSettings.extrasUrl, sourceSettings.extrasKey);
5860 case 'palm':
5961 return getMakerSuiteVector(text, directoriessourceSettings.model, sourceSettings.modelrequest);
62+ case 'vertexai':
63+ return getVertexVector(text, sourceSettings.model, sourceSettings.request);
6064 case 'cohere':
6165 return getCohereVector(text, isQuery, directories, sourceSettings.model);
6266 case 'llamacpp':
@@ -105,7 +109,10 @@ async function getBatchVector(source, sourceSettings, texts, isQuery, directorie
105109 results.push(...await getExtrasBatchVector(batch, sourceSettings.extrasUrl, sourceSettings.extrasKey));
106110 break;
107111 case 'palm':
108112 results.push(...await getMakerSuiteBatchVector(batch, directoriessourceSettings.model, sourceSettings.modelrequest));
113+ break;
114+ case 'vertexai':
115+ results.push(...await getVertexBatchVector(batch, sourceSettings.model, sourceSettings.request));
109116 break;
110117 case 'cohere':
111118 results.push(...await getCohereBatchVector(batch, isQuery, directories, sourceSettings.model));
@@ -178,8 +185,10 @@ function getSourceSettings(source, request) {
178185 model: getConfigValue('extensions.models.embedding', ''),
179186 };
180187 case 'palm':
188+ case 'vertexai':
181189 return {
182190 model: String(request.body.model || 'text-embedding-004'),
191+ request: request, // Pass the request object to get API key and URL
183192 };
184193 case 'mistral':
185194 return {
src/vectors/google-vectors.js+101 -0
@@ -0,0 +1,101 @@
1+import fetch from 'node-fetch';
2+import { getGoogleApiConfig } from '../endpoints/google.js';
3+
4+/**
5+ * Gets the vector for the given text from Google AI Studio
6+ * @param {string[]} texts - The array of texts to get the vector for
7+ * @param {string} model - The model to use for embedding
8+ * @param {import('express').Request} request - The request object to get API key and URL
9+ * @returns {Promise<number[][]>} - The array of vectors for the texts
10+ */
11+export async function getMakerSuiteBatchVector(texts, model, request) {
12+ const { url, headers, apiName } = await getGoogleApiConfig(request, model, 'batchEmbedContents');
13+
14+ const body = {
15+ requests: texts.map(text => ({
16+ model: `models/${model}`,
17+ content: { parts: [{ text }] },
18+ })),
19+ };
20+
21+ const response = await fetch(url, {
22+ body: JSON.stringify(body),
23+ method: 'POST',
24+ headers: headers,
25+ });
26+
27+ if (!response.ok) {
28+ const text = await response.text();
29+ console.warn(`${apiName} batch request failed`, response.statusText, text);
30+ throw new Error(`${apiName} batch request failed`);
31+ }
32+
33+ /** @type {any} */
34+ const data = await response.json();
35+ if (!Array.isArray(data?.embeddings)) {
36+ throw new Error(`${apiName} did not return an array`);
37+ }
38+
39+ const embeddings = data.embeddings.map(embedding => embedding.values);
40+ return embeddings;
41+}
42+
43+/**
44+ * Gets the vector for the given text from Google Vertex AI
45+ * @param {string[]} texts - The array of texts to get the vector for
46+ * @param {string} model - The model to use for embedding
47+ * @param {import('express').Request} request - The request object to get API key and URL
48+ * @returns {Promise<number[][]>} - The array of vectors for the texts
49+ */
50+export async function getVertexBatchVector(texts, model, request) {
51+ const { url, headers, apiName } = await getGoogleApiConfig(request, model, 'predict');
52+
53+ const body = {
54+ instances: texts.map(text => ({ content: text })),
55+ };
56+
57+ const response = await fetch(url, {
58+ body: JSON.stringify(body),
59+ method: 'POST',
60+ headers: headers,
61+ });
62+
63+ if (!response.ok) {
64+ const text = await response.text();
65+ console.warn(`${apiName} batch request failed`, response.statusText, text);
66+ throw new Error(`${apiName} batch request failed`);
67+ }
68+
69+ /** @type {any} */
70+ const data = await response.json();
71+ if (!Array.isArray(data?.predictions)) {
72+ throw new Error(`${apiName} did not return an array`);
73+ }
74+
75+ const embeddings = data.predictions.map(p => p.embeddings.values);
76+ return embeddings;
77+}
78+
79+/**
80+ * Gets the vector for the given text from Google AI Studio
81+ * @param {string} text - The text to get the vector for
82+ * @param {string} model - The model to use for embedding
83+ * @param {import('express').Request} request - The request object to get API key and URL
84+ * @returns {Promise<number[]>} - The vector for the text
85+ */
86+export async function getMakerSuiteVector(text, model, request) {
87+ const [embedding] = await getMakerSuiteBatchVector([text], model, request);
88+ return embedding;
89+}
90+
91+/**
92+ * Gets the vector for the given text from Google Vertex AI
93+ * @param {string} text - The text to get the vector for
94+ * @param {string} model - The model to use for embedding
95+ * @param {import('express').Request} request - The request object to get API key and URL
96+ * @returns {Promise<number[]>} - The vector for the text
97+ */
98+export async function getVertexVector(text, model, request) {
99+ const [embedding] = await getVertexBatchVector([text], model, request);
100+ return embedding;
101+}
src/vectors/makersuite-vectors.js+0 -61
@@ -1,61 +0,0 @@
1-import fetch from 'node-fetch';
2-import { SECRET_KEYS, readSecret } from '../endpoints/secrets.js';
3-import { trimTrailingSlash } from '../util.js';
4-const API_MAKERSUITE = 'https://generativelanguage.googleapis.com';
5-
6-/**
7- * Gets the vector for the given text from gecko model
8- * @param {string[]} texts - The array of texts to get the vector for
9- * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
10- * @param {string} model - The model to use for embedding
11- * @returns {Promise<number[][]>} - The array of vectors for the texts
12- */
13-export async function getMakerSuiteBatchVector(texts, directories, model) {
14- const promises = texts.map(text => getMakerSuiteVector(text, directories, model));
15- return await Promise.all(promises);
16-}
17-
18-/**
19- * Gets the vector for the given text from Gemini API text-embedding-004 model
20- * @param {string} text - The text to get the vector for
21- * @param {import('../users.js').UserDirectoryList} directories - The directories object for the user
22- * @param {string} model - The model to use for embedding (default is 'text-embedding-004')
23- * @returns {Promise<number[]>} - The vector for the text
24- */
25-export async function getMakerSuiteVector(text, directories, model) {
26- const key = readSecret(directories, SECRET_KEYS.MAKERSUITE);
27-
28- if (!key) {
29- console.warn('No Google AI Studio key found');
30- throw new Error('No Google AI Studio key found');
31- }
32-
33- const apiUrl = trimTrailingSlash(API_MAKERSUITE);
34- const url = `${apiUrl}/v1beta/models/${model}:embedContent?key=${key}`;
35- const body = {
36- content: {
37- parts: [
38- { text: text },
39- ],
40- },
41- };
42-
43- const response = await fetch(url, {
44- body: JSON.stringify(body),
45- method: 'POST',
46- headers: {
47- 'Content-Type': 'application/json',
48- },
49- });
50-
51- if (!response.ok) {
52- const text = await response.text();
53- console.warn('Google AI Studio request failed', response.statusText, text);
54- throw new Error('Google AI Studio request failed');
55- }
56-
57- /** @type {any} */
58- const data = await response.json();
59- // noinspection JSValidateTypes
60- return data['embedding']['values'];
61-}