Skip to content

GH-6098: Expose Bedrock Cohere embedding token usage from HTTP headers#6145

Open
wolf8334 wants to merge 1 commit into
spring-projects:mainfrom
wolf8334:GH-6098
Open

GH-6098: Expose Bedrock Cohere embedding token usage from HTTP headers#6145
wolf8334 wants to merge 1 commit into
spring-projects:mainfrom
wolf8334:GH-6098

Conversation

@wolf8334
Copy link
Copy Markdown

Problem

Bedrock Cohere embedding API returns input token count in the `x-amzn-bedrock-input-token-count` HTTP response header, but Spring AI's `AbstractBedrockApi.internalInvocation()` only processes the response body, discarding the headers. This means `EmbeddingResponse` metadata contains no token usage (all zeros).

Changes

  • Added `BedrockResponse` record and `internalInvocationForModel()` to `AbstractBedrockApi` that returns both the parsed response body and the `SdkHttpResponse` headers
  • Added `CohereEmbeddingModelResponse` record and `embeddingForModel()` to `CohereEmbeddingBedrockApi` that extracts the input token count from the `x-amzn-bedrock-input-token-count` header
  • Updated `BedrockCohereEmbeddingModel.call()` to pass token usage via `EmbeddingResponseMetadata` + `DefaultUsage`
  • Added 3 unit tests in `BedrockCohereEmbeddingModelTest` covering: token usage from header present, missing header, and batch embeddings
  • Updated IT tests with token usage assertions

Reference

Same issue was fixed in LangChain4j: langchain4j/langchain4j#5064
Closes GH-6098

@quaff
Copy link
Copy Markdown
Contributor

quaff commented May 26, 2026

You should drop the unrelated commit fd1bda2 for this PR.

…rom HTTP headers

Bedrock Cohere embedding API returns input token count in the
x-amzn-bedrock-input-token-count HTTP response header, but
AbstractBedrockApi.internalInvocation() only processed the response body,
discarding the headers. This caused EmbeddingResponse metadata to contain
no token usage information.
Changes:
- Add BedrockResponse<T> record and internalInvocationForModel() to
  AbstractBedrockApi that returns both the response body and
  SdkHttpResponse headers
- Add CohereEmbeddingModelResponse wrapper and embeddingForModel() to
  CohereEmbeddingBedrockApi that extracts input token count from the
  x-amzn-bedrock-input-token-count header
- Update BedrockCohereEmbeddingModel.call() to pass token usage via
  EmbeddingResponseMetadata + DefaultUsage
- Add unit tests verifying token usage extraction from headers
Fixes spring-projectsGH-6098

Signed-off-by: 钱多多 <2916947+wolf8334@users.noreply.github.com>
@wolf8334
Copy link
Copy Markdown
Author

You should drop the unrelated commit fd1bda2 for this PR.

thank you.

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.

AWS Bedrock Cohere Embedding Model returns no metrics

2 participants