Class LobeGoogleAI

packages/model-runtime/src/providers/google/index.ts:112–605

};

export class LobeGoogleAI implements LobeRuntimeAI {
  private client: GoogleGenAI;
  private isVertexAi: boolean;
  baseURL?: string;
  apiKey?: string;
  provider: string;
  private readonly modelIdMappingOptions: ModelIdMappingOptions;

  constructor({
    apiKey,
    baseURL,
    client,
    isVertexAi,
    id,
    defaultHeaders,
    modelIdMapping,
  }: LobeGoogleAIParams = {}) {
    if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);

    const httpOptions = baseURL
      ? ({ baseUrl: baseURL, headers: defaultHeaders } as HttpOptions)
      : undefined;

    this.apiKey = apiKey;
    this.client = client ?? new GoogleGenAI({ apiKey, httpOptions });
    this.baseURL = client ? undefined : baseURL || DEFAULT_BASE_URL;
    this.isVertexAi = isVertexAi || false;
    this.modelIdMappingOptions = { modelIdMapping };

    this.provider = id || (isVertexAi ? 'vertexai' : 'google');
  }

  async chat(rawPayload: ChatStreamPayload, options?: ChatMethodOptions) {
    try {
      const payload = this.buildPayload(rawPayload);
      const { model, thinkingBudget, thinkingLevel, imageAspectRatio, imageResolution } = payload;

      // https://ai.google.dev/gemini-api/docs/thinking#set-budget
      const thinkingConfig = resolveGoogleThinkingConfig(model, {
        thinkingBudget,
        thinkingLevel,
      }) as ThinkingConfig;

      const contents = await buildGoogleMessages(payload.messages, { model });
      const isImageResponseModel = isGoogleImageResponseModel(model);

      const controller = new AbortController();
      const originalSignal = options?.signal;

      if (originalSignal) {
        if (originalSignal.aborted) {
          controller.abort();
        } else {
          originalSignal.addEventListener('abort', () => {
            controller.abort();
          });
        }
      }