async function getEmbeddingsGemini()

in firestore-semantic-search/functions/src/common/gemini_embeddings.ts [43:76]

• 26 lines of code
• 5 McCabe index (conditional complexity)

async function getEmbeddingsGemini(
  text: string | string[]
): Promise<number[][]> {
  if (!client && (typeof text !== 'string' || text.length !== 0)) {
    await initializeGeminiClient();
  }

  const model = client.getGenerativeModel({model: 'embedding-001'});

  let embeddings: number[][] = [];

  const t0 = performance.now();

  if (Array.isArray(text)) {
    const batches: EmbedContentRequest[] = text.map(doc => ({
      content: {role: 'user', parts: [{text: doc}]},
      taskType: TaskType.RETRIEVAL_QUERY,
    }));

    const result = await model.batchEmbedContents({
      requests: batches,
    });

    embeddings = result.embeddings.map(embedding => embedding.values);
  } else {
    const result = await model.embedContent(text);
    embeddings = [result.embedding.values];
  }

  const duration = performance.now() - t0;
  console.log(`Processed embeddings. This took ${duration}ms`);

  return embeddings;
}