Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions packages/@ant/model-provider/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -67,3 +67,7 @@ export {
anthropicToolChoiceToOpenAI,
} from './shared/openaiConvertTools.js'
export { adaptOpenAIStreamToAnthropic } from './shared/openaiStreamAdapter.js'
export {
normalizeOpenAIUsage,
type AnthropicUsage,
} from './shared/openaiUsage.js'
Original file line number Diff line number Diff line change
Expand Up @@ -34,11 +34,15 @@ function makeChunk(
}

/** Collect all emitted Anthropic events from the stream adapter for assertion */
async function collectEvents(chunks: ChatCompletionChunk[]) {
async function collectEvents(
chunks: ChatCompletionChunk[],
options?: { includeCacheWriteTokens?: boolean },
) {
const events: any[] = []
for await (const event of adaptOpenAIStreamToAnthropic(
mockStream(chunks),
'gpt-4o',
options,
)) {
events.push(event)
}
Expand Down Expand Up @@ -521,6 +525,60 @@ describe('thinking support (reasoning_content)', () => {
})

describe('prompt caching support', () => {
test('maps official OpenAI cache writes when explicitly enabled', async () => {
const events = await collectEvents(
[
makeChunk({
choices: [
{ index: 0, delta: { content: 'hi' }, finish_reason: null },
],
}),
makeChunk({
choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
usage: {
prompt_tokens: 1000,
completion_tokens: 50,
total_tokens: 1050,
prompt_tokens_details: {
cached_tokens: 600,
cache_write_tokens: 250,
},
} as any,
}),
],
{ includeCacheWriteTokens: true },
)

const msgDelta = events.find(e => e.type === 'message_delta') as any
expect(msgDelta.usage.input_tokens).toBe(150)
expect(msgDelta.usage.cache_read_input_tokens).toBe(600)
expect(msgDelta.usage.cache_creation_input_tokens).toBe(250)
})

test('ignores cache writes for compatible providers by default', async () => {
const events = await collectEvents([
makeChunk({
choices: [{ index: 0, delta: { content: 'hi' }, finish_reason: null }],
}),
makeChunk({
choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
usage: {
prompt_tokens: 1000,
completion_tokens: 50,
total_tokens: 1050,
prompt_tokens_details: {
cached_tokens: 600,
cache_write_tokens: 250,
},
} as any,
}),
])

const msgDelta = events.find(e => e.type === 'message_delta') as any
expect(msgDelta.usage.input_tokens).toBe(400)
expect(msgDelta.usage.cache_creation_input_tokens).toBe(0)
})

test('maps cached_tokens to cache_read_input_tokens', async () => {
const events = await collectEvents([
makeChunk({
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
import { describe, expect, test } from 'bun:test'
import { normalizeOpenAIUsage } from '../openaiUsage.js'

describe('normalizeOpenAIUsage', () => {
test('partitions total input into ordinary, cache-read, and cache-write tokens', () => {
expect(
normalizeOpenAIUsage({
totalInputTokens: 1000,
outputTokens: 50,
cacheReadTokens: 600,
cacheWriteTokens: 250,
}),
).toEqual({
input_tokens: 150,
output_tokens: 50,
cache_creation_input_tokens: 250,
cache_read_input_tokens: 600,
})
})

test('clamps overlapping cache segments to the total input', () => {
expect(
normalizeOpenAIUsage({
totalInputTokens: 5000,
outputTokens: 10,
cacheReadTokens: 4000,
cacheWriteTokens: 4000,
}),
).toEqual({
input_tokens: 0,
output_tokens: 10,
cache_creation_input_tokens: 1000,
cache_read_input_tokens: 4000,
})
})

test('clamps negative provider values to zero', () => {
expect(
normalizeOpenAIUsage({
totalInputTokens: -1,
outputTokens: -2,
cacheReadTokens: -3,
cacheWriteTokens: -4,
}),
).toEqual({
input_tokens: 0,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
})
})
})
61 changes: 36 additions & 25 deletions packages/@ant/model-provider/src/shared/openaiStreamAdapter.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import type { BetaRawMessageStreamEvent } from '@anthropic-ai/sdk/resources/beta/messages/messages.mjs'
import type { ChatCompletionChunk } from 'openai/resources/chat/completions/completions.mjs'
import { randomUUID } from 'crypto'
import { normalizeOpenAIUsage } from './openaiUsage.js'

/**
* Adapt an OpenAI streaming response into Anthropic BetaRawMessageStreamEvent.
Expand All @@ -13,10 +14,10 @@ import { randomUUID } from 'crypto'
* finish_reason → message_delta(stop_reason) + message_stop
*
* Usage field mapping (OpenAI → Anthropic):
* prompt_tokens - cached_tokens → input_tokens (non-cached input only)
* prompt_tokens - cached_tokens - cache_write_tokens → input_tokens
* completion_tokens → output_tokens
* prompt_tokens_details.cached_tokens → cache_read_input_tokens
* (no OpenAI equivalent) → cache_creation_input_tokens (always 0)
* prompt_tokens_details.cache_write_tokens → cache_creation_input_tokens
*
* All four fields are emitted in the post-loop message_delta (not message_start)
* so that trailing usage chunks (sent after finish_reason by some
Expand All @@ -35,6 +36,7 @@ import { randomUUID } from 'crypto'
export async function* adaptOpenAIStreamToAnthropic(
stream: AsyncIterable<ChatCompletionChunk>,
model: string,
options?: { includeCacheWriteTokens?: boolean },
): AsyncGenerator<BetaRawMessageStreamEvent, void> {
const messageId = `msg_${randomUUID().replace(/-/g, '').slice(0, 24)}`

Expand All @@ -53,13 +55,13 @@ export async function* adaptOpenAIStreamToAnthropic(
// Track text block state
let textBlockOpen = false

// Track usage — all four Anthropic fields, populated from OpenAI usage fields:
// rawInputTokens tracks the raw prompt_tokens (OpenAI total, including cached).
// inputTokens is the derived Anthropic value (non-cached only = rawInputTokens - cachedReadTokens).
// Track raw OpenAI usage across chunks. The normalized Anthropic fields are
// disjoint: ordinary input + cache reads + cache writes = total input.
let rawInputTokens = 0
let inputTokens = 0
let outputTokens = 0
let cachedReadTokens = 0
let rawCacheReadTokens = 0
let rawCacheWriteTokens = 0
let usage = normalizeOpenAIUsage({ totalInputTokens: 0, outputTokens: 0 })

// Track all open content block indices (for cleanup)
const openBlockIndices = new Set<number>()
Expand All @@ -75,16 +77,32 @@ export async function* adaptOpenAIStreamToAnthropic(
// Extract usage from any chunk that carries it.
if (chunk.usage) {
rawInputTokens = chunk.usage.prompt_tokens ?? rawInputTokens
const rawCached =
((chunk.usage as any).prompt_tokens_details?.cached_tokens as
| number
| undefined) ?? cachedReadTokens
// Anthropic's input_tokens = non-cached input only. OpenAI's prompt_tokens
// includes cached tokens, so subtract. Clamp to 0 in case cached > total
// due to a streaming race.
inputTokens = Math.max(0, rawInputTokens - rawCached)
outputTokens = chunk.usage.completion_tokens ?? outputTokens
cachedReadTokens = rawCached

const usageRecord = chunk.usage as unknown as Record<string, unknown>
const detailsValue = usageRecord.prompt_tokens_details
const details =
detailsValue && typeof detailsValue === 'object'
? (detailsValue as Record<string, unknown>)
: undefined
if (typeof details?.cached_tokens === 'number') {
rawCacheReadTokens = details.cached_tokens
}
if (
options?.includeCacheWriteTokens &&
typeof details?.cache_write_tokens === 'number'
) {
rawCacheWriteTokens = details.cache_write_tokens
} else if (!options?.includeCacheWriteTokens) {
rawCacheWriteTokens = 0
}

usage = normalizeOpenAIUsage({
totalInputTokens: rawInputTokens,
outputTokens,
cacheReadTokens: rawCacheReadTokens,
cacheWriteTokens: rawCacheWriteTokens,
})
}

// Emit message_start on first chunk
Expand All @@ -102,10 +120,8 @@ export async function* adaptOpenAIStreamToAnthropic(
stop_reason: null,
stop_sequence: null,
usage: {
input_tokens: inputTokens,
...usage,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: cachedReadTokens,
},
},
} as unknown as BetaRawMessageStreamEvent
Expand Down Expand Up @@ -312,12 +328,7 @@ export async function* adaptOpenAIStreamToAnthropic(
stop_reason: stopReason,
stop_sequence: null,
},
usage: {
input_tokens: inputTokens,
output_tokens: outputTokens,
cache_read_input_tokens: cachedReadTokens,
cache_creation_input_tokens: 0,
},
usage,
} as BetaRawMessageStreamEvent

yield {
Expand Down
35 changes: 35 additions & 0 deletions packages/@ant/model-provider/src/shared/openaiUsage.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
export type AnthropicUsage = {
input_tokens: number
output_tokens: number
cache_creation_input_tokens: number
cache_read_input_tokens: number
}

/**
* Convert OpenAI's total-input usage into Anthropic's disjoint usage fields.
* Cache reads take priority when malformed provider data makes segments overlap.
*/
export function normalizeOpenAIUsage(params: {
totalInputTokens: number
outputTokens: number
cacheReadTokens?: number
cacheWriteTokens?: number
}): AnthropicUsage {
const totalInput = Math.max(0, params.totalInputTokens)
const cacheRead = Math.min(
Math.max(0, params.cacheReadTokens ?? 0),
totalInput,
)
const remainingAfterRead = Math.max(0, totalInput - cacheRead)
const cacheCreation = Math.min(
Math.max(0, params.cacheWriteTokens ?? 0),
remainingAfterRead,
)

return {
input_tokens: Math.max(0, remainingAfterRead - cacheCreation),
output_tokens: Math.max(0, params.outputTokens),
cache_creation_input_tokens: cacheCreation,
cache_read_input_tokens: cacheRead,
}
}
50 changes: 50 additions & 0 deletions src/services/api/openai/__tests__/openaiShared.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
import { describe, expect, test } from 'bun:test'
import {
getOfficialOpenAIPromptCacheKey,
isOfficialOpenAIBaseURL,
} from '../openaiShared.js'

describe('isOfficialOpenAIBaseURL', () => {
test('treats the SDK default endpoint as official OpenAI', () => {
expect(isOfficialOpenAIBaseURL(undefined)).toBe(true)
expect(isOfficialOpenAIBaseURL('')).toBe(true)
})

test('accepts global and regional official OpenAI endpoints', () => {
expect(isOfficialOpenAIBaseURL('https://api.openai.com/v1')).toBe(true)
expect(isOfficialOpenAIBaseURL('https://eu.api.openai.com/v1')).toBe(true)
expect(isOfficialOpenAIBaseURL('https://api.openai.com:443/v1')).toBe(true)
})

test('rejects OpenAI-compatible and spoofed endpoints', () => {
expect(isOfficialOpenAIBaseURL('https://api.deepseek.com/v1')).toBe(false)
expect(isOfficialOpenAIBaseURL('http://api.openai.com/v1')).toBe(false)
expect(isOfficialOpenAIBaseURL('https://api.openai.com.evil.test/v1')).toBe(
false,
)
expect(isOfficialOpenAIBaseURL('https://api.openai.com:8443/v1')).toBe(
false,
)
expect(isOfficialOpenAIBaseURL('not-a-url')).toBe(false)
})
})

describe('getOfficialOpenAIPromptCacheKey', () => {
test('returns a session key for the SDK default and official endpoint', () => {
expect(getOfficialOpenAIPromptCacheKey(undefined, 'session-1')).toBe(
'ccb:session-1',
)
expect(
getOfficialOpenAIPromptCacheKey('https://api.openai.com/v1', 'session-2'),
).toBe('ccb:session-2')
})

test('returns undefined for compatible endpoints', () => {
expect(
getOfficialOpenAIPromptCacheKey(
'https://api.deepseek.com/v1',
'session-1',
),
).toBeUndefined()
})
})
Loading