This is an automated email from the ASF dual-hosted git repository.
AlinsRan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/apisix.git
The following commit(s) were added to refs/heads/master by this push:
new 96ba388d50 fix(ai-proxy): account for Bedrock prompt cache tokens
(#14044)
96ba388d50 is described below
commit 96ba388d507ea73129a4ee7848c740b6274c5181
Author: AlinsRan <[email protected]>
AuthorDate: Mon Oct 26 16:46:32 2026 +0800
fix(ai-proxy): account for Bedrock prompt cache tokens (#14044)
---
apisix/plugins/ai-protocols/bedrock-converse.lua | 39 ++--
t/fixtures/bedrock/converse-cache-usage-cases.json | 196 +++++++++++++++++++++
.../bedrock/converse-streaming-with-cache.bin | Bin 0 -> 600 bytes
t/fixtures/bedrock/converse-with-cache.json | 1 +
t/lib/server.lua | 6 +-
t/plugin/ai-proxy-bedrock-single.t | 68 +++++++
6 files changed, 296 insertions(+), 14 deletions(-)
diff --git a/apisix/plugins/ai-protocols/bedrock-converse.lua
b/apisix/plugins/ai-protocols/bedrock-converse.lua
index 1d179d01e8..97929d2590 100644
--- a/apisix/plugins/ai-protocols/bedrock-converse.lua
+++ b/apisix/plugins/ai-protocols/bedrock-converse.lua
@@ -30,6 +30,31 @@ local table = table
local _M = {}
+local function cache_token_count(raw, key, alias)
+ local value = raw[key]
+ if type(value) ~= "number" then
+ value = raw[alias]
+ end
+ return type(value) == "number" and value or 0
+end
+
+
+local function normalize_usage(raw)
+ local cache_read = cache_token_count(raw, "cacheReadInputTokens",
"cacheReadInputTokenCount")
+ local cache_write = cache_token_count(raw, "cacheWriteInputTokens",
"cacheWriteInputTokenCount")
+ -- Bedrock inputTokens excludes cached tokens; normalized prompt usage
includes them.
+ local prompt_tokens = (raw.inputTokens or 0) + cache_read + cache_write
+ local completion_tokens = raw.outputTokens or 0
+ return {
+ prompt_tokens = prompt_tokens,
+ completion_tokens = completion_tokens,
+ total_tokens = raw.totalTokens or prompt_tokens + completion_tokens,
+ cache_read_input_tokens = cache_read,
+ cache_creation_input_tokens = cache_write,
+ }
+end
+
+
--- Detect whether the request matches the Bedrock Converse API format.
-- Uses URI suffix (/converse) and body (valid JSON table with messages).
function _M.matches(body, ctx)
@@ -117,12 +142,7 @@ function _M.parse_sse_event(event, ctx, state)
local raw = data.usage
return {
type = "usage_and_done",
- usage = {
- prompt_tokens = raw.inputTokens or 0,
- completion_tokens = raw.outputTokens or 0,
- total_tokens = raw.totalTokens
- or (raw.inputTokens or 0) + (raw.outputTokens or 0),
- },
+ usage = normalize_usage(raw),
raw_usage = raw,
}
end
@@ -140,12 +160,7 @@ function _M.extract_usage(res_body)
return nil, nil
end
local raw = res_body.usage
- return {
- prompt_tokens = raw.inputTokens or 0,
- completion_tokens = raw.outputTokens or 0,
- total_tokens = raw.totalTokens
- or (raw.inputTokens or 0) + (raw.outputTokens or 0),
- }, raw
+ return normalize_usage(raw), raw
end
diff --git a/t/fixtures/bedrock/converse-cache-usage-cases.json
b/t/fixtures/bedrock/converse-cache-usage-cases.json
new file mode 100644
index 0000000000..6ba210f76d
--- /dev/null
+++ b/t/fixtures/bedrock/converse-cache-usage-cases.json
@@ -0,0 +1,196 @@
+[
+ {
+ "name": "no cache",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8
+ },
+ "expected": {
+ "prompt_tokens": 10,
+ "completion_tokens": 8,
+ "total_tokens": 18,
+ "cache_read_input_tokens": 0,
+ "cache_creation_input_tokens": 0
+ }
+ },
+ {
+ "name": "read",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": 20
+ },
+ "expected": {
+ "prompt_tokens": 30,
+ "completion_tokens": 8,
+ "total_tokens": 38,
+ "cache_read_input_tokens": 20,
+ "cache_creation_input_tokens": 0
+ }
+ },
+ {
+ "name": "write",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheWriteInputTokens": 30
+ },
+ "expected": {
+ "prompt_tokens": 40,
+ "completion_tokens": 8,
+ "total_tokens": 48,
+ "cache_read_input_tokens": 0,
+ "cache_creation_input_tokens": 30
+ }
+ },
+ {
+ "name": "both",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": 20,
+ "cacheWriteInputTokens": 30
+ },
+ "expected": {
+ "prompt_tokens": 60,
+ "completion_tokens": 8,
+ "total_tokens": 68,
+ "cache_read_input_tokens": 20,
+ "cache_creation_input_tokens": 30
+ }
+ },
+ {
+ "name": "explicit total",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": 20,
+ "cacheWriteInputTokens": 30,
+ "totalTokens": 68
+ },
+ "expected": {
+ "prompt_tokens": 60,
+ "completion_tokens": 8,
+ "total_tokens": 68,
+ "cache_read_input_tokens": 20,
+ "cache_creation_input_tokens": 30
+ }
+ },
+ {
+ "name": "preserve total",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": 20,
+ "totalTokens": 99
+ },
+ "expected": {
+ "prompt_tokens": 30,
+ "completion_tokens": 8,
+ "total_tokens": 99,
+ "cache_read_input_tokens": 20,
+ "cache_creation_input_tokens": 0
+ }
+ },
+ {
+ "name": "zero total",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "totalTokens": 0
+ },
+ "expected": {
+ "prompt_tokens": 10,
+ "completion_tokens": 8,
+ "total_tokens": 0,
+ "cache_read_input_tokens": 0,
+ "cache_creation_input_tokens": 0
+ }
+ },
+ {
+ "name": "aliases",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokenCount": 20,
+ "cacheWriteInputTokenCount": 30
+ },
+ "expected": {
+ "prompt_tokens": 60,
+ "completion_tokens": 8,
+ "total_tokens": 68,
+ "cache_read_input_tokens": 20,
+ "cache_creation_input_tokens": 30
+ }
+ },
+ {
+ "name": "canonical precedence",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": 4,
+ "cacheReadInputTokenCount": 20,
+ "cacheWriteInputTokens": 6,
+ "cacheWriteInputTokenCount": 30
+ },
+ "expected": {
+ "prompt_tokens": 20,
+ "completion_tokens": 8,
+ "total_tokens": 28,
+ "cache_read_input_tokens": 4,
+ "cache_creation_input_tokens": 6
+ }
+ },
+ {
+ "name": "zero precedence",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": 0,
+ "cacheReadInputTokenCount": 20,
+ "cacheWriteInputTokens": 0,
+ "cacheWriteInputTokenCount": 30
+ },
+ "expected": {
+ "prompt_tokens": 10,
+ "completion_tokens": 8,
+ "total_tokens": 18,
+ "cache_read_input_tokens": 0,
+ "cache_creation_input_tokens": 0
+ }
+ },
+ {
+ "name": "null cache",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": null,
+ "cacheWriteInputTokens": null
+ },
+ "expected": {
+ "prompt_tokens": 10,
+ "completion_tokens": 8,
+ "total_tokens": 18,
+ "cache_read_input_tokens": 0,
+ "cache_creation_input_tokens": 0
+ }
+ },
+ {
+ "name": "null alias fallback",
+ "raw": {
+ "inputTokens": 10,
+ "outputTokens": 8,
+ "cacheReadInputTokens": null,
+ "cacheReadInputTokenCount": 20,
+ "cacheWriteInputTokens": null,
+ "cacheWriteInputTokenCount": 30
+ },
+ "expected": {
+ "prompt_tokens": 60,
+ "completion_tokens": 8,
+ "total_tokens": 68,
+ "cache_read_input_tokens": 20,
+ "cache_creation_input_tokens": 30
+ }
+ }
+]
diff --git a/t/fixtures/bedrock/converse-streaming-with-cache.bin
b/t/fixtures/bedrock/converse-streaming-with-cache.bin
new file mode 100644
index 0000000000..ee4ec14e3a
Binary files /dev/null and
b/t/fixtures/bedrock/converse-streaming-with-cache.bin differ
diff --git a/t/fixtures/bedrock/converse-with-cache.json
b/t/fixtures/bedrock/converse-with-cache.json
new file mode 100644
index 0000000000..154ecbec21
--- /dev/null
+++ b/t/fixtures/bedrock/converse-with-cache.json
@@ -0,0 +1 @@
+{"output":{"message":{"role":"assistant","content":[{"text":"Hello"}]}},"stopReason":"end_turn","usage":{"inputTokens":10,"outputTokens":8,"totalTokens":68,"cacheReadInputTokens":20,"cacheWriteInputTokens":30}}
diff --git a/t/lib/server.lua b/t/lib/server.lua
index b2f057dc4f..33be399269 100644
--- a/t/lib/server.lua
+++ b/t/lib/server.lua
@@ -1480,7 +1480,8 @@ function _M.bedrock_converse()
end
local fixture_loader = require("lib.fixture_loader")
- local content, ferr = fixture_loader.load("bedrock/converse-basic.json")
+ local content, ferr = fixture_loader.load(
+ ngx.req.get_headers()["X-AI-Fixture"] or "bedrock/converse-basic.json")
if not content then
ngx.status = 500
ngx.say(ferr)
@@ -1594,7 +1595,8 @@ function _M.bedrock_converse_stream()
end
local fixture_loader = require("lib.fixture_loader")
- local content, ferr =
fixture_loader.load("bedrock/bedrock-converse-streaming.bin")
+ local content, ferr = fixture_loader.load(
+ ngx.req.get_headers()["X-AI-Fixture"] or
"bedrock/bedrock-converse-streaming.bin")
if not content then
ngx.status = 500
ngx.say(ferr)
diff --git a/t/plugin/ai-proxy-bedrock-single.t
b/t/plugin/ai-proxy-bedrock-single.t
index fc7e7d3b5a..c5f5dacb8b 100644
--- a/t/plugin/ai-proxy-bedrock-single.t
+++ b/t/plugin/ai-proxy-bedrock-single.t
@@ -497,3 +497,71 @@ POST /single-ai/stream/converse
qr/"text"\s*:\s*"Hello!"/
--- error_log eval
qr{\[test\] received uri:
/model/anthropic\.claude-3-5-sonnet-20241022-v2%3A0/converse(?!-stream)}
+
+
+
+=== TEST 19: Bedrock cache usage matches LiteLLM for both response modes
+--- config
+ location /t {
+ content_by_lua_block {
+ local core = require("apisix.core")
+ local bedrock =
require("apisix.plugins.ai-protocols.bedrock-converse")
+ local loader = require("lib.fixture_loader")
+ -- Expected usage was differentialed against LiteLLM 3d526bfe
transform_usage.
+ local cases = core.json.decode(assert(loader.load(
+ "bedrock/converse-cache-usage-cases.json")), { null_as_nil =
true })
+ for _, case in ipairs(cases) do
+ local usage, raw = bedrock.extract_usage({ usage = case.raw })
+ assert(raw == case.raw, case.name .. " raw usage preserved")
+ local result = bedrock.parse_sse_event({
+ headers = { [":event-type"] = "metadata" },
+ payload = core.json.encode({ usage = case.raw }),
+ }, {}, {})
+ assert(result.type == "usage_and_done", case.name)
+ for key, expected in pairs(case.expected) do
+ assert(usage[key] == expected, case.name .. " non-stream "
.. key)
+ assert(result.usage[key] == expected, case.name .. "
stream " .. key)
+ end
+ assert(core.json.stably_encode(result.raw_usage)
+ == core.json.stably_encode(case.raw), case.name .. "
stream raw usage")
+ end
+ local usage = bedrock.extract_usage({ usage = {} })
+ assert(usage.prompt_tokens == 0 and usage.total_tokens == 0,
"empty usage")
+ assert(bedrock.extract_usage({}) == nil, "missing usage")
+ ngx.say("OK")
+ }
+ }
+--- response_body
+OK
+--- no_error_log
+[error]
+
+
+
+=== TEST 20: non-streaming cache tokens reach the LLM logging variables
+--- request
+POST /single-ai/stream/converse
+{"messages":[{"role":"user","content":[{"text":"Say hi"}]}]}
+--- more_headers
+X-AI-Fixture: bedrock/converse-with-cache.json
+--- response_body_like eval
+qr/"cacheReadInputTokens":20/
+--- access_log eval
+qr/ 60 8 68 false false 0 \S* 20 30 0 /
+--- no_error_log
+[error]
+
+
+
+=== TEST 21: streaming cache tokens reach the LLM logging variables
+--- request
+POST /single-ai/stream/converse
+{"stream":true,"messages":[{"role":"user","content":[{"text":"Say hi"}]}]}
+--- more_headers
+X-AI-Fixture: bedrock/converse-streaming-with-cache.bin
+--- response_body_like eval
+qr/metadata/
+--- access_log eval
+qr/ 60 8 68 true false 0 \S* 20 30 0 /
+--- no_error_log
+[error]
