search

spring-projects / spring-ai

An Application Framework for AI Engineering
https://docs.spring.io/spring-ai/reference/index.html
Apache License 2.0
3.32k stars 850 forks source link

Bedrock Titan Chat Missing Usage Information #1769

Open gm2552 opened 2 days ago

gm2552 commented 2 days ago

Bug description Bedrock/Titan chats return token usage information which is deserialized into a TitanChatResponse record. Token usage information is dropped in the BedrockTitanChatModel class when converting TitanChatResponse objects to ChatResponse objects.

    public ChatResponse call(Prompt prompt) {
        TitanChatBedrockApi.TitanChatResponse response = this.chatApi.chatCompletion(this.createRequest(prompt));
        List<Generation> generations = response.results().stream().map((result) -> {
            return new Generation(result.outputText());
        }).toList();
        return new ChatResponse(generations);
    }

Environment All versions of Spring AI supporting Bedrock/Titan. Currently testing with M3-PLATFORM.

Steps to reproduce Execute any valid Chat against a Chat based Bedrock/Titan model using a BedrockTitanChatModel and validate that usage meta data is empty.

Eg:

BedrockTitanChatModel = new BedrockTitanChatModel(...);
ChatClient.builder(chatModel).defaultOptions(options).build();
ChatResponse response = client.prompt(prompt).call().chatResponse();
Usage usage = response.getMetadata().getUsage();
assertThat(usage.getPromptTokens()).isEqualTo(0);

Expected behavior Expecting that the usage in will be non-zero.

assertThat(usage.getPromptTokens()).isGreaterThan(0);