@illuma-ai/agents 1.0.89 → 1.0.93

package/src/llm/bedrock/index.ts

@@ -1,19 +1,25 @@
 /**
- * Optimized ChatBedrockConverse wrapper that fixes contentBlockIndex conflicts
- * and adds support for:
+ * Optimized ChatBedrockConverse wrapper that fixes content block merging for
+ * streaming responses and adds support for latest @langchain/aws features:
  *
  * - Prompt caching support for Bedrock Converse API (Illuma feature)
  * - Application Inference Profiles (PR #9129)
  * - Service Tiers (Priority/Standard/Flex) (PR #9785) - requires AWS SDK 3.966.0+
  *
- * Bedrock sends the same contentBlockIndex for both text and tool_use content blocks,
- * causing LangChain's merge logic to fail with "field[contentBlockIndex] already exists"
- * errors. This wrapper simply strips contentBlockIndex from response_metadata to avoid
- * the conflict.
+ * Bedrock's `@langchain/aws` library does not include an `index` property on content
+ * blocks (unlike Anthropic/OpenAI), which causes LangChain's `_mergeLists` to append
+ * each streaming chunk as a separate array entry instead of merging by index.
  *
- * The contentBlockIndex field is only used internally by Bedrock's streaming protocol
- * and isn't needed by application logic - the index field on tool_call_chunks serves
- * the purpose of tracking tool call ordering.
+ * This wrapper takes full ownership of the stream by directly interfacing with the
+ * AWS SDK client (`this.client`) and using custom handlers from `./utils/` that
+ * include `contentBlockIndex` in response_metadata for every delta type. It then
+ * promotes `contentBlockIndex` to an `index` property on each content block
+ * (mirroring Anthropic's pattern) and strips it from metadata to avoid
+ * `_mergeDicts` conflicts.
+ *
+ * When multiple content block types are present (e.g. reasoning + text), text deltas
+ * are promoted from strings to array form with `index` so they merge correctly once
+ * the accumulated content is already an array.
  *
  * PROMPT CACHING:
  * When promptCache: true is set, this wrapper adds cachePoint markers to the tools array
@@ -21,16 +27,23 @@
  * cached and reused across requests, reducing latency and costs.
  *
  * CACHE TOKEN EXTRACTION:
- * LangChain AWS doesn't extract cacheReadInputTokens/cacheWriteInputTokens from Bedrock's
- * response. This wrapper adds input_token_details to usage_metadata with cache information.
+ * Cache token extraction is handled in `./utils/message_outputs.ts` and added
+ * to usage_metadata with input_token_details (cacheReadInputTokens/cacheWriteInputTokens).
  */
 
 import { ChatBedrockConverse } from '@langchain/aws';
+import { ConverseStreamCommand } from '@aws-sdk/client-bedrock-runtime';
 import { AIMessageChunk } from '@langchain/core/messages';
-import type { BaseMessage, UsageMetadata } from '@langchain/core/messages';
+import type { BaseMessage } from '@langchain/core/messages';
 import { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs';
 import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
 import type { ChatBedrockConverseInput } from '@langchain/aws';
+import {
+  convertToConverseMessages,
+  handleConverseStreamContentBlockStart,
+  handleConverseStreamContentBlockDelta,
+  handleConverseStreamMetadata,
+} from './utils';
 
 /**
  * Service tier type for Bedrock invocations.
@@ -45,8 +58,7 @@ export type ServiceTierType = 'priority' | 'default' | 'flex' | 'reserved';
  * - applicationInferenceProfile: Use an inference profile ARN instead of model ID
  * - serviceTier: Specify service tier (Priority, Standard, Flex, Reserved)
  */
-export interface CustomChatBedrockConverseInput
-  extends ChatBedrockConverseInput {
+export interface IllumaBedrockConverseInput extends ChatBedrockConverseInput {
   /**
    * Enable Bedrock prompt caching for tool definitions.
    * When true, adds cachePoint markers to tools array.
@@ -82,11 +94,11 @@ export interface CustomChatBedrockConverseInput
 /**
  * Extended call options with serviceTier override support.
  */
-export interface CustomChatBedrockConverseCallOptions {
+export interface IllumaBedrockConverseCallOptions {
   serviceTier?: ServiceTierType;
 }
 
-export class CustomChatBedrockConverse extends ChatBedrockConverse {
+export class IllumaBedrockConverse extends ChatBedrockConverse {
   /** Enable Bedrock prompt caching for tool definitions */
   promptCache: boolean;
 
@@ -96,7 +108,7 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
   /** Service tier for model invocation */
   serviceTier?: ServiceTierType;
 
-  constructor(fields?: CustomChatBedrockConverseInput) {
+  constructor(fields?: IllumaBedrockConverseInput) {
     super(fields);
     this.promptCache = fields?.promptCache ?? false;
     this.applicationInferenceProfile = fields?.applicationInferenceProfile;
@@ -108,7 +120,9 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
     // 2. Using different naming conventions (claude-4, claude-opus-4, etc.)
     // We need to ensure tool_choice is properly set for withStructuredOutput to work
     const modelName = (fields?.model ?? '').toLowerCase();
-    const profileName = (fields?.applicationInferenceProfile ?? '').toLowerCase();
+    const profileName = (
+      fields?.applicationInferenceProfile ?? ''
+    ).toLowerCase();
     const isClaudeModel =
       modelName.includes('claude') ||
       modelName.includes('anthropic') ||
@@ -146,7 +160,7 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
    * NOTE: Only Claude models support cachePoint - Nova and other models will reject it.
    */
   override invocationParams(
-    options?: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions
+    options?: this['ParsedCallOptions'] & IllumaBedrockConverseCallOptions
   ): ReturnType<ChatBedrockConverse['invocationParams']> & {
     serviceTier?: { type: ServiceTierType };
   } {
@@ -220,10 +234,9 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
    */
   override async _generateNonStreaming(
     messages: BaseMessage[],
-    options: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions,
+    options: this['ParsedCallOptions'] & IllumaBedrockConverseCallOptions,
     runManager?: CallbackManagerForLLMRun
   ): Promise<ChatResult> {
-    // Temporarily swap model for applicationInferenceProfile support
     const originalModel = this.model;
     if (
       this.applicationInferenceProfile != null &&
@@ -235,156 +248,177 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
     try {
       return await super._generateNonStreaming(messages, options, runManager);
     } finally {
-      // Restore original model
       this.model = originalModel;
     }
   }
 
   /**
-   * Override _streamResponseChunks to:
-   * 1. Use applicationInferenceProfile as modelId (by temporarily swapping this.model)
-   * 2. Strip contentBlockIndex from response_metadata to prevent merge conflicts
-   * 3. Extract cacheReadInputTokens/cacheWriteInputTokens and add to usage_metadata
+   * Own the stream end-to-end so we have direct access to every
+   * `contentBlockDelta.contentBlockIndex` from the AWS SDK.
    *
-   * Note: We delegate to super._streamResponseChunks() to preserve @langchain/aws's
-   * internal chunk handling which correctly preserves array content for reasoning blocks.
+   * This replaces the parent's implementation which strips contentBlockIndex
+   * from text and reasoning deltas, making it impossible to merge correctly.
    */
   override async *_streamResponseChunks(
     messages: BaseMessage[],
-    options: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions,
+    options: this['ParsedCallOptions'] & IllumaBedrockConverseCallOptions,
     runManager?: CallbackManagerForLLMRun
   ): AsyncGenerator<ChatGenerationChunk> {
-    // Temporarily swap model for applicationInferenceProfile support
-    const originalModel = this.model;
-    if (
-      this.applicationInferenceProfile != null &&
-      this.applicationInferenceProfile !== ''
-    ) {
-      this.model = this.applicationInferenceProfile;
+    const { converseMessages, converseSystem } =
+      convertToConverseMessages(messages);
+    const params = this.invocationParams(options);
+
+    let { streamUsage } = this;
+    if ((options as Record<string, unknown>).streamUsage !== undefined) {
+      streamUsage = (options as Record<string, unknown>).streamUsage as boolean;
     }
 
-    try {
-      // Use parent's streaming logic which correctly handles reasoning content
-      const baseStream = super._streamResponseChunks(
-        messages,
-        options,
-        runManager
-      );
+    const modelId = this.getModelId();
 
-      for await (const chunk of baseStream) {
-        // Clean and enhance chunk
-        yield this.processChunk(chunk);
+    const command = new ConverseStreamCommand({
+      modelId,
+      messages: converseMessages,
+      system: converseSystem,
+      ...(params as Record<string, unknown>),
+    });
+
+    const response = await this.client.send(command, {
+      abortSignal: options.signal,
+    });
+
+    if (!response.stream) {
+      return;
+    }
+
+    const seenBlockIndices = new Set<number>();
+
+    for await (const event of response.stream) {
+      if (event.contentBlockStart != null) {
+        const startChunk = handleConverseStreamContentBlockStart(
+          event.contentBlockStart
+        );
+        if (startChunk != null) {
+          const idx = event.contentBlockStart.contentBlockIndex;
+          if (idx != null) {
+            seenBlockIndices.add(idx);
+          }
+          yield this.enrichChunk(startChunk, seenBlockIndices);
+        }
+      } else if (event.contentBlockDelta != null) {
+        const deltaChunk = handleConverseStreamContentBlockDelta(
+          event.contentBlockDelta
+        );
+
+        const idx = event.contentBlockDelta.contentBlockIndex;
+        if (idx != null) {
+          seenBlockIndices.add(idx);
+        }
+
+        yield this.enrichChunk(deltaChunk, seenBlockIndices);
+
+        await runManager?.handleLLMNewToken(
+          deltaChunk.text,
+          undefined,
+          undefined,
+          undefined,
+          undefined,
+          { chunk: deltaChunk }
+        );
+      } else if (event.metadata != null) {
+        yield handleConverseStreamMetadata(event.metadata, { streamUsage });
+      } else if (event.contentBlockStop != null) {
+        const stopIdx = event.contentBlockStop.contentBlockIndex;
+        if (stopIdx != null) {
+          seenBlockIndices.add(stopIdx);
+        }
+      } else {
+        yield new ChatGenerationChunk({
+          text: '',
+          message: new AIMessageChunk({
+            content: '',
+            response_metadata: event,
+          }),
+        });
       }
-    } finally {
-      // Restore original model
-      this.model = originalModel;
     }
   }
 
   /**
-   * Process a chunk by:
-   * 1. Removing contentBlockIndex from response_metadata
-   * 2. Extracting cache token information from Bedrock's usage data
+   * Inject `index` on content blocks for proper merge behaviour, then strip
+   * `contentBlockIndex` from response_metadata to prevent `_mergeDicts` conflicts.
+   *
+   * Text string content is promoted to array form only when the stream contains
+   * multiple content block indices (e.g. reasoning at index 0, text at index 1),
+   * ensuring text merges correctly with the already-array accumulated content.
    */
-  private processChunk(chunk: ChatGenerationChunk): ChatGenerationChunk {
+  enrichChunk(
+    chunk: ChatGenerationChunk,
+    seenBlockIndices: Set<number>
+  ): ChatGenerationChunk {
     const message = chunk.message;
     if (!(message instanceof AIMessageChunk)) {
       return chunk;
     }
 
-    const responseMetadata = message.response_metadata as Record<
-      string,
-      unknown
-    >;
-    let needsModification = false;
-    let cleanedMetadata = responseMetadata;
-
-    // Check if contentBlockIndex exists anywhere in response_metadata
-    const hasContentBlockIndex = this.hasContentBlockIndex(responseMetadata);
-    if (hasContentBlockIndex) {
-      cleanedMetadata = this.removeContentBlockIndex(
-        responseMetadata
-      ) as Record<string, unknown>;
-      needsModification = true;
-    }
+    const metadata = message.response_metadata as Record<string, unknown>;
+    const blockIndex = this.extractContentBlockIndex(metadata);
+    const hasMetadataIndex = blockIndex != null;
 
-    // Extract cache tokens from metadata.usage (Bedrock streaming format)
-    // The metadata chunk contains usage with cacheReadInputTokens/cacheWriteInputTokens
-    const metadata = responseMetadata.metadata as
-      | Record<string, unknown>
-      | undefined;
-    const usage = (metadata?.usage ?? responseMetadata.usage) as
-      | Record<string, unknown>
-      | undefined;
-
-    let enhancedUsageMetadata: UsageMetadata | undefined =
-      message.usage_metadata;
-
-    if (usage) {
-      const cacheRead = (usage.cacheReadInputTokens as number | undefined) ?? 0;
-      const cacheWrite =
-        (usage.cacheWriteInputTokens as number | undefined) ?? 0;
-      const inputTokens = (usage.inputTokens as number | undefined) ?? 0;
-      const outputTokens = (usage.outputTokens as number | undefined) ?? 0;
-
-      if (cacheRead > 0 || cacheWrite > 0) {
-        needsModification = true;
-        enhancedUsageMetadata = {
-          input_tokens: message.usage_metadata?.input_tokens ?? inputTokens,
-          output_tokens: message.usage_metadata?.output_tokens ?? outputTokens,
-          total_tokens:
-            message.usage_metadata?.total_tokens ??
-            (usage.totalTokens as number | undefined) ??
-            0,
-          input_token_details: {
-            cache_read: cacheRead,
-            cache_creation: cacheWrite,
-          },
-        };
-      }
+    let content: AIMessageChunk['content'] = message.content;
+    let contentModified = false;
+
+    if (Array.isArray(content) && blockIndex != null) {
+      content = content.map((block) =>
+        typeof block === 'object' && !('index' in block)
+          ? { ...block, index: blockIndex }
+          : block
+      );
+      contentModified = true;
+    } else if (
+      typeof content === 'string' &&
+      content !== '' &&
+      blockIndex != null &&
+      seenBlockIndices.size > 1
+    ) {
+      content = [{ type: 'text', text: content, index: blockIndex }];
+      contentModified = true;
     }
 
-    if (needsModification) {
-      return new ChatGenerationChunk({
-        text: chunk.text,
-        message: new AIMessageChunk({
-          ...message,
-          response_metadata: cleanedMetadata,
-          usage_metadata: enhancedUsageMetadata,
-        }),
-        generationInfo: chunk.generationInfo,
-      });
+    if (!contentModified && !hasMetadataIndex) {
+      return chunk;
     }
 
-    return chunk;
+    const cleanedMetadata = hasMetadataIndex
+      ? (this.removeContentBlockIndex(metadata) as Record<string, unknown>)
+      : metadata;
+
+    return new ChatGenerationChunk({
+      text: chunk.text,
+      message: new AIMessageChunk({
+        ...message,
+        content,
+        response_metadata: cleanedMetadata,
+      }),
+      generationInfo: chunk.generationInfo,
+    });
   }
 
   /**
-   * Check if contentBlockIndex exists at any level in the object
+   * Extract `contentBlockIndex` from the top level of response_metadata.
+   * Our custom handlers always place it at the top level.
    */
-  private hasContentBlockIndex(obj: unknown): boolean {
-    if (obj === null || obj === undefined || typeof obj !== 'object') {
-      return false;
-    }
-
-    if ('contentBlockIndex' in obj) {
-      return true;
-    }
-
-    for (const value of Object.values(obj)) {
-      if (typeof value === 'object' && value !== null) {
-        if (this.hasContentBlockIndex(value)) {
-          return true;
-        }
-      }
+  private extractContentBlockIndex(
+    metadata: Record<string, unknown>
+  ): number | undefined {
+    if (
+      'contentBlockIndex' in metadata &&
+      typeof metadata.contentBlockIndex === 'number'
+    ) {
+      return metadata.contentBlockIndex;
     }
-
-    return false;
+    return undefined;
   }
 
-  /**
-   * Recursively remove contentBlockIndex from all levels of an object
-   */
   private removeContentBlockIndex(obj: unknown): unknown {
     if (obj === null || obj === undefined) {
       return obj;
@@ -408,4 +442,7 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
   }
 }
 
+/** @deprecated Use IllumaBedrockConverse. Alias kept for backward compatibility with existing imports. */
+export const CustomChatBedrockConverse = IllumaBedrockConverse;
+
 export type { ChatBedrockConverseInput };