@librechat/agents 3.0.776 → 3.1.1

package/dist/types/tools/search/schema.d.ts

@@ -1,4 +1,3 @@
-import { z } from 'zod';
 export declare enum DATE_RANGE {
     PAST_HOUR = "h",
     PAST_24_HOURS = "d",
@@ -8,9 +7,28 @@ export declare enum DATE_RANGE {
 }
 export declare const DEFAULT_QUERY_DESCRIPTION: string;
 export declare const DEFAULT_COUNTRY_DESCRIPTION: string;
-export declare const querySchema: z.ZodString;
-export declare const dateSchema: z.ZodOptional<z.ZodNativeEnum<typeof DATE_RANGE>>;
-export declare const countrySchema: z.ZodOptional<z.ZodString>;
-export declare const imagesSchema: z.ZodOptional<z.ZodBoolean>;
-export declare const videosSchema: z.ZodOptional<z.ZodBoolean>;
-export declare const newsSchema: z.ZodOptional<z.ZodBoolean>;
+export declare const querySchema: {
+    readonly type: "string";
+    readonly description: string;
+};
+export declare const dateSchema: {
+    readonly type: "string";
+    readonly enum: DATE_RANGE[];
+    readonly description: "Date range for search results.";
+};
+export declare const countrySchema: {
+    readonly type: "string";
+    readonly description: string;
+};
+export declare const imagesSchema: {
+    readonly type: "boolean";
+    readonly description: "Whether to also run an image search.";
+};
+export declare const videosSchema: {
+    readonly type: "boolean";
+    readonly description: "Whether to also run a video search.";
+};
+export declare const newsSchema: {
+    readonly type: "boolean";
+    readonly description: "Whether to also run a news search.";
+};

package/dist/types/tools/search/tool.d.ts

@@ -1,54 +1,3 @@
-import { z } from 'zod';
 import { DynamicStructuredTool } from '@langchain/core/tools';
 import type * as t from './types';
-import { DATE_RANGE } from './schema';
-/**
- * Creates a search tool with a schema that dynamically includes the country field
- * only when the searchProvider is 'serper'.
- *
- * Supports multiple scraper providers:
- * - Firecrawl (default): Full-featured web scraping with multiple formats
- * - Serper: Lightweight scraping using Serper's scrape API
- *
- * @example
- * ```typescript
- * // Using Firecrawl scraper (default)
- * const searchTool = createSearchTool({
- *   searchProvider: 'serper',
- *   scraperProvider: 'firecrawl',
- *   firecrawlApiKey: 'your-firecrawl-key'
- * });
- *
- * // Using Serper scraper
- * const searchTool = createSearchTool({
- *   searchProvider: 'serper',
- *   scraperProvider: 'serper',
- *   serperApiKey: 'your-serper-key'
- * });
- * ```
- *
- * @param config - The search tool configuration
- * @returns A DynamicStructuredTool with a schema that depends on the searchProvider
- */
-export declare const createSearchTool: (config?: t.SearchToolConfig) => DynamicStructuredTool<z.ZodObject<{
-    query: z.ZodString;
-    date: z.ZodOptional<z.ZodNativeEnum<typeof DATE_RANGE>>;
-    country?: z.ZodOptional<z.ZodString>;
-    images: z.ZodOptional<z.ZodBoolean>;
-    videos: z.ZodOptional<z.ZodBoolean>;
-    news: z.ZodOptional<z.ZodBoolean>;
-}, "strip", z.ZodTypeAny, {
-    query: string;
-    videos?: boolean | undefined;
-    images?: boolean | undefined;
-    news?: boolean | undefined;
-    date?: DATE_RANGE | undefined;
-    country?: unknown;
-}, {
-    query: string;
-    videos?: boolean | undefined;
-    images?: boolean | undefined;
-    news?: boolean | undefined;
-    date?: DATE_RANGE | undefined;
-    country?: unknown;
-}>>;
+export declare const createSearchTool: (config?: t.SearchToolConfig) => DynamicStructuredTool;

package/dist/types/tools/search/types.d.ts

@@ -1,4 +1,3 @@
-import { z } from 'zod';
 import type { Logger as WinstonLogger } from 'winston';
 import type { RunnableConfig } from '@langchain/core/runnables';
 import type { BaseReranker } from './rerankers';
@@ -567,25 +566,8 @@ export type ProcessSourcesFields = {
     proMode: boolean;
     onGetHighlights: SearchToolConfig['onGetHighlights'];
 };
-export type SearchToolSchema = z.ZodObject<{
-    query: z.ZodString;
-    date: z.ZodOptional<z.ZodNativeEnum<typeof DATE_RANGE>>;
-    country?: z.ZodOptional<z.ZodString>;
-    images: z.ZodOptional<z.ZodBoolean>;
-    videos: z.ZodOptional<z.ZodBoolean>;
-    news: z.ZodOptional<z.ZodBoolean>;
-}, 'strip', z.ZodTypeAny, {
-    query: string;
-    date?: DATE_RANGE;
-    country?: unknown;
-    images?: boolean;
-    videos?: boolean;
-    news?: boolean;
-}, {
-    query: string;
-    date?: DATE_RANGE;
-    country?: unknown;
-    images?: boolean;
-    videos?: boolean;
-    news?: boolean;
-}>;
+export interface SearchToolSchema {
+    type: 'object';
+    properties: Record<string, unknown>;
+    required: string[];
+}

package/dist/types/types/tools.d.ts

@@ -57,6 +57,8 @@ export type FileRef = {
     id: string;
     name: string;
     path?: string;
+    /** Session ID this file belongs to (for multi-session file tracking) */
+    session_id?: string;
 };
 export type FileRefs = FileRef[];
 export type ExecuteResult = {

package/dist/types/utils/index.d.ts

@@ -4,3 +4,4 @@ export * from './misc';
 export * from './handlers';
 export * from './run';
 export * from './tokens';
+export * from './schema';

package/dist/types/utils/schema.d.ts

@@ -0,0 +1,8 @@
+import type { ZodTypeAny } from 'zod';
+/** Checks if a schema is a Zod schema by looking for the _def property */
+export declare function isZodSchema(schema: unknown): schema is ZodTypeAny;
+/**
+ * Converts a schema to JSON schema format.
+ * Handles both Zod schemas (converts) and JSON schemas (passthrough).
+ */
+export declare function toJsonSchema(schema: unknown, name?: string, description?: string): Record<string, unknown>;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@librechat/agents",
-  "version": "3.0.776",
+  "version": "3.1.01",
   "main": "./dist/cjs/main.cjs",
   "module": "./dist/esm/main.mjs",
   "types": "./dist/types/index.d.ts",
@@ -51,12 +51,14 @@
     "image": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/image.ts --provider 'google' --name 'Jo' --location 'New York, NY'",
     "code_exec_files": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/code_exec_files.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
     "code_exec_session": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/code_exec_session.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
+    "code_exec_multi_session": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/code_exec_multi_session.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
     "code_exec_simple": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/code_exec_simple.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
     "simple": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/simple.ts --provider 'bedrock' --name 'Jo' --location 'New York, NY'",
     "caching": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/caching.ts --name 'Jo' --location 'New York, NY'",
     "thinking": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/thinking.ts --name 'Jo' --location 'New York, NY'",
+    "thinking:bedrock": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/thinking-bedrock.ts --name 'Jo' --location 'New York, NY'",
     "memory": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/memory.ts --provider 'openAI' --name 'Jo' --location 'New York, NY'",
-    "tool": "node --trace-warnings -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/tools.ts --provider 'openrouter' --name 'Jo' --location 'New York, NY'",
+    "tool": "node --trace-warnings -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/tools.ts --provider 'bedrock' --name 'Jo' --location 'New York, NY'",
     "search": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/search.ts --provider 'bedrock' --name 'Jo' --location 'New York, NY'",
     "tool_search": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/tool_search.ts",
     "programmatic_exec": "node -r dotenv/config --loader ./tsconfig-paths-bootstrap.mjs --experimental-specifier-resolution=node ./src/scripts/programmatic_exec.ts",
@@ -110,6 +112,7 @@
     }
   },
   "dependencies": {
+    "@aws-sdk/client-bedrock-runtime": "^3.970.0",
     "@langchain/anthropic": "^0.3.26",
     "@langchain/aws": "^0.1.15",
     "@langchain/core": "^0.3.80",

package/src/agents/AgentContext.ts

@@ -1,6 +1,5 @@
 /* eslint-disable no-console */
 // src/agents/AgentContext.ts
-import { zodToJsonSchema } from 'zod-to-json-schema';
 import { SystemMessage } from '@langchain/core/messages';
 import { RunnableLambda } from '@langchain/core/runnables';
 import type {
@@ -12,6 +11,7 @@ import type { RunnableConfig, Runnable } from '@langchain/core/runnables';
 import type * as t from '@/types';
 import type { createPruneMessages } from '@/messages';
 import { ContentTypes, Providers } from '@/common';
+import { toJsonSchema } from '@/utils/schema';
 
 /**
  * Encapsulates agent-specific state that can vary between agents in a multi-agent system
@@ -479,15 +479,10 @@ export class AgentContext {
           genericTool.schema != null &&
           typeof genericTool.schema === 'object'
         ) {
-          const schema = genericTool.schema as {
-            describe: (desc: string) => unknown;
-          };
-          const describedSchema = schema.describe(
-            (genericTool.description as string) || ''
-          );
-          const jsonSchema = zodToJsonSchema(
-            describedSchema as Parameters<typeof zodToJsonSchema>[0],
-            (genericTool.name as string) || ''
+          const jsonSchema = toJsonSchema(
+            genericTool.schema,
+            (genericTool.name as string | undefined) ?? '',
+            (genericTool.description as string | undefined) ?? ''
           );
           toolTokens += tokenCounter(
             new SystemMessage(JSON.stringify(jsonSchema))
@@ -496,7 +491,6 @@ export class AgentContext {
       }
     }
 
-    // Add tool tokens to existing instruction tokens (which may already include system message tokens)
     this.instructionTokens += toolTokens;
   }
 

package/src/graphs/Graph.ts

@@ -1042,8 +1042,8 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
 
     /**
      * Extract and store code execution session context from artifacts.
-     * Only update session_id when files are generated - this ensures we don't
-     * lose the original session that contains the files.
+     * Each file is stamped with its source session_id to support multi-session file tracking.
+     * When the same filename appears in a later execution, the newer version replaces the old.
      */
     const toolName = output.name;
     if (
@@ -1060,17 +1060,35 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         artifact.session_id !== ''
       ) {
         /**
-         * Files were generated - update session with the new session_id.
-         * The new session_id is the one that contains these files.
+         * Stamp each new file with its source session_id.
+         * This enables files from different executions (parallel or sequential)
+         * to be tracked and passed to subsequent calls.
          */
+        const filesWithSession: t.FileRefs = newFiles.map((file) => ({
+          ...file,
+          session_id: artifact.session_id,
+        }));
+
         const existingSession = this.sessions.get(Constants.EXECUTE_CODE) as
           | t.CodeSessionContext
           | undefined;
         const existingFiles = existingSession?.files ?? [];
 
+        /**
+         * Merge files, preferring latest versions by name.
+         * If a file with the same name exists, replace it with the new version.
+         * This handles cases where files are edited/recreated in subsequent executions.
+         */
+        const newFileNames = new Set(filesWithSession.map((f) => f.name));
+        const filteredExisting = existingFiles.filter(
+          (f) => !newFileNames.has(f.name)
+        );
+
         this.sessions.set(Constants.EXECUTE_CODE, {
+          /** Keep latest session_id for reference/fallback */
           session_id: artifact.session_id,
-          files: [...existingFiles, ...newFiles],
+          /** Accumulated files with latest versions preferred */
+          files: [...filteredExisting, ...filesWithSession],
           lastUpdated: Date.now(),
         });
       }

package/src/graphs/MultiAgentGraph.ts

@@ -1,4 +1,3 @@
-import { z } from 'zod';
 import { tool } from '@langchain/core/tools';
 import { PromptTemplate } from '@langchain/core/prompts';
 import {
@@ -292,7 +291,8 @@ export class MultiAgentGraph extends StandardGraph {
 
       tools.push(
         tool(
-          async (input: Record<string, unknown>, config) => {
+          async (rawInput, config) => {
+            const input = rawInput as Record<string, unknown>;
             const state = getCurrentTaskInput() as t.BaseGraphState;
             const toolCallId =
               (config as ToolRunnableConfig | undefined)?.toolCall?.id ??
@@ -343,13 +343,17 @@ export class MultiAgentGraph extends StandardGraph {
           {
             name: toolName,
             schema: hasHandoffInput
-              ? z.object({
-                [promptKey]: z
-                  .string()
-                  .optional()
-                  .describe(handoffInputDescription as string),
-              })
-              : z.object({}),
+              ? {
+                type: 'object',
+                properties: {
+                  [promptKey]: {
+                    type: 'string',
+                    description: handoffInputDescription as string,
+                  },
+                },
+                required: [],
+              }
+              : { type: 'object', properties: {}, required: [] },
             description: toolDescription,
           }
         )
@@ -371,7 +375,8 @@ export class MultiAgentGraph extends StandardGraph {
 
         tools.push(
           tool(
-            async (input: Record<string, unknown>, config) => {
+            async (rawInput, config) => {
+              const input = rawInput as Record<string, unknown>;
               const toolCallId =
                 (config as ToolRunnableConfig | undefined)?.toolCall?.id ??
                 'unknown';
@@ -470,13 +475,17 @@ export class MultiAgentGraph extends StandardGraph {
             {
               name: toolName,
               schema: hasHandoffInput
-                ? z.object({
-                  [promptKey]: z
-                    .string()
-                    .optional()
-                    .describe(handoffInputDescription as string),
-                })
-                : z.object({}),
+                ? {
+                  type: 'object',
+                  properties: {
+                    [promptKey]: {
+                      type: 'string',
+                      description: handoffInputDescription as string,
+                    },
+                  },
+                  required: [],
+                }
+                : { type: 'object', properties: {}, required: [] },
               description: toolDescription,
             }
           )

package/src/llm/bedrock/index.ts

@@ -1,5 +1,9 @@
 /**
  * Optimized ChatBedrockConverse wrapper that fixes contentBlockIndex conflicts
+ * and adds support for latest @langchain/aws features:
+ *
+ * - 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"
@@ -12,15 +16,74 @@
  */
 
 import { ChatBedrockConverse } from '@langchain/aws';
-import type { ChatBedrockConverseInput } from '@langchain/aws';
 import { AIMessageChunk } from '@langchain/core/messages';
-import type { BaseMessage } from '@langchain/core/messages';
-import { ChatGenerationChunk } from '@langchain/core/outputs';
+import { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs';
 import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
+import type { ChatBedrockConverseInput } from '@langchain/aws';
+import type { BaseMessage } from '@langchain/core/messages';
+
+/**
+ * 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:
+ * - applicationInferenceProfile: Use an inference profile ARN instead of model ID
+ * - serviceTier: Specify service tier (Priority, Standard, Flex, Reserved)
+ */
+export interface CustomChatBedrockConverseInput
+  extends ChatBedrockConverseInput {
+  /**
+   * 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 {
-  constructor(fields?: ChatBedrockConverseInput) {
+  /**
+   * 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.applicationInferenceProfile = fields?.applicationInferenceProfile;
+    this.serviceTier = fields?.serviceTier;
   }
 
   static lc_name(): string {
@@ -28,52 +91,126 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
   }
 
   /**
-   * Override _streamResponseChunks to strip contentBlockIndex from response_metadata
-   * This prevents LangChain's merge conflicts when the same index is used for
-   * different content types (text vs tool calls)
+   * 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 add serviceTier support.
+   */
+  override invocationParams(
+    options?: this['ParsedCallOptions'] & CustomChatBedrockConverseCallOptions
+  ): ReturnType<ChatBedrockConverse['invocationParams']> & {
+    serviceTier?: { type: ServiceTierType };
+  } {
+    const baseParams = super.invocationParams(options);
+
+    /** Service tier from options or fall back to class-level setting */
+    const serviceTierType = options?.serviceTier ?? this.serviceTier;
+
+    return {
+      ...baseParams,
+      serviceTier: serviceTierType ? { type: serviceTierType } : undefined,
+    };
+  }
+
+  /**
+   * Override _generateNonStreaming to use applicationInferenceProfile as modelId.
+   * Uses the same model-swapping pattern as streaming for consistency.
    */
-  async *_streamResponseChunks(
+  override async _generateNonStreaming(
     messages: BaseMessage[],
-    options: this['ParsedCallOptions'],
+    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
+   *
+   * 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> {
-    const baseStream = super._streamResponseChunks(
-      messages,
-      options,
-      runManager
-    );
-
-    for await (const chunk of baseStream) {
-      // Only process if we have response_metadata
-      if (
-        chunk.message instanceof AIMessageChunk &&
-        (chunk.message as Partial<AIMessageChunk>).response_metadata &&
-        typeof chunk.message.response_metadata === 'object'
-      ) {
-        // Check if contentBlockIndex exists anywhere in response_metadata (top level or nested)
-        const hasContentBlockIndex = this.hasContentBlockIndex(
-          chunk.message.response_metadata
-        );
-
-        if (hasContentBlockIndex) {
-          const cleanedMetadata = this.removeContentBlockIndex(
-            chunk.message.response_metadata
-          ) as Record<string, unknown>;
-
-          yield new ChatGenerationChunk({
-            text: chunk.text,
-            message: new AIMessageChunk({
-              ...chunk.message,
-              response_metadata: cleanedMetadata,
-            }),
-            generationInfo: chunk.generationInfo,
-          });
-          continue;
-        }
+    // 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 contentBlockIndex from response_metadata to prevent merge conflicts
+        yield this.cleanChunk(chunk);
       }
+    } finally {
+      // Restore original model
+      this.model = originalModel;
+    }
+  }
 
-      yield chunk;
+  /**
+   * Clean a chunk by removing contentBlockIndex from response_metadata.
+   */
+  private cleanChunk(chunk: ChatGenerationChunk): ChatGenerationChunk {
+    const message = chunk.message;
+    if (!(message instanceof AIMessageChunk)) {
+      return chunk;
     }
+
+    const metadata = message.response_metadata as Record<string, unknown>;
+    const hasContentBlockIndex = this.hasContentBlockIndex(metadata);
+    if (!hasContentBlockIndex) {
+      return chunk;
+    }
+
+    const cleanedMetadata = this.removeContentBlockIndex(metadata) as Record<
+      string,
+      unknown
+    >;
+
+    return new ChatGenerationChunk({
+      text: chunk.text,
+      message: new AIMessageChunk({
+        ...message,
+        response_metadata: cleanedMetadata,
+      }),
+      generationInfo: chunk.generationInfo,
+    });
   }
 
   /**