Re-implementing Amazon Bedrock Chat model with Amazon Bedrock Converse API.

Author: wmz7year

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/BedrockChatResponseMetadata.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2023 - 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.bedrock;
+
+import java.util.HashMap;
+
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.EmptyUsage;
+import org.springframework.ai.chat.metadata.Usage;
+
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseMetrics;
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamMetadataEvent;
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamMetrics;
+
+/**
+ * {@link ChatResponseMetadata} implementation for {@literal Amazon Bedrock}.
+ *
+ * @author Wei Jiang
+ * @since 1.0.0
+ */
+public class BedrockChatResponseMetadata extends HashMap<String, Object> implements ChatResponseMetadata {
+
+	protected static final String AI_METADATA_STRING = "{ @type: %1$s, id: %2$s, usage: %3$s, latency: %4$sms}";
+
+	private final String id;
+
+	private final Usage usage;
+
+	private final Long latencyMs;
+
+	public static BedrockChatResponseMetadata from(ConverseResponse response) {
+		String requestId = response.responseMetadata().requestId();
+
+		BedrockUsage usage = BedrockUsage.from(response.usage());
+
+		ConverseMetrics metrics = response.metrics();
+
+		return new BedrockChatResponseMetadata(requestId, usage, metrics.latencyMs());
+	}
+
+	public static BedrockChatResponseMetadata from(ConverseStreamMetadataEvent converseStreamMetadataEvent) {
+		BedrockUsage usage = BedrockUsage.from(converseStreamMetadataEvent.usage());
+
+		ConverseStreamMetrics metrics = converseStreamMetadataEvent.metrics();
+
+		return new BedrockChatResponseMetadata(null, usage, metrics.latencyMs());
+	}
+
+	protected BedrockChatResponseMetadata(String id, BedrockUsage usage, Long latencyMs) {
+		this.id = id;
+		this.usage = usage;
+		this.latencyMs = latencyMs;
+	}
+
+	public String getId() {
+		return this.id;
+	}
+
+	public Long getLatencyMs() {
+		return latencyMs;
+	}
+
+	@Override
+	public Usage getUsage() {
+		Usage usage = this.usage;
+		return usage != null ? usage : new EmptyUsage();
+	}
+
+	@Override
+	public String toString() {
+		return AI_METADATA_STRING.formatted(getClass().getName(), getId(), getUsage(), getLatencyMs());
+	}
+
+}

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/BedrockUsage.java

@@ -19,42 +19,49 @@
 import org.springframework.ai.chat.metadata.Usage;
 import org.springframework.util.Assert;
 
+import software.amazon.awssdk.services.bedrockruntime.model.TokenUsage;
+
 /**
  * {@link Usage} implementation for Bedrock API.
  *
  * @author Christian Tzolov
+ * @author Wei Jiang
  * @since 0.8.0
  */
 public class BedrockUsage implements Usage {
 
 	public static BedrockUsage from(AmazonBedrockInvocationMetrics usage) {
-		return new BedrockUsage(usage);
+		return new BedrockUsage(usage.inputTokenCount().longValue(), usage.outputTokenCount().longValue());
 	}
 
-	private final AmazonBedrockInvocationMetrics usage;
+	public static BedrockUsage from(TokenUsage usage) {
+		Assert.notNull(usage, "'TokenUsage' must not be null.");
 
-	protected BedrockUsage(AmazonBedrockInvocationMetrics usage) {
-		Assert.notNull(usage, "OpenAI Usage must not be null");
-		this.usage = usage;
+		return new BedrockUsage(usage.inputTokens().longValue(), usage.outputTokens().longValue());
 	}
 
-	protected AmazonBedrockInvocationMetrics getUsage() {
-		return this.usage;
+	private final Long inputTokens;
+
+	private final Long outputTokens;
+
+	protected BedrockUsage(Long inputTokens, Long outputTokens) {
+		this.inputTokens = inputTokens;
+		this.outputTokens = outputTokens;
 	}
 
 	@Override
 	public Long getPromptTokens() {
-		return getUsage().inputTokenCount().longValue();
+		return inputTokens;
 	}
 
 	@Override
 	public Long getGenerationTokens() {
-		return getUsage().outputTokenCount().longValue();
+		return outputTokens;
 	}
 
 	@Override
 	public String toString() {
-		return getUsage().toString();
+		return "BedrockUsage [inputTokens=" + inputTokens + ", outputTokens=" + outputTokens + "]";
 	}
 
 }

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/anthropic/AnthropicChatOptions.java

@@ -25,7 +25,11 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 /**
+ * Java {@link ChatOptions} for the Bedrock Anthropic chat generative model chat options.
+ * https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-text-completion.html
+ *
  * @author Christian Tzolov
+ * @author Wei Jiang
  */
 @JsonInclude(Include.NON_NULL)
 public class AnthropicChatOptions implements ChatOptions {
@@ -44,7 +48,7 @@ public class AnthropicChatOptions implements ChatOptions {
 	 * reaching this maximum. This parameter only specifies the absolute maximum number of tokens to generate. We
 	 * recommend a limit of 4,000 tokens for optimal performance.
 	 */
-	private @JsonProperty("max_tokens_to_sample") Integer maxTokensToSample;
+	private @JsonProperty("max_tokens") Integer maxTokens;
 
 	/**
 	 * Specify the number of token choices the generative uses to generate the next token.
@@ -62,11 +66,6 @@ public class AnthropicChatOptions implements ChatOptions {
 	 * generating further tokens. The returned text doesn't contain the stop sequence.
 	 */
 	private @JsonProperty("stop_sequences") List<String> stopSequences;
-
-	/**
-	 * The version of the generative to use. The default value is bedrock-2023-05-31.
-	 */
-	private @JsonProperty("anthropic_version") String anthropicVersion;
 	// @formatter:on
 
 	public static Builder builder() {
@@ -82,8 +81,8 @@ public Builder withTemperature(Float temperature) {
 			return this;
 		}
 
-		public Builder withMaxTokensToSample(Integer maxTokensToSample) {
-			this.options.setMaxTokensToSample(maxTokensToSample);
+		public Builder withMaxTokens(Integer maxTokens) {
+			this.options.setMaxTokens(maxTokens);
 			return this;
 		}
 
@@ -102,11 +101,6 @@ public Builder withStopSequences(List<String> stopSequences) {
 			return this;
 		}
 
-		public Builder withAnthropicVersion(String anthropicVersion) {
-			this.options.setAnthropicVersion(anthropicVersion);
-			return this;
-		}
-
 		public AnthropicChatOptions build() {
 			return this.options;
 		}
@@ -122,12 +116,12 @@ public void setTemperature(Float temperature) {
 		this.temperature = temperature;
 	}
 
-	public Integer getMaxTokensToSample() {
-		return this.maxTokensToSample;
+	public Integer getMaxTokens() {
+		return maxTokens;
 	}
 
-	public void setMaxTokensToSample(Integer maxTokensToSample) {
-		this.maxTokensToSample = maxTokensToSample;
+	public void setMaxTokens(Integer maxTokens) {
+		this.maxTokens = maxTokens;
 	}
 
 	@Override
@@ -156,21 +150,12 @@ public void setStopSequences(List<String> stopSequences) {
 		this.stopSequences = stopSequences;
 	}
 
-	public String getAnthropicVersion() {
-		return this.anthropicVersion;
-	}
-
-	public void setAnthropicVersion(String anthropicVersion) {
-		this.anthropicVersion = anthropicVersion;
-	}
-
 	public static AnthropicChatOptions fromOptions(AnthropicChatOptions fromOptions) {
 		return builder().withTemperature(fromOptions.getTemperature())
-			.withMaxTokensToSample(fromOptions.getMaxTokensToSample())
+			.withMaxTokens(fromOptions.getMaxTokens())
 			.withTopK(fromOptions.getTopK())
 			.withTopP(fromOptions.getTopP())
 			.withStopSequences(fromOptions.getStopSequences())
-			.withAnthropicVersion(fromOptions.getAnthropicVersion())
 			.build();
 	}