spring-ai
spring-ai copied to clipboard
spring-projects
Bedrock AI models add usage information.
This PR closes: https://github.com/spring-projects/spring-ai/issues/600
When Bedrock invokes model API, Bedrock will return input, output token count, and latency information from the SDKHttpResponse header.
extract this information as ChatResponseMetadata to ChatResponse
@wmz7year thank you for looking at this. It would be a valuable contribution to add more, structured usage metadata.
From the PR code, I have the impression that we can implement (most of) this inside AnthropicChatBedrockApi.java without having to expose it via the AmazonBedrockInvocationContex and reduce some necessary code repetition.
For example It looks like the BedrockAnthropicChatResponseMetadata , BedrockAnthropic3ChatResponseMetadata, CohereChatResponseMetadata, BedrockAi21Jurassic2ChatResponseMetadata, BedrockLlama2ChatResponseMetadata and BedrockTitanChatResponseMetadata are the same? Then why not create one instance and use it inside the AnthropicChatBedrockApi to generate the metadata inside the AnthropicChatBedrockApi?
What do you think about this approach?
The idea behind AmazonBedrockInvocationContext is that the Amazon Bedrock invoke model API will return input/output tokens and latency information from HTTP headers, not only from the response body. So I wrap it up as an AmazonBedrockInvocationContext to include more information besides the API response.
Check this out: https://github.com/spring-projects/spring-ai/blob/65e6880fa0f8db22ebcf780c3c840dc5e063b37f/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java#L220
And also, the reason why I create multiple metadata classes like BedrockAnthropicChatResponseMetadata is because each model has a different response structure. Currently, it seems to look similar, but for example, the response message ID is different. So I decided to keep this up for future changes in each model.
How do you think about this? btw I can reduce metadata classes to a single one maybe we can add a dynamic field like Map<String, Object> attribute fields to cover the different metadata fields?
Thanks, @wmz7year I need to look into this further, but wont be able to do it in coming 2 weeks.
Thanks, @wmz7year I need to look into this further, but wont be able to do it in coming 2 weeks.
Sure, if you have any thoughts just let me know
re-implement with: https://github.com/spring-projects/spring-ai/pull/813