Skip to content

fix(ai-proxy): account for Bedrock prompt cache tokens - #14044

Open
AlinsRan wants to merge 1 commit into
apache:masterfrom
AlinsRan:fix/bedrock-cache-usage
Open

AlinsRan wants to merge 1 commit into
apache:masterfrom
AlinsRan:fix/bedrock-cache-usage

Conversation

@AlinsRan

Copy link
Copy Markdown
Contributor

Description

Bedrock Converse ordinary responses and streaming metadata currently omit cache-read and cache-write tokens. Normalize both paths consistently with LiteLLM: add cache tokens to prompt_tokens, expose the cache counters, include them in the fallback total, and preserve an explicit upstream totalTokens and raw usage. Support the InvokeModel cache counter aliases, giving canonical Converse fields precedence even when zero.

For inputTokens=10, cache reads of 20, cache writes of 30, and outputTokens=8, prompt usage now reports 60 instead of 10, with both cache logging variables populated.

Add 12 differential cases and end-to-end assertions for both response modes, preserving all existing test cases. Let the existing Bedrock mock select fixtures through X-AI-Fixture.

Validation:

  • prove -I. -v t/plugin/ai-proxy-bedrock-single.t: all 65 assertions pass locally.
  • luacheck apisix/plugins/ai-protocols/bedrock-converse.lua t/lib/server.lua: no warnings or errors.
  • Differential against LiteLLM 3d526bfe3f6b9cbe06e80156d84d5d1aae5633e3: executed the source AST's AmazonConverseConfig.transform_usage and cache-count methods with passive output containers. All 120 selected fields across 12 cases and both Lua response modes match. The Lua harness uses the real apisix.core.json module.

The same fix was submitted to EE first: https://github.com/api7/api7-ee-3-gateway/pull/2331.

Which issue(s) this PR fixes:

Related: https://github.com/api7/api7-ee-3-gateway/issues/2323

Checklist

  • I have explained the need for this PR and the problem it solves
  • I have explained the changes or the new features added to this PR
  • I have added tests corresponding to this change
  • I have updated the documentation to reflect this change (no configuration or API changes)
  • I have verified that this change is backward compatible

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant