Skip to content

Commit c26a892

Browse files
seeMseeM
committed
cleanup and better cache breakpoint logging
1 parent 48876a6 commit c26a892

File tree

5 files changed

+219
-133
lines changed

5 files changed

+219
-133
lines changed

‎extensions/positron-assistant/src/anthropic.ts

Lines changed: 128 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -8,18 +8,26 @@ import * as vscode from 'vscode';
88
import Anthropic from '@anthropic-ai/sdk';
99
import { ModelConfig } from './config';
1010
import { isLanguageModelImagePart, LanguageModelImagePart } from './languageModelParts.js';
11-
import { isChatImagePart, isLanguageModelJsonDataPart, parseJsonLanguageModelDataPart, processMessages } from './utils.js';
11+
import { isChatImagePart, isCacheBreakpointDataPart, parseCacheBreakpoint, processMessages } from './utils.js';
1212
import { DEFAULT_MAX_TOKEN_OUTPUT } from './constants.js';
1313
import { log } from './extension.js';
1414

1515
/**
1616
* Options for controlling cache behavior in the Anthropic language model.
1717
*/
1818
export interface CacheControlOptions {
19-
/** Add a cache control point to the system prompt (default: true). */
19+
/** Add a cache breakpoint to the system prompt (default: true). */
2020
system?: boolean;
2121
}
2222

23+
/**
24+
* Block params that set cache breakpoints.
25+
*/
26+
type CacheControllableBlockParam = Anthropic.TextBlockParam |
27+
Anthropic.ImageBlockParam |
28+
Anthropic.ToolUseBlockParam |
29+
Anthropic.ToolResultBlockParam;
30+
2331
export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProvider {
2432
name: string;
2533
provider: string;
@@ -200,11 +208,10 @@ export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProv
200208
});
201209
} else {
202210
// For LanguageModelChatMessage, ensure it has non-empty message content
203-
const processedMessages = processMessages([text]);
204-
if (processedMessages.length === 0) {
211+
messages.push(...toAnthropicMessages([text]));
212+
if (messages.length === 0) {
205213
return 0;
206214
}
207-
messages.push(...processedMessages.map(toAnthropicMessage));
208215
}
209216
const result = await this._client.messages.countTokens({
210217
model: this._config.model,
@@ -223,30 +230,37 @@ export class AnthropicLanguageModel implements positron.ai.LanguageModelChatProv
223230
}
224231

225232
function toAnthropicMessages(messages: vscode.LanguageModelChatMessage2[]): Anthropic.MessageParam[] {
226-
const anthropicMessages = processMessages(messages).map(toAnthropicMessage);
233+
let userMessageIndex = 0;
234+
let assistantMessageIndex = 0;
235+
const anthropicMessages = processMessages(messages).map((message) => {
236+
const source = message.role === vscode.LanguageModelChatMessageRole.User ?
237+
`User message ${userMessageIndex++}` :
238+
`Assistant message ${assistantMessageIndex++}`;
239+
return toAnthropicMessage(message, source);
240+
});
227241
return anthropicMessages;
228242
}
229243

230-
function toAnthropicMessage(message: vscode.LanguageModelChatMessage2): Anthropic.MessageParam {
244+
function toAnthropicMessage(message: vscode.LanguageModelChatMessage2, source: string): Anthropic.MessageParam {
231245
switch (message.role) {
232246
case vscode.LanguageModelChatMessageRole.Assistant:
233-
return toAnthropicAssistantMessage(message);
247+
return toAnthropicAssistantMessage(message, source);
234248
case vscode.LanguageModelChatMessageRole.User:
235-
return toAnthropicUserMessage(message);
249+
return toAnthropicUserMessage(message, source);
236250
default:
237251
throw new Error(`Unsupported message role: ${message.role}`);
238252
}
239253
}
240254

241-
function toAnthropicAssistantMessage(message: vscode.LanguageModelChatMessage2): Anthropic.MessageParam {
255+
function toAnthropicAssistantMessage(message: vscode.LanguageModelChatMessage2, source: string): Anthropic.MessageParam {
242256
const content: Anthropic.ContentBlockParam[] = [];
243257
for (let i = 0; i < message.content.length; i++) {
244258
const [part, nextPart] = [message.content[i], message.content[i + 1]];
245259
const dataPart = nextPart instanceof vscode.LanguageModelDataPart ? nextPart : undefined;
246260
if (part instanceof vscode.LanguageModelTextPart) {
247-
content.push(toAnthropicTextBlock(part, dataPart));
261+
content.push(toAnthropicTextBlock(part, source, dataPart));
248262
} else if (part instanceof vscode.LanguageModelToolCallPart) {
249-
content.push(toAnthropicToolUseBlock(part, dataPart));
263+
content.push(toAnthropicToolUseBlock(part, source, dataPart));
250264
} else if (part instanceof vscode.LanguageModelDataPart) {
251265
// Skip extra data parts. They're handled in part conversion.
252266
} else {
@@ -259,20 +273,20 @@ function toAnthropicAssistantMessage(message: vscode.LanguageModelChatMessage2):
259273
};
260274
}
261275

262-
function toAnthropicUserMessage(message: vscode.LanguageModelChatMessage2): Anthropic.MessageParam {
276+
function toAnthropicUserMessage(message: vscode.LanguageModelChatMessage2, source: string): Anthropic.MessageParam {
263277
const content: Anthropic.ContentBlockParam[] = [];
264278
for (let i = 0; i < message.content.length; i++) {
265279
const [part, nextPart] = [message.content[i], message.content[i + 1]];
266280
const dataPart = nextPart instanceof vscode.LanguageModelDataPart ? nextPart : undefined;
267281
if (part instanceof vscode.LanguageModelTextPart) {
268-
content.push(toAnthropicTextBlock(part, dataPart));
282+
content.push(toAnthropicTextBlock(part, source, dataPart));
269283
} else if (part instanceof vscode.LanguageModelToolResultPart) {
270-
content.push(toAnthropicToolResultBlock(part, dataPart));
284+
content.push(toAnthropicToolResultBlock(part, source, dataPart));
271285
} else if (part instanceof vscode.LanguageModelToolResultPart2) {
272-
content.push(toAnthropicToolResultBlock(part, dataPart));
286+
content.push(toAnthropicToolResultBlock(part, source, dataPart));
273287
} else if (part instanceof vscode.LanguageModelDataPart) {
274288
if (isChatImagePart(part)) {
275-
content.push(chatImagePartToAnthropicImageBlock(part, dataPart));
289+
content.push(chatImagePartToAnthropicImageBlock(part, source, dataPart));
276290
} else {
277291
// Skip other data parts.
278292
}
@@ -286,66 +300,106 @@ function toAnthropicUserMessage(message: vscode.LanguageModelChatMessage2): Anth
286300
};
287301
}
288302

289-
function toAnthropicTextBlock(part: vscode.LanguageModelTextPart, dataPart?: vscode.LanguageModelDataPart): Anthropic.TextBlockParam {
290-
return withCacheControl({
291-
type: 'text',
292-
text: part.value,
293-
}, dataPart);
303+
function toAnthropicTextBlock(
304+
part: vscode.LanguageModelTextPart,
305+
source: string,
306+
dataPart?: vscode.LanguageModelDataPart,
307+
): Anthropic.TextBlockParam {
308+
return withCacheControl(
309+
{
310+
type: 'text',
311+
text: part.value,
312+
},
313+
source,
314+
dataPart,
315+
);
294316
}
295317

296-
function toAnthropicToolUseBlock(part: vscode.LanguageModelToolCallPart, dataPart?: vscode.LanguageModelDataPart): Anthropic.ToolUseBlockParam {
297-
return withCacheControl({
298-
type: 'tool_use',
299-
id: part.callId,
300-
name: part.name,
301-
input: part.input,
302-
}, dataPart);
318+
function toAnthropicToolUseBlock(
319+
part: vscode.LanguageModelToolCallPart,
320+
source: string,
321+
dataPart?: vscode.LanguageModelDataPart,
322+
): Anthropic.ToolUseBlockParam {
323+
return withCacheControl(
324+
{
325+
type: 'tool_use',
326+
id: part.callId,
327+
name: part.name,
328+
input: part.input,
329+
},
330+
source,
331+
dataPart,
332+
);
303333
}
304334

305-
function toAnthropicToolResultBlock(part: vscode.LanguageModelToolResultPart, dataPart?: vscode.LanguageModelDataPart): Anthropic.ToolResultBlockParam {
335+
function toAnthropicToolResultBlock(
336+
part: vscode.LanguageModelToolResultPart,
337+
source: string,
338+
dataPart?: vscode.LanguageModelDataPart,
339+
): Anthropic.ToolResultBlockParam {
306340
const content: Anthropic.ToolResultBlockParam['content'] = [];
307341
for (let i = 0; i < part.content.length; i++) {
308342
const [resultPart, resultNextPart] = [part.content[i], part.content[i + 1]];
309343
const resultDataPart = resultNextPart instanceof vscode.LanguageModelDataPart ? resultNextPart : undefined;
310344
if (resultPart instanceof vscode.LanguageModelTextPart) {
311-
content.push(toAnthropicTextBlock(resultPart, resultDataPart));
345+
content.push(toAnthropicTextBlock(resultPart, source, resultDataPart));
312346
} else if (isLanguageModelImagePart(resultPart)) {
313-
content.push(languageModelImagePartToAnthropicImageBlock(resultPart, resultDataPart));
347+
content.push(languageModelImagePartToAnthropicImageBlock(resultPart, source, resultDataPart));
314348
} else if (resultPart instanceof vscode.LanguageModelDataPart) {
315349
// Skip data parts.
316350
} else {
317351
throw new Error('Unsupported part type on tool result part content');
318352
}
319353
}
320-
return withCacheControl({
321-
type: 'tool_result',
322-
tool_use_id: part.callId,
323-
content,
324-
}, dataPart);
354+
return withCacheControl(
355+
{
356+
type: 'tool_result',
357+
tool_use_id: part.callId,
358+
content,
359+
},
360+
source,
361+
dataPart,
362+
);
325363
}
326364

327-
function chatImagePartToAnthropicImageBlock(part: vscode.LanguageModelDataPart, dataPart?: vscode.LanguageModelDataPart): Anthropic.ImageBlockParam {
328-
return withCacheControl({
329-
type: 'image',
330-
source: {
331-
type: 'base64',
332-
// We may pass an unsupported mime type; let Anthropic throw the error.
333-
media_type: part.mimeType as Anthropic.Base64ImageSource['media_type'],
334-
data: Buffer.from(part.data).toString('base64'),
365+
function chatImagePartToAnthropicImageBlock(
366+
part: vscode.LanguageModelDataPart,
367+
source: string,
368+
dataPart?: vscode.LanguageModelDataPart,
369+
): Anthropic.ImageBlockParam {
370+
return withCacheControl(
371+
{
372+
type: 'image',
373+
source: {
374+
type: 'base64',
375+
// We may pass an unsupported mime type; let Anthropic throw the error.
376+
media_type: part.mimeType as Anthropic.Base64ImageSource['media_type'],
377+
data: Buffer.from(part.data).toString('base64'),
378+
},
335379
},
336-
}, dataPart);
380+
source,
381+
dataPart,
382+
);
337383
}
338384

339-
function languageModelImagePartToAnthropicImageBlock(part: LanguageModelImagePart, dataPart?: vscode.LanguageModelDataPart): Anthropic.ImageBlockParam {
340-
return withCacheControl({
341-
type: 'image',
342-
source: {
343-
type: 'base64',
344-
// We may pass an unsupported mime type; let Anthropic throw the error.
345-
media_type: part.value.mimeType as Anthropic.Base64ImageSource['media_type'],
346-
data: part.value.base64,
385+
function languageModelImagePartToAnthropicImageBlock(
386+
part: LanguageModelImagePart,
387+
source: string,
388+
dataPart?: vscode.LanguageModelDataPart,
389+
): Anthropic.ImageBlockParam {
390+
return withCacheControl(
391+
{
392+
type: 'image',
393+
source: {
394+
type: 'base64',
395+
// We may pass an unsupported mime type; let Anthropic throw the error.
396+
media_type: part.value.mimeType as Anthropic.Base64ImageSource['media_type'],
397+
data: part.value.base64,
398+
},
347399
},
348-
}, dataPart);
400+
source,
401+
dataPart,
402+
);
349403
}
350404

351405
function toAnthropicTools(tools: vscode.LanguageModelChatTool[]): Anthropic.ToolUnion[] {
@@ -396,10 +450,10 @@ function toAnthropicSystem(system: unknown, cacheSystem = true): Anthropic.Messa
396450
}];
397451

398452
if (cacheSystem) {
399-
// Add a cache control point to the last system prompt block.
453+
// Add a cache breakpoint to the last system prompt block.
400454
const lastSystemBlock = anthropicSystem[anthropicSystem.length - 1];
401455
lastSystemBlock.cache_control = { type: 'ephemeral' };
402-
log.debug(`[anthropic] Adding cache control point to system prompt`);
456+
log.debug(`[anthropic] Adding cache breakpoint to system prompt`);
403457
}
404458

405459
return anthropicSystem;
@@ -413,7 +467,7 @@ function toAnthropicSystem(system: unknown, cacheSystem = true): Anthropic.Messa
413467
const [part, nextPart] = [system[i], system[i + 1]];
414468
const dataPart = nextPart instanceof vscode.LanguageModelDataPart ? nextPart : undefined;
415469
if (part instanceof vscode.LanguageModelTextPart) {
416-
anthropicSystem.push(toAnthropicTextBlock(part, dataPart));
470+
anthropicSystem.push(toAnthropicTextBlock(part, 'System prompt', dataPart));
417471
}
418472
}
419473
return anthropicSystem;
@@ -430,42 +484,27 @@ function isCacheControlOptions(options: unknown): options is CacheControlOptions
430484
return cacheControlOptions.system === undefined || typeof cacheControlOptions.system === 'boolean';
431485
}
432486

433-
function withCacheControl<T>(part: T, dataPart: vscode.LanguageModelDataPart | undefined): T {
434-
if (!isLanguageModelJsonDataPart(dataPart)) {
487+
function withCacheControl<T extends CacheControllableBlockParam>(
488+
part: T,
489+
source: string,
490+
dataPart: vscode.LanguageModelDataPart | undefined,
491+
): T {
492+
if (!isCacheBreakpointDataPart(dataPart)) {
435493
return part;
436494
}
437495

438-
let data: unknown;
439496
try {
440-
data = parseJsonLanguageModelDataPart(dataPart);
497+
const cachBreakpoint = parseCacheBreakpoint(dataPart);
498+
log.debug(`[anthropic] Adding cache breakpoint to ${part.type} part. Source: ${source}`);
499+
return {
500+
...part,
501+
// Pass the data through without validation, let the Anthropic API handle errors.
502+
cache_control: {
503+
type: cachBreakpoint.type,
504+
},
505+
};
441506
} catch (error) {
442-
log.error(`[anthropic] Failed to parse language model json data part: ${error}`);
507+
log.error(`[anthropic] Failed to parse cache breakpoint: ${error}`);
443508
return part;
444509
}
445-
446-
if (!isCacheControlData(data)) {
447-
return part;
448-
}
449-
450-
log.debug(`[anthropic] Adding cache control point via LanguageModelDataPart: ${JSON.stringify(data.data)}`);
451-
return {
452-
...part,
453-
// Pass the data through without validation, let the Anthropic API handle errors.
454-
cache_control: data.data as Anthropic.CacheControlEphemeral,
455-
};
456-
}
457-
458-
function isCacheControlData(data: unknown): data is { kind: 'cacheControl'; data: unknown } {
459-
return typeof data === 'object' && data !== null && 'kind' in data && data.kind === 'cacheControl' && 'data' in data;
460-
}
461-
462-
/**
463-
* Create a language model part that represents a cache control point.
464-
* @returns A language model part representing the cache control point.
465-
*/
466-
export function languageModelCacheControlPart(): vscode.LanguageModelDataPart {
467-
return vscode.LanguageModelDataPart.json({
468-
kind: 'cacheControl',
469-
data: { type: 'ephemeral' } satisfies Anthropic.CacheControlEphemeral,
470-
});
471510
}

0 commit comments

Comments
 (0)