illuma-agents 1.0.7 → 1.0.9

package/src/llm/fake.ts

@@ -1,133 +1,133 @@
-import { ChatGenerationChunk } from '@langchain/core/outputs';
-import { AIMessageChunk } from '@langchain/core/messages';
-import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
-import type { BaseMessage } from '@langchain/core/messages';
-import { FakeListChatModel } from '@langchain/core/utils/testing';
-import { ToolCall, ToolCallChunk } from '@langchain/core/messages/tool';
-
-type SplitStrategy = {
-  type: 'regex' | 'fixed';
-  value: RegExp | number;
-};
-
-export class FakeChatModel extends FakeListChatModel {
-  private splitStrategy: SplitStrategy;
-  private toolCalls: ToolCall[] = [];
-  private addedToolCalls: boolean = false;
-
-  constructor({
-    responses,
-    sleep,
-    emitCustomEvent,
-    splitStrategy = { type: 'regex', value: /(?<=\s+)|(?=\s+)/ },
-    toolCalls = []
-  }: {
-    responses: string[];
-    sleep?: number;
-    emitCustomEvent?: boolean;
-    splitStrategy?: SplitStrategy;
-    toolCalls?: ToolCall[];
-  }) {
-    super({ responses, sleep, emitCustomEvent });
-    this.splitStrategy = splitStrategy;
-    this.toolCalls = toolCalls;
-  }
-
-  private splitText(text: string): string[] {
-    if (this.splitStrategy.type === 'regex') {
-      return text.split(this.splitStrategy.value as RegExp);
-    } else {
-      const chunkSize = this.splitStrategy.value as number;
-      const chunks: string[] = [];
-      for (let i = 0; i < text.length; i += chunkSize) {
-        chunks.push(text.slice(i, i + chunkSize));
-      }
-      return chunks;
-    }
-  }
-  _createResponseChunk(text: string, tool_call_chunks?: ToolCallChunk[]): ChatGenerationChunk {
-    return new ChatGenerationChunk({
-      text,
-      generationInfo: {},
-      message: new AIMessageChunk({
-        content: text,
-        tool_call_chunks,
-        additional_kwargs: tool_call_chunks ? {
-          tool_calls: tool_call_chunks.map((toolCall) => ({
-            index: toolCall.index ?? 0,
-            id: toolCall.id ?? '',
-            type: 'function',
-            function: {
-              name: toolCall.name ?? '',
-              arguments: toolCall.args ?? '',
-            },
-          })),
-        } : undefined,
-      })});
-  }
-
-  async *_streamResponseChunks(
-    _messages: BaseMessage[],
-    options: this['ParsedCallOptions'],
-    runManager?: CallbackManagerForLLMRun
-  ): AsyncGenerator<ChatGenerationChunk> {
-    const response = this._currentResponse();
-    this._incrementResponse();
-
-    if (this.emitCustomEvent) {
-      await runManager?.handleCustomEvent('some_test_event', {
-        someval: true,
-      });
-    }
-
-    const chunks = this.splitText(response);
-    for await (const chunk of chunks) {
-      await this._sleepIfRequested();
-
-      if (options.thrownErrorString != null && options.thrownErrorString) {
-        throw new Error(options.thrownErrorString);
-      }
-
-      const responseChunk = super._createResponseChunk(chunk);
-      yield responseChunk;
-      void runManager?.handleLLMNewToken(chunk);
-    }
-
-    await this._sleepIfRequested();
-    if (this.toolCalls.length > 0 && !this.addedToolCalls) {
-      this.addedToolCalls = true;
-      const toolCallChunks = this.toolCalls.map((toolCall) => {;
-        return {
-          name: toolCall.name,
-          args: JSON.stringify(toolCall.args),
-          id: toolCall.id,
-          type: 'tool_call_chunk',
-        } as ToolCallChunk;
-      });
-      const responseChunk = this._createResponseChunk('', toolCallChunks);
-      yield responseChunk;
-      void runManager?.handleLLMNewToken('');
-    }
-  }
-}
-
-export function createFakeStreamingLLM({
-  responses,
-  sleep,
-  splitStrategy,
-  toolCalls,
-} : {
-  responses: string[],
-  sleep?: number,
-  splitStrategy?: SplitStrategy,
-  toolCalls?: ToolCall[]
-}
-): FakeChatModel {
-  return new FakeChatModel({
-    sleep,
-    responses,
-    emitCustomEvent: true,
-    splitStrategy,
-    toolCalls,
-  });
-}
+import { ChatGenerationChunk } from '@langchain/core/outputs';
+import { AIMessageChunk } from '@langchain/core/messages';
+import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
+import type { BaseMessage } from '@langchain/core/messages';
+import { FakeListChatModel } from '@langchain/core/utils/testing';
+import { ToolCall, ToolCallChunk } from '@langchain/core/messages/tool';
+
+type SplitStrategy = {
+  type: 'regex' | 'fixed';
+  value: RegExp | number;
+};
+
+export class FakeChatModel extends FakeListChatModel {
+  private splitStrategy: SplitStrategy;
+  private toolCalls: ToolCall[] = [];
+  private addedToolCalls: boolean = false;
+
+  constructor({
+    responses,
+    sleep,
+    emitCustomEvent,
+    splitStrategy = { type: 'regex', value: /(?<=\s+)|(?=\s+)/ },
+    toolCalls = []
+  }: {
+    responses: string[];
+    sleep?: number;
+    emitCustomEvent?: boolean;
+    splitStrategy?: SplitStrategy;
+    toolCalls?: ToolCall[];
+  }) {
+    super({ responses, sleep, emitCustomEvent });
+    this.splitStrategy = splitStrategy;
+    this.toolCalls = toolCalls;
+  }
+
+  private splitText(text: string): string[] {
+    if (this.splitStrategy.type === 'regex') {
+      return text.split(this.splitStrategy.value as RegExp);
+    } else {
+      const chunkSize = this.splitStrategy.value as number;
+      const chunks: string[] = [];
+      for (let i = 0; i < text.length; i += chunkSize) {
+        chunks.push(text.slice(i, i + chunkSize));
+      }
+      return chunks;
+    }
+  }
+  _createResponseChunk(text: string, tool_call_chunks?: ToolCallChunk[]): ChatGenerationChunk {
+    return new ChatGenerationChunk({
+      text,
+      generationInfo: {},
+      message: new AIMessageChunk({
+        content: text,
+        tool_call_chunks,
+        additional_kwargs: tool_call_chunks ? {
+          tool_calls: tool_call_chunks.map((toolCall) => ({
+            index: toolCall.index ?? 0,
+            id: toolCall.id ?? '',
+            type: 'function',
+            function: {
+              name: toolCall.name ?? '',
+              arguments: toolCall.args ?? '',
+            },
+          })),
+        } : undefined,
+      })});
+  }
+
+  async *_streamResponseChunks(
+    _messages: BaseMessage[],
+    options: this['ParsedCallOptions'],
+    runManager?: CallbackManagerForLLMRun
+  ): AsyncGenerator<ChatGenerationChunk> {
+    const response = this._currentResponse();
+    this._incrementResponse();
+
+    if (this.emitCustomEvent) {
+      await runManager?.handleCustomEvent('some_test_event', {
+        someval: true,
+      });
+    }
+
+    const chunks = this.splitText(response);
+    for await (const chunk of chunks) {
+      await this._sleepIfRequested();
+
+      if (options.thrownErrorString != null && options.thrownErrorString) {
+        throw new Error(options.thrownErrorString);
+      }
+
+      const responseChunk = super._createResponseChunk(chunk);
+      yield responseChunk;
+      void runManager?.handleLLMNewToken(chunk);
+    }
+
+    await this._sleepIfRequested();
+    if (this.toolCalls.length > 0 && !this.addedToolCalls) {
+      this.addedToolCalls = true;
+      const toolCallChunks = this.toolCalls.map((toolCall) => {;
+        return {
+          name: toolCall.name,
+          args: JSON.stringify(toolCall.args),
+          id: toolCall.id,
+          type: 'tool_call_chunk',
+        } as ToolCallChunk;
+      });
+      const responseChunk = this._createResponseChunk('', toolCallChunks);
+      yield responseChunk;
+      void runManager?.handleLLMNewToken('');
+    }
+  }
+}
+
+export function createFakeStreamingLLM({
+  responses,
+  sleep,
+  splitStrategy,
+  toolCalls,
+} : {
+  responses: string[],
+  sleep?: number,
+  splitStrategy?: SplitStrategy,
+  toolCalls?: ToolCall[]
+}
+): FakeChatModel {
+  return new FakeChatModel({
+    sleep,
+    responses,
+    emitCustomEvent: true,
+    splitStrategy,
+    toolCalls,
+  });
+}

package/src/llm/google/index.ts

@@ -11,15 +11,30 @@ import type {
 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 } from './types';
+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);
 
@@ -111,6 +126,59 @@ export class CustomChatGoogleGenerativeAI extends ChatGoogleGenerativeAI {
     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'> {
@@ -127,6 +195,60 @@ export class CustomChatGoogleGenerativeAI extends ChatGoogleGenerativeAI {
     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'],
@@ -135,7 +257,8 @@ export class CustomChatGoogleGenerativeAI extends ChatGoogleGenerativeAI {
     const prompt = convertBaseMessagesToContent(
       messages,
       this._isMultimodalModel,
-      this.useSystemInstruction
+      this.useSystemInstruction,
+      this.model
     );
     let actualPrompt = prompt;
     if (prompt?.[0].role === 'system') {
@@ -166,18 +289,10 @@ export class CustomChatGoogleGenerativeAI extends ChatGoogleGenerativeAI {
         this.streamUsage !== false &&
         options.streamUsage !== false
       ) {
-        const genAIUsageMetadata = response.usageMetadata as
-          | GeminiApiUsageMetadata
-          | undefined;
-
-        const output_tokens =
-          (genAIUsageMetadata?.candidatesTokenCount ?? 0) +
-          (genAIUsageMetadata?.thoughtsTokenCount ?? 0);
-        lastUsageMetadata = {
-          input_tokens: genAIUsageMetadata?.promptTokenCount ?? 0,
-          output_tokens,
-          total_tokens: genAIUsageMetadata?.totalTokenCount ?? 0,
-        };
+        lastUsageMetadata = this._convertToUsageMetadata(
+          response.usageMetadata as GeminiApiUsageMetadata | undefined,
+          this.model
+        );
       }
 
       const chunk = convertResponseContentToChatGenerationChunk(response, {