Amazon Bedrock Chat adds tool support.

Author: wmz7year

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

@@ -0,0 +1,69 @@
+/*
+ * 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.List;
+
+import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
+
+import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
+import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock.Type;
+import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
+
+/**
+ * Amazon Bedrock Chat model converse interface generation metadata, encapsulating
+ * information on the completion.
+ *
+ * @author Wei Jiang
+ * @since 1.0.0
+ */
+public class BedrockConverseChatGenerationMetadata implements ChatGenerationMetadata {
+
+	private final String stopReason;
+
+	private final List<ContentBlock> contentBlocks;
+
+	public BedrockConverseChatGenerationMetadata(String stopReason, List<ContentBlock> contentBlocks) {
+		super();
+
+		this.stopReason = stopReason;
+		this.contentBlocks = contentBlocks;
+	}
+
+	public static BedrockConverseChatGenerationMetadata from(ConverseResponse response, List<ContentBlock> contents) {
+		return new BedrockConverseChatGenerationMetadata(response.stopReasonAsString(), contents);
+	}
+
+	@Override
+	public <T> T getContentFilterMetadata() {
+		return null;
+	}
+
+	@Override
+	public String getFinishReason() {
+		return stopReason;
+	}
+
+	public List<ContentBlock> getContentBlocks() {
+		return contentBlocks;
+	}
+
+	public boolean isToolUse() {
+		return contentBlocks.stream().anyMatch(content -> content.type() == Type.TOOL_USE);
+
+	}
+
+}

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

@@ -15,12 +15,20 @@
  */
 package org.springframework.ai.bedrock.anthropic3;
 
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonInclude.Include;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.model.function.FunctionCallback;
+import org.springframework.ai.model.function.FunctionCallingOptions;
+import org.springframework.boot.context.properties.NestedConfigurationProperty;
+import org.springframework.util.Assert;
 
+import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 /**
  * Java {@link ChatOptions} for the Bedrock Anthropic chat generative model chat options.
@@ -31,7 +39,7 @@
  * @since 1.0.0
  */
 @JsonInclude(Include.NON_NULL)
-public class Anthropic3ChatOptions implements ChatOptions {
+public class Anthropic3ChatOptions implements ChatOptions, FunctionCallingOptions {
 
 	// @formatter:off
 	/**
@@ -66,6 +74,31 @@ public class Anthropic3ChatOptions implements ChatOptions {
 	 */
 	private @JsonProperty("stop_sequences") List<String> stopSequences;
 
+	/**
+	 * Tool Function Callbacks to register with the ChatModel. For Prompt
+	 * Options the functionCallbacks are automatically enabled for the duration of the
+	 * prompt execution. For Default Options the functionCallbacks are registered but
+	 * disabled by default. Use the enableFunctions to set the functions from the registry
+	 * to be used by the ChatModel chat completion requests.
+	 */
+	@NestedConfigurationProperty
+	@JsonIgnore
+	private List<FunctionCallback> functionCallbacks = new ArrayList<>();
+
+	/**
+	 * List of functions, identified by their names, to configure for function calling in
+	 * the chat completion requests. Functions with those names must exist in the
+	 * functionCallbacks registry. The {@link #functionCallbacks} from the PromptOptions
+	 * are automatically enabled for the duration of the prompt execution.
+	 *
+	 * Note that function enabled with the default options are enabled for all chat
+	 * completion requests. This could impact the token count and the billing. If the
+	 * functions is set in a prompt options, then the enabled functions are only active
+	 * for the duration of this prompt execution.
+	 */
+	@NestedConfigurationProperty
+	@JsonIgnore
+	private Set<String> functions = new HashSet<>();
 	// @formatter:on
 
 	public static Builder builder() {
@@ -101,6 +134,23 @@ public Builder withStopSequences(List<String> stopSequences) {
 			return this;
 		}
 
+		public Builder withFunctionCallbacks(List<FunctionCallback> functionCallbacks) {
+			this.options.functionCallbacks = functionCallbacks;
+			return this;
+		}
+
+		public Builder withFunctions(Set<String> functionNames) {
+			Assert.notNull(functionNames, "Function names must not be null");
+			this.options.functions = functionNames;
+			return this;
+		}
+
+		public Builder withFunction(String functionName) {
+			Assert.hasText(functionName, "Function name must not be empty");
+			this.options.functions.add(functionName);
+			return this;
+		}
+
 		public Anthropic3ChatOptions build() {
 			return this.options;
 		}
@@ -150,12 +200,36 @@ public void setStopSequences(List<String> stopSequences) {
 		this.stopSequences = stopSequences;
 	}
 
+	@Override
+	public List<FunctionCallback> getFunctionCallbacks() {
+		return this.functionCallbacks;
+	}
+
+	@Override
+	public void setFunctionCallbacks(List<FunctionCallback> functionCallbacks) {
+		Assert.notNull(functionCallbacks, "FunctionCallbacks must not be null");
+		this.functionCallbacks = functionCallbacks;
+	}
+
+	@Override
+	public Set<String> getFunctions() {
+		return this.functions;
+	}
+
+	@Override
+	public void setFunctions(Set<String> functions) {
+		Assert.notNull(functions, "Function must not be null");
+		this.functions = functions;
+	}
+
 	public static Anthropic3ChatOptions fromOptions(Anthropic3ChatOptions fromOptions) {
 		return builder().withTemperature(fromOptions.getTemperature())
 			.withMaxTokens(fromOptions.getMaxTokens())
 			.withTopK(fromOptions.getTopK())
 			.withTopP(fromOptions.getTopP())
 			.withStopSequences(fromOptions.getStopSequences())
+			.withFunctionCallbacks(fromOptions.getFunctionCallbacks())
+			.withFunctions(fromOptions.getFunctions())
 			.build();
 	}
 

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

@@ -16,18 +16,33 @@
 package org.springframework.ai.bedrock.anthropic3;
 
 import reactor.core.publisher.Flux;
-import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
+import software.amazon.awssdk.services.bedrockruntime.model.Tool;
+import software.amazon.awssdk.services.bedrockruntime.model.ToolConfiguration;
+import software.amazon.awssdk.services.bedrockruntime.model.ToolInputSchema;
+import software.amazon.awssdk.services.bedrockruntime.model.ToolSpecification;
 
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import org.springframework.ai.bedrock.BedrockConverseChatGenerationMetadata;
 import org.springframework.ai.bedrock.api.BedrockConverseApi;
+import org.springframework.ai.bedrock.api.BedrockConverseApi.BedrockConverseRequest;
 import org.springframework.ai.bedrock.api.BedrockConverseApiUtils;
+import org.springframework.ai.chat.messages.Message;
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
 import org.springframework.ai.chat.model.StreamingChatModel;
 import org.springframework.ai.chat.prompt.ChatOptions;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.model.ModelDescription;
+import org.springframework.ai.model.ModelOptionsUtils;
+import org.springframework.ai.model.function.AbstractFunctionCallSupport;
+import org.springframework.ai.model.function.FunctionCallbackContext;
 import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
 
 /**
  * Java {@link ChatModel} and {@link StreamingChatModel} for the Bedrock Anthropic3 chat
@@ -38,7 +53,9 @@
  * @author Wei Jiang
  * @since 1.0.0
  */
-public class BedrockAnthropic3ChatModel implements ChatModel, StreamingChatModel {
+public class BedrockAnthropic3ChatModel
+		extends AbstractFunctionCallSupport<Message, BedrockConverseRequest, ChatResponse>
+		implements ChatModel, StreamingChatModel {
 
 	private final String modelId;
 
@@ -56,6 +73,13 @@ public BedrockAnthropic3ChatModel(BedrockConverseApi converseApi, Anthropic3Chat
 	}
 
 	public BedrockAnthropic3ChatModel(String modelId, BedrockConverseApi converseApi, Anthropic3ChatOptions options) {
+		this(modelId, converseApi, options, null);
+	}
+
+	public BedrockAnthropic3ChatModel(String modelId, BedrockConverseApi converseApi, Anthropic3ChatOptions options,
+			FunctionCallbackContext functionCallbackContext) {
+		super(functionCallbackContext);
+
 		Assert.notNull(modelId, "modelId must not be null.");
 		Assert.notNull(converseApi, "BedrockConverseApi must not be null.");
 		Assert.notNull(options, "Anthropic3ChatOptions must not be null.");
@@ -69,29 +93,125 @@ public BedrockAnthropic3ChatModel(String modelId, BedrockConverseApi converseApi
 	public ChatResponse call(Prompt prompt) {
 		Assert.notNull(prompt, "Prompt must not be null.");
 
-		var request = BedrockConverseApiUtils.createConverseRequest(modelId, prompt, defaultOptions);
-
-		ConverseResponse response = this.converseApi.converse(request);
+		var request = createBedrockConverseRequest(prompt);
 
-		return BedrockConverseApiUtils.convertConverseResponse(response);
+		return this.callWithFunctionSupport(request);
 	}
 
 	@Override
 	public Flux<ChatResponse> stream(Prompt prompt) {
 		Assert.notNull(prompt, "Prompt must not be null.");
 
+		// TODO
 		var request = BedrockConverseApiUtils.createConverseStreamRequest(modelId, prompt, defaultOptions);
 
 		Flux<ConverseStreamOutput> fluxResponse = this.converseApi.converseStream(request);
 
 		return fluxResponse.map(output -> BedrockConverseApiUtils.convertConverseStreamOutput(output));
 	}
 
+	private BedrockConverseRequest createBedrockConverseRequest(Prompt prompt) {
+		var request = BedrockConverseApiUtils.createBedrockConverseRequest(modelId, prompt, defaultOptions);
+
+		ToolConfiguration toolConfiguration = createToolConfiguration(prompt);
+		request.setToolConfiguration(toolConfiguration);
+
+		return request;
+	}
+
+	private ToolConfiguration createToolConfiguration(Prompt prompt) {
+		Set<String> functionsForThisRequest = new HashSet<>();
+
+		if (this.defaultOptions != null) {
+			Set<String> promptEnabledFunctions = this.handleFunctionCallbackConfigurations(this.defaultOptions,
+					!IS_RUNTIME_CALL);
+			functionsForThisRequest.addAll(promptEnabledFunctions);
+		}
+
+		if (prompt.getOptions() != null) {
+			if (prompt.getOptions() instanceof ChatOptions runtimeOptions) {
+				Anthropic3ChatOptions updatedRuntimeOptions = ModelOptionsUtils.copyToTarget(runtimeOptions,
+						ChatOptions.class, Anthropic3ChatOptions.class);
+
+				Set<String> defaultEnabledFunctions = this.handleFunctionCallbackConfigurations(updatedRuntimeOptions,
+						IS_RUNTIME_CALL);
+				functionsForThisRequest.addAll(defaultEnabledFunctions);
+			}
+			else {
+				throw new IllegalArgumentException("Prompt options are not of type ChatOptions: "
+						+ prompt.getOptions().getClass().getSimpleName());
+			}
+		}
+
+		if (CollectionUtils.isEmpty(functionsForThisRequest)) {
+			return null;
+		}
+		else {
+			return ToolConfiguration.builder().tools(getFunctionTools(functionsForThisRequest)).build();
+		}
+	}
+
+	private List<Tool> getFunctionTools(Set<String> functionNames) {
+		return this.resolveFunctionCallbacks(functionNames).stream().map(functionCallback -> {
+			var description = functionCallback.getDescription();
+			var name = functionCallback.getName();
+			String inputSchema = functionCallback.getInputTypeSchema();
+
+			return Tool.builder()
+				.toolSpec(ToolSpecification.builder()
+					.name(name)
+					.description(description)
+					.inputSchema(ToolInputSchema.builder()
+						.json(BedrockConverseApiUtils.convertObjectToDocument(ModelOptionsUtils.jsonToMap(inputSchema)))
+						.build())
+					.build())
+				.build();
+		}).toList();
+	}
+
 	@Override
 	public ChatOptions getDefaultOptions() {
 		return Anthropic3ChatOptions.fromOptions(this.defaultOptions);
 	}
 
+	@Override
+	protected BedrockConverseRequest doCreateToolResponseRequest(BedrockConverseRequest previousRequest,
+			Message responseMessage, List<Message> conversationHistory) {
+		// TODO
+		return null;
+	}
+
+	@Override
+	protected List<Message> doGetUserMessages(BedrockConverseRequest request) {
+		return BedrockConverseApiUtils.getMessagesInstructions(request.getMessages());
+	}
+
+	@Override
+	protected Message doGetToolResponseMessage(ChatResponse response) {
+		return response.getResult().getOutput();
+	}
+
+	@Override
+	protected ChatResponse doChatCompletion(BedrockConverseRequest request) {
+		return converseApi.converse(request);
+	}
+
+	@Override
+	protected Flux<ChatResponse> doChatCompletionStream(BedrockConverseRequest request) {
+		return converseApi.converseStream(request);
+	}
+
+	@Override
+	protected boolean isToolFunctionCall(ChatResponse response) {
+		Generation result = response.getResult();
+
+		if (result.getMetadata() instanceof BedrockConverseChatGenerationMetadata metadata) {
+			return metadata.isToolUse();
+		}
+
+		return false;
+	}
+
 	/**
 	 * Anthropic3 models version.
 	 */