Add metadata to BedrockAnthropic3ChatModel response

This commit enhances the Bedrock Anthropic model's output:

- Add response ID, model name, and usage data to ChatResponseMetadata
- Introduce DefaultUsage class for token usage information
- Update BedrockAnthropic3ChatModel to include new metadata
- Add Jackson annotations for serialization/deserialization
- Implement unit tests for DefaultUsage

These changes provide structured, serializable metadata in the
ChatResponse, improving the model's output with additional
information.
This commit is contained in:
dafriz
2024-10-10 19:47:04 +11:00
committed by Mark Pollack
parent 6d38c85ff6
commit c9f50da12a
3 changed files with 246 additions and 4 deletions

View File

@@ -18,11 +18,14 @@ package org.springframework.ai.bedrock.anthropic3;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.DefaultUsage;
import org.springframework.ai.chat.metadata.Usage;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.anthropic3.api.Anthropic3ChatBedrockApi;
@@ -82,11 +85,17 @@ public class BedrockAnthropic3ChatModel implements ChatModel, StreamingChatModel
AnthropicChatResponse response = this.anthropicChatApi.chatCompletion(request);
List<Generation> generations = response.content().stream().map(content -> {
return new Generation(content.text(), Map.of())
.withGenerationMetadata(ChatGenerationMetadata.from(response.stopReason(), null));
return new Generation(new AssistantMessage(content.text()),
ChatGenerationMetadata.from(response.stopReason(), null));
}).toList();
return new ChatResponse(generations);
ChatResponseMetadata metadata = ChatResponseMetadata.builder()
.withId(response.id())
.withModel(response.model())
.withUsage(extractUsage(response))
.build();
return new ChatResponse(generations, metadata);
}
@Override
@@ -116,6 +125,11 @@ public class BedrockAnthropic3ChatModel implements ChatModel, StreamingChatModel
});
}
protected Usage extractUsage(AnthropicChatResponse response) {
return new DefaultUsage(response.usage().inputTokens().longValue(),
response.usage().outputTokens().longValue());
}
/**
* Accessible for testing.
*/