Merge pull request continuedev#5371 from chezsmithy/feat-cache-tools-config

feat(config-yaml): Add promptCaching to Default Completion Options and enable Bedrock Tools Caching

Author: Patrick-Erichsen

core/index.d.ts

@@ -1023,6 +1023,7 @@ export interface BaseCompletionOptions {
   toolChoice?: ToolChoice;
   reasoning?: boolean;
   reasoningBudgetTokens?: number;
+  promptCaching?: boolean;
 }
 
 export interface ModelCapability {

core/llm/llms/Bedrock.ts

@@ -5,6 +5,7 @@ import {
   ConverseStreamCommandOutput,
   InvokeModelCommand,
   Message,
+  ToolConfiguration,
 } from "@aws-sdk/client-bedrock-runtime";
 import { fromNodeProviderChain } from "@aws-sdk/credential-providers";
 
@@ -291,9 +292,10 @@ class Bedrock extends BaseLLM {
     const convertedMessages = this._convertMessages(messages);
 
     const shouldCacheSystemMessage =
-      !!systemMessage && this.cacheBehavior?.cacheSystemMessage;
+      !!systemMessage && this.cacheBehavior?.cacheSystemMessage || this.completionOptions.promptCaching;
     const enablePromptCaching =
-      shouldCacheSystemMessage || this.cacheBehavior?.cacheConversation;
+      shouldCacheSystemMessage || this.cacheBehavior?.cacheConversation || this.completionOptions.promptCaching;
+    const shouldCacheToolsConfig = this.completionOptions.promptCaching;
 
     // Add header for prompt caching
     if (enablePromptCaching) {
@@ -305,28 +307,34 @@ class Bedrock extends BaseLLM {
 
     const supportsTools =
       PROVIDER_TOOL_SUPPORT.bedrock?.(options.model || "") ?? false;
+
+    let toolConfig = supportsTools && options.tools
+    ? {
+        tools: options.tools.map((tool) => ({
+          toolSpec: {
+            name: tool.function.name,
+            description: tool.function.description,
+            inputSchema: {
+              json: tool.function.parameters,
+            },
+          },
+        })),
+      } as ToolConfiguration
+    : undefined;
+
+    if (toolConfig?.tools && shouldCacheToolsConfig) {
+      toolConfig.tools.push({ cachePoint: { type: "default" } });
+    }
+
     return {
       modelId: options.model,
-      messages: convertedMessages,
       system: systemMessage
         ? shouldCacheSystemMessage
           ? [{ text: systemMessage }, { cachePoint: { type: "default" } }]
           : [{ text: systemMessage }]
         : undefined,
-      toolConfig:
-        supportsTools && options.tools
-          ? {
-              tools: options.tools.map((tool) => ({
-                toolSpec: {
-                  name: tool.function.name,
-                  description: tool.function.description,
-                  inputSchema: {
-                    json: tool.function.parameters,
-                  },
-                },
-              })),
-            }
-          : undefined,
+      toolConfig: toolConfig,
+      messages: convertedMessages,
       inferenceConfig: {
         maxTokens: options.maxTokens,
         temperature: options.temperature,

packages/config-yaml/src/schemas/models.ts

@@ -49,6 +49,7 @@ export const completionOptionsSchema = z.object({
   n: z.number().optional(),
   reasoning: z.boolean().optional(),
   reasoningBudgetTokens: z.number().optional(),
+  promptCaching: z.boolean().optional(),
 });
 export type CompletionOptions = z.infer<typeof completionOptionsSchema>;