Skip to content

Commit 6d4c75e

Browse files
committed
Return ChatResponseMetadata with BedrockAnthropic3ChatModel call prompt method.
Include id, model, inputTokens and outputTokens usage data.
1 parent c5e9e0a commit 6d4c75e

File tree

1 file changed

+25
-3
lines changed

1 file changed

+25
-3
lines changed

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323

2424
import org.springframework.ai.chat.messages.AssistantMessage;
2525
import org.springframework.ai.chat.messages.UserMessage;
26+
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
27+
import org.springframework.ai.chat.metadata.Usage;
2628
import reactor.core.publisher.Flux;
2729

2830
import org.springframework.ai.bedrock.anthropic3.api.Anthropic3ChatBedrockApi;
@@ -83,11 +85,16 @@ public ChatResponse call(Prompt prompt) {
8385

8486
List<Generation> generations = response.content().stream().map(content -> {
8587
return new Generation(new AssistantMessage(content.text()),
86-
ChatGenerationMetadata.from(response.stopReason(), new Anthropic3ChatBedrockApi.AnthropicUsage(
87-
response.usage().inputTokens(), response.usage().outputTokens())));
88+
ChatGenerationMetadata.from(response.stopReason(), null));
8889
}).toList();
8990

90-
return new ChatResponse(generations);
91+
ChatResponseMetadata metadata = ChatResponseMetadata.builder()
92+
.withId(response.id())
93+
.withModel(response.model())
94+
.withUsage(extractUsage(response))
95+
.build();
96+
97+
return new ChatResponse(generations, metadata);
9198
}
9299

93100
@Override
@@ -117,6 +124,21 @@ public Flux<ChatResponse> stream(Prompt prompt) {
117124
});
118125
}
119126

127+
private Usage extractUsage(AnthropicChatResponse response) {
128+
return new Usage() {
129+
130+
@Override
131+
public Long getPromptTokens() {
132+
return response.usage().inputTokens().longValue();
133+
}
134+
135+
@Override
136+
public Long getGenerationTokens() {
137+
return response.usage().outputTokens().longValue();
138+
}
139+
};
140+
}
141+
120142
/**
121143
* Accessible for testing.
122144
*/

0 commit comments

Comments
 (0)