All files / src/providers/gemini Models.ts

73.33% Statements 11/15
88.23% Branches 15/17
50% Functions 3/6
73.33% Lines 11/15

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 57 58 59 60 61 62 63 64 65 66 67              13x 13x       2x 2x             2x 2x   2x 51x 51x   51x                                       51x                                      
import { ModelInfo } from "../Provider.js";
import { Capabilities } from "./Capabilities.js";
import { ModelRegistry } from "../../models/ModelRegistry.js";
import { GeminiListModelsResponse } from "./types.js";
 
export class GeminiModels {
  constructor(
    private readonly baseUrl: string,
    private readonly apiKey: string
  ) {}
 
  async execute(): Promise<ModelInfo[]> {
    try {
      const response = await fetch(`${this.baseUrl}/models?key=${this.apiKey}`, {
        method: "GET",
        headers: {
          "Content-Type": "application/json"
        }
      });
 
      Eif (response.ok) {
        const { models } = (await response.json()) as GeminiListModelsResponse;
 
        return models.map((m) => {
          const modelId = m.name.replace("models/", "");
          const registryModel = ModelRegistry.find(modelId, "gemini");
 
          const info: ModelInfo = {
            id: modelId,
            name: registryModel?.name || m.displayName || modelId,
            provider: "gemini",
            family: registryModel?.family || modelId,
            context_window: registryModel?.context_window || Capabilities.getContextWindow(modelId),
            max_output_tokens:
              registryModel?.max_output_tokens || Capabilities.getMaxOutputTokens(modelId),
            modalities: registryModel?.modalities || Capabilities.getModalities(modelId),
            capabilities: Capabilities.getCapabilities(modelId),
            pricing: registryModel?.pricing || Capabilities.getPricing(modelId),
            metadata: {
              ...(registryModel?.metadata || {}),
              description: m.description,
              input_token_limit: m.inputTokenLimit,
              output_token_limit: m.outputTokenLimit,
              supported_generation_methods: m.supportedGenerationMethods
            }
          };
 
          return info;
        });
      }
    } catch (_error) {
      // Fallback
    }
 
    return ModelRegistry.all()
      .filter((m) => m.provider === "gemini")
      .map((m) => ({
        ...m,
        capabilities: Capabilities.getCapabilities(m.id)
      })) as ModelInfo[];
  }
 
  find(modelId: string) {
    return ModelRegistry.find(modelId, "gemini");
  }
}