fix(ai): normalize Bedrock cache usage (#38427)
Co-authored-by: Aiden Cline <63023139+rekram1-node@users.noreply.github.com>
This commit is contained in:
parent
f1f0f47ee2
commit
6e8aefcfa0
5 changed files with 114 additions and 20 deletions
|
|
@ -436,21 +436,22 @@ const mapFinishReason = (reason: string): FinishReason => {
|
||||||
return "unknown"
|
return "unknown"
|
||||||
}
|
}
|
||||||
|
|
||||||
// AWS Bedrock Converse reports `inputTokens` (inclusive total) with
|
// AWS reports inputTokens separately from cache reads and writes.
|
||||||
// `cacheReadInputTokens` and `cacheWriteInputTokens` as subsets. Pass
|
// Bedrock does not break reasoning out of outputTokens for current models.
|
||||||
// the total through and derive the non-cached breakdown. Bedrock does
|
|
||||||
// not break reasoning out of `outputTokens` for any current model.
|
|
||||||
const mapUsage = (usage: BedrockUsageSchema | undefined): Usage | undefined => {
|
const mapUsage = (usage: BedrockUsageSchema | undefined): Usage | undefined => {
|
||||||
if (!usage) return undefined
|
if (!usage) return undefined
|
||||||
const cacheTotal = (usage.cacheReadInputTokens ?? 0) + (usage.cacheWriteInputTokens ?? 0)
|
const inputTokens = ProviderShared.sumTokens(
|
||||||
const nonCached = ProviderShared.subtractTokens(usage.inputTokens, cacheTotal)
|
usage.inputTokens,
|
||||||
|
usage.cacheReadInputTokens,
|
||||||
|
usage.cacheWriteInputTokens,
|
||||||
|
)
|
||||||
return new Usage({
|
return new Usage({
|
||||||
inputTokens: usage.inputTokens,
|
inputTokens,
|
||||||
outputTokens: usage.outputTokens,
|
outputTokens: usage.outputTokens,
|
||||||
nonCachedInputTokens: nonCached,
|
nonCachedInputTokens: usage.inputTokens,
|
||||||
cacheReadInputTokens: usage.cacheReadInputTokens,
|
cacheReadInputTokens: usage.cacheReadInputTokens,
|
||||||
cacheWriteInputTokens: usage.cacheWriteInputTokens,
|
cacheWriteInputTokens: usage.cacheWriteInputTokens,
|
||||||
totalTokens: ProviderShared.totalTokens(usage.inputTokens, usage.outputTokens, usage.totalTokens),
|
totalTokens: ProviderShared.totalTokens(inputTokens, usage.outputTokens, usage.totalTokens),
|
||||||
providerMetadata: { bedrock: usage },
|
providerMetadata: { bedrock: usage },
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -34,11 +34,12 @@ import { ProviderFailureClassification } from "./errors"
|
||||||
*
|
*
|
||||||
* **Semantics by provider**:
|
* **Semantics by provider**:
|
||||||
*
|
*
|
||||||
* - OpenAI Chat / Responses / Gemini / Bedrock: provider reports inclusive
|
* - OpenAI Chat / Responses / Gemini: provider reports inclusive
|
||||||
* `inputTokens` and an inclusive `outputTokens`; mapper subtracts to
|
* `inputTokens` and an inclusive `outputTokens`; mapper subtracts to
|
||||||
* derive the breakdown.
|
* derive the breakdown.
|
||||||
* - Anthropic: provider reports the breakdown natively (`input_tokens` is
|
* - Anthropic and Bedrock report the input breakdown natively: Anthropic's
|
||||||
* non-cached only); mapper sums to derive the inclusive `inputTokens`.
|
* `input_tokens` and Bedrock's `inputTokens` are non-cached only. Their
|
||||||
|
* mappers sum the breakdown to derive the inclusive `inputTokens`.
|
||||||
* Anthropic does *not* break extended-thinking out of `output_tokens`, so
|
* Anthropic does *not* break extended-thinking out of `output_tokens`, so
|
||||||
* `reasoningTokens` is `undefined` and `outputTokens` carries the
|
* `reasoningTokens` is `undefined` and `outputTokens` carries the
|
||||||
* combined total — a documented limitation of the Anthropic API.
|
* combined total — a documented limitation of the Anthropic API.
|
||||||
|
|
|
||||||
File diff suppressed because one or more lines are too long
|
|
@ -13,12 +13,8 @@ const RECORDING_REGION = process.env.BEDROCK_RECORDING_REGION ?? "us-east-1"
|
||||||
// call wouldn't deterministically prove cache mapping works. Override with
|
// call wouldn't deterministically prove cache mapping works. Override with
|
||||||
// BEDROCK_CACHE_MODEL_ID if your account has access elsewhere.
|
// BEDROCK_CACHE_MODEL_ID if your account has access elsewhere.
|
||||||
const model = AmazonBedrock.configure({
|
const model = AmazonBedrock.configure({
|
||||||
credentials: {
|
apiKey: process.env.AWS_BEARER_TOKEN_BEDROCK ?? "fixture",
|
||||||
region: RECORDING_REGION,
|
region: RECORDING_REGION,
|
||||||
accessKeyId: process.env.AWS_ACCESS_KEY_ID ?? "fixture",
|
|
||||||
secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY ?? "fixture",
|
|
||||||
sessionToken: process.env.AWS_SESSION_TOKEN,
|
|
||||||
},
|
|
||||||
}).model(process.env.BEDROCK_CACHE_MODEL_ID ?? "us.anthropic.claude-haiku-4-5-20251001-v1:0")
|
}).model(process.env.BEDROCK_CACHE_MODEL_ID ?? "us.anthropic.claude-haiku-4-5-20251001-v1:0")
|
||||||
|
|
||||||
const cacheRequest = LLM.request({
|
const cacheRequest = LLM.request({
|
||||||
|
|
@ -36,7 +32,7 @@ const recorded = recordedTests({
|
||||||
prefix: "bedrock-converse-cache",
|
prefix: "bedrock-converse-cache",
|
||||||
provider: "amazon-bedrock",
|
provider: "amazon-bedrock",
|
||||||
protocol: "bedrock-converse",
|
protocol: "bedrock-converse",
|
||||||
requires: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
|
requires: ["AWS_BEARER_TOKEN_BEDROCK"],
|
||||||
// Two identical requests in one cassette — replay walks the cassette in
|
// Two identical requests in one cassette — replay walks the cassette in
|
||||||
// recording order so the second call replays the cached-hit interaction.
|
// recording order so the second call replays the cached-hit interaction.
|
||||||
})
|
})
|
||||||
|
|
@ -45,10 +41,20 @@ describe("Bedrock Converse cache recorded", () => {
|
||||||
recorded.effect.with("writes then reads cachePoint on identical second call", { tags: ["cache"] }, () =>
|
recorded.effect.with("writes then reads cachePoint on identical second call", { tags: ["cache"] }, () =>
|
||||||
Effect.gen(function* () {
|
Effect.gen(function* () {
|
||||||
const first = yield* LLMClient.generate(cacheRequest)
|
const first = yield* LLMClient.generate(cacheRequest)
|
||||||
expect(first.usage?.cacheReadInputTokens ?? 0).toBeGreaterThanOrEqual(0)
|
expect(first.usage?.cacheWriteInputTokens ?? 0).toBeGreaterThan(0)
|
||||||
|
expect(first.usage?.inputTokens).toBe(
|
||||||
|
(first.usage?.nonCachedInputTokens ?? 0) +
|
||||||
|
(first.usage?.cacheReadInputTokens ?? 0) +
|
||||||
|
(first.usage?.cacheWriteInputTokens ?? 0),
|
||||||
|
)
|
||||||
|
|
||||||
const second = yield* LLMClient.generate(cacheRequest)
|
const second = yield* LLMClient.generate(cacheRequest)
|
||||||
expect(second.usage?.cacheReadInputTokens ?? 0).toBeGreaterThan(0)
|
expect(second.usage?.cacheReadInputTokens ?? 0).toBeGreaterThan(0)
|
||||||
|
expect(second.usage?.inputTokens).toBe(
|
||||||
|
(second.usage?.nonCachedInputTokens ?? 0) +
|
||||||
|
(second.usage?.cacheReadInputTokens ?? 0) +
|
||||||
|
(second.usage?.cacheWriteInputTokens ?? 0),
|
||||||
|
)
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
})
|
})
|
||||||
|
|
|
||||||
|
|
@ -269,6 +269,39 @@ describe("Bedrock Converse route", () => {
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
it.effect("adds cache reads and writes to Bedrock input usage", () =>
|
||||||
|
Effect.gen(function* () {
|
||||||
|
const body = eventStreamBody(
|
||||||
|
["messageStart", { role: "assistant" }],
|
||||||
|
["contentBlockDelta", { contentBlockIndex: 0, delta: { text: "Hello" } }],
|
||||||
|
["contentBlockStop", { contentBlockIndex: 0 }],
|
||||||
|
["messageStop", { stopReason: "end_turn" }],
|
||||||
|
[
|
||||||
|
"metadata",
|
||||||
|
{
|
||||||
|
usage: {
|
||||||
|
inputTokens: 5,
|
||||||
|
outputTokens: 2,
|
||||||
|
totalTokens: 12,
|
||||||
|
cacheReadInputTokens: 3,
|
||||||
|
cacheWriteInputTokens: 2,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
|
||||||
|
|
||||||
|
expect(response.usage).toMatchObject({
|
||||||
|
inputTokens: 10,
|
||||||
|
nonCachedInputTokens: 5,
|
||||||
|
cacheReadInputTokens: 3,
|
||||||
|
cacheWriteInputTokens: 2,
|
||||||
|
outputTokens: 2,
|
||||||
|
totalTokens: 12,
|
||||||
|
})
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
it.effect("assembles streamed tool call input", () =>
|
it.effect("assembles streamed tool call input", () =>
|
||||||
Effect.gen(function* () {
|
Effect.gen(function* () {
|
||||||
const body = eventStreamBody(
|
const body = eventStreamBody(
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue