Skip to content

Commit ff2974f

Browse files
hntrlFilipZmijewski
authored andcommitted
fix(aws): map Bedrock prompt cache usage metadata (langchain-ai#10839)
## Summary `ChatBedrockConverse` currently drops Bedrock prompt-caching usage fields (`cacheReadInputTokens` and `cacheWriteInputTokens`) when building LangChain usage metadata. This change maps those fields into `usage_metadata.input_token_details` for both invoke and stream metadata paths so cache behavior is visible in LangSmith and downstream cost/usage analytics. ## Changes - `@langchain/aws` - Updated `convertConverseMessageToLangChainMessage` to map: - `cacheReadInputTokens` -> `input_token_details.cache_read` - `cacheWriteInputTokens` -> `input_token_details.cache_creation` - Updated `handleConverseStreamMetadata` to emit the same cache token details in streaming usage metadata. - Added focused unit tests covering: - non-stream cache token mapping - stream cache token mapping - omission of `input_token_details` when Bedrock cache fields are absent
1 parent 43f75b1 commit ff2974f

3 files changed

Lines changed: 116 additions & 0 deletions

File tree

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
---
2+
"@langchain/aws": patch
3+
---
4+
5+
fix(aws): map Bedrock prompt cache usage metadata to input token details
6+
7+
Include `cacheReadInputTokens` and `cacheWriteInputTokens` from Bedrock Converse
8+
responses in `usage_metadata.input_token_details` for both invoke and stream
9+
metadata handling.

libs/providers/langchain-aws/src/utils/message_outputs.ts

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,24 @@ export function convertConverseMessageToLangChainMessage(
3333
}
3434
let tokenUsage: UsageMetadata | undefined;
3535
if (responseMetadata.usage) {
36+
const inputTokenDetails = {
37+
...(responseMetadata.usage.cacheReadInputTokens !== undefined && {
38+
cache_read: responseMetadata.usage.cacheReadInputTokens,
39+
}),
40+
...(responseMetadata.usage.cacheWriteInputTokens !== undefined && {
41+
cache_creation: responseMetadata.usage.cacheWriteInputTokens,
42+
}),
43+
};
3644
const input_tokens = responseMetadata.usage.inputTokens ?? 0;
3745
const output_tokens = responseMetadata.usage.outputTokens ?? 0;
3846
tokenUsage = {
3947
input_tokens,
4048
output_tokens,
4149
total_tokens:
4250
responseMetadata.usage.totalTokens ?? input_tokens + output_tokens,
51+
input_token_details: Object.keys(inputTokenDetails).length
52+
? inputTokenDetails
53+
: undefined,
4354
};
4455
}
4556

@@ -193,12 +204,23 @@ export function handleConverseStreamMetadata(
193204
streamUsage: boolean;
194205
}
195206
): ChatGenerationChunk {
207+
const inputTokenDetails = {
208+
...(metadata.usage?.cacheReadInputTokens !== undefined && {
209+
cache_read: metadata.usage.cacheReadInputTokens,
210+
}),
211+
...(metadata.usage?.cacheWriteInputTokens !== undefined && {
212+
cache_creation: metadata.usage.cacheWriteInputTokens,
213+
}),
214+
};
196215
const inputTokens = metadata.usage?.inputTokens ?? 0;
197216
const outputTokens = metadata.usage?.outputTokens ?? 0;
198217
const usage_metadata: UsageMetadata = {
199218
input_tokens: inputTokens,
200219
output_tokens: outputTokens,
201220
total_tokens: metadata.usage?.totalTokens ?? inputTokens + outputTokens,
221+
input_token_details: Object.keys(inputTokenDetails).length
222+
? inputTokenDetails
223+
: undefined,
202224
};
203225
return new ChatGenerationChunk({
204226
text: "",
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { describe, expect, test } from "vitest";
2+
import type * as Bedrock from "@aws-sdk/client-bedrock-runtime";
3+
import {
4+
convertConverseMessageToLangChainMessage,
5+
handleConverseStreamMetadata,
6+
} from "../message_outputs.js";
7+
8+
describe("message output usage metadata conversion", () => {
9+
test("maps Bedrock prompt cache tokens for non-stream responses", () => {
10+
const message: Bedrock.Message = {
11+
role: "assistant",
12+
content: [{ text: "Hello" }],
13+
};
14+
const responseMetadata = {
15+
usage: {
16+
inputTokens: 10,
17+
outputTokens: 5,
18+
totalTokens: 15,
19+
cacheReadInputTokens: 7,
20+
cacheWriteInputTokens: 3,
21+
},
22+
} as Omit<Bedrock.ConverseResponse, "output">;
23+
24+
const result = convertConverseMessageToLangChainMessage(
25+
message,
26+
responseMetadata
27+
);
28+
29+
expect(result.usage_metadata).toEqual({
30+
input_tokens: 10,
31+
output_tokens: 5,
32+
total_tokens: 15,
33+
input_token_details: {
34+
cache_read: 7,
35+
cache_creation: 3,
36+
},
37+
});
38+
});
39+
40+
test("does not add input_token_details when Bedrock cache fields are absent", () => {
41+
const message: Bedrock.Message = {
42+
role: "assistant",
43+
content: [{ text: "Hello" }],
44+
};
45+
const responseMetadata = {
46+
usage: {
47+
inputTokens: 10,
48+
outputTokens: 5,
49+
totalTokens: 15,
50+
},
51+
} as Omit<Bedrock.ConverseResponse, "output">;
52+
53+
const result = convertConverseMessageToLangChainMessage(
54+
message,
55+
responseMetadata
56+
);
57+
58+
expect(result.usage_metadata?.input_token_details).toBeUndefined();
59+
});
60+
61+
test("maps Bedrock prompt cache tokens for stream metadata", () => {
62+
const result = handleConverseStreamMetadata(
63+
{
64+
usage: {
65+
inputTokens: 20,
66+
outputTokens: 4,
67+
totalTokens: 24,
68+
cacheReadInputTokens: 9,
69+
cacheWriteInputTokens: 6,
70+
},
71+
},
72+
{ streamUsage: true }
73+
);
74+
75+
expect(result.message.usage_metadata).toEqual({
76+
input_tokens: 20,
77+
output_tokens: 4,
78+
total_tokens: 24,
79+
input_token_details: {
80+
cache_read: 9,
81+
cache_creation: 6,
82+
},
83+
});
84+
});
85+
});

0 commit comments

Comments
 (0)