anthropic prompt caching

Author: seeM

extensions/positron-assistant/src/anthropic.ts

@@ -10,6 +10,22 @@ import { ModelConfig } from './config';
 import { isLanguageModelImagePart, LanguageModelImagePart } from './languageModelParts.js';
 import { isChatImagePart, processMessages } from './utils.js';
 import { DEFAULT_MAX_TOKEN_OUTPUT } from './constants.js';
+import { log } from './extension.js';
+
+/**
+ * Options for controlling cache behavior in the Anthropic language model.
+ */
+interface CacheControlOptions {
+	/** Add a cache control point to the last tool description (default: true). */
+	lastTool?: boolean;
+
+	/** Add a cache control point to the system prompt (default: true). */
+	system?: boolean;
+
+	/** Add a cache control point to the last user message (default: false). */
+	lastUserMessage?: boolean;
+}
+
 
 export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProvider {
 	name: string;
@@ -56,16 +72,45 @@ export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProv
 		progress: vscode.Progress<vscode.ChatResponseFragment2>,
 		token: vscode.CancellationToken
 	) {
-		const anthropicMessages = processMessages(messages).map(toAnthropicMessage);
-		const tools = options.tools?.map(tool => toAnthropicTool(tool));
+		const cacheControlOptions = isCacheControlOptions(options.modelOptions?.cacheControl)
+			? options.modelOptions.cacheControl
+			: undefined;
+		const tools = options.tools && toAnthropicTools(options.tools, cacheControlOptions?.lastTool);
 		const tool_choice = options.toolMode && toAnthropicToolChoice(options.toolMode);
-		const stream = this._client.messages.stream({
+		const system = options.modelOptions?.system &&
+			toAnthropicSystem(options.modelOptions.system, cacheControlOptions?.system);
+		const anthropicMessages = toAnthropicMessages(messages, cacheControlOptions?.lastUserMessage);
+
+		const body: Anthropic.MessageStreamParams = {
 			model: this._config.model,
 			max_tokens: options.modelOptions?.maxTokens ?? this.maxOutputTokens,
-			messages: anthropicMessages,
-			tool_choice,
 			tools,
-			system: options.modelOptions?.system,
+			tool_choice,
+			system,
+			messages: anthropicMessages,
+		};
+		const stream = this._client.messages.stream(body);
+
+		// Log request information - the request ID is only available upon connection.
+		stream.on('connect', () => {
+			if (log.logLevel <= vscode.LogLevel.Trace) {
+				log.trace(`[anthropic] SEND messages.stream [${stream.request_id}]: ${JSON.stringify(body)}`);
+			} else {
+				const userMessages = body.messages.filter(m => m.role === 'user');
+				const assistantMessages = body.messages.filter(m => m.role === 'assistant');
+				log.debug(
+					`[anthropic] SEND messages.stream [${stream.request_id}]: ` +
+					`model: ${body.model}; ` +
+					`cache options: ${cacheControlOptions ? JSON.stringify(cacheControlOptions) : 'default'}; ` +
+					`tools: ${body.tools?.map(t => t.name).sort().join(', ') ?? 'none'}; ` +
+					`tool choice: ${body.tool_choice ? JSON.stringify(body.tool_choice) : 'default'}; ` +
+					`system chars: ${body.system ? JSON.stringify(body.system).length : 0}; ` +
+					`user messages: ${userMessages.length}; ` +
+					`user message characters: ${JSON.stringify(userMessages).length}; ` +
+					`assistant messages: ${assistantMessages.length}; ` +
+					`assistant message characters: ${JSON.stringify(assistantMessages).length}`
+				);
+			}
 		});
 
 		token.onCancellationRequested(() => {
@@ -102,20 +147,31 @@ export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProv
 			}
 			throw error;
 		}
+
+		// Log usage information.
+		const message = await stream.finalMessage();
+		if (log.logLevel <= vscode.LogLevel.Trace) {
+			log.trace(`[anthropic] RECV messages.stream [${stream.request_id}]: ${JSON.stringify(message)}`);
+		} else {
+			log.debug(
+				`[anthropic] RECV messages.stream [${stream.request_id}]: ` +
+				`usage: ${JSON.stringify(message.usage)}`
+			);
+		}
 	}
 
 	get providerName(): string {
 		return AnthropicLanguageModel.source.provider.displayName;
 	}
 
-	private onContentBlock(block: Anthropic.Messages.ContentBlock, progress: vscode.Progress<vscode.ChatResponseFragment2>): void {
+	private onContentBlock(block: Anthropic.ContentBlock, progress: vscode.Progress<vscode.ChatResponseFragment2>): void {
 		switch (block.type) {
 			case 'tool_use':
 				return this.onToolUseBlock(block, progress);
 		}
 	}
 
-	private onToolUseBlock(block: Anthropic.Messages.ToolUseBlock, progress: vscode.Progress<vscode.ChatResponseFragment2>): void {
+	private onToolUseBlock(block: Anthropic.ToolUseBlock, progress: vscode.Progress<vscode.ChatResponseFragment2>): void {
 		progress.report({
 			index: 0,
 			part: new vscode.LanguageModelToolCallPart(block.id, block.name, block.input as any),
@@ -170,6 +226,49 @@ export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProv
 	}
 }
 
+function toAnthropicMessages(messages: vscode.LanguageModelChatMessage2[], cacheLastUserMessage = false): Anthropic.MessageParam[] {
+	const anthropicMessages = processMessages(messages).map(toAnthropicMessage);
+
+	if (cacheLastUserMessage) {
+		// Add a cache control point to the last valid user message.
+		for (let i = anthropicMessages.length - 1; i >= 0; i--) {
+			const message = anthropicMessages[i];
+
+			// Skip non-user messages.
+			if (message.role !== 'user') {
+				continue;
+			}
+
+			if (typeof message.content === 'string') {
+				// Content is a single string, make it a text block with a cache control point.
+				const text = message.content;
+				message.content = [{
+					type: 'text',
+					text,
+					cache_control: { type: 'ephemeral' },
+				}];
+				log.debug(`[anthropic] Adding cache control point to last user message block`);
+				break;
+			} else {
+				// Content is an array, try to add a cache control point to the last content block.
+				const lastContentBlock = message.content[message.content.length - 1];
+
+				// Thinking blocks cannot be cache control points.
+				if (lastContentBlock.type === 'thinking'
+					|| lastContentBlock.type === 'redacted_thinking') {
+					continue;
+				}
+
+				lastContentBlock.cache_control = { type: 'ephemeral' };
+				log.debug(`[anthropic] Adding cache control point to last user message block`);
+				break;
+			}
+		}
+	}
+
+	return anthropicMessages;
+}
+
 function toAnthropicMessage(message: vscode.LanguageModelChatMessage2): Anthropic.MessageParam {
 	switch (message.role) {
 		case vscode.LanguageModelChatMessageRole.Assistant:
@@ -281,6 +380,25 @@ function languageModelImagePartToAnthropicImageBlock(part: LanguageModelImagePar
 	};
 }
 
+function toAnthropicTools(tools: vscode.LanguageModelChatTool[], cacheLastTool = true): Anthropic.ToolUnion[] {
+	if (tools.length === 0) {
+		return [];
+	}
+	const anthropicTools = tools.map(tool => toAnthropicTool(tool));
+
+	// Ensure a stable sort order for prompt caching.
+	anthropicTools.sort((a, b) => a.name.localeCompare(b.name));
+
+	if (cacheLastTool) {
+		// Add a cache control point to the last tool description.
+		const lastTool = anthropicTools[anthropicTools.length - 1];
+		log.debug(`[anthropic] Adding cache control point to last tool: ${lastTool.name}`);
+		lastTool.cache_control = { type: 'ephemeral' };
+	}
+
+	return anthropicTools;
+}
+
 function toAnthropicTool(tool: vscode.LanguageModelChatTool): Anthropic.ToolUnion {
 	const input_schema = tool.inputSchema as Anthropic.Tool.InputSchema ?? {
 		type: 'object',
@@ -308,3 +426,36 @@ function toAnthropicToolChoice(toolMode: vscode.LanguageModelChatToolMode): Anth
 			throw new Error(`Unsupported tool mode: ${toolMode}`);
 	}
 }
+
+function toAnthropicSystem(system: unknown, cacheSystem = true): Anthropic.MessageCreateParams['system'] {
+	if (typeof system === 'string') {
+		const anthropicSystem: Anthropic.MessageCreateParams['system'] = [{
+			type: 'text',
+			text: system,
+		}];
+
+		if (cacheSystem) {
+			// Add a cache control point to the last system prompt block.
+			const lastSystemBlock = anthropicSystem[anthropicSystem.length - 1];
+			lastSystemBlock.cache_control = { type: 'ephemeral' };
+			log.debug(`[anthropic] Adding cache control point to system prompt`);
+		}
+
+		return anthropicSystem;
+	}
+	// Pass the system prompt through as-is.
+	// We may pass an invalid system prompt; let Anthropic throw the error.
+	return system as Anthropic.MessageCreateParams['system'];
+}
+
+function isCacheControlOptions(options: unknown): options is CacheControlOptions {
+	if (typeof options !== 'object' || options === null) {
+		return false;
+	}
+	const cacheControlOptions = options as CacheControlOptions;
+	return (
+		(cacheControlOptions.lastTool === undefined || typeof cacheControlOptions.lastTool === 'boolean') &&
+		(cacheControlOptions.system === undefined || typeof cacheControlOptions.system === 'boolean') &&
+		(cacheControlOptions.lastUserMessage === undefined || typeof cacheControlOptions.lastUserMessage === 'boolean')
+	);
+}