diff --git a/apisix/plugins/ai-protocols/bedrock-converse.lua b/apisix/plugins/ai-protocols/bedrock-converse.lua index 1d179d01e8ae..97929d2590f0 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 000000000000..6ba210f76df1 --- /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 000000000000..ee4ec14e3a0f 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 000000000000..154ecbec2148 --- /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 b2f057dc4f1c..33be39926909 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 fc7e7d3b5a97..c5f5dacb8bd0 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]