From c5e9e0a3f3d7436fc324d78d8dd24370c8def3e8 Mon Sep 17 00:00:00 2001 From: dafriz Date: Thu, 10 Oct 2024 19:47:04 +1100 Subject: [PATCH 1/3] Add missing usage collection to BedrockAnthropic3ChatModel call with prompt and additionally replace use of deprecated methods. --- .../ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 95a74f933c1..4409f516c52 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,10 +18,10 @@ 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 reactor.core.publisher.Flux; @@ -82,8 +82,9 @@ public ChatResponse call(Prompt prompt) { 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(), new Anthropic3ChatBedrockApi.AnthropicUsage( + response.usage().inputTokens(), response.usage().outputTokens()))); }).toList(); return new ChatResponse(generations); From 6d4c75e93ff37d19209b7004a688d13698850977 Mon Sep 17 00:00:00 2001 From: dafriz Date: Thu, 10 Oct 2024 23:52:39 +1100 Subject: [PATCH 2/3] Return ChatResponseMetadata with BedrockAnthropic3ChatModel call prompt method. Include id, model, inputTokens and outputTokens usage data. --- .../BedrockAnthropic3ChatModel.java | 28 +++++++++++++++++-- 1 file changed, 25 insertions(+), 3 deletions(-) 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 4409f516c52..3992d34d3bd 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 @@ -23,6 +23,8 @@ 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.Usage; import reactor.core.publisher.Flux; import org.springframework.ai.bedrock.anthropic3.api.Anthropic3ChatBedrockApi; @@ -83,11 +85,16 @@ public ChatResponse call(Prompt prompt) { List generations = response.content().stream().map(content -> { return new Generation(new AssistantMessage(content.text()), - ChatGenerationMetadata.from(response.stopReason(), new Anthropic3ChatBedrockApi.AnthropicUsage( - response.usage().inputTokens(), response.usage().outputTokens()))); + 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 @@ -117,6 +124,21 @@ public Flux stream(Prompt prompt) { }); } + private Usage extractUsage(AnthropicChatResponse response) { + return new Usage() { + + @Override + public Long getPromptTokens() { + return response.usage().inputTokens().longValue(); + } + + @Override + public Long getGenerationTokens() { + return response.usage().outputTokens().longValue(); + } + }; + } + /** * Accessible for testing. */ From 4b1c6d2f8a129ce8f796812d6725b6604fb886eb Mon Sep 17 00:00:00 2001 From: dafriz Date: Sat, 12 Oct 2024 20:59:27 +1100 Subject: [PATCH 3/3] Optimize BedrockAnthropic3ChatModel.extractUsage --- .../ai/bedrock/anthropic3/BedrockAnthropic3ChatModel.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 3992d34d3bd..4ae4643ffdc 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 @@ -91,7 +91,7 @@ public ChatResponse call(Prompt prompt) { ChatResponseMetadata metadata = ChatResponseMetadata.builder() .withId(response.id()) .withModel(response.model()) - .withUsage(extractUsage(response)) + .withUsage(extractUsage(response.usage())) .build(); return new ChatResponse(generations, metadata); @@ -124,17 +124,17 @@ public Flux stream(Prompt prompt) { }); } - private Usage extractUsage(AnthropicChatResponse response) { + private Usage extractUsage(Anthropic3ChatBedrockApi.AnthropicUsage usage) { return new Usage() { @Override public Long getPromptTokens() { - return response.usage().inputTokens().longValue(); + return usage.inputTokens().longValue(); } @Override public Long getGenerationTokens() { - return response.usage().outputTokens().longValue(); + return usage.outputTokens().longValue(); } }; }