@illuma-ai/agents 1.0.81

package/src/llm/google/index.ts

@@ -0,0 +1,337 @@
+/* eslint-disable @typescript-eslint/ban-ts-comment */
+import { AIMessageChunk } from '@langchain/core/messages';
+import { ChatGenerationChunk } from '@langchain/core/outputs';
+import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
+import { getEnvironmentVariable } from '@langchain/core/utils/env';
+import { GoogleGenerativeAI as GenerativeAI } from '@google/generative-ai';
+import type {
+  GenerateContentRequest,
+  SafetySetting,
+} from '@google/generative-ai';
+import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
+import type { BaseMessage, UsageMetadata } from '@langchain/core/messages';
+import type { GeminiGenerationConfig } from '@langchain/google-common';
+import type { GeminiApiUsageMetadata, InputTokenDetails } from './types';
+import type { GoogleClientOptions } from '@/types';
+import {
+  convertResponseContentToChatGenerationChunk,
+  convertBaseMessagesToContent,
+  mapGenerateContentResultToChatResult,
+} from './utils/common';
+
+export class CustomChatGoogleGenerativeAI extends ChatGoogleGenerativeAI {
+  thinkingConfig?: GeminiGenerationConfig['thinkingConfig'];
+
+  /**
+   * Override to add gemini-3 model support for multimodal and function calling thought signatures
+   */
+  get _isMultimodalModel(): boolean {
+    return (
+      this.model.startsWith('gemini-1.5') ||
+      this.model.startsWith('gemini-2') ||
+      (this.model.startsWith('gemma-3-') &&
+        !this.model.startsWith('gemma-3-1b')) ||
+      this.model.startsWith('gemini-3')
+    );
+  }
+
+  constructor(fields: GoogleClientOptions) {
+    super(fields);
+
+    this.model = fields.model.replace(/^models\//, '');
+
+    this.maxOutputTokens = fields.maxOutputTokens ?? this.maxOutputTokens;
+
+    if (this.maxOutputTokens != null && this.maxOutputTokens < 0) {
+      throw new Error('`maxOutputTokens` must be a positive integer');
+    }
+
+    this.temperature = fields.temperature ?? this.temperature;
+    if (
+      this.temperature != null &&
+      (this.temperature < 0 || this.temperature > 2)
+    ) {
+      throw new Error('`temperature` must be in the range of [0.0,2.0]');
+    }
+
+    this.topP = fields.topP ?? this.topP;
+    if (this.topP != null && this.topP < 0) {
+      throw new Error('`topP` must be a positive integer');
+    }
+
+    if (this.topP != null && this.topP > 1) {
+      throw new Error('`topP` must be below 1.');
+    }
+
+    this.topK = fields.topK ?? this.topK;
+    if (this.topK != null && this.topK < 0) {
+      throw new Error('`topK` must be a positive integer');
+    }
+
+    this.stopSequences = fields.stopSequences ?? this.stopSequences;
+
+    this.apiKey = fields.apiKey ?? getEnvironmentVariable('GOOGLE_API_KEY');
+    if (this.apiKey == null || this.apiKey === '') {
+      throw new Error(
+        'Please set an API key for Google GenerativeAI ' +
+          'in the environment variable GOOGLE_API_KEY ' +
+          'or in the `apiKey` field of the ' +
+          'ChatGoogleGenerativeAI constructor'
+      );
+    }
+
+    this.safetySettings = fields.safetySettings ?? this.safetySettings;
+    if (this.safetySettings && this.safetySettings.length > 0) {
+      const safetySettingsSet = new Set(
+        this.safetySettings.map((s) => s.category)
+      );
+      if (safetySettingsSet.size !== this.safetySettings.length) {
+        throw new Error(
+          'The categories in `safetySettings` array must be unique'
+        );
+      }
+    }
+
+    this.thinkingConfig = fields.thinkingConfig ?? this.thinkingConfig;
+
+    this.streaming = fields.streaming ?? this.streaming;
+    this.json = fields.json;
+
+    // @ts-ignore - Accessing private property from parent class
+    this.client = new GenerativeAI(this.apiKey).getGenerativeModel(
+      {
+        model: this.model,
+        safetySettings: this.safetySettings as SafetySetting[],
+        generationConfig: {
+          stopSequences: this.stopSequences,
+          maxOutputTokens: this.maxOutputTokens,
+          temperature: this.temperature,
+          topP: this.topP,
+          topK: this.topK,
+          ...(this.json != null
+            ? { responseMimeType: 'application/json' }
+            : {}),
+        },
+      },
+      {
+        apiVersion: fields.apiVersion,
+        baseUrl: fields.baseUrl,
+        customHeaders: fields.customHeaders,
+      }
+    );
+    this.streamUsage = fields.streamUsage ?? this.streamUsage;
+  }
+
+  static lc_name(): 'IllumaGoogleGenerativeAI' {
+    return 'IllumaGoogleGenerativeAI';
+  }
+
+  /**
+   * Helper function to convert Gemini API usage metadata to LangChain format
+   * Includes support for cached tokens and tier-based tracking for gemini-3-pro-preview
+   */
+  private _convertToUsageMetadata(
+    usageMetadata: GeminiApiUsageMetadata | undefined,
+    model: string
+  ): UsageMetadata | undefined {
+    if (!usageMetadata) {
+      return undefined;
+    }
+
+    const output: UsageMetadata = {
+      input_tokens: usageMetadata.promptTokenCount ?? 0,
+      output_tokens:
+        (usageMetadata.candidatesTokenCount ?? 0) +
+        (usageMetadata.thoughtsTokenCount ?? 0),
+      total_tokens: usageMetadata.totalTokenCount ?? 0,
+    };
+
+    if (usageMetadata.cachedContentTokenCount) {
+      output.input_token_details ??= {};
+      output.input_token_details.cache_read =
+        usageMetadata.cachedContentTokenCount;
+    }
+
+    // gemini-3-pro-preview has bracket based tracking of tokens per request
+    if (model === 'gemini-3-pro-preview') {
+      const over200k = Math.max(
+        0,
+        (usageMetadata.promptTokenCount ?? 0) - 200000
+      );
+      const cachedOver200k = Math.max(
+        0,
+        (usageMetadata.cachedContentTokenCount ?? 0) - 200000
+      );
+      if (over200k) {
+        output.input_token_details = {
+          ...output.input_token_details,
+          over_200k: over200k,
+        } as InputTokenDetails;
+      }
+      if (cachedOver200k) {
+        output.input_token_details = {
+          ...output.input_token_details,
+          cache_read_over_200k: cachedOver200k,
+        } as InputTokenDetails;
+      }
+    }
+
+    return output;
+  }
+
+  invocationParams(
+    options?: this['ParsedCallOptions']
+  ): Omit<GenerateContentRequest, 'contents'> {
+    const params = super.invocationParams(options);
+    if (this.thinkingConfig) {
+      /** @ts-ignore */
+      this.client.generationConfig = {
+        /** @ts-ignore */
+        ...this.client.generationConfig,
+        /** @ts-ignore */
+        thinkingConfig: this.thinkingConfig,
+      };
+    }
+    return params;
+  }
+
+  async _generate(
+    messages: BaseMessage[],
+    options: this['ParsedCallOptions'],
+    runManager?: CallbackManagerForLLMRun
+  ): Promise<import('@langchain/core/outputs').ChatResult> {
+    const prompt = convertBaseMessagesToContent(
+      messages,
+      this._isMultimodalModel,
+      this.useSystemInstruction,
+      this.model
+    );
+    let actualPrompt = prompt;
+    if (prompt?.[0].role === 'system') {
+      const [systemInstruction] = prompt;
+      /** @ts-ignore */
+      this.client.systemInstruction = systemInstruction;
+      actualPrompt = prompt.slice(1);
+    }
+    const parameters = this.invocationParams(options);
+    const request = {
+      ...parameters,
+      contents: actualPrompt,
+    };
+
+    const res = await this.caller.callWithOptions(
+      { signal: options.signal },
+      async () =>
+        /** @ts-ignore */
+        this.client.generateContent(request)
+    );
+
+    const response = res.response;
+    const usageMetadata = this._convertToUsageMetadata(
+      /** @ts-ignore */
+      response.usageMetadata,
+      this.model
+    );
+
+    /** @ts-ignore */
+    const generationResult = mapGenerateContentResultToChatResult(response, {
+      usageMetadata,
+    });
+
+    await runManager?.handleLLMNewToken(
+      generationResult.generations[0].text || '',
+      undefined,
+      undefined,
+      undefined,
+      undefined,
+      undefined
+    );
+    return generationResult;
+  }
+
+  async *_streamResponseChunks(
+    messages: BaseMessage[],
+    options: this['ParsedCallOptions'],
+    runManager?: CallbackManagerForLLMRun
+  ): AsyncGenerator<ChatGenerationChunk> {
+    const prompt = convertBaseMessagesToContent(
+      messages,
+      this._isMultimodalModel,
+      this.useSystemInstruction,
+      this.model
+    );
+    let actualPrompt = prompt;
+    if (prompt?.[0].role === 'system') {
+      const [systemInstruction] = prompt;
+      /** @ts-ignore */
+      this.client.systemInstruction = systemInstruction;
+      actualPrompt = prompt.slice(1);
+    }
+    const parameters = this.invocationParams(options);
+    const request = {
+      ...parameters,
+      contents: actualPrompt,
+    };
+    const stream = await this.caller.callWithOptions(
+      { signal: options.signal },
+      async () => {
+        /** @ts-ignore */
+        const { stream } = await this.client.generateContentStream(request);
+        return stream;
+      }
+    );
+
+    let index = 0;
+    let lastUsageMetadata: UsageMetadata | undefined;
+    for await (const response of stream) {
+      if (
+        'usageMetadata' in response &&
+        this.streamUsage !== false &&
+        options.streamUsage !== false
+      ) {
+        lastUsageMetadata = this._convertToUsageMetadata(
+          response.usageMetadata as GeminiApiUsageMetadata | undefined,
+          this.model
+        );
+      }
+
+      const chunk = convertResponseContentToChatGenerationChunk(response, {
+        usageMetadata: undefined,
+        index,
+      });
+      index += 1;
+      if (!chunk) {
+        continue;
+      }
+
+      yield chunk;
+      await runManager?.handleLLMNewToken(
+        chunk.text || '',
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        { chunk }
+      );
+    }
+
+    if (lastUsageMetadata) {
+      const finalChunk = new ChatGenerationChunk({
+        text: '',
+        message: new AIMessageChunk({
+          content: '',
+          usage_metadata: lastUsageMetadata,
+        }),
+      });
+      yield finalChunk;
+      await runManager?.handleLLMNewToken(
+        finalChunk.text || '',
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        { chunk: finalChunk }
+      );
+    }
+  }
+}