feat: add virtual id for extended output capabilities

This commit is contained in:
Tom X Nguyen 2025-03-03 07:47:12 +07:00
parent 8e8079d528
commit aec03d71c3
2 changed files with 92 additions and 21 deletions

View file

@ -29,7 +29,7 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
let stream: AnthropicStream<Anthropic.Messages.RawMessageStreamEvent>
const cacheControl: CacheControlEphemeral = { type: "ephemeral" }
let { id: modelId, maxTokens, thinking, temperature } = this.getModel()
let { id: modelId, maxTokens, thinking, temperature, virtualId } = this.getModel()
switch (modelId) {
case "claude-3-7-sonnet-20250219":
@ -82,6 +82,15 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
// prompt caching: https://x.com/alexalbert__/status/1823751995901272068
// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
// https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393
// Check for the thinking-128k variant first
if (virtualId === "claude-3-7-sonnet-20250219:thinking-128k") {
return {
headers: { "anthropic-beta": "output-128k-2025-02-19" },
}
}
// Then check for models that support prompt caching
switch (modelId) {
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
@ -184,31 +193,58 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
let id = modelId && modelId in anthropicModels ? (modelId as AnthropicModelId) : anthropicDefaultModelId
const info: ModelInfo = anthropicModels[id]
// Track the original model ID for special variant handling
const virtualId = id
// The `:thinking` variant is a virtual identifier for the
// `claude-3-7-sonnet-20250219` model with a thinking budget.
// We can handle this more elegantly in the future.
if (id === "claude-3-7-sonnet-20250219:thinking") {
if (id === "claude-3-7-sonnet-20250219:thinking" || id === "claude-3-7-sonnet-20250219:thinking-128k") {
id = "claude-3-7-sonnet-20250219"
}
return {
id,
info,
virtualId, // Include the original ID to use for header selection
...getModelParams({ options: this.options, model: info, defaultMaxTokens: ANTHROPIC_DEFAULT_MAX_TOKENS }),
}
}
async completePrompt(prompt: string) {
let { id: modelId, maxTokens, thinking, temperature } = this.getModel()
let { id: modelId, maxTokens, thinking, temperature, virtualId } = this.getModel()
const message = await this.client.messages.create({
model: modelId,
max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
thinking,
temperature,
messages: [{ role: "user", content: prompt }],
stream: false,
})
const message = await this.client.messages.create(
{
model: modelId,
max_tokens: maxTokens ?? ANTHROPIC_DEFAULT_MAX_TOKENS,
thinking,
temperature,
messages: [{ role: "user", content: prompt }],
stream: false,
},
(() => {
// Check for the thinking-128k variant first
if (virtualId === "claude-3-7-sonnet-20250219:thinking-128k") {
return {
headers: { "anthropic-beta": "output-128k-2025-02-19" },
}
}
// Then check for models that support prompt caching
switch (modelId) {
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307":
return {
headers: { "anthropic-beta": "prompt-caching-2024-07-31" },
}
default:
return undefined
}
})(),
)
const content = message.content.find(({ type }) => type === "text")
return content?.type === "text" ? content.text : ""
@ -223,17 +259,40 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
try {
// Use the current model
const actualModelId = this.getModel().id
const { id: actualModelId, virtualId } = this.getModel()
const response = await this.client.messages.countTokens({
model: actualModelId,
messages: [
{
role: "user",
content: content,
},
],
})
const response = await this.client.messages.countTokens(
{
model: actualModelId,
messages: [
{
role: "user",
content: content,
},
],
},
(() => {
// Check for the thinking-128k variant first
if (virtualId === "claude-3-7-sonnet-20250219:thinking-128k") {
return {
headers: { "anthropic-beta": "output-128k-2025-02-19" },
}
}
// Then check for models that support prompt caching
switch (actualModelId) {
case "claude-3-5-sonnet-20241022":
case "claude-3-5-haiku-20241022":
case "claude-3-opus-20240229":
case "claude-3-haiku-20240307":
return {
headers: { "anthropic-beta": "prompt-caching-2024-07-31" },
}
default:
return undefined
}
})(),
)
return response.input_tokens
} catch (error) {

View file

@ -98,6 +98,18 @@ export interface ModelInfo {
export type AnthropicModelId = keyof typeof anthropicModels
export const anthropicDefaultModelId: AnthropicModelId = "claude-3-7-sonnet-20250219"
export const anthropicModels = {
"claude-3-7-sonnet-20250219:thinking-128k": {
maxTokens: 128_000,
contextWindow: 200_000,
supportsImages: true,
supportsComputerUse: true,
supportsPromptCache: true,
inputPrice: 3.0, // $3 per million input tokens
outputPrice: 15.0, // $15 per million output tokens
cacheWritesPrice: 3.75, // $3.75 per million tokens
cacheReadsPrice: 0.3, // $0.30 per million tokens
thinking: true,
},
"claude-3-7-sonnet-20250219:thinking": {
maxTokens: 64_000,
contextWindow: 200_000,