Google Embedding Models

npm install @niledatabase/server @google/generative-ai

const { GoogleGenerativeAI } = require('@google/generative-ai');
const GOOGLE_API_KEY = 'my-google-api-key';
const model = 'models/text-embedding-004'; // or your favorite Google model

const genAI = new GoogleGenerativeAI(GOOGLE_API_KEY);

const embed = genAI.getGenerativeModel({
  model: model,
  outputDimensionality: 768, // optional, default 768 but can be scaled down
});

const input_text = `The future belongs to those who believe in the beauty of their 
                    dreams.`;

// Note, if you don't want to specify taskType and title, you can simply use:
// const doc_vec_resp = await model.embedContent(input_text);
const doc_vec_resp = await embed.embedContent({
  content: { parts: [{ text: input_text }] },
  taskType: 'RETRIEVAL_DOCUMENT', // for embeddings of documents stored in pg_vector
  // title of the document. can be used with RETRIEVAL_DOCUMENT and
  // possibly improve the quality of the embeddings
  // title: ""
});

const question = 'Who does the future belong to?';

const question_vec_resp = await embed.embedContent({
  content: { parts: [{ text: question }] },
  taskType: 'RETRIEVAL_QUERY', //embeddings of questions used to search documents
});

// Google returns a response object with an embedding field that
// contains the embeddings in the value field
const doc_vec = doc_vec_resp.embedding.values;
const question_vec = question_vec_resp.embedding.values;

// set up Nile as the vector store
const { Nile } = await import('@niledatabase/server');
const NILEDB_USER = 'you need a nile user with access to a nile database';
const NILEDB_PASSWORD = 'and a password for that user';
const nile = Nile({
  user: NILEDB_USER,
  password: NILEDB_PASSWORD,
});

// create table to store vectors - vector size must match the model dimensions
await nile.db.query(
  'CREATE TABLE IF NOT EXISTS embeddings (embedding vector(768))',
);

// store vector in a table
await nile.db.query('INSERT INTO embeddings (embedding) values ($1)', [
  JSON.stringify(doc_vec.map((v) => Number(v))),
]);

// search for similar vectors
let db_resp = await nile.db.query(
  'select embedding from embeddings order by embedding<=>$1 limit 1',
  [JSON.stringify(question_vec.map((v) => Number(v)))],
);

// Postgres returns an object, with array of rows each column is a
// property of the row the vector is represented as a string
let similar_str = db_resp.rows[0].embedding;
let similar = similar_str
  .substring(1, similar_str.length - 1)
  .split(',')
  .map((v) => parseFloat(v));

// check that we got the same vector back
let same = true;

for (let i = 0; i < similar.length; i++) {
  if (Math.abs(similar[i] - doc_vec[i]) > 0.000001) {
    same = false;
  }
}

console.log('got same vector? ' + same);