@librechat/agents 3.1.97 → 3.1.99

package/src/llm/bedrock/llm.spec.ts

@@ -18,14 +18,16 @@ import {
   ConverseCommand,
   ConverseStreamCommand,
 } from '@aws-sdk/client-bedrock-runtime';
-import type { ConverseResponse } from '@aws-sdk/client-bedrock-runtime';
+import type { ConverseResponse, Tool } from '@aws-sdk/client-bedrock-runtime';
 import {
   convertConverseMessageToLangChainMessage,
   handleConverseStreamMetadata,
   convertToConverseMessages,
 } from './utils';
+import type { GraphTools } from '@/types';
 import { toLangChainContent } from '@/messages/langchain';
 import { CustomChatBedrockConverse, ServiceTierType } from './index';
+import { partitionAndMarkBedrockToolCache } from './toolCache';
 
 jest.setTimeout(120000);
 
@@ -85,6 +87,17 @@ function humanMessageWithContent(
   return new HumanMessage({ content: toLangChainContent(content) });
 }
 
+function getBedrockToolName(entry: Tool): string {
+  if ('cachePoint' in entry) {
+    return 'cachePoint';
+  }
+  return entry.toolSpec?.name ?? 'missing';
+}
+
+function getBedrockToolNames(tools: Tool[] | undefined): string[] {
+  return (tools ?? []).map(getBedrockToolName);
+}
+
 describe('CustomChatBedrockConverse', () => {
   describe('applicationInferenceProfile parameter', () => {
     test('should initialize applicationInferenceProfile from constructor', () => {
@@ -240,6 +253,51 @@ describe('CustomChatBedrockConverse', () => {
     });
   });
 
+  describe('guardrailConfig configuration', () => {
+    test('should pass guardrailConfig through to ConverseStreamCommand', async () => {
+      const guardrailConfig = {
+        guardrailIdentifier: 'test-guardrail-id',
+        guardrailVersion: 'DRAFT',
+        trace: 'enabled_full' as const,
+        streamProcessingMode: 'sync' as const,
+      };
+      const mockSend = jest.fn<any>().mockResolvedValue({
+        stream: (async function* streamChunks() {
+          yield {
+            contentBlockDelta: {
+              contentBlockIndex: 0,
+              delta: { text: 'Guardrail response' },
+            },
+          };
+        })(),
+      });
+
+      const mockClient = {
+        send: mockSend,
+      } as unknown as BedrockRuntimeClient;
+
+      const model = new CustomChatBedrockConverse({
+        ...baseConstructorArgs,
+        model: 'anthropic.claude-3-haiku-20240307-v1:0',
+        guardrailConfig,
+        client: mockClient,
+      });
+
+      let chunks = 0;
+      for await (const _chunk of await model.stream([
+        new HumanMessage('Hello'),
+      ])) {
+        chunks += 1;
+      }
+
+      expect(mockSend).toHaveBeenCalledTimes(1);
+      expect(chunks).toBe(1);
+      const commandArg = mockSend.mock.calls[0][0] as ConverseStreamCommand;
+      expect(commandArg).toBeInstanceOf(ConverseStreamCommand);
+      expect(commandArg.input.guardrailConfig).toEqual(guardrailConfig);
+    });
+  });
+
   describe('serviceTier configuration', () => {
     test('should set serviceTier in constructor', () => {
       const model = new CustomChatBedrockConverse({
@@ -313,6 +371,101 @@ describe('CustomChatBedrockConverse', () => {
     });
   });
 
+  describe('promptCache tool configuration', () => {
+    test('adds a Bedrock cache point for directly-bound tools', () => {
+      const model = new CustomChatBedrockConverse({
+        ...baseConstructorArgs,
+        promptCache: true,
+      });
+
+      const params = model.invocationParams({
+        tools: [
+          {
+            type: 'function',
+            function: {
+              name: 'direct_tool',
+              description: 'Direct tool',
+              parameters: { type: 'object', properties: {} },
+            },
+          },
+        ],
+      });
+
+      expect(getBedrockToolNames(params.toolConfig?.tools)).toEqual([
+        'direct_tool',
+        'cachePoint',
+      ]);
+    });
+
+    test('adds the Bedrock cache point before deferred tools', () => {
+      const model = new CustomChatBedrockConverse({
+        ...baseConstructorArgs,
+        promptCache: true,
+      });
+      const tools = partitionAndMarkBedrockToolCache(
+        [
+          {
+            type: 'function',
+            function: {
+              name: 'static_tool',
+              description: 'Static tool',
+              parameters: { type: 'object', properties: {} },
+            },
+          },
+          {
+            type: 'function',
+            function: {
+              name: 'deferred_tool',
+              description: 'Deferred tool',
+              parameters: { type: 'object', properties: {} },
+            },
+          },
+        ] as GraphTools,
+        (name) => name === 'deferred_tool'
+      );
+
+      const params = model.invocationParams({ tools });
+
+      expect(getBedrockToolNames(params.toolConfig?.tools)).toEqual([
+        'static_tool',
+        'cachePoint',
+        'deferred_tool',
+      ]);
+      expect(JSON.stringify(params.toolConfig?.tools)).not.toContain(
+        '__lc_bedrock_cache_point_after'
+      );
+    });
+
+    test('does not fall back to caching when Graph marks all tools deferred', () => {
+      const model = new CustomChatBedrockConverse({
+        ...baseConstructorArgs,
+        promptCache: true,
+      });
+      const tools = partitionAndMarkBedrockToolCache(
+        [
+          {
+            type: 'function',
+            function: {
+              name: 'deferred_tool',
+              description: 'Deferred tool',
+              parameters: { type: 'object', properties: {} },
+            },
+          },
+        ] as GraphTools,
+        () => true
+      );
+
+      const params = model.invocationParams({ tools });
+
+      expect(getBedrockToolNames(params.toolConfig?.tools)).toEqual([
+        'deferred_tool',
+      ]);
+      expect(JSON.stringify(params.toolConfig?.tools)).not.toContain(
+        '__lc_bedrock_skip_tool_cache'
+      );
+    });
+  });
+
   describe('contentBlockIndex cleanup', () => {
     // Access private methods for testing via any cast
     function getModelWithCleanMethods() {

package/src/llm/bedrock/toolCache.test.ts

@@ -0,0 +1,131 @@
+import { tool } from '@langchain/core/tools';
+import type { Tool } from '@aws-sdk/client-bedrock-runtime';
+import type { GraphTools } from '@/types';
+import {
+  insertBedrockToolCachePoint,
+  partitionAndMarkBedrockToolCache,
+} from './toolCache';
+
+type OpenAITool = {
+  type: 'function';
+  function: {
+    name: string;
+    description?: string;
+    parameters?: object;
+  };
+};
+
+function createOpenAITool(name: string): OpenAITool {
+  return {
+    type: 'function',
+    function: {
+      name,
+      description: `${name} description`,
+      parameters: {
+        type: 'object',
+        properties: {},
+      },
+    },
+  };
+}
+
+function toolName(entry: Tool): string {
+  if ('cachePoint' in entry) {
+    return 'cachePoint';
+  }
+  return entry.toolSpec?.name ?? 'missing';
+}
+
+function toolNames(tools: Tool[] | undefined): string[] {
+  return (tools ?? []).map(toolName);
+}
+
+describe('partitionAndMarkBedrockToolCache', () => {
+  it('inserts the Bedrock cache point after the last static tool', () => {
+    const tools = [
+      createOpenAITool('static_one'),
+      createOpenAITool('static_two'),
+      createOpenAITool('dynamic_one'),
+    ] as GraphTools;
+
+    const marked = partitionAndMarkBedrockToolCache(
+      tools,
+      (name) => name === 'dynamic_one'
+    ) as Tool[];
+    const result = insertBedrockToolCachePoint({ tools: marked }, false);
+
+    expect(toolNames(result?.tools)).toEqual([
+      'static_one',
+      'static_two',
+      'cachePoint',
+      'dynamic_one',
+    ]);
+    expect(JSON.stringify(result?.tools)).not.toContain(
+      '__lc_bedrock_cache_point_after'
+    );
+  });
+
+  it('converts LangChain tools to Bedrock tool specs before marking', () => {
+    const staticTool = tool(async () => 'static', {
+      name: 'static_tool',
+      description: 'Static tool',
+      schema: {
+        type: 'object',
+        properties: {},
+      },
+    });
+    const dynamicTool = tool(async () => 'dynamic', {
+      name: 'dynamic_tool',
+      description: 'Dynamic tool',
+      schema: {
+        type: 'object',
+        properties: {},
+      },
+    });
+
+    const marked = partitionAndMarkBedrockToolCache(
+      [dynamicTool, staticTool] as GraphTools,
+      (name) => name === 'dynamic_tool'
+    ) as Tool[];
+    const result = insertBedrockToolCachePoint({ tools: marked }, false);
+
+    expect(toolNames(result?.tools)).toEqual([
+      'static_tool',
+      'cachePoint',
+      'dynamic_tool',
+    ]);
+  });
+
+  it('does not add a cache point when every tool is deferred', () => {
+    const tools = [createOpenAITool('dynamic_one')] as GraphTools;
+    const marked = partitionAndMarkBedrockToolCache(
+      tools,
+      () => true
+    ) as Tool[];
+    const result = insertBedrockToolCachePoint({ tools: marked }, false);
+
+    expect(toolNames(result?.tools)).toEqual(['dynamic_one']);
+    expect(JSON.stringify(result?.tools)).not.toContain(
+      '__lc_bedrock_skip_tool_cache'
+    );
+  });
+
+  it('can fall back to caching all directly-bound tools', () => {
+    const result = insertBedrockToolCachePoint(
+      {
+        tools: [
+          {
+            toolSpec: {
+              name: 'direct_tool',
+              description: 'Direct tool',
+              inputSchema: { json: { type: 'object', properties: {} } },
+            },
+          },
+        ],
+      },
+      true
+    );
+
+    expect(toolNames(result?.tools)).toEqual(['direct_tool', 'cachePoint']);
+  });
+});

package/src/llm/bedrock/toolCache.ts

@@ -0,0 +1,191 @@
+import type { Tool, ToolConfiguration } from '@aws-sdk/client-bedrock-runtime';
+import type { BindToolsInput } from '@langchain/core/language_models/chat_models';
+import type { OpenAIClient } from '@langchain/openai';
+import type { DocumentType } from '@smithy/types';
+import type { GraphTools } from '@/types';
+import { _convertToOpenAITool } from '@/llm/openai';
+
+const CACHE_POINT: Tool.CachePointMember = {
+  cachePoint: { type: 'default' },
+};
+
+const BEDROCK_TOOL_CACHE_MARKER = '__lc_bedrock_cache_point_after';
+const BEDROCK_TOOL_CACHE_DISABLED_MARKER = '__lc_bedrock_skip_tool_cache';
+
+type BedrockToolWithCacheMarker = Tool & {
+  [BEDROCK_TOOL_CACHE_MARKER]?: true;
+  [BEDROCK_TOOL_CACHE_DISABLED_MARKER]?: true;
+};
+
+type OpenAIFunctionTool = Extract<
+  OpenAIClient.ChatCompletionTool,
+  { type: 'function' }
+>;
+
+type ToolNameCandidate = {
+  name?: unknown;
+  function?: {
+    name?: unknown;
+  };
+  toolSpec?: {
+    name?: unknown;
+  };
+};
+
+function isBedrockToolSpec(tool: unknown): tool is Tool.ToolSpecMember {
+  return (
+    typeof tool === 'object' &&
+    tool != null &&
+    'toolSpec' in tool &&
+    typeof (tool as ToolNameCandidate).toolSpec?.name === 'string'
+  );
+}
+
+function isBedrockCachePoint(tool: Tool): tool is Tool.CachePointMember {
+  return 'cachePoint' in tool && tool.cachePoint != null;
+}
+
+function getToolName(tool: unknown): string | undefined {
+  const candidate = tool as ToolNameCandidate;
+  if (typeof candidate.toolSpec?.name === 'string') {
+    return candidate.toolSpec.name;
+  }
+  if (typeof candidate.name === 'string') {
+    return candidate.name;
+  }
+  if (typeof candidate.function?.name === 'string') {
+    return candidate.function.name;
+  }
+  return undefined;
+}
+
+function openAIToBedrockTool(tool: OpenAIFunctionTool): Tool.ToolSpecMember {
+  return {
+    toolSpec: {
+      name: tool.function.name,
+      description: tool.function.description,
+      inputSchema: { json: tool.function.parameters as DocumentType },
+    },
+  };
+}
+
+function toBedrockTool(tool: unknown): BedrockToolWithCacheMarker {
+  if (isBedrockToolSpec(tool)) {
+    return { ...tool };
+  }
+
+  return openAIToBedrockTool(
+    _convertToOpenAITool(tool as BindToolsInput) as OpenAIFunctionTool
+  ) as BedrockToolWithCacheMarker;
+}
+
+function markCachePointAfter(
+  tool: BedrockToolWithCacheMarker
+): BedrockToolWithCacheMarker {
+  return {
+    ...tool,
+    [BEDROCK_TOOL_CACHE_MARKER]: true,
+  };
+}
+
+function markToolCacheDisabled(
+  tool: BedrockToolWithCacheMarker
+): BedrockToolWithCacheMarker {
+  return {
+    ...tool,
+    [BEDROCK_TOOL_CACHE_DISABLED_MARKER]: true,
+  };
+}
+
+function stripCachePointMarker(tool: BedrockToolWithCacheMarker): Tool {
+  const {
+    [BEDROCK_TOOL_CACHE_MARKER]: _marker,
+    [BEDROCK_TOOL_CACHE_DISABLED_MARKER]: _disabled,
+    ...rest
+  } = tool;
+  return rest as Tool;
+}
+
+export function partitionAndMarkBedrockToolCache(
+  tools: GraphTools | undefined,
+  isDeferred: (toolName: string) => boolean
+): GraphTools | undefined {
+  if (tools == null || tools.length === 0) {
+    return tools;
+  }
+
+  const staticTools: BedrockToolWithCacheMarker[] = [];
+  const deferredTools: BedrockToolWithCacheMarker[] = [];
+
+  for (const tool of tools as readonly unknown[]) {
+    const converted = toBedrockTool(tool);
+    const name = getToolName(converted) ?? getToolName(tool);
+
+    if (name != null && isDeferred(name)) {
+      deferredTools.push(converted);
+      continue;
+    }
+
+    staticTools.push(converted);
+  }
+
+  if (staticTools.length === 0) {
+    deferredTools[0] = markToolCacheDisabled(deferredTools[0]);
+    return [...deferredTools] as GraphTools;
+  }
+
+  staticTools[staticTools.length - 1] = markCachePointAfter(
+    staticTools[staticTools.length - 1]
+  );
+
+  return [...staticTools, ...deferredTools] as GraphTools;
+}
+
+export function insertBedrockToolCachePoint(
+  toolConfig: ToolConfiguration | undefined,
+  fallbackToEnd: boolean
+): ToolConfiguration | undefined {
+  const tools = toolConfig?.tools as BedrockToolWithCacheMarker[] | undefined;
+  if (tools == null || tools.length === 0) {
+    return toolConfig;
+  }
+
+  let markerIndex = -1;
+  let hasCachePoint = false;
+  let hasDisabledMarker = false;
+  const cleanedTools: Tool[] = [];
+
+  for (let i = 0; i < tools.length; i++) {
+    const tool = tools[i];
+    if (isBedrockCachePoint(tool)) {
+      hasCachePoint = true;
+      cleanedTools.push(tool);
+      continue;
+    }
+    if (tool[BEDROCK_TOOL_CACHE_MARKER] === true) {
+      markerIndex = cleanedTools.length;
+    }
+    if (tool[BEDROCK_TOOL_CACHE_DISABLED_MARKER] === true) {
+      hasDisabledMarker = true;
+    }
+    cleanedTools.push(stripCachePointMarker(tool));
+  }
+
+  if (hasCachePoint || hasDisabledMarker) {
+    return { ...toolConfig, tools: cleanedTools };
+  }
+
+  const insertionIndex = markerIndex >= 0 ? markerIndex : tools.length - 1;
+  if (markerIndex < 0 && !fallbackToEnd) {
+    return { ...toolConfig, tools: cleanedTools };
+  }
+
+  return {
+    ...toolConfig,
+    tools: [
+      ...cleanedTools.slice(0, insertionIndex + 1),
+      CACHE_POINT,
+      ...cleanedTools.slice(insertionIndex + 1),
+    ],
+  };
+}