@illuma-ai/agents 1.0.81

package/src/llm/bedrock/index.ts

@@ -0,0 +1,411 @@
+/**
+ * Optimized ChatBedrockConverse wrapper that fixes contentBlockIndex conflicts
+ * and adds support for:
+ *
+ * - 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.
+ *
+ * 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.
+ *
+ * PROMPT CACHING:
+ * When promptCache: true is set, this wrapper adds cachePoint markers to the tools array
+ * to enable Bedrock prompt caching for tool definitions. This allows tool schemas to be
+ * 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.
+ */
+
+import { ChatBedrockConverse } from '@langchain/aws';
+import { AIMessageChunk } from '@langchain/core/messages';
+import type { BaseMessage, UsageMetadata } 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';
+
+/**
+ * Service tier type for Bedrock invocations.
+ * Requires AWS SDK >= 3.966.0 to actually work.
+ * @see https://docs.aws.amazon.com/bedrock/latest/userguide/service-tiers-inference.html
+ */
+export type ServiceTierType = 'priority' | 'default' | 'flex' | 'reserved';
+
+/**
+ * Extended input interface with additional features:
+ * - promptCache: Enable Bedrock prompt caching for tool definitions
+ * - applicationInferenceProfile: Use an inference profile ARN instead of model ID
+ * - serviceTier: Specify service tier (Priority, Standard, Flex, Reserved)
+ */
+export interface CustomChatBedrockConverseInput
+  extends ChatBedrockConverseInput {
+  /**
+   * Enable Bedrock prompt caching for tool definitions.
+   * When true, adds cachePoint markers to tools array.
+   */
+  promptCache?: boolean;
+
+  /**
+   * Application Inference Profile ARN to use for the model.
+   * For example, "arn:aws:bedrock:eu-west-1:123456789102:application-inference-profile/fm16bt65tzgx"
+   * When provided, this ARN will be used for the actual inference calls instead of the model ID.
+   * Must still provide `model` as normal modelId to benefit from all the metadata.
+   * @see https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-create.html
+   */
+  applicationInferenceProfile?: string;
+
+  /**
+   * Service tier for model invocation.
+   * Specifies the processing tier type used for serving the request.
+   * Supported values are 'priority', 'default', 'flex', and 'reserved'.
+   *
+   * - 'priority': Prioritized processing for lower latency
+   * - 'default': Standard processing tier
+   * - 'flex': Flexible processing tier with lower cost
+   * - 'reserved': Reserved capacity for consistent performance
+   *
+   * If not provided, AWS uses the default tier.
+   * Note: Requires AWS SDK >= 3.966.0 to work.
+   * @see https://docs.aws.amazon.com/bedrock/latest/userguide/service-tiers-inference.html
+   */
+  serviceTier?: ServiceTierType;
+}
+
+/**
+ * Extended call options with serviceTier override support.
+ */
+export interface CustomChatBedrockConverseCallOptions {
+  serviceTier?: ServiceTierType;
+}
+
+export class CustomChatBedrockConverse extends ChatBedrockConverse {
+  /** Enable Bedrock prompt caching for tool definitions */
+  promptCache: boolean;
+
+  /** Application Inference Profile ARN to use instead of model ID */
+  applicationInferenceProfile?: string;
+
+  /** Service tier for model invocation */
+  serviceTier?: ServiceTierType;
+
+  constructor(fields?: CustomChatBedrockConverseInput) {
+    super(fields);
+    this.promptCache = fields?.promptCache ?? false;
+    this.applicationInferenceProfile = fields?.applicationInferenceProfile;
+    this.serviceTier = fields?.serviceTier;
+
+    // Fix: Force supportsToolChoiceValues for Claude models
+    // The parent constructor checks `model.includes('claude-3')` but this fails when:
+    // 1. Using applicationInferenceProfile ARNs (arn:aws:bedrock:...)
+    // 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 isClaudeModel =
+      modelName.includes('claude') ||
+      modelName.includes('anthropic') ||
+      profileName.includes('claude') ||
+      profileName.includes('anthropic');
+
+    if (isClaudeModel && !this.supportsToolChoiceValues?.length) {
+      // Claude models support all tool choice values
+      this.supportsToolChoiceValues = ['auto', 'any', 'tool'];
+    }
+  }
+
+  static lc_name(): string {
+    return 'IllumaBedrockConverse';
+  }
+
+  /**
+   * Get the model ID to use for API calls.
+   * Returns applicationInferenceProfile if set, otherwise returns this.model.
+   */
+  protected getModelId(): string {
+    return this.applicationInferenceProfile ?? this.model;
+  }
+
+  /**
+   * Override invocationParams to:
+   * 1. Add cachePoint to tools when promptCache is enabled
+   * 2. Add serviceTier support
+   *
+   * CACHING STRATEGY: Separate cachePoints for core tools and MCP tools
+   * - Core tools (web_search, execute_code, etc.) are stable → cache first
+   * - MCP tools (have '_mcp_' in name) are dynamic → cache separately after
+   * - This allows core tools to stay cached when MCP selection changes
+   *
+   * NOTE: Only Claude models support cachePoint - Nova and other models will reject it.
+   */
+  override invocationParams(
+    options?: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions
+  ): ReturnType<ChatBedrockConverse['invocationParams']> & {
+    serviceTier?: { type: ServiceTierType };
+  } {
+    const params = super.invocationParams(options);
+
+    // Add cachePoint to tools array if promptCache is enabled and tools exist
+    // Only Claude models support cachePoint - check model name
+    const modelId = this.model.toLowerCase();
+    const isClaudeModel =
+      modelId.includes('claude') || modelId.includes('anthropic');
+
+    if (
+      this.promptCache &&
+      isClaudeModel &&
+      params.toolConfig?.tools &&
+      Array.isArray(params.toolConfig.tools) &&
+      params.toolConfig.tools.length > 0
+    ) {
+      // Separate core tools from MCP tools
+      // MCP tools have '_mcp_' in their name (e.g., 'search_emails_mcp_Google-Workspace')
+      const coreTools: typeof params.toolConfig.tools = [];
+      const mcpTools: typeof params.toolConfig.tools = [];
+
+      for (const tool of params.toolConfig.tools) {
+        // Check if tool has a name property with '_mcp_' pattern
+        const toolName =
+          (tool as { toolSpec?: { name?: string } }).toolSpec?.name ?? '';
+        if (toolName.includes('_mcp_')) {
+          mcpTools.push(tool);
+        } else {
+          coreTools.push(tool);
+        }
+      }
+
+      // Build tools array with strategic cachePoints:
+      // [CoreTool1, CoreTool2, cachePoint] + [MCPTool1, MCPTool2, cachePoint]
+      const toolsWithCache: typeof params.toolConfig.tools = [];
+
+      // Add core tools with cachePoint (if any)
+      if (coreTools.length > 0) {
+        toolsWithCache.push(...coreTools);
+        toolsWithCache.push({ cachePoint: { type: 'default' } });
+      }
+
+      // Add MCP tools with their own cachePoint (if any)
+      if (mcpTools.length > 0) {
+        toolsWithCache.push(...mcpTools);
+        toolsWithCache.push({ cachePoint: { type: 'default' } });
+      }
+
+      // If no tools at all (shouldn't happen but safety check)
+      if (toolsWithCache.length === 0) {
+        toolsWithCache.push({ cachePoint: { type: 'default' } });
+      }
+
+      params.toolConfig.tools = toolsWithCache;
+    }
+
+    // Add serviceTier support
+    const serviceTierType = options?.serviceTier ?? this.serviceTier;
+
+    return {
+      ...params,
+      serviceTier: serviceTierType ? { type: serviceTierType } : undefined,
+    };
+  }
+
+  /**
+   * Override _generateNonStreaming to use applicationInferenceProfile as modelId.
+   * Uses the same model-swapping pattern as streaming for consistency.
+   */
+  override async _generateNonStreaming(
+    messages: BaseMessage[],
+    options: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions,
+    runManager?: CallbackManagerForLLMRun
+  ): Promise<ChatResult> {
+    // Temporarily swap model for applicationInferenceProfile support
+    const originalModel = this.model;
+    if (
+      this.applicationInferenceProfile != null &&
+      this.applicationInferenceProfile !== ''
+    ) {
+      this.model = this.applicationInferenceProfile;
+    }
+
+    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
+   *
+   * Note: We delegate to super._streamResponseChunks() to preserve @langchain/aws's
+   * internal chunk handling which correctly preserves array content for reasoning blocks.
+   */
+  override async *_streamResponseChunks(
+    messages: BaseMessage[],
+    options: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions,
+    runManager?: CallbackManagerForLLMRun
+  ): AsyncGenerator<ChatGenerationChunk> {
+    // Temporarily swap model for applicationInferenceProfile support
+    const originalModel = this.model;
+    if (
+      this.applicationInferenceProfile != null &&
+      this.applicationInferenceProfile !== ''
+    ) {
+      this.model = this.applicationInferenceProfile;
+    }
+
+    try {
+      // Use parent's streaming logic which correctly handles reasoning content
+      const baseStream = super._streamResponseChunks(
+        messages,
+        options,
+        runManager
+      );
+
+      for await (const chunk of baseStream) {
+        // Clean and enhance chunk
+        yield this.processChunk(chunk);
+      }
+    } 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
+   */
+  private processChunk(chunk: ChatGenerationChunk): 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;
+    }
+
+    // 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,
+          },
+        };
+      }
+    }
+
+    if (needsModification) {
+      return new ChatGenerationChunk({
+        text: chunk.text,
+        message: new AIMessageChunk({
+          ...message,
+          response_metadata: cleanedMetadata,
+          usage_metadata: enhancedUsageMetadata,
+        }),
+        generationInfo: chunk.generationInfo,
+      });
+    }
+
+    return chunk;
+  }
+
+  /**
+   * Check if contentBlockIndex exists at any level in the object
+   */
+  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;
+        }
+      }
+    }
+
+    return false;
+  }
+
+  /**
+   * Recursively remove contentBlockIndex from all levels of an object
+   */
+  private removeContentBlockIndex(obj: unknown): unknown {
+    if (obj === null || obj === undefined) {
+      return obj;
+    }
+
+    if (Array.isArray(obj)) {
+      return obj.map((item) => this.removeContentBlockIndex(item));
+    }
+
+    if (typeof obj === 'object') {
+      const cleaned: Record<string, unknown> = {};
+      for (const [key, value] of Object.entries(obj)) {
+        if (key !== 'contentBlockIndex') {
+          cleaned[key] = this.removeContentBlockIndex(value);
+        }
+      }
+      return cleaned;
+    }
+
+    return obj;
+  }
+}
+
+export type { ChatBedrockConverseInput };