From c9f50da12a50bd3ccba176ffdbefce429d744b7c Mon Sep 17 00:00:00 2001 From: dafriz Date: Thu, 10 Oct 2024 19:47:04 +1100 Subject: [PATCH] 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. --- .../BedrockAnthropic3ChatModel.java | 22 +++- .../ai/chat/metadata/DefaultUsage.java | 111 +++++++++++++++++ .../ai/chat/metadata/DefaultUsageTests.java | 117 ++++++++++++++++++ 3 files changed, 246 insertions(+), 4 deletions(-) create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/DefaultUsage.java create mode 100644 spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java index 95a74f933..4bef201bb 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java @@ -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 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. */ diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/DefaultUsage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/DefaultUsage.java new file mode 100644 index 000000000..5f5ee9c51 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/DefaultUsage.java @@ -0,0 +1,111 @@ +/* + * Copyright 2024 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.chat.metadata; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Objects; + +/** + * Default implementation of the {@link Usage} interface. + * + * @author Mark Pollack + * @since 1.0.0 + */ +public class DefaultUsage implements Usage { + + private final Long promptTokens; + + private final Long generationTokens; + + private final Long totalTokens; + + /** + * Create a new DefaultUsage with promptTokens and generationTokens. + * @param promptTokens the number of tokens in the prompt, or {@code null} if not + * available + * @param generationTokens the number of tokens in the generation, or {@code null} if + * not available + */ + public DefaultUsage(Long promptTokens, Long generationTokens) { + this(promptTokens, generationTokens, null); + } + + /** + * Create a new DefaultUsage with promptTokens, generationTokens, and totalTokens. + * @param promptTokens the number of tokens in the prompt, or {@code null} if not + * available + * @param generationTokens the number of tokens in the generation, or {@code null} if + * not available + * @param totalTokens the total number of tokens, or {@code null} to calculate from + * promptTokens and generationTokens + */ + @JsonCreator + public DefaultUsage(@JsonProperty("promptTokens") Long promptTokens, + @JsonProperty("generationTokens") Long generationTokens, @JsonProperty("totalTokens") Long totalTokens) { + this.promptTokens = promptTokens != null ? promptTokens : 0L; + this.generationTokens = generationTokens != null ? generationTokens : 0L; + this.totalTokens = totalTokens != null ? totalTokens + : calculateTotalTokens(this.promptTokens, this.generationTokens); + } + + @Override + @JsonProperty("promptTokens") + public Long getPromptTokens() { + return promptTokens; + } + + @Override + @JsonProperty("generationTokens") + public Long getGenerationTokens() { + return generationTokens; + } + + @Override + @JsonProperty("totalTokens") + public Long getTotalTokens() { + return totalTokens; + } + + private Long calculateTotalTokens(Long promptTokens, Long generationTokens) { + return promptTokens + generationTokens; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (o == null || getClass() != o.getClass()) + return false; + DefaultUsage that = (DefaultUsage) o; + return Objects.equals(promptTokens, that.promptTokens) + && Objects.equals(generationTokens, that.generationTokens) + && Objects.equals(totalTokens, that.totalTokens); + } + + @Override + public int hashCode() { + return Objects.hash(promptTokens, generationTokens, totalTokens); + } + + @Override + public String toString() { + return "DefaultUsage{" + "promptTokens=" + promptTokens + ", generationTokens=" + generationTokens + + ", totalTokens=" + totalTokens + '}'; + } + +} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java new file mode 100644 index 000000000..68985f67b --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java @@ -0,0 +1,117 @@ +/* + * Copyright 2024 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.chat.metadata; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; +import static org.junit.jupiter.api.Assertions.*; + +public class DefaultUsageTests { + + private final ObjectMapper objectMapper = new ObjectMapper(); + + @Test + void testSerializationWithAllFields() throws Exception { + DefaultUsage usage = new DefaultUsage(100L, 50L, 150L); + String json = objectMapper.writeValueAsString(usage); + assertEquals("{\"promptTokens\":100,\"generationTokens\":50,\"totalTokens\":150}", json); + } + + @Test + void testDeserializationWithAllFields() throws Exception { + String json = "{\"promptTokens\":100,\"generationTokens\":50,\"totalTokens\":150}"; + DefaultUsage usage = objectMapper.readValue(json, DefaultUsage.class); + assertEquals(100L, usage.getPromptTokens()); + assertEquals(50L, usage.getGenerationTokens()); + assertEquals(150L, usage.getTotalTokens()); + } + + @Test + void testSerializationWithNullFields() throws Exception { + DefaultUsage usage = new DefaultUsage(null, null, null); + String json = objectMapper.writeValueAsString(usage); + assertEquals("{\"promptTokens\":0,\"generationTokens\":0,\"totalTokens\":0}", json); + } + + @Test + void testDeserializationWithMissingFields() throws Exception { + String json = "{\"promptTokens\":100}"; + DefaultUsage usage = objectMapper.readValue(json, DefaultUsage.class); + assertEquals(100L, usage.getPromptTokens()); + assertEquals(0L, usage.getGenerationTokens()); + assertEquals(100L, usage.getTotalTokens()); + } + + @Test + void testDeserializationWithNullFields() throws Exception { + String json = "{\"promptTokens\":null,\"generationTokens\":null,\"totalTokens\":null}"; + DefaultUsage usage = objectMapper.readValue(json, DefaultUsage.class); + assertEquals(0L, usage.getPromptTokens()); + assertEquals(0L, usage.getGenerationTokens()); + assertEquals(0L, usage.getTotalTokens()); + } + + @Test + void testRoundTripSerialization() throws Exception { + DefaultUsage original = new DefaultUsage(100L, 50L, 150L); + String json = objectMapper.writeValueAsString(original); + DefaultUsage deserialized = objectMapper.readValue(json, DefaultUsage.class); + assertEquals(original.getPromptTokens(), deserialized.getPromptTokens()); + assertEquals(original.getGenerationTokens(), deserialized.getGenerationTokens()); + assertEquals(original.getTotalTokens(), deserialized.getTotalTokens()); + } + + @Test + void testTwoArgumentConstructorAndSerialization() throws Exception { + DefaultUsage usage = new DefaultUsage(100L, 50L); + + // Test that the fields are set correctly + assertEquals(100L, usage.getPromptTokens()); + assertEquals(50L, usage.getGenerationTokens()); + assertEquals(150L, usage.getTotalTokens()); // 100 + 50 = 150 + + // Test serialization + String json = objectMapper.writeValueAsString(usage); + assertEquals("{\"promptTokens\":100,\"generationTokens\":50,\"totalTokens\":150}", json); + + // Test deserialization + DefaultUsage deserializedUsage = objectMapper.readValue(json, DefaultUsage.class); + assertEquals(100L, deserializedUsage.getPromptTokens()); + assertEquals(50L, deserializedUsage.getGenerationTokens()); + assertEquals(150L, deserializedUsage.getTotalTokens()); + } + + @Test + void testTwoArgumentConstructorWithNullValues() throws Exception { + DefaultUsage usage = new DefaultUsage(null, null); + + // Test that null values are converted to 0 + assertEquals(0L, usage.getPromptTokens()); + assertEquals(0L, usage.getGenerationTokens()); + assertEquals(0L, usage.getTotalTokens()); + + // Test serialization + String json = objectMapper.writeValueAsString(usage); + assertEquals("{\"promptTokens\":0,\"generationTokens\":0,\"totalTokens\":0}", json); + + // Test deserialization + DefaultUsage deserializedUsage = objectMapper.readValue(json, DefaultUsage.class); + assertEquals(0L, deserializedUsage.getPromptTokens()); + assertEquals(0L, deserializedUsage.getGenerationTokens()); + assertEquals(0L, deserializedUsage.getTotalTokens()); + } + +}