Update makersuite-vectors.js to use Gemini API text-embedding-004 model

4acd7296118aafd2e02ea966d68cd604fdb6447a

ceruleandeep <83318388+ceruleandeep@users.noreply.github.com>

Signed
1 files changed, +18 -12Showing whitespace changes
src/vectors/makersuite-vectors.js+18 -12
@@ -1,5 +1,6 @@
1const fetch = require('node-fetch').default;1const fetch = require('node-fetch').default;
2const { SECRET_KEYS, readSecret } = require('../endpoints/secrets');2const { SECRET_KEYS, readSecret } = require('../endpoints/secrets');
3const API_MAKERSUITE = 'https://generativelanguage.googleapis.com';
34
4/**5/**
5 * Gets the vector for the given text from gecko model6 * Gets the vector for the given text from gecko model
@@ -9,12 +10,11 @@ const { SECRET_KEYS, readSecret } = require('../endpoints/secrets');
9 */10 */
10async function getMakerSuiteBatchVector(texts, directories) {11async function getMakerSuiteBatchVector(texts, directories) {
11 const promises = texts.map(text => getMakerSuiteVector(text, directories));12 const promises = texts.map(text => getMakerSuiteVector(text, directories));
12 const vectors = await Promise.all(promises);13 return await Promise.all(promises);
13 return vectors;
14}14}
1515
16/**16/**
17 * Gets the vector for the given text from PaLM gecko model17 * Gets the vector for the given text from Gemini API text-embedding-004 model
18 * @param {string} text - The text to get the vector for18 * @param {string} text - The text to get the vector for
19 * @param {import('../users').UserDirectoryList} directories - The directories object for the user19 * @param {import('../users').UserDirectoryList} directories - The directories object for the user
20 * @returns {Promise<number[]>} - The vector for the text20 * @returns {Promise<number[]>} - The vector for the text
@@ -27,14 +27,23 @@ async function getMakerSuiteVector(text, directories) {
27 throw new Error('No Google AI Studio key found');27 throw new Error('No Google AI Studio key found');
28 }28 }
2929
30 const response = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/embedding-gecko-001:embedText?key=${key}`, {30 const apiUrl = new URL(API_MAKERSUITE);
31 const model = 'text-embedding-004';
32 const url = `${apiUrl.origin}/v1beta/models/${model}:embedContent?key=${key}`;
33 const body = {
34 content: {
35 parts: [
36 { text: text },
37 ],
38 },
39 };
40
41 const response = await fetch(url, {
42 body: JSON.stringify(body),
31 method: 'POST',43 method: 'POST',
32 headers: {44 headers: {
33 'Content-Type': 'application/json',45 'Content-Type': 'application/json',
34 },46 },
35 body: JSON.stringify({
36 text: text,
37 }),
38 });47 });
3948
40 if (!response.ok) {49 if (!response.ok) {
@@ -44,11 +53,8 @@ async function getMakerSuiteVector(text, directories) {
44 }53 }
4554
46 const data = await response.json();55 const data = await response.json();
4756 // noinspection JSValidateTypes
48 // Access the "value" dictionary57 return data['embedding']['values'];
49 const vector = data.embedding.value;
50
51 return vector;
52}58}
5359
54module.exports = {60module.exports = {