All files / src/providers/gemini Embeddings.ts

89.47% Statements 17/19
58.33% Branches 7/12
100% Functions 4/4
88.88% Lines 16/18

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56              12x 12x       2x 2x 2x   2x   3x           3x     3x       2x   2x           2x       2x 2x 3x   2x                
import { EmbeddingRequest, EmbeddingResponse } from "../Provider.js";
import { GeminiBatchEmbedRequest, GeminiBatchEmbedResponse, GeminiEmbedRequest } from "./types.js";
import { handleGeminiError } from "./Errors.js";
import { logger } from "../../utils/logger.js";
 
export class GeminiEmbeddings {
  constructor(
    private readonly baseUrl: string,
    private readonly apiKey: string
  ) {}
 
  async execute(request: EmbeddingRequest): Promise<EmbeddingResponse> {
    const modelId = request.model || "text-embedding-004";
    const url = `${this.baseUrl}/models/${modelId}:batchEmbedContents?key=${this.apiKey}`;
    const inputs = Array.isArray(request.input) ? request.input : [request.input];
 
    const payload: GeminiBatchEmbedRequest = {
      requests: inputs.map((text) => {
        const item: GeminiEmbedRequest = {
          model: `models/${modelId}`,
          content: {
            parts: [{ text: String(text) }]
          }
        };
        Iif (request.dimensions) {
          item.outputDimensionality = request.dimensions;
        }
        return item;
      })
    };
 
    logger.logRequest("Gemini", "POST", url, payload);
 
    const response = await fetch(url, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(payload)
    });
 
    Iif (!response.ok) {
      await handleGeminiError(response, modelId);
    }
 
    const json = (await response.json()) as GeminiBatchEmbedResponse;
    logger.logResponse("Gemini", response.status, response.statusText, json);
    const vectors = json.embeddings?.map((e) => e.values) || [];
 
    return {
      model: modelId,
      vectors: vectors,
      input_tokens: 0,
      dimensions: vectors[0]?.length || 0
    };
  }
}