illuma-agents 1.0.10 → 1.0.12

package/src/llm/bedrock/index.ts

@@ -1,5 +1,6 @@
 /**
  * Optimized ChatBedrockConverse wrapper that fixes contentBlockIndex conflicts
+ * and adds prompt caching support for Bedrock Converse API.
  *
  * 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"
@@ -9,18 +10,35 @@
  * 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 type { ChatBedrockConverseInput } from '@langchain/aws';
 import { AIMessageChunk } from '@langchain/core/messages';
-import type { BaseMessage } from '@langchain/core/messages';
+import type { BaseMessage, UsageMetadata } from '@langchain/core/messages';
 import { ChatGenerationChunk } from '@langchain/core/outputs';
 import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
 
+/** Extended input type with promptCache option */
+export interface CustomChatBedrockConverseInput extends ChatBedrockConverseInput {
+  promptCache?: boolean;
+}
+
 export class CustomChatBedrockConverse extends ChatBedrockConverse {
-  constructor(fields?: ChatBedrockConverseInput) {
+  promptCache: boolean;
+
+  constructor(fields?: CustomChatBedrockConverseInput) {
     super(fields);
+    this.promptCache = fields?.promptCache ?? false;
   }
 
   static lc_name(): string {
@@ -28,9 +46,92 @@ 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)
+   * Override invocationParams to add cachePoint to tools when promptCache is enabled.
+   * This enables Bedrock prompt caching for tool definitions.
+   * 
+   * 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.
+   */
+  invocationParams(
+    options?: this['ParsedCallOptions']
+  ): ReturnType<ChatBedrockConverse['invocationParams']> {
+    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 = [];
+      const coreToolNames: string[] = [];
+      const mcpToolNames: string[] = [];
+      
+      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);
+          mcpToolNames.push(toolName);
+        } else {
+          coreTools.push(tool);
+          coreToolNames.push(toolName);
+        }
+      }
+
+      // Always log cache structure (INFO level for tracking)
+      console.log(`[Cache] 🔧 Tools | Core: [${coreToolNames.join(', ')}] (${coreTools.length}) | MCP: [${mcpToolNames.join(', ')}] (${mcpTools.length})`);
+
+      // Build tools array with strategic cachePoints:
+      // [CoreTool1, CoreTool2, cachePoint] + [MCPTool1, MCPTool2, cachePoint]
+      const toolsWithCache: typeof params.toolConfig.tools = [];
+      let cachePointCount = 0;
+      
+      // Add core tools with cachePoint (if any)
+      if (coreTools.length > 0) {
+        toolsWithCache.push(...coreTools);
+        toolsWithCache.push({ cachePoint: { type: 'default' } });
+        cachePointCount++;
+      }
+      
+      // Add MCP tools with their own cachePoint (if any)
+      if (mcpTools.length > 0) {
+        toolsWithCache.push(...mcpTools);
+        toolsWithCache.push({ cachePoint: { type: 'default' } });
+        cachePointCount++;
+      }
+      
+      // If no tools at all (shouldn't happen but safety check)
+      if (toolsWithCache.length === 0) {
+        toolsWithCache.push({ cachePoint: { type: 'default' } });
+        cachePointCount++;
+      }
+      
+      console.log(`[Cache] 📍 Tool cachePoints: ${cachePointCount} | Order: [${coreToolNames.length > 0 ? 'CoreTools→CP' : ''}${mcpToolNames.length > 0 ? '→MCPTools→CP' : ''}]`);
+      
+      params.toolConfig.tools = toolsWithCache;
+    }
+
+    return params;
+  }
+
+  /**
+   * Override _streamResponseChunks to:
+   * 1. Strip contentBlockIndex from response_metadata to prevent merge conflicts
+   * 2. Extract cacheReadInputTokens/cacheWriteInputTokens and add to usage_metadata
    */
   async *_streamResponseChunks(
     messages: BaseMessage[],
@@ -50,21 +151,57 @@ export class CustomChatBedrockConverse extends ChatBedrockConverse {
         (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
-        );
+        const responseMetadata = chunk.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) {
-          const cleanedMetadata = this.removeContentBlockIndex(
-            chunk.message.response_metadata
-          ) as Record<string, unknown>;
+          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 = chunk.message.usage_metadata;
+        
+        if (usage) {
+          const cacheRead = (usage.cacheReadInputTokens as number) ?? 0;
+          const cacheWrite = (usage.cacheWriteInputTokens as number) ?? 0;
+          const inputTokens = (usage.inputTokens as number) ?? 0;
+          const outputTokens = (usage.outputTokens as number) ?? 0;
+          
+          if (cacheRead > 0 || cacheWrite > 0) {
+            // Always log cache results for tracking
+            const cacheStatus = cacheRead > 0 && cacheWrite === 0 ? '✅ HIT' : 
+                               cacheWrite > 0 && cacheRead === 0 ? '📝 WRITE' :
+                               cacheRead > 0 && cacheWrite > 0 ? '🔄 PARTIAL' : '❌ MISS';
+            console.log(`[Cache] ${cacheStatus} | read=${cacheRead} | write=${cacheWrite} | input=${inputTokens} | output=${outputTokens}`);
+            
+            needsModification = true;
+            enhancedUsageMetadata = {
+              input_tokens: chunk.message.usage_metadata?.input_tokens ?? inputTokens,
+              output_tokens: chunk.message.usage_metadata?.output_tokens ?? outputTokens,
+              total_tokens: chunk.message.usage_metadata?.total_tokens ?? (usage.totalTokens as number) ?? 0,
+              input_token_details: {
+                cache_read: cacheRead,
+                cache_creation: cacheWrite,
+              },
+            };
+          }
+        }
 
+        if (needsModification) {
           yield new ChatGenerationChunk({
             text: chunk.text,
             message: new AIMessageChunk({
               ...chunk.message,
               response_metadata: cleanedMetadata,
+              usage_metadata: enhancedUsageMetadata,
             }),
             generationInfo: chunk.generationInfo,
           });

package/src/messages/__tests__/tools.test.ts

@@ -0,0 +1,473 @@
+// src/messages/__tests__/tools.test.ts
+import {
+  AIMessageChunk,
+  ToolMessage,
+  HumanMessage,
+} from '@langchain/core/messages';
+import type { BaseMessage } from '@langchain/core/messages';
+import { extractToolDiscoveries, hasToolSearchInCurrentTurn } from '../tools';
+import { Constants } from '@/common';
+
+describe('Tool Discovery Functions', () => {
+  /**
+   * Helper to create an AIMessageChunk with tool calls
+   */
+  const createAIMessage = (
+    content: string,
+    toolCalls: Array<{
+      id: string;
+      name: string;
+      args: Record<string, unknown>;
+    }>
+  ): AIMessageChunk => {
+    return new AIMessageChunk({
+      content,
+      tool_calls: toolCalls.map((tc) => ({
+        id: tc.id,
+        name: tc.name,
+        args: tc.args,
+        type: 'tool_call' as const,
+      })),
+    });
+  };
+
+  /**
+   * Helper to create a ToolMessage (tool search result)
+   */
+  const createToolSearchResult = (
+    toolCallId: string,
+    discoveredTools: string[]
+  ): ToolMessage => {
+    return new ToolMessage({
+      content: `Found ${discoveredTools.length} tools`,
+      tool_call_id: toolCallId,
+      name: Constants.TOOL_SEARCH_REGEX,
+      artifact: {
+        tool_references: discoveredTools.map((name) => ({
+          tool_name: name,
+          match_score: 0.9,
+          matched_field: 'description',
+          snippet: 'Test snippet',
+        })),
+        metadata: {
+          total_searched: 10,
+          pattern: 'test',
+        },
+      },
+    });
+  };
+
+  /**
+   * Helper to create a regular ToolMessage (non-search)
+   */
+  const createRegularToolMessage = (
+    toolCallId: string,
+    name: string,
+    content: string
+  ): ToolMessage => {
+    return new ToolMessage({
+      content,
+      tool_call_id: toolCallId,
+      name,
+    });
+  };
+
+  describe('extractToolDiscoveries', () => {
+    it('extracts tool names from a single tool search result', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search for database tools'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'database' },
+          },
+        ]),
+        createToolSearchResult('call_1', ['database_query', 'database_insert']),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual(['database_query', 'database_insert']);
+    });
+
+    it('extracts tool names from multiple tool search results in same turn', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search for tools'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'database' },
+          },
+          {
+            id: 'call_2',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'file' },
+          },
+        ]),
+        createToolSearchResult('call_1', ['database_query']),
+        createToolSearchResult('call_2', ['file_read', 'file_write']),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual(['database_query', 'file_read', 'file_write']);
+    });
+
+    it('returns empty array when no messages', () => {
+      const result = extractToolDiscoveries([]);
+      expect(result).toEqual([]);
+    });
+
+    it('returns empty array when last message is not a ToolMessage', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        createAIMessage('Hi there!', []),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual([]);
+    });
+
+    it('returns empty array when no AI message with tool calls found', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        new ToolMessage({
+          content: 'Some result',
+          tool_call_id: 'orphan_call',
+          name: Constants.TOOL_SEARCH_REGEX,
+        }),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual([]);
+    });
+
+    it('ignores tool search results from previous turns', () => {
+      const messages: BaseMessage[] = [
+        // Turn 1: Previous search
+        new HumanMessage('Search for old tools'),
+        createAIMessage('Searching...', [
+          {
+            id: 'old_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'old' },
+          },
+        ]),
+        createToolSearchResult('old_call', ['old_tool_1', 'old_tool_2']),
+        // Turn 2: Current turn
+        new HumanMessage('Search for new tools'),
+        createAIMessage('Searching again...', [
+          {
+            id: 'new_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'new' },
+          },
+        ]),
+        createToolSearchResult('new_call', ['new_tool']),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      // Should only return tools from current turn
+      expect(result).toEqual(['new_tool']);
+    });
+
+    it('ignores non-search tool results', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Do some work'),
+        createAIMessage('Working...', [
+          {
+            id: 'search_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+          { id: 'other_call', name: 'get_weather', args: { city: 'NYC' } },
+        ]),
+        createToolSearchResult('search_call', ['found_tool']),
+        createRegularToolMessage('other_call', 'get_weather', '{"temp": 72}'),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual(['found_tool']);
+    });
+
+    it('handles empty tool_references in artifact', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'xyz' },
+          },
+        ]),
+        new ToolMessage({
+          content: 'No tools found',
+          tool_call_id: 'call_1',
+          name: Constants.TOOL_SEARCH_REGEX,
+          artifact: {
+            tool_references: [],
+            metadata: { total_searched: 10, pattern: 'xyz' },
+          },
+        }),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual([]);
+    });
+
+    it('handles missing artifact', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+        ]),
+        new ToolMessage({
+          content: 'Error occurred',
+          tool_call_id: 'call_1',
+          name: Constants.TOOL_SEARCH_REGEX,
+          // No artifact
+        }),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual([]);
+    });
+
+    it('ignores tool messages with wrong tool_call_id', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+        ]),
+        // This has a different tool_call_id that doesn't match the AI message
+        createToolSearchResult('wrong_id', ['should_not_appear']),
+        createToolSearchResult('call_1', ['correct_tool']),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      expect(result).toEqual(['correct_tool']);
+    });
+
+    it('only looks at messages after the latest AI parent', () => {
+      const messages: BaseMessage[] = [
+        // First AI message with search
+        createAIMessage('First search', [
+          {
+            id: 'first_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'first' },
+          },
+        ]),
+        createToolSearchResult('first_call', ['first_tool']),
+        // Second AI message - this is the "latest AI parent" for the last tool message
+        createAIMessage('Second search', [
+          {
+            id: 'second_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'second' },
+          },
+        ]),
+        createToolSearchResult('second_call', ['second_tool']),
+      ];
+
+      const result = extractToolDiscoveries(messages);
+
+      // Should only find second_tool (from the turn of the latest AI parent)
+      expect(result).toEqual(['second_tool']);
+    });
+  });
+
+  describe('hasToolSearchInCurrentTurn', () => {
+    it('returns true when current turn has tool search results', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+        ]),
+        createToolSearchResult('call_1', ['found_tool']),
+      ];
+
+      const result = hasToolSearchInCurrentTurn(messages);
+
+      expect(result).toBe(true);
+    });
+
+    it('returns false when no messages', () => {
+      const result = hasToolSearchInCurrentTurn([]);
+      expect(result).toBe(false);
+    });
+
+    it('returns false when last message is not a ToolMessage', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        createAIMessage('Hi!', []),
+      ];
+
+      const result = hasToolSearchInCurrentTurn(messages);
+
+      expect(result).toBe(false);
+    });
+
+    it('returns false when no AI parent found', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        new ToolMessage({
+          content: 'Result',
+          tool_call_id: 'orphan',
+          name: Constants.TOOL_SEARCH_REGEX,
+        }),
+      ];
+
+      const result = hasToolSearchInCurrentTurn(messages);
+
+      expect(result).toBe(false);
+    });
+
+    it('returns false when current turn has no tool search (only other tools)', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Get weather'),
+        createAIMessage('Getting weather...', [
+          { id: 'call_1', name: 'get_weather', args: { city: 'NYC' } },
+        ]),
+        createRegularToolMessage('call_1', 'get_weather', '{"temp": 72}'),
+      ];
+
+      const result = hasToolSearchInCurrentTurn(messages);
+
+      expect(result).toBe(false);
+    });
+
+    it('returns true even with mixed tool types in current turn', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search and get weather'),
+        createAIMessage('Working...', [
+          {
+            id: 'search_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+          { id: 'weather_call', name: 'get_weather', args: { city: 'NYC' } },
+        ]),
+        createRegularToolMessage('weather_call', 'get_weather', '{"temp": 72}'),
+        createToolSearchResult('search_call', ['found_tool']),
+      ];
+
+      const result = hasToolSearchInCurrentTurn(messages);
+
+      expect(result).toBe(true);
+    });
+
+    it('returns false for tool search from previous turn', () => {
+      const messages: BaseMessage[] = [
+        // Previous turn with search
+        createAIMessage('Searching...', [
+          {
+            id: 'old_call',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'old' },
+          },
+        ]),
+        createToolSearchResult('old_call', ['old_tool']),
+        // Current turn without search
+        new HumanMessage('Get weather now'),
+        createAIMessage('Getting weather...', [
+          { id: 'weather_call', name: 'get_weather', args: { city: 'NYC' } },
+        ]),
+        createRegularToolMessage('weather_call', 'get_weather', '{"temp": 72}'),
+      ];
+
+      const result = hasToolSearchInCurrentTurn(messages);
+
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('Integration: extractToolDiscoveries + hasToolSearchInCurrentTurn', () => {
+    it('hasToolSearchInCurrentTurn is true when extractToolDiscoveries returns results', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Search'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+        ]),
+        createToolSearchResult('call_1', ['tool_a', 'tool_b']),
+      ];
+
+      const hasSearch = hasToolSearchInCurrentTurn(messages);
+      const discoveries = extractToolDiscoveries(messages);
+
+      expect(hasSearch).toBe(true);
+      expect(discoveries.length).toBeGreaterThan(0);
+    });
+
+    it('both return empty/false for non-search turns', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Get weather'),
+        createAIMessage('Getting...', [
+          { id: 'call_1', name: 'get_weather', args: { city: 'NYC' } },
+        ]),
+        createRegularToolMessage('call_1', 'get_weather', '{"temp": 72}'),
+      ];
+
+      const hasSearch = hasToolSearchInCurrentTurn(messages);
+      const discoveries = extractToolDiscoveries(messages);
+
+      expect(hasSearch).toBe(false);
+      expect(discoveries).toEqual([]);
+    });
+
+    it('hasToolSearchInCurrentTurn can be used as quick check before extraction', () => {
+      const messagesWithSearch: BaseMessage[] = [
+        new HumanMessage('Search'),
+        createAIMessage('Searching...', [
+          {
+            id: 'call_1',
+            name: Constants.TOOL_SEARCH_REGEX,
+            args: { pattern: 'test' },
+          },
+        ]),
+        createToolSearchResult('call_1', ['tool_a']),
+      ];
+
+      const messagesWithoutSearch: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        createAIMessage('Hi!', []),
+      ];
+
+      // Pattern: quick check first, then extract only if needed
+      if (hasToolSearchInCurrentTurn(messagesWithSearch)) {
+        const discoveries = extractToolDiscoveries(messagesWithSearch);
+        expect(discoveries).toEqual(['tool_a']);
+      }
+
+      if (hasToolSearchInCurrentTurn(messagesWithoutSearch)) {
+        // This should not execute
+        expect(true).toBe(false);
+      }
+    });
+  });
+});