Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add metadata to BedrockAnthropic3ChatModel response

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

22
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java
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.
*/

111
spring-ai-core/src/main/java/org/springframework/ai/chat/metadata/DefaultUsage.java Normal file
View File

@@ -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 + '}';
}
}

117
spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java Normal file
View File

@@ -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());
}
}