Add Bedrock Converse API chat model support
Introduces support for Amazon Bedrock Converse API through a new BedrockProxyChatModel implementation. This enables integration with Bedrock's conversation models with features including: - Support for sync/async chat completions - Stream response handling - Tool/function calling capabilities - System message support - Image input support - Observation and metrics integration - Configurable model parameters and AWS credentials Adds core support classes: - BedrockUsage: Implements Usage interface for token tracking - ConverseApiUtils: Utility class for handling Bedrock API responses including: - Tool use event aggregation and processing - Chat response transformation from stream outputs - Model options conversion - Support for metadata aggregation - URLValidator: Utility for URL validation and normalization with support for: - Basic and strict URL validation - URL normalization - Multimodal input handling - Enhanced FunctionCallingOptionsBuilder with merge capabilities for both ChatOptions and FunctionCallingOptions - Added BEDROCK_CONVERSE to AiProvider enum for metrics tracking - Extended AWS credentials support with session token capability - Added configurable session token property to BedrockAwsConnectionProperties Adds new auto-configuration support: - BedrockConverseProxyChatAutoConfiguration for automatic setup of the Bedrock Converse chat model - BedrockConverseProxyChatProperties for configuration including: - Model selection (defaults to Claude 3 Sonnet) - Timeout settings (defaults to 5 minutes) - Temperature and token control - Top-K and Top-P sampling parameters - Integration with existing BedrockAwsConnectionConfiguration for AWS credentials Updates to testing infrastructure: - Adds comprehensive test suite for Bedrock Converse properties and auto-configuration - Integration tests for chat completion and streaming scenarios - Property validation tests for configuration options - Temporarily disabled other Bedrock tests due to AWS quota limitations - Added ObjectMapper configuration for proper JSON handling Added new spring-ai-bedrock-converse-spring-boot-starter module Updates module configuration in parent POM and BOM to include new bedrock-converse modules and starters. Adds necessary auto-configuration imports for seamless integration with Spring Boot applications. Unrelated changes: - Disabled several Bedrock model tests (Jurassic2, Llama, Titan) due to AWS quota limitations - Disabled PaLM2 tests due to API decommissioning by Google Resolves #809, #802 Add docs and fix configs - Move timeout configuration from chat properties to connection properties - Add comprehensive documentation for Bedrock Converse API usage and configuration - Update tests to reflect configuration changes Co-authored-by: maxjiang153 <maxjiang153@users.noreply.github.com> Standardize AWS credential handling in integration tests - Improve how we manage AWS credentials across our integration test suite and ensures consistent test configuration. We're replacing individual environment variable checks with @RequiresAwsCredentials annotation and standardizing the use of BedrockTestUtils for context creation in tests We also align all AWS regions to US_EAST_1 for consistency and add missing dependency versioning for Oracle Free. These changes make our AWS tests more easier to maintain. Key changes: - Replace @EnabledIfEnvironmentVariable with @RequiresAwsCredentials - Standardize context creation via BedrockTestUtils - Set AWS region to US_EAST_1 - Add Oracle Free dependency version in pom.xml
This commit is contained in:
committed by
Mark Pollack
parent
932fc873f1
commit
0d2d4b7385
84
models/spring-ai-bedrock-converse/pom.xml
Normal file
84
models/spring-ai-bedrock-converse/pom.xml
Normal file
@@ -0,0 +1,84 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai</artifactId>
|
||||
<version>1.0.0-SNAPSHOT</version>
|
||||
<relativePath>../../pom.xml</relativePath>
|
||||
</parent>
|
||||
<artifactId>spring-ai-bedrock-converse</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>Spring AI Model - Amazon Bedrock Converse API</name>
|
||||
<description>Amazon Bedrock models support using the Converse API</description>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
|
||||
<scm>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
<connection>git://github.com/spring-projects/spring-ai.git</connection>
|
||||
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
|
||||
</scm>
|
||||
|
||||
<properties>
|
||||
<aws.sdk.version>2.29.3</aws.sdk.version>
|
||||
</properties>
|
||||
|
||||
<dependencies>
|
||||
|
||||
<!-- production dependencies -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-core</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-retry</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>software.amazon.awssdk</groupId>
|
||||
<artifactId>bedrockruntime</artifactId>
|
||||
<version>${aws.sdk.version}</version>
|
||||
<exclusions>
|
||||
<exclusion>
|
||||
<groupId>commons-logging</groupId>
|
||||
<artifactId>commons-logging</artifactId>
|
||||
</exclusion>
|
||||
</exclusions>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>software.amazon.awssdk</groupId>
|
||||
<artifactId>sts</artifactId>
|
||||
<version>${aws.sdk.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>software.amazon.awssdk</groupId>
|
||||
<artifactId>netty-nio-client</artifactId>
|
||||
<version>${aws.sdk.version}</version>
|
||||
</dependency>
|
||||
|
||||
<!-- test dependencies -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-test</artifactId>
|
||||
<version>${project.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>io.micrometer</groupId>
|
||||
<artifactId>micrometer-observation-test</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
@@ -0,0 +1,710 @@
|
||||
/*
|
||||
* 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.bedrock.converse;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.net.URL;
|
||||
import java.net.URLConnection;
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.ai.bedrock.converse.api.ConverseApiUtils;
|
||||
import org.springframework.ai.bedrock.converse.api.URLValidator;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.MessageType;
|
||||
import org.springframework.ai.chat.messages.ToolResponseMessage;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
|
||||
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
||||
import org.springframework.ai.chat.metadata.DefaultUsage;
|
||||
import org.springframework.ai.chat.model.AbstractToolCallSupport;
|
||||
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.MessageAggregator;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationContext;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
|
||||
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
|
||||
import org.springframework.ai.chat.prompt.ChatOptions;
|
||||
import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.model.ModelOptionsUtils;
|
||||
import org.springframework.ai.model.function.FunctionCallback;
|
||||
import org.springframework.ai.model.function.FunctionCallbackContext;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
|
||||
import org.springframework.ai.observation.conventions.AiProvider;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.CollectionUtils;
|
||||
import org.springframework.util.StreamUtils;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
import io.micrometer.observation.Observation;
|
||||
import io.micrometer.observation.ObservationRegistry;
|
||||
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.core.publisher.Sinks.EmitFailureHandler;
|
||||
import reactor.core.publisher.Sinks.EmitResult;
|
||||
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
|
||||
import software.amazon.awssdk.core.SdkBytes;
|
||||
import software.amazon.awssdk.core.document.Document;
|
||||
import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient;
|
||||
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseMetrics;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamRequest;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ImageBlock;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ImageSource;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.InferenceConfiguration;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.Message;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.SystemContentBlock;
|
||||
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.ToolResultBlock;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ToolResultContentBlock;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ToolSpecification;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ToolUseBlock;
|
||||
|
||||
/**
|
||||
* A {@link ChatModel} implementation that uses the Amazon Bedrock Converse API to
|
||||
* interact with the <a href=
|
||||
* "https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html">Supported
|
||||
* models</a>. <br/>
|
||||
* <br/>
|
||||
* The Converse API doesn't support any embedding models (such as Titan Embeddings G1 -
|
||||
* Text) or image generation models (such as Stability AI).
|
||||
*
|
||||
* <p>
|
||||
* https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html
|
||||
* <p>
|
||||
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
|
||||
* <p>
|
||||
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStream.html
|
||||
* <p>
|
||||
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
|
||||
* <p>
|
||||
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Wei Jiang
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class BedrockProxyChatModel extends AbstractToolCallSupport implements ChatModel {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(BedrockProxyChatModel.class);
|
||||
|
||||
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
|
||||
|
||||
private final BedrockRuntimeClient bedrockRuntimeClient;
|
||||
|
||||
private final BedrockRuntimeAsyncClient bedrockRuntimeAsyncClient;
|
||||
|
||||
private FunctionCallingOptions defaultOptions;
|
||||
|
||||
/**
|
||||
* Observation registry used for instrumentation.
|
||||
*/
|
||||
private final ObservationRegistry observationRegistry;
|
||||
|
||||
/**
|
||||
* Conventions to use for generating observations.
|
||||
*/
|
||||
private ChatModelObservationConvention observationConvention;
|
||||
|
||||
public BedrockProxyChatModel(BedrockRuntimeClient bedrockRuntimeClient,
|
||||
BedrockRuntimeAsyncClient bedrockRuntimeAsyncClient, FunctionCallingOptions defaultOptions,
|
||||
FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks,
|
||||
ObservationRegistry observationRegistry) {
|
||||
|
||||
super(functionCallbackContext, defaultOptions, toolFunctionCallbacks);
|
||||
|
||||
Assert.notNull(bedrockRuntimeClient, "bedrockRuntimeClient must not be null");
|
||||
Assert.notNull(bedrockRuntimeAsyncClient, "bedrockRuntimeAsyncClient must not be null");
|
||||
|
||||
this.bedrockRuntimeClient = bedrockRuntimeClient;
|
||||
this.bedrockRuntimeAsyncClient = bedrockRuntimeAsyncClient;
|
||||
this.defaultOptions = defaultOptions;
|
||||
this.observationRegistry = observationRegistry;
|
||||
}
|
||||
|
||||
/**
|
||||
* Invoke the model and return the response.
|
||||
*
|
||||
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
|
||||
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
|
||||
* https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/bedrockruntime/BedrockRuntimeClient.html#converse
|
||||
* @param bedrockConverseRequest Model invocation request.
|
||||
* @return The model invocation response.
|
||||
*/
|
||||
@Override
|
||||
public ChatResponse call(Prompt prompt) {
|
||||
|
||||
ConverseRequest converseRequest = this.createRequest(prompt);
|
||||
|
||||
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
|
||||
.prompt(prompt)
|
||||
.provider(AiProvider.BEDROCK_CONVERSE.value())
|
||||
.requestOptions(buildRequestOptions(converseRequest))
|
||||
.build();
|
||||
|
||||
ChatResponse chatResponse = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
|
||||
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
|
||||
this.observationRegistry)
|
||||
.observe(() -> {
|
||||
|
||||
ConverseResponse converseResponse = this.bedrockRuntimeClient.converse(converseRequest);
|
||||
|
||||
var response = this.toChatResponse(converseResponse);
|
||||
|
||||
observationContext.setResponse(response);
|
||||
|
||||
return response;
|
||||
});
|
||||
|
||||
if (!this.isProxyToolCalls(prompt, this.defaultOptions) && chatResponse != null
|
||||
&& this.isToolCall(chatResponse, Set.of("tool_use"))) {
|
||||
var toolCallConversation = this.handleToolCalls(prompt, chatResponse);
|
||||
return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
|
||||
}
|
||||
|
||||
return chatResponse;
|
||||
}
|
||||
|
||||
private ChatOptions buildRequestOptions(ConverseRequest request) {
|
||||
return ChatOptionsBuilder.builder()
|
||||
.withModel(request.modelId())
|
||||
.withMaxTokens(request.inferenceConfig().maxTokens())
|
||||
.withStopSequences(request.inferenceConfig().stopSequences())
|
||||
.withTemperature(request.inferenceConfig().temperature() != null
|
||||
? request.inferenceConfig().temperature().doubleValue() : null)
|
||||
.withTopP(request.inferenceConfig().topP() != null ? request.inferenceConfig().topP().doubleValue() : null)
|
||||
.build();
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatOptions getDefaultOptions() {
|
||||
return this.defaultOptions;
|
||||
}
|
||||
|
||||
public ConverseStreamRequest createStreamRequest(Prompt prompt) {
|
||||
|
||||
ConverseRequest converseRequest = this.createRequest(prompt);
|
||||
|
||||
return ConverseStreamRequest.builder()
|
||||
.modelId(converseRequest.modelId())
|
||||
.messages(converseRequest.messages())
|
||||
.system(converseRequest.system())
|
||||
.additionalModelRequestFields(converseRequest.additionalModelRequestFields())
|
||||
.toolConfig(converseRequest.toolConfig())
|
||||
.build();
|
||||
}
|
||||
|
||||
ConverseRequest createRequest(Prompt prompt) {
|
||||
|
||||
Set<String> functionsForThisRequest = new HashSet<>();
|
||||
|
||||
List<Message> instructionMessages = prompt.getInstructions()
|
||||
.stream()
|
||||
.filter(message -> message.getMessageType() != MessageType.SYSTEM)
|
||||
.map(message -> {
|
||||
if (message.getMessageType() == MessageType.USER) {
|
||||
List<ContentBlock> contents = new ArrayList<>();
|
||||
if (message instanceof UserMessage) {
|
||||
var userMessage = (UserMessage) message;
|
||||
contents.add(ContentBlock.fromText(userMessage.getContent()));
|
||||
|
||||
if (!CollectionUtils.isEmpty(userMessage.getMedia())) {
|
||||
List<ContentBlock> mediaContent = userMessage.getMedia().stream().map(media -> {
|
||||
ContentBlock cb = ContentBlock.fromImage(ImageBlock.builder()
|
||||
.format(media.getMimeType().getSubtype())
|
||||
.source(ImageSource
|
||||
.fromBytes(SdkBytes.fromByteArray(getContentMediaData(media.getData()))))
|
||||
.build());
|
||||
return cb;
|
||||
}).toList();
|
||||
contents.addAll(mediaContent);
|
||||
}
|
||||
}
|
||||
return Message.builder().content(contents).role(ConversationRole.USER).build();
|
||||
}
|
||||
else if (message.getMessageType() == MessageType.ASSISTANT) {
|
||||
AssistantMessage assistantMessage = (AssistantMessage) message;
|
||||
List<ContentBlock> contentBlocks = new ArrayList<>();
|
||||
if (StringUtils.hasText(message.getContent())) {
|
||||
contentBlocks.add(ContentBlock.fromText(message.getContent()));
|
||||
}
|
||||
if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
|
||||
for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) {
|
||||
|
||||
var argumentsDocument = ConverseApiUtils
|
||||
.convertObjectToDocument(ModelOptionsUtils.jsonToMap(toolCall.arguments()));
|
||||
|
||||
contentBlocks.add(ContentBlock.fromToolUse(ToolUseBlock.builder()
|
||||
.toolUseId(toolCall.id())
|
||||
.name(toolCall.name())
|
||||
.input(argumentsDocument)
|
||||
.build()));
|
||||
|
||||
}
|
||||
}
|
||||
return Message.builder().content(contentBlocks).role(ConversationRole.ASSISTANT).build();
|
||||
}
|
||||
else if (message.getMessageType() == MessageType.TOOL) {
|
||||
List<ContentBlock> contentBlocks = ((ToolResponseMessage) message).getResponses()
|
||||
.stream()
|
||||
.map(toolResponse -> {
|
||||
ToolResultBlock toolResultBlock = ToolResultBlock.builder()
|
||||
.toolUseId(toolResponse.id())
|
||||
.content(ToolResultContentBlock.builder().text(toolResponse.responseData()).build())
|
||||
.build();
|
||||
return ContentBlock.fromToolResult(toolResultBlock);
|
||||
})
|
||||
.toList();
|
||||
return Message.builder().content(contentBlocks).role(ConversationRole.USER).build();
|
||||
}
|
||||
else {
|
||||
throw new IllegalArgumentException("Unsupported message type: " + message.getMessageType());
|
||||
}
|
||||
})
|
||||
.toList();
|
||||
|
||||
List<SystemContentBlock> systemMessages = prompt.getInstructions()
|
||||
.stream()
|
||||
.filter(m -> m.getMessageType() == MessageType.SYSTEM)
|
||||
.map(sysMessage -> SystemContentBlock.builder().text(sysMessage.getContent()).build())
|
||||
.toList();
|
||||
|
||||
FunctionCallingOptions updatedRuntimeOptions = (FunctionCallingOptions) this.defaultOptions.copy();
|
||||
|
||||
if (prompt.getOptions() != null) {
|
||||
if (prompt.getOptions() instanceof FunctionCallingOptions) {
|
||||
var functionCallingOptions = (FunctionCallingOptions) prompt.getOptions();
|
||||
updatedRuntimeOptions = ((PortableFunctionCallingOptions) updatedRuntimeOptions)
|
||||
.merge(functionCallingOptions);
|
||||
}
|
||||
else if (prompt.getOptions() instanceof ChatOptions) {
|
||||
var chatOptions = (ChatOptions) prompt.getOptions();
|
||||
updatedRuntimeOptions = ((PortableFunctionCallingOptions) updatedRuntimeOptions).merge(chatOptions);
|
||||
}
|
||||
}
|
||||
|
||||
functionsForThisRequest.addAll(this.runtimeFunctionCallbackConfigurations(updatedRuntimeOptions));
|
||||
|
||||
ToolConfiguration toolConfiguration = null;
|
||||
|
||||
if (!CollectionUtils.isEmpty(functionsForThisRequest)) {
|
||||
toolConfiguration = ToolConfiguration.builder().tools(getFunctionTools(functionsForThisRequest)).build();
|
||||
}
|
||||
|
||||
InferenceConfiguration inferenceConfiguration = InferenceConfiguration.builder()
|
||||
.maxTokens(updatedRuntimeOptions.getMaxTokens())
|
||||
.stopSequences(updatedRuntimeOptions.getStopSequences())
|
||||
.temperature(updatedRuntimeOptions.getTemperature() != null
|
||||
? updatedRuntimeOptions.getTemperature().floatValue() : null)
|
||||
.topP(updatedRuntimeOptions.getTopP() != null ? updatedRuntimeOptions.getTopP().floatValue() : null)
|
||||
.build();
|
||||
Document additionalModelRequestFields = ConverseApiUtils
|
||||
.getChatOptionsAdditionalModelRequestFields(defaultOptions, prompt.getOptions());
|
||||
|
||||
return ConverseRequest.builder()
|
||||
.modelId(updatedRuntimeOptions.getModel())
|
||||
.inferenceConfig(inferenceConfiguration)
|
||||
.messages(instructionMessages)
|
||||
.system(systemMessages)
|
||||
.additionalModelRequestFields(additionalModelRequestFields)
|
||||
.toolConfig(toolConfiguration)
|
||||
.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
|
||||
.fromJson(ConverseApiUtils.convertObjectToDocument(ModelOptionsUtils.jsonToMap(inputSchema))))
|
||||
.build())
|
||||
.build();
|
||||
}).toList();
|
||||
}
|
||||
|
||||
private static byte[] getContentMediaData(Object mediaData) {
|
||||
if (mediaData instanceof byte[] bytes) {
|
||||
return bytes;
|
||||
}
|
||||
else if (mediaData instanceof String text) {
|
||||
if (URLValidator.isValidURLBasic(text)) {
|
||||
try {
|
||||
URL url = new URL(text);
|
||||
URLConnection connection = url.openConnection();
|
||||
try (InputStream is = connection.getInputStream()) {
|
||||
return StreamUtils.copyToByteArray(is);
|
||||
}
|
||||
}
|
||||
catch (IOException e) {
|
||||
throw new RuntimeException("Failed to read media data from URL: " + text, e);
|
||||
}
|
||||
}
|
||||
return text.getBytes();
|
||||
}
|
||||
else if (mediaData instanceof URL url) {
|
||||
try (InputStream is = url.openConnection().getInputStream()) {
|
||||
return StreamUtils.copyToByteArray(is);
|
||||
}
|
||||
catch (IOException e) {
|
||||
throw new RuntimeException("Failed to read media data from URL: " + url, e);
|
||||
}
|
||||
}
|
||||
else {
|
||||
throw new IllegalArgumentException("Unsupported media data type: " + mediaData.getClass().getSimpleName());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert {@link ConverseResponse} to {@link ChatResponse} includes model output,
|
||||
* stopReason, usage, metrics etc.
|
||||
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html#API_runtime_Converse_ResponseSyntax
|
||||
* @param response The Bedrock Converse response.
|
||||
* @return The ChatResponse entity.
|
||||
*/
|
||||
private ChatResponse toChatResponse(ConverseResponse response) {
|
||||
|
||||
Assert.notNull(response, "'response' must not be null.");
|
||||
|
||||
Message message = response.output().message();
|
||||
|
||||
List<Generation> generations = message.content()
|
||||
.stream()
|
||||
.filter(content -> content.type() != ContentBlock.Type.TOOL_USE)
|
||||
.map(content -> {
|
||||
return new Generation(new AssistantMessage(content.text(), Map.of()),
|
||||
ChatGenerationMetadata.from(response.stopReasonAsString(), null));
|
||||
})
|
||||
.toList();
|
||||
|
||||
List<Generation> allGenerations = new ArrayList<>(generations);
|
||||
|
||||
if (response.stopReasonAsString() != null && generations.isEmpty()) {
|
||||
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
|
||||
ChatGenerationMetadata.from(response.stopReasonAsString(), null));
|
||||
allGenerations.add(generation);
|
||||
}
|
||||
|
||||
List<ContentBlock> toolUseContentBlocks = message.content()
|
||||
.stream()
|
||||
.filter(c -> c.type() == ContentBlock.Type.TOOL_USE)
|
||||
.toList();
|
||||
|
||||
if (!CollectionUtils.isEmpty(toolUseContentBlocks)) {
|
||||
|
||||
List<AssistantMessage.ToolCall> toolCalls = new ArrayList<>();
|
||||
|
||||
for (ContentBlock toolUseContentBlock : toolUseContentBlocks) {
|
||||
|
||||
var functionCallId = toolUseContentBlock.toolUse().toolUseId();
|
||||
var functionName = toolUseContentBlock.toolUse().name();
|
||||
var functionArguments = toolUseContentBlock.toolUse().input().toString();
|
||||
|
||||
toolCalls
|
||||
.add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
|
||||
}
|
||||
|
||||
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
|
||||
Generation toolCallGeneration = new Generation(assistantMessage,
|
||||
ChatGenerationMetadata.from(response.stopReasonAsString(), null));
|
||||
allGenerations.add(toolCallGeneration);
|
||||
}
|
||||
|
||||
DefaultUsage usage = new DefaultUsage(response.usage().inputTokens().longValue(),
|
||||
response.usage().outputTokens().longValue(), response.usage().totalTokens().longValue());
|
||||
|
||||
Document modelResponseFields = response.additionalModelResponseFields();
|
||||
|
||||
ConverseMetrics metrics = response.metrics();
|
||||
|
||||
var chatResponseMetaData = ChatResponseMetadata.builder()
|
||||
.withId(response.responseMetadata().requestId())
|
||||
.withUsage(usage)
|
||||
.build();
|
||||
|
||||
return new ChatResponse(allGenerations, chatResponseMetaData);
|
||||
}
|
||||
|
||||
/**
|
||||
* Invoke the model and return the response stream.
|
||||
*
|
||||
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
|
||||
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
|
||||
* https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/bedrockruntime/BedrockRuntimeAsyncClient.html#converseStream
|
||||
* @param bedrockConverseRequest Model invocation request.
|
||||
* @return The model invocation response stream.
|
||||
*/
|
||||
@Override
|
||||
public Flux<ChatResponse> stream(Prompt prompt) {
|
||||
Assert.notNull(prompt, "'prompt' must not be null");
|
||||
|
||||
return Flux.deferContextual(contextView -> {
|
||||
|
||||
ConverseRequest converseRequest = this.createRequest(prompt);
|
||||
|
||||
// System.out.println(">>>>> CONVERSE REQUEST: " + converseRequest);
|
||||
|
||||
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
|
||||
.prompt(prompt)
|
||||
.provider(AiProvider.BEDROCK_CONVERSE.value())
|
||||
.requestOptions(buildRequestOptions(converseRequest))
|
||||
.build();
|
||||
|
||||
Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(
|
||||
this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
|
||||
this.observationRegistry);
|
||||
|
||||
observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();
|
||||
|
||||
ConverseStreamRequest converseStreamRequest = ConverseStreamRequest.builder()
|
||||
.modelId(converseRequest.modelId())
|
||||
.messages(converseRequest.messages())
|
||||
.system(converseRequest.system())
|
||||
.additionalModelRequestFields(converseRequest.additionalModelRequestFields())
|
||||
.toolConfig(converseRequest.toolConfig())
|
||||
.build();
|
||||
|
||||
Flux<ConverseStreamOutput> response = converseStream(converseStreamRequest);
|
||||
|
||||
// @formatter:off
|
||||
Flux<ChatResponse> chatResponses = ConverseApiUtils.toChatResponse(response);
|
||||
|
||||
Flux<ChatResponse> chatResponseFlux = chatResponses.switchMap(chatResponse -> {
|
||||
if (!this.isProxyToolCalls(prompt, this.defaultOptions) && chatResponse != null
|
||||
&& this.isToolCall(chatResponse, Set.of("tool_use"))) {
|
||||
var toolCallConversation = this.handleToolCalls(prompt, chatResponse);
|
||||
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
|
||||
}
|
||||
return Mono.just(chatResponse);
|
||||
})
|
||||
.doOnError(observation::error)
|
||||
.doFinally(s -> observation.stop())
|
||||
.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
|
||||
// @formatter:on
|
||||
|
||||
return new MessageAggregator().aggregate(chatResponseFlux, observationContext::setResponse);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Invoke the model and return the response stream.
|
||||
*
|
||||
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
|
||||
* https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
|
||||
* https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/bedrockruntime/BedrockRuntimeAsyncClient.html#converseStream
|
||||
* @param converseStreamRequest Model invocation request.
|
||||
* @return The model invocation response stream.
|
||||
*/
|
||||
public Flux<ConverseStreamOutput> converseStream(ConverseStreamRequest converseStreamRequest) {
|
||||
Assert.notNull(converseStreamRequest, "'converseStreamRequest' must not be null");
|
||||
|
||||
Sinks.Many<ConverseStreamOutput> eventSink = Sinks.many().multicast().onBackpressureBuffer();
|
||||
|
||||
ConverseStreamResponseHandler.Visitor visitor = ConverseStreamResponseHandler.Visitor.builder()
|
||||
.onDefault((output) -> {
|
||||
logger.debug("Received converse stream output:{}", output);
|
||||
eventSink.tryEmitNext(output);
|
||||
})
|
||||
.build();
|
||||
|
||||
ConverseStreamResponseHandler responseHandler = ConverseStreamResponseHandler.builder()
|
||||
.onEventStream(stream -> stream.subscribe((e) -> e.accept(visitor)))
|
||||
.onComplete(() -> {
|
||||
EmitResult emitResult = eventSink.tryEmitComplete();
|
||||
|
||||
while (!emitResult.isSuccess()) {
|
||||
logger.info("Emitting complete:{}", emitResult);
|
||||
emitResult = eventSink.tryEmitComplete();
|
||||
}
|
||||
|
||||
eventSink.emitComplete(EmitFailureHandler.busyLooping(Duration.ofSeconds(3)));
|
||||
logger.info("Completed streaming response.");
|
||||
})
|
||||
.onError((error) -> {
|
||||
logger.error("Error handling Bedrock converse stream response", error);
|
||||
eventSink.tryEmitError(error);
|
||||
})
|
||||
.build();
|
||||
|
||||
this.bedrockRuntimeAsyncClient.converseStream(converseStreamRequest, responseHandler);
|
||||
|
||||
return eventSink.asFlux();
|
||||
|
||||
}
|
||||
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
public static class Builder {
|
||||
|
||||
private AwsCredentialsProvider credentialsProvider;
|
||||
|
||||
private Region region = Region.US_EAST_1;
|
||||
|
||||
private Duration timeout = Duration.ofMinutes(10);
|
||||
|
||||
private FunctionCallingOptions defaultOptions = new FunctionCallingOptionsBuilder().build();
|
||||
|
||||
private FunctionCallbackContext functionCallbackContext;
|
||||
|
||||
private List<FunctionCallback> toolFunctionCallbacks;
|
||||
|
||||
private ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
|
||||
|
||||
private ChatModelObservationConvention customObservationConvention;
|
||||
|
||||
private BedrockRuntimeClient bedrockRuntimeClient;
|
||||
|
||||
private BedrockRuntimeAsyncClient bedrockRuntimeAsyncClient;
|
||||
|
||||
private Builder() {
|
||||
}
|
||||
|
||||
public Builder withCredentialsProvider(AwsCredentialsProvider credentialsProvider) {
|
||||
Assert.notNull(credentialsProvider, "'credentialsProvider' must not be null.");
|
||||
this.credentialsProvider = credentialsProvider;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withRegion(Region region) {
|
||||
Assert.notNull(region, "'region' must not be null.");
|
||||
this.region = region;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withTimeout(Duration timeout) {
|
||||
Assert.notNull(timeout, "'timeout' must not be null.");
|
||||
this.timeout = timeout;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withDefaultOptions(FunctionCallingOptions defaultOptions) {
|
||||
Assert.notNull(defaultOptions, "'defaultOptions' must not be null.");
|
||||
this.defaultOptions = defaultOptions;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withFunctionCallbackContext(FunctionCallbackContext functionCallbackContext) {
|
||||
this.functionCallbackContext = functionCallbackContext;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withToolFunctionCallbacks(List<FunctionCallback> toolFunctionCallbacks) {
|
||||
this.toolFunctionCallbacks = toolFunctionCallbacks;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withObservationRegistry(ObservationRegistry observationRegistry) {
|
||||
Assert.notNull(observationRegistry, "'observationRegistry' must not be null.");
|
||||
this.observationRegistry = observationRegistry;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withCustomObservationConvention(ChatModelObservationConvention observationConvention) {
|
||||
Assert.notNull(observationConvention, "'observationConvention' must not be null.");
|
||||
this.customObservationConvention = observationConvention;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withBedrockRuntimeClient(BedrockRuntimeClient bedrockRuntimeClient) {
|
||||
this.bedrockRuntimeClient = bedrockRuntimeClient;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withBedrockRuntimeAsyncClient(BedrockRuntimeAsyncClient bedrockRuntimeAsyncClient) {
|
||||
this.bedrockRuntimeAsyncClient = bedrockRuntimeAsyncClient;
|
||||
return this;
|
||||
}
|
||||
|
||||
public BedrockProxyChatModel build() {
|
||||
|
||||
if (this.bedrockRuntimeClient == null) {
|
||||
this.bedrockRuntimeClient = BedrockRuntimeClient.builder()
|
||||
.region(this.region)
|
||||
.httpClientBuilder(null)
|
||||
.credentialsProvider(this.credentialsProvider)
|
||||
.overrideConfiguration(c -> c.apiCallTimeout(this.timeout))
|
||||
.build();
|
||||
}
|
||||
|
||||
if (this.bedrockRuntimeAsyncClient == null) {
|
||||
|
||||
// TODO: Is it ok to configure the NettyNioAsyncHttpClient explicitly???
|
||||
var httpClientBuilder = NettyNioAsyncHttpClient.builder()
|
||||
.tcpKeepAlive(true)
|
||||
.connectionAcquisitionTimeout(Duration.ofSeconds(30))
|
||||
.maxConcurrency(200);
|
||||
|
||||
var builder = BedrockRuntimeAsyncClient.builder()
|
||||
.region(this.region)
|
||||
.httpClientBuilder(httpClientBuilder)
|
||||
.credentialsProvider(this.credentialsProvider)
|
||||
.overrideConfiguration(c -> c.apiCallTimeout(this.timeout));
|
||||
this.bedrockRuntimeAsyncClient = builder.build();
|
||||
}
|
||||
|
||||
var bedrockProxyChatModel = new BedrockProxyChatModel(this.bedrockRuntimeClient,
|
||||
this.bedrockRuntimeAsyncClient, this.defaultOptions, this.functionCallbackContext,
|
||||
this.toolFunctionCallbacks, this.observationRegistry);
|
||||
|
||||
if (this.customObservationConvention != null) {
|
||||
bedrockProxyChatModel.setObservationConvention(this.customObservationConvention);
|
||||
}
|
||||
|
||||
return bedrockProxyChatModel;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Use the provided convention for reporting observation data
|
||||
* @param observationConvention The provided convention
|
||||
*/
|
||||
public void setObservationConvention(ChatModelObservationConvention observationConvention) {
|
||||
Assert.notNull(observationConvention, "observationConvention cannot be null");
|
||||
this.observationConvention = observationConvention;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
/*
|
||||
* 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.bedrock.converse.api;
|
||||
|
||||
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 Converse API.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Wei Jiang
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class BedrockUsage implements Usage {
|
||||
|
||||
public static BedrockUsage from(TokenUsage usage) {
|
||||
Assert.notNull(usage, "'TokenUsage' must not be null.");
|
||||
|
||||
return new BedrockUsage(usage.inputTokens().longValue(), usage.outputTokens().longValue());
|
||||
}
|
||||
|
||||
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 inputTokens;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Long getGenerationTokens() {
|
||||
return outputTokens;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return "BedrockUsage [inputTokens=" + inputTokens + ", outputTokens=" + outputTokens + "]";
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,499 @@
|
||||
/*
|
||||
* 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.bedrock.converse.api;
|
||||
|
||||
import java.math.BigDecimal;
|
||||
import java.math.BigInteger;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
|
||||
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
||||
import org.springframework.ai.chat.metadata.DefaultUsage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.ChatOptions;
|
||||
import org.springframework.ai.model.ModelOptions;
|
||||
import org.springframework.ai.model.ModelOptionsUtils;
|
||||
import org.springframework.util.CollectionUtils;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import software.amazon.awssdk.core.SdkField;
|
||||
import software.amazon.awssdk.core.document.Document;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockDelta;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockDeltaEvent;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStart;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStartEvent;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlockStopEvent;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamMetadataEvent;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamMetrics;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput.EventType;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler.Visitor;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamTrace;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.MessageStartEvent;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.MessageStopEvent;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.TokenUsage;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ToolUseBlockStart;
|
||||
|
||||
/**
|
||||
* Amazon Bedrock Converse API utils.
|
||||
*
|
||||
* @author Wei Jiang
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class ConverseApiUtils {
|
||||
|
||||
public static boolean isToolUseStart(ConverseStreamOutput event) {
|
||||
if (event == null || event.sdkEventType() == null || event.sdkEventType() != EventType.CONTENT_BLOCK_START) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return ContentBlockStart.Type.TOOL_USE == ((ContentBlockStartEvent) event).start().type();
|
||||
}
|
||||
|
||||
public static boolean isToolUseFinish(ConverseStreamOutput event) {
|
||||
if (event == null || event.sdkEventType() == null || event.sdkEventType() != EventType.METADATA) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
public record Aggregation(MetadataAggregation metadataAggregation, ChatResponse chatResponse) {
|
||||
public Aggregation() {
|
||||
this(MetadataAggregation.builder().build(), EMPTY_CHAT_RESPONSE);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Special event used to aggregate multiple tool use events into a single event with
|
||||
* list of aggregated ContentBlockToolUse.
|
||||
*/
|
||||
public static class ToolUseAggregationEvent implements ConverseStreamOutput {
|
||||
|
||||
public record ToolUseEntry(Integer index, String id, String name, String input) {
|
||||
}
|
||||
|
||||
private Integer index;
|
||||
|
||||
private String id;
|
||||
|
||||
private String name;
|
||||
|
||||
private String partialJson = "";
|
||||
|
||||
private List<ToolUseEntry> toolUseEntries = new ArrayList<>();
|
||||
|
||||
private DefaultUsage usage;
|
||||
|
||||
public List<ToolUseEntry> toolUseEntries() {
|
||||
return this.toolUseEntries;
|
||||
}
|
||||
|
||||
public boolean isEmpty() {
|
||||
return (this.index == null || this.id == null || this.name == null
|
||||
|| !StringUtils.hasText(this.partialJson));
|
||||
}
|
||||
|
||||
ToolUseAggregationEvent withIndex(Integer index) {
|
||||
this.index = index;
|
||||
return this;
|
||||
}
|
||||
|
||||
ToolUseAggregationEvent withId(String id) {
|
||||
this.id = id;
|
||||
return this;
|
||||
}
|
||||
|
||||
ToolUseAggregationEvent withName(String name) {
|
||||
this.name = name;
|
||||
return this;
|
||||
}
|
||||
|
||||
ToolUseAggregationEvent withUsage(DefaultUsage usage) {
|
||||
this.usage = usage;
|
||||
return this;
|
||||
}
|
||||
|
||||
ToolUseAggregationEvent appendPartialJson(String partialJson) {
|
||||
this.partialJson = this.partialJson + partialJson;
|
||||
return this;
|
||||
}
|
||||
|
||||
void squashIntoContentBlock() {
|
||||
this.toolUseEntries.add(new ToolUseEntry(this.index, this.id, this.name, this.partialJson));
|
||||
this.index = null;
|
||||
this.id = null;
|
||||
this.name = null;
|
||||
this.partialJson = "";
|
||||
this.usage = null;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return "EventToolUseBuilder [index=" + this.index + ", id=" + this.id + ", name=" + this.name
|
||||
+ ", partialJson=" + this.partialJson + ", toolUseMap=" + "]";
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<SdkField<?>> sdkFields() {
|
||||
return List.of();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void accept(Visitor visitor) {
|
||||
throw new UnsupportedOperationException();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public static ConverseStreamOutput mergeToolUseEvents(ConverseStreamOutput previousEvent,
|
||||
ConverseStreamOutput event) {
|
||||
|
||||
ToolUseAggregationEvent toolUseEventAggregator = (ToolUseAggregationEvent) previousEvent;
|
||||
|
||||
if (event.sdkEventType() == EventType.CONTENT_BLOCK_START) {
|
||||
|
||||
ContentBlockStartEvent contentBlockStart = (ContentBlockStartEvent) event;
|
||||
|
||||
if (ContentBlockStart.Type.TOOL_USE.equals(contentBlockStart.start().type())) {
|
||||
ToolUseBlockStart cbToolUse = contentBlockStart.start().toolUse();
|
||||
|
||||
return toolUseEventAggregator.withIndex(contentBlockStart.contentBlockIndex())
|
||||
.withId(cbToolUse.toolUseId())
|
||||
.withName(cbToolUse.name())
|
||||
.appendPartialJson(""); // CB START always has empty JSON.
|
||||
}
|
||||
}
|
||||
else if (event.sdkEventType() == EventType.CONTENT_BLOCK_DELTA) {
|
||||
ContentBlockDeltaEvent contentBlockDelta = (ContentBlockDeltaEvent) event;
|
||||
if (ContentBlockDelta.Type.TOOL_USE == contentBlockDelta.delta().type()) {
|
||||
return toolUseEventAggregator.appendPartialJson(contentBlockDelta.delta().toolUse().input());
|
||||
}
|
||||
}
|
||||
else if (event.sdkEventType() == EventType.CONTENT_BLOCK_STOP) {
|
||||
return toolUseEventAggregator;
|
||||
}
|
||||
else if (event.sdkEventType() == EventType.MESSAGE_STOP) {
|
||||
return toolUseEventAggregator;
|
||||
}
|
||||
else if (event.sdkEventType() == EventType.METADATA) {
|
||||
ConverseStreamMetadataEvent metadataEvent = (ConverseStreamMetadataEvent) event;
|
||||
DefaultUsage usage = new DefaultUsage(metadataEvent.usage().inputTokens().longValue(),
|
||||
metadataEvent.usage().outputTokens().longValue(), metadataEvent.usage().totalTokens().longValue());
|
||||
toolUseEventAggregator.withUsage(usage);
|
||||
// TODO
|
||||
if (!toolUseEventAggregator.isEmpty()) {
|
||||
toolUseEventAggregator.squashIntoContentBlock();
|
||||
return toolUseEventAggregator;
|
||||
}
|
||||
}
|
||||
|
||||
return event;
|
||||
}
|
||||
|
||||
public static Flux<ChatResponse> toChatResponse(Flux<ConverseStreamOutput> responses) {
|
||||
|
||||
AtomicBoolean isInsideTool = new AtomicBoolean(false);
|
||||
|
||||
return responses.map(event -> {
|
||||
if (ConverseApiUtils.isToolUseStart(event)) {
|
||||
isInsideTool.set(true);
|
||||
}
|
||||
return event;
|
||||
}).windowUntil(event -> { // Group all chunks belonging to the same function call.
|
||||
if (isInsideTool.get() && ConverseApiUtils.isToolUseFinish(event)) {
|
||||
isInsideTool.set(false);
|
||||
return true;
|
||||
}
|
||||
return !isInsideTool.get();
|
||||
}).concatMapIterable(window -> {// Merging the window chunks into a single chunk.
|
||||
Mono<ConverseStreamOutput> monoChunk = window.reduce(new ToolUseAggregationEvent(),
|
||||
ConverseApiUtils::mergeToolUseEvents);
|
||||
return List.of(monoChunk);
|
||||
}).flatMap(mono -> mono).scanWith(() -> new Aggregation(), (lastAggregation, nextEvent) -> {
|
||||
|
||||
// System.out.println(nextEvent);
|
||||
if (nextEvent instanceof ToolUseAggregationEvent toolUseAggregationEvent) {
|
||||
|
||||
if (CollectionUtils.isEmpty(toolUseAggregationEvent.toolUseEntries())) {
|
||||
return new Aggregation();
|
||||
}
|
||||
|
||||
List<AssistantMessage.ToolCall> toolCalls = new ArrayList<>();
|
||||
|
||||
for (ToolUseAggregationEvent.ToolUseEntry toolUseEntry : toolUseAggregationEvent.toolUseEntries()) {
|
||||
var functionCallId = toolUseEntry.id();
|
||||
var functionName = toolUseEntry.name();
|
||||
var functionArguments = toolUseEntry.input();
|
||||
toolCalls.add(
|
||||
new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
|
||||
}
|
||||
|
||||
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
|
||||
Generation toolCallGeneration = new Generation(assistantMessage,
|
||||
ChatGenerationMetadata.from("tool_use", null));
|
||||
|
||||
var chatResponseMetaData = ChatResponseMetadata.builder()
|
||||
.withUsage(toolUseAggregationEvent.usage)
|
||||
.build();
|
||||
|
||||
return new Aggregation(
|
||||
MetadataAggregation.builder().copy(lastAggregation.metadataAggregation()).build(),
|
||||
new ChatResponse(List.of(toolCallGeneration), chatResponseMetaData));
|
||||
|
||||
}
|
||||
else if (nextEvent instanceof MessageStartEvent messageStartEvent) {
|
||||
var newMeta = MetadataAggregation.builder()
|
||||
.copy(lastAggregation.metadataAggregation())
|
||||
.withRole(messageStartEvent.role().toString())
|
||||
.build();
|
||||
return new Aggregation(newMeta, ConverseApiUtils.EMPTY_CHAT_RESPONSE);
|
||||
}
|
||||
else if (nextEvent instanceof MessageStopEvent messageStopEvent) {
|
||||
var newMeta = MetadataAggregation.builder()
|
||||
.copy(lastAggregation.metadataAggregation())
|
||||
.withStopReason(messageStopEvent.stopReasonAsString())
|
||||
.withAdditionalModelResponseFields(messageStopEvent.additionalModelResponseFields())
|
||||
.build();
|
||||
return new Aggregation(newMeta, ConverseApiUtils.EMPTY_CHAT_RESPONSE);
|
||||
}
|
||||
else if (nextEvent instanceof ContentBlockStartEvent contentBlockStartEvent) {
|
||||
// TODO ToolUse support
|
||||
return new Aggregation();
|
||||
}
|
||||
else if (nextEvent instanceof ContentBlockDeltaEvent contentBlockDeltaEvent) {
|
||||
if (contentBlockDeltaEvent.delta().type().equals(ContentBlockDelta.Type.TEXT)) {
|
||||
|
||||
var generation = new Generation(
|
||||
new AssistantMessage(contentBlockDeltaEvent.delta().text(), Map.of()),
|
||||
ChatGenerationMetadata.from(lastAggregation.metadataAggregation().stopReason(), null));
|
||||
|
||||
return new Aggregation(
|
||||
MetadataAggregation.builder().copy(lastAggregation.metadataAggregation()).build(),
|
||||
new ChatResponse(List.of(generation)));
|
||||
}
|
||||
else if (contentBlockDeltaEvent.delta().type().equals(ContentBlockDelta.Type.TOOL_USE)) {
|
||||
// TODO ToolUse support
|
||||
}
|
||||
return new Aggregation();
|
||||
}
|
||||
else if (nextEvent instanceof ContentBlockStopEvent contentBlockStopEvent) {
|
||||
// TODO ToolUse support
|
||||
return new Aggregation();
|
||||
}
|
||||
else if (nextEvent instanceof ConverseStreamMetadataEvent metadataEvent) {
|
||||
// return new Aggregation();
|
||||
var newMeta = MetadataAggregation.builder()
|
||||
.copy(lastAggregation.metadataAggregation())
|
||||
.withTokenUsage(metadataEvent.usage())
|
||||
.withMetrics(metadataEvent.metrics())
|
||||
.withTrace(metadataEvent.trace())
|
||||
.build();
|
||||
|
||||
DefaultUsage usage = new DefaultUsage(metadataEvent.usage().inputTokens().longValue(),
|
||||
metadataEvent.usage().outputTokens().longValue(),
|
||||
metadataEvent.usage().totalTokens().longValue());
|
||||
|
||||
// TODO
|
||||
Document modelResponseFields = lastAggregation.metadataAggregation().additionalModelResponseFields();
|
||||
ConverseStreamMetrics metrics = metadataEvent.metrics();
|
||||
|
||||
var chatResponseMetaData = ChatResponseMetadata.builder().withUsage(usage).build();
|
||||
|
||||
return new Aggregation(newMeta, new ChatResponse(List.of(), chatResponseMetaData));
|
||||
}
|
||||
else {
|
||||
return new Aggregation();
|
||||
}
|
||||
})
|
||||
// .skip(1)
|
||||
.map(aggregation -> aggregation.chatResponse())
|
||||
.filter(chatResponse -> chatResponse != ConverseApiUtils.EMPTY_CHAT_RESPONSE);
|
||||
}
|
||||
|
||||
public static final ChatResponse EMPTY_CHAT_RESPONSE = ChatResponse.builder()
|
||||
.withGenerations(List.of())
|
||||
.withMetadata("empty", true)
|
||||
.build();
|
||||
|
||||
public record MetadataAggregation(String role, String stopReason, Document additionalModelResponseFields,
|
||||
TokenUsage tokenUsage, ConverseStreamMetrics metrics, ConverseStreamTrace trace) {
|
||||
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
public final static class Builder {
|
||||
|
||||
private String role;
|
||||
|
||||
private String stopReason;
|
||||
|
||||
private Document additionalModelResponseFields;
|
||||
|
||||
private TokenUsage tokenUsage;
|
||||
|
||||
private ConverseStreamMetrics metrics;
|
||||
|
||||
private ConverseStreamTrace trace;
|
||||
|
||||
private Builder() {
|
||||
}
|
||||
|
||||
public Builder copy(MetadataAggregation metadataAggregation) {
|
||||
this.role = metadataAggregation.role;
|
||||
this.stopReason = metadataAggregation.stopReason;
|
||||
this.additionalModelResponseFields = metadataAggregation.additionalModelResponseFields;
|
||||
this.tokenUsage = metadataAggregation.tokenUsage;
|
||||
this.metrics = metadataAggregation.metrics;
|
||||
this.trace = metadataAggregation.trace;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withRole(String role) {
|
||||
this.role = role;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withStopReason(String stopReason) {
|
||||
this.stopReason = stopReason;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withAdditionalModelResponseFields(Document additionalModelResponseFields) {
|
||||
this.additionalModelResponseFields = additionalModelResponseFields;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withTokenUsage(TokenUsage tokenUsage) {
|
||||
this.tokenUsage = tokenUsage;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withMetrics(ConverseStreamMetrics metrics) {
|
||||
this.metrics = metrics;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withTrace(ConverseStreamTrace trace) {
|
||||
this.trace = trace;
|
||||
return this;
|
||||
}
|
||||
|
||||
public MetadataAggregation build() {
|
||||
return new MetadataAggregation(role, stopReason, additionalModelResponseFields, tokenUsage, metrics,
|
||||
trace);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
public static Document getChatOptionsAdditionalModelRequestFields(ChatOptions defaultOptions,
|
||||
ModelOptions promptOptions) {
|
||||
if (defaultOptions == null && promptOptions == null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
Map<String, Object> attributes = new HashMap<>();
|
||||
|
||||
if (defaultOptions != null) {
|
||||
attributes.putAll(ModelOptionsUtils.objectToMap(defaultOptions));
|
||||
}
|
||||
|
||||
if (promptOptions != null) {
|
||||
if (promptOptions instanceof ChatOptions runtimeOptions) {
|
||||
attributes.putAll(ModelOptionsUtils.objectToMap(runtimeOptions));
|
||||
}
|
||||
else {
|
||||
throw new IllegalArgumentException(
|
||||
"Prompt options are not of type ChatOptions:" + promptOptions.getClass().getSimpleName());
|
||||
}
|
||||
}
|
||||
|
||||
attributes.remove("model");
|
||||
attributes.remove("proxyToolCalls");
|
||||
attributes.remove("functions");
|
||||
attributes.remove("toolContext");
|
||||
attributes.remove("functionCallbacks");
|
||||
|
||||
attributes.remove("temperature");
|
||||
attributes.remove("topK");
|
||||
attributes.remove("stopSequences");
|
||||
attributes.remove("maxTokens");
|
||||
attributes.remove("topP");
|
||||
|
||||
return convertObjectToDocument(attributes);
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
public static Document convertObjectToDocument(Object value) {
|
||||
if (value == null) {
|
||||
return Document.fromNull();
|
||||
}
|
||||
else if (value instanceof String stringValue) {
|
||||
return Document.fromString(stringValue);
|
||||
}
|
||||
else if (value instanceof Boolean booleanValue) {
|
||||
return Document.fromBoolean(booleanValue);
|
||||
}
|
||||
else if (value instanceof Integer integerValue) {
|
||||
return Document.fromNumber(integerValue);
|
||||
}
|
||||
else if (value instanceof Long longValue) {
|
||||
return Document.fromNumber(longValue);
|
||||
}
|
||||
else if (value instanceof Float floatValue) {
|
||||
return Document.fromNumber(floatValue);
|
||||
}
|
||||
else if (value instanceof Double doubleValue) {
|
||||
return Document.fromNumber(doubleValue);
|
||||
}
|
||||
else if (value instanceof BigDecimal bigDecimalValue) {
|
||||
return Document.fromNumber(bigDecimalValue);
|
||||
}
|
||||
else if (value instanceof BigInteger bigIntegerValue) {
|
||||
return Document.fromNumber(bigIntegerValue);
|
||||
}
|
||||
else if (value instanceof List listValue) {
|
||||
return Document.fromList(listValue.stream().map(v -> convertObjectToDocument(v)).toList());
|
||||
}
|
||||
else if (value instanceof Map mapValue) {
|
||||
return convertMapToDocument(mapValue);
|
||||
}
|
||||
else {
|
||||
throw new IllegalArgumentException("Unsupported value type:" + value.getClass().getSimpleName());
|
||||
}
|
||||
}
|
||||
|
||||
private static Document convertMapToDocument(Map<String, Object> value) {
|
||||
Map<String, Document> attr = value.entrySet()
|
||||
.stream()
|
||||
.collect(Collectors.toMap(e -> e.getKey(), e -> convertObjectToDocument(e.getValue())));
|
||||
|
||||
return Document.fromMap(attr);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
/*
|
||||
* 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.bedrock.converse.api;
|
||||
|
||||
import java.net.MalformedURLException;
|
||||
import java.net.URISyntaxException;
|
||||
import java.net.URL;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
/**
|
||||
* Utility class for detecting and normalizing URLs. Intended for use with multimodal user
|
||||
* inputs.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class URLValidator {
|
||||
|
||||
// Basic URL regex pattern
|
||||
// Protocol (http:// or https://)
|
||||
private static final Pattern URL_PATTERN = Pattern.compile("^(https?://)" +
|
||||
|
||||
"((([a-zA-Z0-9-]+\\.)+[a-zA-Z]{2,6})|" + // Domain name
|
||||
"(localhost))" + // OR localhost
|
||||
"(:[0-9]{1,5})?" + // Optional port
|
||||
"(/[\\w\\-./]*)*" + // Optional path
|
||||
"(\\?[\\w=&\\-.]*)?" + // Optional query parameters
|
||||
"(#[\\w-]*)?" + // Optional fragment
|
||||
"$");
|
||||
|
||||
/**
|
||||
* Quick validation using regex pattern Good for basic checks but may not catch all
|
||||
* edge cases
|
||||
*/
|
||||
public static boolean isValidURLBasic(String urlString) {
|
||||
if (urlString == null || urlString.trim().isEmpty()) {
|
||||
return false;
|
||||
}
|
||||
return URL_PATTERN.matcher(urlString).matches();
|
||||
}
|
||||
|
||||
/**
|
||||
* Thorough validation using URL class More comprehensive but might be slower
|
||||
* Validates protocol, host, port, and basic structure
|
||||
*/
|
||||
public static boolean isValidURLStrict(String urlString) {
|
||||
if (urlString == null || urlString.trim().isEmpty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
try {
|
||||
URL url = new URL(urlString);
|
||||
// Additional validation by attempting to convert to URI
|
||||
url.toURI();
|
||||
|
||||
// Ensure protocol is http or https
|
||||
String protocol = url.getProtocol().toLowerCase();
|
||||
if (!protocol.equals("http") && !protocol.equals("https")) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Validate host (not empty and contains at least one dot, unless it's
|
||||
// localhost)
|
||||
String host = url.getHost();
|
||||
if (host == null || host.isEmpty()) {
|
||||
return false;
|
||||
}
|
||||
if (!host.equals("localhost") && !host.contains(".")) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Validate port (if specified)
|
||||
int port = url.getPort();
|
||||
if (port != -1 && (port < 1 || port > 65535)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
catch (MalformedURLException | URISyntaxException e) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Attempts to fix common URL issues Adds protocol if missing, removes extra spaces
|
||||
*/
|
||||
public static String normalizeURL(String urlString) {
|
||||
if (urlString == null || urlString.trim().isEmpty()) {
|
||||
return null;
|
||||
}
|
||||
|
||||
String normalized = urlString.trim();
|
||||
|
||||
// Add protocol if missing
|
||||
if (!normalized.toLowerCase().startsWith("http://") && !normalized.toLowerCase().startsWith("https://")) {
|
||||
normalized = "https://" + normalized;
|
||||
}
|
||||
|
||||
// Remove multiple forward slashes in path (except after protocol)
|
||||
normalized = normalized.replaceAll("(?<!:)/{2,}", "/");
|
||||
|
||||
// Remove trailing slash (unless it's just protocol://host)
|
||||
if (normalized.matches("https?://[^/]+/+$")) {
|
||||
normalized = normalized.replaceAll("/+$", "");
|
||||
}
|
||||
|
||||
return normalized;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,341 @@
|
||||
/*
|
||||
* 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.bedrock.converse;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.URL;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.converter.BeanOutputConverter;
|
||||
import org.springframework.ai.converter.ListOutputConverter;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.core.ParameterizedTypeReference;
|
||||
import org.springframework.core.convert.support.DefaultConversionService;
|
||||
import org.springframework.core.io.ClassPathResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.util.MimeTypeUtils;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
@SpringBootTest(classes = BedrockConverseTestConfiguration.class)
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
class BedrockConverseChatClientIT {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(BedrockConverseChatClientIT.class);
|
||||
|
||||
@Value("classpath:/prompts/system-message.st")
|
||||
private Resource systemTextResource;
|
||||
|
||||
@Autowired
|
||||
private ChatModel chatModel;
|
||||
|
||||
@Test
|
||||
void call() {
|
||||
|
||||
// @formatter:off
|
||||
ChatResponse response = ChatClient.create(this.chatModel).prompt()
|
||||
.advisors(new SimpleLoggerAdvisor())
|
||||
.system(s -> s.text(this.systemTextResource)
|
||||
.param("name", "Bob")
|
||||
.param("voice", "pirate"))
|
||||
.user("Tell me about 3 famous pirates from the Golden Age of Piracy and what they did")
|
||||
.call()
|
||||
.chatResponse();
|
||||
// @formatter:on
|
||||
|
||||
logger.info("" + response);
|
||||
assertThat(response.getResults()).hasSize(1);
|
||||
assertThat(response.getResults().get(0).getOutput().getContent()).contains("Blackbeard");
|
||||
}
|
||||
|
||||
@Test
|
||||
void listOutputConverterString() {
|
||||
// @formatter:off
|
||||
List<String> collection = ChatClient.create(this.chatModel).prompt()
|
||||
.user(u -> u.text("List five {subject}")
|
||||
.param("subject", "ice cream flavors"))
|
||||
.call()
|
||||
.entity(new ParameterizedTypeReference<List<String>>() {});
|
||||
// @formatter:on
|
||||
|
||||
logger.info(collection.toString());
|
||||
assertThat(collection).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void listOutputConverterBean() {
|
||||
|
||||
// @formatter:off
|
||||
List<ActorsFilms> actorsFilms = ChatClient.create(this.chatModel).prompt()
|
||||
.user("Generate the filmography of 5 movies for Tom Hanks and Bill Murray.")
|
||||
.call()
|
||||
.entity(new ParameterizedTypeReference<List<ActorsFilms>>() {
|
||||
});
|
||||
// @formatter:on
|
||||
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms).hasSize(2);
|
||||
}
|
||||
|
||||
@Test
|
||||
void customOutputConverter() {
|
||||
|
||||
var toStringListConverter = new ListOutputConverter(new DefaultConversionService());
|
||||
|
||||
// @formatter:off
|
||||
List<String> flavors = ChatClient.create(this.chatModel).prompt()
|
||||
.user(u -> u.text("List five {subject}")
|
||||
.param("subject", "ice cream flavors"))
|
||||
.call()
|
||||
.entity(toStringListConverter);
|
||||
// @formatter:on
|
||||
|
||||
logger.info("ice cream flavors" + flavors);
|
||||
assertThat(flavors).hasSize(5);
|
||||
assertThat(flavors).containsAnyOf("Vanilla", "vanilla");
|
||||
}
|
||||
|
||||
@Test
|
||||
void mapOutputConverter() {
|
||||
// @formatter:off
|
||||
Map<String, Object> result = ChatClient.create(this.chatModel).prompt()
|
||||
.user(u -> u.text("Provide me a List of {subject}")
|
||||
.param("subject", "an array of numbers from 1 to 9 under they key name 'numbers'"))
|
||||
.call()
|
||||
.entity(new ParameterizedTypeReference<Map<String, Object>>() {
|
||||
});
|
||||
// @formatter:on
|
||||
|
||||
assertThat(result.get("numbers")).isEqualTo(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanOutputConverter() {
|
||||
|
||||
// @formatter:off
|
||||
ActorsFilms actorsFilms = ChatClient.create(this.chatModel).prompt()
|
||||
.user("Generate the filmography for a random actor.")
|
||||
.call()
|
||||
.entity(ActorsFilms.class);
|
||||
// @formatter:on
|
||||
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isNotBlank();
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanOutputConverterRecords() {
|
||||
|
||||
// @formatter:off
|
||||
ActorsFilms actorsFilms = ChatClient.create(this.chatModel).prompt()
|
||||
.user("Generate the filmography of 5 movies for Tom Hanks.")
|
||||
.call()
|
||||
.entity(ActorsFilms.class);
|
||||
// @formatter:on
|
||||
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
|
||||
assertThat(actorsFilms.movies()).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanStreamOutputConverterRecords() {
|
||||
|
||||
BeanOutputConverter<ActorsFilms> outputConverter = new BeanOutputConverter<>(ActorsFilms.class);
|
||||
|
||||
// @formatter:off
|
||||
Flux<ChatResponse> chatResponse = ChatClient.create(this.chatModel)
|
||||
.prompt()
|
||||
.advisors(new SimpleLoggerAdvisor())
|
||||
.user(u -> u
|
||||
.text("Generate the filmography of 5 movies for Tom Hanks. " + System.lineSeparator()
|
||||
+ "{format}")
|
||||
.param("format", outputConverter.getFormat()))
|
||||
.stream()
|
||||
.chatResponse();
|
||||
|
||||
List<ChatResponse> chatResponses = chatResponse.collectList()
|
||||
.block()
|
||||
.stream()
|
||||
.toList();
|
||||
|
||||
String generationTextFromStream = chatResponses
|
||||
.stream()
|
||||
.filter(cr -> cr.getResult() != null)
|
||||
.map(cr -> cr.getResult().getOutput().getContent())
|
||||
.collect(Collectors.joining());
|
||||
// @formatter:on
|
||||
|
||||
ActorsFilms actorsFilms = outputConverter.convert(generationTextFromStream);
|
||||
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
|
||||
assertThat(actorsFilms.movies()).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void functionCallTest() {
|
||||
|
||||
// @formatter:off
|
||||
String response = ChatClient.create(this.chatModel)
|
||||
.prompt("What's the weather like in San Francisco, Tokyo, and Paris? Return the temperature in Celsius.")
|
||||
.function("getCurrentWeather", "Get the weather in location", new MockWeatherService())
|
||||
.call()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
logger.info("Response: {}", response);
|
||||
|
||||
assertThat(response).contains("30", "10", "15");
|
||||
}
|
||||
|
||||
@Test
|
||||
void defaultFunctionCallTest() {
|
||||
|
||||
// @formatter:off
|
||||
String response = ChatClient.builder(this.chatModel)
|
||||
.defaultFunction("getCurrentWeather", "Get the weather in location", new MockWeatherService())
|
||||
.defaultUser(u -> u.text("What's the weather like in San Francisco, Tokyo, and Paris? Return the temperature in Celsius."))
|
||||
.build()
|
||||
.prompt()
|
||||
.call()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
logger.info("Response: {}", response);
|
||||
|
||||
assertThat(response).contains("30", "10", "15");
|
||||
}
|
||||
|
||||
@Test
|
||||
void streamFunctionCallTest() {
|
||||
|
||||
// @formatter:off
|
||||
Flux<String> response = ChatClient.create(this.chatModel).prompt()
|
||||
.user("What's the weather like in San Francisco, Tokyo, and Paris? Return the temperature in Celsius.")
|
||||
.function("getCurrentWeather", "Get the weather in location", new MockWeatherService())
|
||||
.stream()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
String content = response.collectList().block().stream().collect(Collectors.joining());
|
||||
logger.info("Response: {}", content);
|
||||
|
||||
assertThat(content).contains("30", "10", "15");
|
||||
}
|
||||
|
||||
@Test
|
||||
void singularStreamFunctionCallTest() {
|
||||
|
||||
// @formatter:off
|
||||
Flux<String> response = ChatClient.create(this.chatModel).prompt()
|
||||
.user("What's the weather like in Paris? Return the temperature in Celsius.")
|
||||
.function("getCurrentWeather", "Get the weather in location", new MockWeatherService())
|
||||
.stream()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
String content = response.collectList().block().stream().collect(Collectors.joining());
|
||||
logger.info("Response: {}", content);
|
||||
|
||||
assertThat(content).contains("15");
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "anthropic.claude-3-5-sonnet-20240620-v1:0" })
|
||||
void multiModalityEmbeddedImage(String modelName) throws IOException {
|
||||
|
||||
// @formatter:off
|
||||
String response = ChatClient.create(this.chatModel).prompt()
|
||||
.options(FunctionCallingOptions.builder().withModel(modelName).build())
|
||||
.user(u -> u.text("Explain what do you see on this picture?")
|
||||
.media(MimeTypeUtils.IMAGE_PNG, new ClassPathResource("/test.png")))
|
||||
.call()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
logger.info(response);
|
||||
assertThat(response).contains("bananas", "apple");
|
||||
assertThat(response).containsAnyOf("bowl", "basket");
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "anthropic.claude-3-5-sonnet-20240620-v1:0" })
|
||||
void multiModalityImageUrl(String modelName) throws IOException {
|
||||
|
||||
// TODO: add url method that wrapps the checked exception.
|
||||
URL url = new URL("https://docs.spring.io/spring-ai/reference/_images/multimodal.test.png");
|
||||
|
||||
// @formatter:off
|
||||
String response = ChatClient.create(this.chatModel).prompt()
|
||||
// TODO consider adding model(...) method to ChatClient as a shortcut to
|
||||
.options(FunctionCallingOptions.builder().withModel(modelName).build())
|
||||
.user(u -> u.text("Explain what do you see on this picture?").media(MimeTypeUtils.IMAGE_PNG, url))
|
||||
.call()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
logger.info(response);
|
||||
assertThat(response).contains("bananas", "apple");
|
||||
assertThat(response).containsAnyOf("bowl", "basket");
|
||||
}
|
||||
|
||||
@Test
|
||||
void streamingMultiModalityImageUrl() throws IOException {
|
||||
|
||||
// TODO: add url method that wrapps the checked exception.
|
||||
URL url = new URL("https://docs.spring.io/spring-ai/reference/_images/multimodal.test.png");
|
||||
|
||||
// @formatter:off
|
||||
Flux<String> response = ChatClient.create(this.chatModel).prompt()
|
||||
.user(u -> u.text("Explain what do you see on this picture?")
|
||||
.media(MimeTypeUtils.IMAGE_PNG, url))
|
||||
.stream()
|
||||
.content();
|
||||
// @formatter:on
|
||||
|
||||
String content = response.collectList().block().stream().collect(Collectors.joining());
|
||||
|
||||
logger.info("Response: {}", content);
|
||||
assertThat(content).contains("bananas", "apple");
|
||||
assertThat(content).containsAnyOf("bowl", "basket");
|
||||
}
|
||||
|
||||
record ActorsFilms(String actor, List<String> movies) {
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* 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.converse;
|
||||
|
||||
import java.time.Duration;
|
||||
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.boot.SpringBootConfiguration;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
|
||||
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
@SpringBootConfiguration
|
||||
public class BedrockConverseTestConfiguration {
|
||||
|
||||
@Bean
|
||||
public BedrockProxyChatModel bedrockConverseChatModel() {
|
||||
|
||||
// String modelId = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
// String modelId = "anthropic.claude-3-5-sonnet-20241022-v2:0";
|
||||
// String modelId = "meta.llama3-8b-instruct-v1:0";
|
||||
// String modelId = "ai21.jamba-1-5-large-v1:0";
|
||||
String modelId = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
|
||||
return BedrockProxyChatModel.builder()
|
||||
.withCredentialsProvider(EnvironmentVariableCredentialsProvider.create())
|
||||
.withRegion(Region.US_EAST_1)
|
||||
.withTimeout(Duration.ofSeconds(120))
|
||||
// .withRegion(Region.US_EAST_1)
|
||||
.withDefaultOptions(FunctionCallingOptions.builder().withModel(modelId).build())
|
||||
.build();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,339 @@
|
||||
/*
|
||||
* 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.converse;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.Message;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
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.Prompt;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
|
||||
import org.springframework.ai.converter.BeanOutputConverter;
|
||||
import org.springframework.ai.converter.ListOutputConverter;
|
||||
import org.springframework.ai.converter.MapOutputConverter;
|
||||
import org.springframework.ai.model.Media;
|
||||
import org.springframework.ai.model.function.FunctionCallbackWrapper;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.core.convert.support.DefaultConversionService;
|
||||
import org.springframework.core.io.ClassPathResource;
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.util.MimeTypeUtils;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
@SpringBootTest(classes = BedrockConverseTestConfiguration.class, properties = "spring.ai.retry.on-http-codes=429")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
class BedrockProxyChatModelIT {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(BedrockProxyChatModelIT.class);
|
||||
|
||||
@Autowired
|
||||
protected ChatModel chatModel;
|
||||
|
||||
@Autowired
|
||||
protected StreamingChatModel streamingChatModel;
|
||||
|
||||
@Value("classpath:/prompts/system-message.st")
|
||||
private Resource systemResource;
|
||||
|
||||
private static void validateChatResponseMetadata(ChatResponse response, String model) {
|
||||
// assertThat(response.getMetadata().getId()).isNotEmpty();
|
||||
// assertThat(response.getMetadata().getModel()).containsIgnoringCase(model);
|
||||
assertThat(response.getMetadata().getUsage().getPromptTokens()).isPositive();
|
||||
assertThat(response.getMetadata().getUsage().getGenerationTokens()).isPositive();
|
||||
assertThat(response.getMetadata().getUsage().getTotalTokens()).isPositive();
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"anthropic.claude-3-5-sonnet-20240620-v1:0" })
|
||||
void roleTest(String modelName) {
|
||||
UserMessage userMessage = new UserMessage(
|
||||
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
|
||||
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(this.systemResource);
|
||||
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate"));
|
||||
Prompt prompt = new Prompt(List.of(userMessage, systemMessage),
|
||||
FunctionCallingOptions.builder().withModel(modelName).build());
|
||||
ChatResponse response = this.chatModel.call(prompt);
|
||||
assertThat(response.getResults()).hasSize(1);
|
||||
assertThat(response.getMetadata().getUsage().getGenerationTokens()).isGreaterThan(0);
|
||||
assertThat(response.getMetadata().getUsage().getPromptTokens()).isGreaterThan(0);
|
||||
assertThat(response.getMetadata().getUsage().getTotalTokens())
|
||||
.isEqualTo(response.getMetadata().getUsage().getPromptTokens()
|
||||
+ response.getMetadata().getUsage().getGenerationTokens());
|
||||
Generation generation = response.getResults().get(0);
|
||||
assertThat(generation.getOutput().getContent()).contains("Blackbeard");
|
||||
assertThat(generation.getMetadata().getFinishReason()).isEqualTo("end_turn");
|
||||
logger.info(response.toString());
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled
|
||||
void testMessageHistory() {
|
||||
UserMessage userMessage = new UserMessage(
|
||||
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
|
||||
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(this.systemResource);
|
||||
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate"));
|
||||
Prompt prompt = new Prompt(List.of(userMessage, systemMessage));
|
||||
|
||||
ChatResponse response = this.chatModel.call(prompt);
|
||||
assertThat(response.getResult().getOutput().getContent()).containsAnyOf("Blackbeard", "Bartholomew");
|
||||
|
||||
var promptWithMessageHistory = new Prompt(List.of(new UserMessage("Dummy"), response.getResult().getOutput(),
|
||||
new UserMessage("Repeat the last assistant message.")));
|
||||
|
||||
response = this.chatModel.call(promptWithMessageHistory);
|
||||
|
||||
assertThat(response.getResult().getOutput().getContent()).containsAnyOf("Blackbeard", "Bartholomew");
|
||||
}
|
||||
|
||||
@Test
|
||||
void streamingWithTokenUsage() {
|
||||
var promptOptions = FunctionCallingOptions.builder().withTemperature(0.0).build();
|
||||
|
||||
var prompt = new Prompt("List two colors of the Polish flag. Be brief.", promptOptions);
|
||||
var streamingTokenUsage = this.chatModel.stream(prompt).blockLast().getMetadata().getUsage();
|
||||
var referenceTokenUsage = this.chatModel.call(prompt).getMetadata().getUsage();
|
||||
|
||||
assertThat(streamingTokenUsage.getPromptTokens()).isGreaterThan(0);
|
||||
assertThat(streamingTokenUsage.getGenerationTokens()).isGreaterThan(0);
|
||||
assertThat(streamingTokenUsage.getTotalTokens()).isGreaterThan(0);
|
||||
|
||||
assertThat(streamingTokenUsage.getPromptTokens()).isEqualTo(referenceTokenUsage.getPromptTokens());
|
||||
assertThat(streamingTokenUsage.getGenerationTokens()).isEqualTo(referenceTokenUsage.getGenerationTokens());
|
||||
assertThat(streamingTokenUsage.getTotalTokens()).isEqualTo(referenceTokenUsage.getTotalTokens());
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
void listOutputConverter() {
|
||||
DefaultConversionService conversionService = new DefaultConversionService();
|
||||
ListOutputConverter listOutputConverter = new ListOutputConverter(conversionService);
|
||||
|
||||
String format = listOutputConverter.getFormat();
|
||||
String template = """
|
||||
List five {subject}
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template,
|
||||
Map.of("subject", "ice cream flavors", "format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
Generation generation = this.chatModel.call(prompt).getResult();
|
||||
|
||||
List<String> list = listOutputConverter.convert(generation.getOutput().getContent());
|
||||
assertThat(list).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void mapOutputConverter() {
|
||||
MapOutputConverter mapOutputConverter = new MapOutputConverter();
|
||||
|
||||
String format = mapOutputConverter.getFormat();
|
||||
String template = """
|
||||
Provide me a List of {subject}
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template,
|
||||
Map.of("subject", "an array of numbers from 1 to 9 under they key name 'numbers'", "format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
Generation generation = this.chatModel.call(prompt).getResult();
|
||||
|
||||
Map<String, Object> result = mapOutputConverter.convert(generation.getOutput().getContent());
|
||||
assertThat(result.get("numbers")).isEqualTo(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanOutputConverterRecords() {
|
||||
|
||||
BeanOutputConverter<ActorsFilmsRecord> beanOutputConverter = new BeanOutputConverter<>(ActorsFilmsRecord.class);
|
||||
|
||||
String format = beanOutputConverter.getFormat();
|
||||
String template = """
|
||||
Generate the filmography of 5 movies for Tom Hanks.
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
Generation generation = this.chatModel.call(prompt).getResult();
|
||||
|
||||
ActorsFilmsRecord actorsFilms = beanOutputConverter.convert(generation.getOutput().getContent());
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
|
||||
assertThat(actorsFilms.movies()).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanStreamOutputConverterRecords() {
|
||||
|
||||
BeanOutputConverter<ActorsFilmsRecord> beanOutputConverter = new BeanOutputConverter<>(ActorsFilmsRecord.class);
|
||||
|
||||
String format = beanOutputConverter.getFormat();
|
||||
String template = """
|
||||
Generate the filmography of 5 movies for Tom Hanks.
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
|
||||
String generationTextFromStream = this.streamingChatModel.stream(prompt)
|
||||
.collectList()
|
||||
.block()
|
||||
.stream()
|
||||
.map(ChatResponse::getResults)
|
||||
.flatMap(List::stream)
|
||||
.map(Generation::getOutput)
|
||||
.map(AssistantMessage::getContent)
|
||||
.collect(Collectors.joining());
|
||||
|
||||
ActorsFilmsRecord actorsFilms = beanOutputConverter.convert(generationTextFromStream);
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
|
||||
assertThat(actorsFilms.movies()).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void multiModalityTest() throws IOException {
|
||||
|
||||
var imageData = new ClassPathResource("/test.png");
|
||||
|
||||
var userMessage = new UserMessage("Explain what do you see on this picture?",
|
||||
List.of(new Media(MimeTypeUtils.IMAGE_PNG, imageData)));
|
||||
|
||||
var response = this.chatModel.call(new Prompt(List.of(userMessage)));
|
||||
|
||||
logger.info(response.getResult().getOutput().getContent());
|
||||
assertThat(response.getResult().getOutput().getContent()).contains("banan", "apple", "basket");
|
||||
}
|
||||
|
||||
@Test
|
||||
void functionCallTest() {
|
||||
|
||||
UserMessage userMessage = new UserMessage(
|
||||
"What's the weather like in San Francisco, Tokyo and Paris? Return the result in Celsius.");
|
||||
|
||||
List<Message> messages = new ArrayList<>(List.of(userMessage));
|
||||
|
||||
var promptOptions = FunctionCallingOptions.builder()
|
||||
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
|
||||
.withName("getCurrentWeather")
|
||||
.withDescription(
|
||||
"Get the weather in location. Return temperature in 36°F or 36°C format. Use multi-turn if needed.")
|
||||
.build()))
|
||||
.build();
|
||||
|
||||
ChatResponse response = this.chatModel.call(new Prompt(messages, promptOptions));
|
||||
|
||||
logger.info("Response: {}", response);
|
||||
|
||||
Generation generation = response.getResult();
|
||||
assertThat(generation.getOutput().getContent()).contains("30", "10", "15");
|
||||
}
|
||||
|
||||
@Test
|
||||
void streamFunctionCallTest() {
|
||||
|
||||
UserMessage userMessage = new UserMessage(
|
||||
// "What's the weather like in San Francisco? Return the result in
|
||||
// Celsius.");
|
||||
"What's the weather like in San Francisco, Tokyo and Paris? Return the result in Celsius.");
|
||||
|
||||
List<Message> messages = new ArrayList<>(List.of(userMessage));
|
||||
|
||||
var promptOptions = FunctionCallingOptions.builder()
|
||||
.withModel("anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
|
||||
.withName("getCurrentWeather")
|
||||
.withDescription(
|
||||
"Get the weather in location. Return temperature in 36°F or 36°C format. Use multi-turn if needed.")
|
||||
.build()))
|
||||
.build();
|
||||
|
||||
Flux<ChatResponse> response = this.chatModel.stream(new Prompt(messages, promptOptions));
|
||||
|
||||
String content = response.collectList()
|
||||
.block()
|
||||
.stream()
|
||||
.filter(cr -> cr.getResult() != null)
|
||||
.map(cr -> cr.getResult().getOutput().getContent())
|
||||
.collect(Collectors.joining());
|
||||
|
||||
logger.info("Response: {}", content);
|
||||
assertThat(content).contains("30", "10", "15");
|
||||
}
|
||||
|
||||
@Test
|
||||
void validateCallResponseMetadata() {
|
||||
String model = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
// @formatter:off
|
||||
ChatResponse response = ChatClient.create(this.chatModel).prompt()
|
||||
.options(FunctionCallingOptions.builder().withModel(model).build())
|
||||
.user("Tell me about 3 famous pirates from the Golden Age of Piracy and what they did")
|
||||
.call()
|
||||
.chatResponse();
|
||||
// @formatter:on
|
||||
|
||||
logger.info(response.toString());
|
||||
validateChatResponseMetadata(response, model);
|
||||
}
|
||||
|
||||
@Test
|
||||
void validateStreamCallResponseMetadata() {
|
||||
String model = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
// @formatter:off
|
||||
ChatResponse response = ChatClient.create(this.chatModel).prompt()
|
||||
.options(FunctionCallingOptions.builder().withModel(model).build())
|
||||
.user("Tell me about 3 famous pirates from the Golden Age of Piracy and what they did")
|
||||
.stream()
|
||||
.chatResponse()
|
||||
.blockLast();
|
||||
// @formatter:on
|
||||
|
||||
logger.info(response.toString());
|
||||
validateChatResponseMetadata(response, model);
|
||||
}
|
||||
|
||||
record ActorsFilmsRecord(String actor, List<String> movies) {
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
/*
|
||||
* 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.converse;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import io.micrometer.observation.ObservationRegistry;
|
||||
import io.micrometer.observation.tck.TestObservationRegistry;
|
||||
import io.micrometer.observation.tck.TestObservationRegistryAssert;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation.HighCardinalityKeyNames;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation.LowCardinalityKeyNames;
|
||||
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
|
||||
import org.springframework.ai.observation.conventions.AiOperationType;
|
||||
import org.springframework.ai.observation.conventions.AiProvider;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.SpringBootConfiguration;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
|
||||
/**
|
||||
* Integration tests for observation instrumentation in {@link BedrockProxyChatModel}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@SpringBootTest(classes = BedrockProxyChatModelObservationIT.Config.class,
|
||||
properties = "spring.ai.retry.on-http-codes=429")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockProxyChatModelObservationIT {
|
||||
|
||||
@Autowired
|
||||
TestObservationRegistry observationRegistry;
|
||||
|
||||
@Autowired
|
||||
BedrockProxyChatModel chatModel;
|
||||
|
||||
@BeforeEach
|
||||
void beforeEach() {
|
||||
this.observationRegistry.clear();
|
||||
}
|
||||
|
||||
@Test
|
||||
void observationForChatOperation() {
|
||||
var options = PortableFunctionCallingOptions.builder()
|
||||
.withModel("anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.withMaxTokens(2048)
|
||||
.withStopSequences(List.of("this-is-the-end"))
|
||||
.withTemperature(0.7)
|
||||
// .withTopK(1)
|
||||
.withTopP(1.0)
|
||||
.build();
|
||||
|
||||
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
|
||||
|
||||
ChatResponse chatResponse = this.chatModel.call(prompt);
|
||||
assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty();
|
||||
|
||||
ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
|
||||
assertThat(responseMetadata).isNotNull();
|
||||
|
||||
validate(responseMetadata, "[\"end_turn\"]");
|
||||
}
|
||||
|
||||
@Test
|
||||
void observationForStreamingChatOperation() {
|
||||
var options = PortableFunctionCallingOptions.builder()
|
||||
.withModel("anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.withMaxTokens(2048)
|
||||
.withStopSequences(List.of("this-is-the-end"))
|
||||
.withTemperature(0.7)
|
||||
.withTopP(1.0)
|
||||
.build();
|
||||
|
||||
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
|
||||
|
||||
Flux<ChatResponse> chatResponseFlux = this.chatModel.stream(prompt);
|
||||
|
||||
List<ChatResponse> responses = chatResponseFlux.collectList().block();
|
||||
assertThat(responses).isNotEmpty();
|
||||
assertThat(responses).hasSizeGreaterThan(3);
|
||||
|
||||
String aggregatedResponse = responses.subList(0, responses.size() - 1)
|
||||
.stream()
|
||||
.filter(r -> r.getResult() != null)
|
||||
.map(r -> r.getResult().getOutput().getContent())
|
||||
.collect(Collectors.joining());
|
||||
assertThat(aggregatedResponse).isNotEmpty();
|
||||
|
||||
ChatResponse lastChatResponse = responses.get(responses.size() - 1);
|
||||
|
||||
ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata();
|
||||
assertThat(responseMetadata).isNotNull();
|
||||
|
||||
validate(responseMetadata, "[\"end_turn\"]");
|
||||
}
|
||||
|
||||
private void validate(ChatResponseMetadata responseMetadata, String finishReasons) {
|
||||
TestObservationRegistryAssert.assertThat(this.observationRegistry)
|
||||
.doesNotHaveAnyRemainingCurrentObservation()
|
||||
.hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME)
|
||||
.that()
|
||||
.hasContextualNameEqualTo("chat " + "anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
|
||||
AiOperationType.CHAT.value())
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(),
|
||||
AiProvider.BEDROCK_CONVERSE.value())
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(),
|
||||
"anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
// .hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(),
|
||||
// responseMetadata.getModel())
|
||||
.doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.REQUEST_FREQUENCY_PENALTY.asString())
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
|
||||
.doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.REQUEST_PRESENCE_PENALTY.asString())
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(),
|
||||
"[\"this-is-the-end\"]")
|
||||
// .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(),
|
||||
// "0.7")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0")
|
||||
// .hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_ID.asString(),
|
||||
// responseMetadata.getId())
|
||||
// .hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(),
|
||||
// finishReasons)
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
|
||||
String.valueOf(responseMetadata.getUsage().getPromptTokens()))
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(),
|
||||
String.valueOf(responseMetadata.getUsage().getGenerationTokens()))
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),
|
||||
String.valueOf(responseMetadata.getUsage().getTotalTokens()))
|
||||
.hasBeenStarted()
|
||||
.hasBeenStopped();
|
||||
}
|
||||
|
||||
@SpringBootConfiguration
|
||||
static class Config {
|
||||
|
||||
@Bean
|
||||
public TestObservationRegistry observationRegistry() {
|
||||
return TestObservationRegistry.create();
|
||||
}
|
||||
|
||||
@Bean
|
||||
public BedrockProxyChatModel bedrockConverseChatModel(ObservationRegistry observationRegistry) {
|
||||
|
||||
String modelId = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
|
||||
return BedrockProxyChatModel.builder()
|
||||
.withCredentialsProvider(EnvironmentVariableCredentialsProvider.create())
|
||||
.withRegion(Region.US_EAST_1)
|
||||
.withObservationRegistry(observationRegistry)
|
||||
.withDefaultOptions(FunctionCallingOptions.builder().withModel(modelId).build())
|
||||
.build();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
/*
|
||||
* 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.bedrock.converse;
|
||||
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonClassDescription;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude.Include;
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
public class MockWeatherService implements Function<MockWeatherService.Request, MockWeatherService.Response> {
|
||||
|
||||
@Override
|
||||
public Response apply(Request request) {
|
||||
|
||||
double temperature = 0;
|
||||
if (request.location().contains("Paris")) {
|
||||
temperature = 15;
|
||||
}
|
||||
else if (request.location().contains("Tokyo")) {
|
||||
temperature = 10;
|
||||
}
|
||||
else if (request.location().contains("San Francisco")) {
|
||||
temperature = 30;
|
||||
}
|
||||
|
||||
return new Response(temperature, Unit.C);
|
||||
}
|
||||
|
||||
/**
|
||||
* Temperature units.
|
||||
*/
|
||||
public enum Unit {
|
||||
|
||||
/**
|
||||
* Celsius.
|
||||
*/
|
||||
C("metric"),
|
||||
/**
|
||||
* Fahrenheit.
|
||||
*/
|
||||
F("imperial");
|
||||
|
||||
/**
|
||||
* Human readable unit name.
|
||||
*/
|
||||
public final String unitName;
|
||||
|
||||
private Unit(String text) {
|
||||
this.unitName = text;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Weather Function request.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
@JsonClassDescription("Weather API request")
|
||||
public record Request(@JsonProperty(required = true,
|
||||
value = "location") @JsonPropertyDescription("The city and state e.g. San Francisco, CA") String location,
|
||||
@JsonProperty(required = true, value = "unit") @JsonPropertyDescription("Temperature unit") Unit unit) {
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Weather Function response.
|
||||
*/
|
||||
public record Response(double temp, Unit unit) {
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
/*
|
||||
* 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.bedrock.converse;
|
||||
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
|
||||
import java.lang.annotation.ElementType;
|
||||
import java.lang.annotation.Retention;
|
||||
import java.lang.annotation.RetentionPolicy;
|
||||
import java.lang.annotation.Target;
|
||||
|
||||
@Target({ ElementType.TYPE, ElementType.METHOD })
|
||||
@Retention(RetentionPolicy.RUNTIME)
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SESSION_TOKEN", matches = ".*")
|
||||
public @interface RequiresAwsCredentials {
|
||||
|
||||
// You can add custom properties here if needed
|
||||
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* 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.bedrock.converse.experiements;
|
||||
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
|
||||
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
/**
|
||||
* Used for reverse engineering the protocol.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
|
||||
public class BedrockConverseChatModelMain {
|
||||
|
||||
public static void main(String[] args) {
|
||||
|
||||
// String modelId = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
String modelId = "ai21.jamba-1-5-large-v1:0";
|
||||
var prompt = new Prompt("Tell me a joke?", ChatOptionsBuilder.builder().withModel(modelId).build());
|
||||
|
||||
var chatModel = BedrockProxyChatModel.builder()
|
||||
.withCredentialsProvider(EnvironmentVariableCredentialsProvider.create())
|
||||
.withRegion(Region.US_EAST_1)
|
||||
.build();
|
||||
|
||||
var chatResponse = chatModel.call(prompt);
|
||||
System.out.println(chatResponse);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
/*
|
||||
* 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.bedrock.converse.experiements;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.bedrock.converse.MockWeatherService;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.model.function.FunctionCallbackWrapper;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamOutput;
|
||||
|
||||
/**
|
||||
* Used for reverse engineering the protocol
|
||||
*/
|
||||
public class BedrockConverseChatModelMain2 {
|
||||
|
||||
public static void main(String[] args) {
|
||||
|
||||
// String modelId = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
// String modelId = "ai21.jamba-1-5-large-v1:0";
|
||||
String modelId = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
|
||||
// var prompt = new Prompt("Tell me a joke?",
|
||||
// ChatOptionsBuilder.builder().withModel(modelId).build());
|
||||
var prompt = new Prompt(
|
||||
// "What's the weather like in San Francisco, Tokyo, and Paris? Return the
|
||||
// temperature in Celsius.",
|
||||
"What's the weather like in Paris? Return the temperature in Celsius.",
|
||||
PortableFunctionCallingOptions.builder()
|
||||
.withModel(modelId)
|
||||
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
|
||||
.withName("getCurrentWeather")
|
||||
.withDescription("Get the weather in location")
|
||||
.build()))
|
||||
.build());
|
||||
|
||||
BedrockProxyChatModel chatModel = BedrockProxyChatModel.builder()
|
||||
.withCredentialsProvider(EnvironmentVariableCredentialsProvider.create())
|
||||
.withRegion(Region.US_EAST_1)
|
||||
.build();
|
||||
|
||||
var streamRequest = chatModel.createStreamRequest(prompt);
|
||||
|
||||
Flux<ConverseStreamOutput> responses = chatModel.converseStream(streamRequest);
|
||||
List<ConverseStreamOutput> responseList = responses.collectList().block();
|
||||
System.out.println(responseList);
|
||||
|
||||
// Flux<ChatResponse> responses2 = ConverseApiUtils.toChatResponse(responses);
|
||||
// List<ChatResponse> responseList2 = responses2.collectList().block();
|
||||
// System.out.println(responseList2);
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
"You are a helpful AI assistant. Your name is {name}.
|
||||
You are an AI assistant that helps people find information.
|
||||
Your name is {name}
|
||||
You should reply to the user's request with your name and also in the style of a {voice}.
|
||||
BIN
models/spring-ai-bedrock-converse/src/test/resources/test.png
Normal file
BIN
models/spring-ai-bedrock-converse/src/test/resources/test.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 164 KiB |
@@ -34,7 +34,7 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
public class BedrockAnthropic3CreateRequestTests {
|
||||
|
||||
private Anthropic3ChatBedrockApi anthropicChatApi = new Anthropic3ChatBedrockApi(AnthropicChatModel.CLAUDE_V2.id(),
|
||||
Region.EU_CENTRAL_1.id(), Duration.ofMillis(1000L));
|
||||
Region.US_EAST_1.id(), Duration.ofMillis(1000L));
|
||||
|
||||
@Test
|
||||
public void createRequestWithChatOptions() {
|
||||
|
||||
@@ -42,7 +42,7 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
public class TitanChatBedrockApiIT {
|
||||
|
||||
TitanChatBedrockApi titanBedrockApi = new TitanChatBedrockApi(TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
|
||||
EnvironmentVariableCredentialsProvider.create(), Region.EU_CENTRAL_1.id(), ModelOptionsUtils.OBJECT_MAPPER,
|
||||
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), ModelOptionsUtils.OBJECT_MAPPER,
|
||||
Duration.ofMinutes(2));
|
||||
|
||||
TitanChatRequest titanChatRequest = TitanChatRequest.builder("Give me the names of 3 famous pirates?")
|
||||
|
||||
4
pom.xml
4
pom.xml
@@ -82,6 +82,7 @@
|
||||
<module>models/spring-ai-anthropic</module>
|
||||
<module>models/spring-ai-azure-openai</module>
|
||||
<module>models/spring-ai-bedrock</module>
|
||||
<module>models/spring-ai-bedrock-converse</module>
|
||||
<module>models/spring-ai-huggingface</module>
|
||||
<module>models/spring-ai-minimax</module>
|
||||
<module>models/spring-ai-mistral-ai</module>
|
||||
@@ -102,6 +103,7 @@
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-aws-opensearch-store</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-azure-openai</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-bedrock-converse</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-huggingface</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-minimax</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-mistral-ai</module>
|
||||
@@ -201,6 +203,7 @@
|
||||
<postgresql.version>42.7.2</postgresql.version>
|
||||
<elasticsearch-java.version>8.13.3</elasticsearch-java.version>
|
||||
<milvus.version>2.3.5</milvus.version>
|
||||
<oracle.free.version>1.19.8</oracle.free.version>
|
||||
<gemfire.testcontainers.version>2.3.0</gemfire.testcontainers.version>
|
||||
<pinecone.version>0.8.0</pinecone.version>
|
||||
<fastjson.version>2.0.46</fastjson.version>
|
||||
@@ -297,6 +300,7 @@
|
||||
checkstyle.build.directory=${project.build.directory}
|
||||
checkstyle.suppressions.file=${project.basedir}/src/checkstyle/checkstyle-suppressions.xml
|
||||
checkstyle.additional.suppressions.file=${project.basedir}/src/checkstyle/checkstyle-suppressions.xml
|
||||
checkstyle.header.file=${project.basedir}/src/checkstyle/checkstyle-header.txt
|
||||
</propertyExpansion>
|
||||
<consoleOutput>true</consoleOutput>
|
||||
<failsOnError>${maven-checkstyle-plugin.failsOnError}
|
||||
|
||||
1
response.json
Normal file
1
response.json
Normal file
@@ -0,0 +1 @@
|
||||
{"id":"msg_bdrk_01Q7dCYZucVfVy6vJej5f4Wo","type":"message","role":"assistant","model":"claude-3-5-sonnet-20240620","content":[{"type":"text","text":"Hello! How can I assist you today? Feel free to ask any questions or let me know if you need help with anything."}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":8,"output_tokens":29}}
|
||||
@@ -86,6 +86,12 @@
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bedrock-converse</artifactId>
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-huggingface</artifactId>
|
||||
@@ -323,6 +329,12 @@
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bedrock-converse-spring-boot-starter</artifactId>
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-chroma-store-spring-boot-starter</artifactId>
|
||||
|
||||
@@ -26,6 +26,8 @@ import java.util.Set;
|
||||
|
||||
import org.springframework.ai.chat.prompt.ChatOptions;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.CollectionUtils;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
/**
|
||||
* Builder for {@link FunctionCallingOptions}. Using the {@link FunctionCallingOptions}
|
||||
@@ -291,6 +293,69 @@ public class FunctionCallingOptionsBuilder {
|
||||
.build();
|
||||
}
|
||||
|
||||
public PortableFunctionCallingOptions merge(FunctionCallingOptions options) {
|
||||
|
||||
var builder = PortableFunctionCallingOptions.builder()
|
||||
.withModel(StringUtils.hasText(options.getModel()) ? options.getModel() : this.model)
|
||||
.withFrequencyPenalty(
|
||||
options.getFrequencyPenalty() != null ? options.getFrequencyPenalty() : this.frequencyPenalty)
|
||||
.withMaxTokens(options.getMaxTokens() != null ? options.getMaxTokens() : this.maxTokens)
|
||||
.withPresencePenalty(
|
||||
options.getPresencePenalty() != null ? options.getPresencePenalty() : this.presencePenalty)
|
||||
.withStopSequences(options.getStopSequences() != null ? options.getStopSequences() : this.stopSequences)
|
||||
.withTemperature(options.getTemperature() != null ? options.getTemperature() : this.temperature)
|
||||
.withTopK(options.getTopK() != null ? options.getTopK() : this.topK)
|
||||
.withTopP(options.getTopP() != null ? options.getTopP() : this.topP)
|
||||
.withProxyToolCalls(
|
||||
options.getProxyToolCalls() != null ? options.getProxyToolCalls() : this.proxyToolCalls);
|
||||
|
||||
Set<String> functions = new HashSet<>();
|
||||
if (!CollectionUtils.isEmpty(this.functions)) {
|
||||
functions.addAll(this.functions);
|
||||
}
|
||||
if (!CollectionUtils.isEmpty(options.getFunctions())) {
|
||||
functions.addAll(options.getFunctions());
|
||||
}
|
||||
builder.withFunctions(functions);
|
||||
|
||||
List<FunctionCallback> functionCallbacks = new ArrayList<>();
|
||||
if (!CollectionUtils.isEmpty(this.functionCallbacks)) {
|
||||
functionCallbacks.addAll(this.functionCallbacks);
|
||||
}
|
||||
if (!CollectionUtils.isEmpty(options.getFunctionCallbacks())) {
|
||||
functionCallbacks.addAll(options.getFunctionCallbacks());
|
||||
}
|
||||
builder.withFunctionCallbacks(functionCallbacks);
|
||||
|
||||
Map<String, Object> context = new HashMap<>();
|
||||
if (!CollectionUtils.isEmpty(this.context)) {
|
||||
context.putAll(this.context);
|
||||
}
|
||||
if (!CollectionUtils.isEmpty(options.getToolContext())) {
|
||||
context.putAll(options.getToolContext());
|
||||
}
|
||||
builder.withToolContext(context);
|
||||
|
||||
return builder.build();
|
||||
}
|
||||
|
||||
public PortableFunctionCallingOptions merge(ChatOptions options) {
|
||||
|
||||
var builder = PortableFunctionCallingOptions.builder()
|
||||
.withModel(StringUtils.hasText(options.getModel()) ? options.getModel() : this.model)
|
||||
.withFrequencyPenalty(
|
||||
options.getFrequencyPenalty() != null ? options.getFrequencyPenalty() : this.frequencyPenalty)
|
||||
.withMaxTokens(options.getMaxTokens() != null ? options.getMaxTokens() : this.maxTokens)
|
||||
.withPresencePenalty(
|
||||
options.getPresencePenalty() != null ? options.getPresencePenalty() : this.presencePenalty)
|
||||
.withStopSequences(options.getStopSequences() != null ? options.getStopSequences() : this.stopSequences)
|
||||
.withTemperature(options.getTemperature() != null ? options.getTemperature() : this.temperature)
|
||||
.withTopK(options.getTopK() != null ? options.getTopK() : this.topK)
|
||||
.withTopP(options.getTopP() != null ? options.getTopP() : this.topP);
|
||||
|
||||
return builder.build();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -43,6 +43,7 @@ public enum AiProvider {
|
||||
ZHIPUAI("zhipuai"),
|
||||
SPRING_AI("spring_ai"),
|
||||
VERTEX_AI("vertex_ai"),
|
||||
BEDROCK_CONVERSE("bedrock_converse"),
|
||||
ONNX("onnx");
|
||||
|
||||
private final String value;
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
* xref:api/index.adoc[AI Models]
|
||||
** xref:api/chatmodel.adoc[Chat Models]
|
||||
*** xref:api/chat/comparison.adoc[Chat Models Comparison]
|
||||
*** xref:api/bedrock-converse.adoc[Amazon Bedrock Converse]
|
||||
*** xref:api/bedrock-chat.adoc[Amazon Bedrock]
|
||||
**** xref:api/chat/bedrock/bedrock-anthropic3.adoc[Anthropic3]
|
||||
**** xref:api/chat/bedrock/bedrock-anthropic.adoc[Anthropic2]
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
= Bedrock Converse API
|
||||
|
||||
link:https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html[Amazon Bedrock] Converse API provides a unified interface for conversational AI models with enhanced capabilities including function/tool calling, multimodal inputs, and streaming responses.
|
||||
|
||||
The Bedrock Converse API has the following high-level features:
|
||||
|
||||
* Tool/Function Calling: Support for function definitions and tool use during conversations
|
||||
* Multimodal Input: Ability to process both text and image inputs in conversations
|
||||
* Streaming Support: Real-time streaming of model responses
|
||||
* System Messages: Support for system-level instructions and context setting
|
||||
* Metrics Integration: Built-in support for observation and metrics tracking
|
||||
|
||||
The https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html[Amazon Bedrock User Guide] contains detailed information on how to use the AWS hosted service.
|
||||
|
||||
TIP: The Bedrock Converse API provides a unified interface across multiple model providers while handling AWS-specific authentication and infrastructure concerns.
|
||||
|
||||
== Prerequisites
|
||||
|
||||
Refer to the xref:api/bedrock.adoc[Spring AI documentation on Amazon Bedrock] for setting up API access.
|
||||
|
||||
* Obtain AWS credentials: If you don't have an AWS account and AWS CLI configured yet, this video guide can help you configure it: link:https://youtu.be/gswVHTrRX8I?si=buaY7aeI0l3-bBVb[AWS CLI & SDK Setup in Less Than 4 Minutes!]. You should be able to obtain your access and security keys.
|
||||
|
||||
* Enable the Models to use: Go to link:https://us-east-1.console.aws.amazon.com/bedrock/home[Amazon Bedrock] and from the link:https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess[Model Access] menu on the left, configure access to the models you are going to use.
|
||||
|
||||
|
||||
== Auto-configuration
|
||||
|
||||
Add the `spring-ai-bedrock-converse-spring-boot-starter` dependency to your project's Maven `pom.xml` or Gradle `build.gradle` build files:
|
||||
|
||||
[tabs]
|
||||
======
|
||||
Maven::
|
||||
+
|
||||
[source,xml]
|
||||
----
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bedrock-converse-spring-boot-starter</artifactId>
|
||||
</dependency>
|
||||
----
|
||||
|
||||
Gradle::
|
||||
+
|
||||
[source,gradle]
|
||||
----
|
||||
dependencies {
|
||||
implementation 'org.springframework.ai:spring-ai-bedrock-converse-spring-boot-starter'
|
||||
}
|
||||
----
|
||||
======
|
||||
|
||||
TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Management] section to add the Spring AI BOM to your build file.
|
||||
|
||||
|
||||
=== Chat Properties
|
||||
|
||||
The prefix `spring.ai.bedrock.aws` is the property prefix to configure the connection to AWS Bedrock.
|
||||
|
||||
[cols="3,3,1", stripes=even]
|
||||
|====
|
||||
| Property | Description | Default
|
||||
|
||||
| spring.ai.bedrock.aws.region | AWS region to use. | us-east-1
|
||||
| spring.ai.bedrock.aws.timeout | AWS timeout to use. | 5m
|
||||
| spring.ai.bedrock.aws.access-key | AWS access key. | -
|
||||
| spring.ai.bedrock.aws.secret-key | AWS secret key. | -
|
||||
| spring.ai.bedrock.aws.session-token | AWS session token for temporary credentials. | -
|
||||
|====
|
||||
|
||||
The prefix `spring.ai.bedrock.converse.chat` is the property prefix that configures the chat model implementation for the Converse API.
|
||||
|
||||
[cols="3,5,1", stripes=even]
|
||||
|====
|
||||
| Property | Description | Default
|
||||
|
||||
| spring.ai.bedrock.converse.chat.enabled | Enable Bedrock Converse chat model. | true
|
||||
| spring.ai.bedrock.converse.chat.options.model | The model ID to use. You can use the https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html[Supported models and model features] | anthropic.claude-3-sonnet-20240229-v1:0
|
||||
| spring.ai.bedrock.converse.chat.options.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.8
|
||||
| spring.ai.bedrock.converse.chat.options.top-p | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default
|
||||
| spring.ai.bedrock.converse.chat.options.top-k | Number of token choices for generating the next token. | AWS Bedrock default
|
||||
| spring.ai.bedrock.converse.chat.options.max-tokens | Maximum number of tokens in the generated response. | 500
|
||||
|====
|
||||
|
||||
== Runtime Options [[chat-options]]
|
||||
|
||||
Use the portable `ChatOptions` or `FunctionCallingOptions` portable builders to create model configurations, such as temperature, maxToken, topP, etc.
|
||||
|
||||
On start-up, the default options can be configured with the `BedrockConverseProxyChatModel(api, options)` constructor or the `spring.ai.bedrock.converse.chat.options.*` properties.
|
||||
|
||||
At run-time you can override the default options by adding new, request specific, options to the `Prompt` call:
|
||||
|
||||
[source,java]
|
||||
----
|
||||
var options = FunctionCallingOptions.builder()
|
||||
.withModel("anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.withTemperature(0.6)
|
||||
.withMaxTokens(300)
|
||||
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new WeatherService())
|
||||
.withName("getCurrentWeather")
|
||||
.withDescription("Get the weather in location. Return temperature in 36°F or 36°C format. Use multi-turn if needed.")
|
||||
.build()))
|
||||
.build();
|
||||
|
||||
ChatResponse response = chatModel.call(new Prompt("What is current weather in Amsterdam?", options));
|
||||
----
|
||||
|
||||
== Tool/Function Calling
|
||||
|
||||
The Bedrock Converse API supports function calling capabilities, allowing models to use tools during conversations. Here's an example of how to define and use functions:
|
||||
|
||||
[source,java]
|
||||
----
|
||||
@Bean
|
||||
@Description("Get the weather in location. Return temperature in 36°F or 36°C format.")
|
||||
public Function<Request, Response> weatherFunction() {
|
||||
return new MockWeatherService();
|
||||
}
|
||||
|
||||
String response = ChatClient.create(this.chatModel)
|
||||
.prompt("What's the weather like in Boston?")
|
||||
.function("weatherFunction")
|
||||
.call()
|
||||
.content();
|
||||
----
|
||||
|
||||
== Sample Controller
|
||||
|
||||
Create a new Spring Boot project and add the `spring-ai-bedrock-converse-spring-boot-starter` to your dependencies.
|
||||
|
||||
Add an `application.properties` file under `src/main/resources`:
|
||||
|
||||
[source,properties]
|
||||
----
|
||||
spring.ai.bedrock.aws.region=eu-central-1
|
||||
spring.ai.bedrock.aws.timeout=10m
|
||||
spring.ai.bedrock.aws.access-key=${AWS_ACCESS_KEY_ID}
|
||||
spring.ai.bedrock.aws.secret-key=${AWS_SECRET_ACCESS_KEY}
|
||||
|
||||
spring.ai.bedrock.converse.chat.options.temperature=0.8
|
||||
spring.ai.bedrock.converse.chat.options.top-k=15
|
||||
----
|
||||
|
||||
Here's an example controller using the chat model:
|
||||
|
||||
[source,java]
|
||||
----
|
||||
@RestController
|
||||
public class ChatController {
|
||||
|
||||
private final ChatClient chatClient;
|
||||
|
||||
@Autowired
|
||||
public ChatController(ChatClient.Builder builder) {
|
||||
this.chatClient = builder.build();
|
||||
}
|
||||
|
||||
@GetMapping("/ai/generate")
|
||||
public Map generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
|
||||
return Map.of("generation", this.chatClient.prompt(message).call().content());
|
||||
}
|
||||
|
||||
@GetMapping("/ai/generateStream")
|
||||
public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
|
||||
return this.chatClient.prompt(message).stream().content();
|
||||
}
|
||||
}
|
||||
----
|
||||
|
||||
@@ -198,7 +198,7 @@ Next, create an https://github.com/spring-projects/spring-ai/blob/main/models/sp
|
||||
AnthropicChatBedrockApi anthropicApi = new AnthropicChatBedrockApi(
|
||||
AnthropicChatBedrockApi.AnthropicModel.CLAUDE_V2.id(),
|
||||
EnvironmentVariableCredentialsProvider.create(),
|
||||
Region.EU_CENTRAL_1.id(),
|
||||
Region.US_EAST_1.id(),
|
||||
new ObjectMapper(),
|
||||
Duration.ofMillis(1000L));
|
||||
|
||||
@@ -234,7 +234,7 @@ Here is a simple snippet how to use the api programmatically:
|
||||
[source,java]
|
||||
----
|
||||
AnthropicChatBedrockApi anthropicChatApi = new AnthropicChatBedrockApi(
|
||||
AnthropicModel.CLAUDE_V2.id(), Region.EU_CENTRAL_1.id(), Duration.ofMillis(1000L));
|
||||
AnthropicModel.CLAUDE_V2.id(), Region.US_EAST_1.id(), Duration.ofMillis(1000L));
|
||||
|
||||
AnthropicChatRequest request = AnthropicChatRequest
|
||||
.builder(String.format(AnthropicChatBedrockApi.PROMPT_TEMPLATE, "Name 3 famous pirates"))
|
||||
|
||||
@@ -271,7 +271,7 @@ Here is a simple snippet how to use the api programmatically:
|
||||
[source,java]
|
||||
----
|
||||
Anthropic3ChatBedrockApi anthropicChatApi = new Anthropic3ChatBedrockApi(
|
||||
AnthropicModel.CLAUDE_V2.id(), Region.EU_CENTRAL_1.id(), Duration.ofMillis(1000L));
|
||||
AnthropicModel.CLAUDE_V2.id(), Region.US_EAST_1.id(), Duration.ofMillis(1000L));
|
||||
|
||||
AnthropicChatRequest request = AnthropicChatRequest
|
||||
.builder(String.format(Anthropic3ChatBedrockApi.PROMPT_TEMPLATE, "Name 3 famous pirates"))
|
||||
|
||||
@@ -220,7 +220,7 @@ Here is a simple snippet how to use the api programmatically:
|
||||
[source,java]
|
||||
----
|
||||
TitanChatBedrockApi titanBedrockApi = new TitanChatBedrockApi(TitanChatCompletionModel.TITAN_TEXT_EXPRESS_V1.id(),
|
||||
Region.EU_CENTRAL_1.id(), Duration.ofMillis(1000L));
|
||||
Region.US_EAST_1.id(), Duration.ofMillis(1000L));
|
||||
|
||||
TitanChatRequest titanChatRequest = TitanChatRequest.builder("Give me the names of 3 famous pirates?")
|
||||
.withTemperature(0.5)
|
||||
|
||||
@@ -281,6 +281,14 @@
|
||||
<optional>true</optional>
|
||||
</dependency>
|
||||
|
||||
<!-- Bedrock Converse -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bedrock-converse</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
<optional>true</optional>
|
||||
</dependency>
|
||||
|
||||
<!-- Mistral AI LLM -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
@@ -447,7 +455,7 @@
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>oracle-free</artifactId>
|
||||
<version>1.19.8</version>
|
||||
<version>${oracle.free.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
@@ -463,12 +471,6 @@
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>oracle-free</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>com.redis</groupId>
|
||||
<artifactId>testcontainers-redis</artifactId>
|
||||
|
||||
@@ -18,6 +18,7 @@ package org.springframework.ai.autoconfigure.bedrock;
|
||||
|
||||
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
|
||||
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
|
||||
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
@@ -43,6 +44,12 @@ public class BedrockAwsConnectionConfiguration {
|
||||
public AwsCredentialsProvider credentialsProvider(BedrockAwsConnectionProperties properties) {
|
||||
|
||||
if (StringUtils.hasText(properties.getAccessKey()) && StringUtils.hasText(properties.getSecretKey())) {
|
||||
|
||||
if (StringUtils.hasText(properties.getSessionToken())) {
|
||||
return StaticCredentialsProvider.create(AwsSessionCredentials.create(properties.getAccessKey(),
|
||||
properties.getSecretKey(), properties.getSessionToken()));
|
||||
}
|
||||
|
||||
return StaticCredentialsProvider
|
||||
.create(AwsBasicCredentials.create(properties.getAccessKey(), properties.getSecretKey()));
|
||||
}
|
||||
@@ -61,9 +68,6 @@ public class BedrockAwsConnectionConfiguration {
|
||||
return DefaultAwsRegionProviderChain.builder().build();
|
||||
}
|
||||
|
||||
/**
|
||||
* @author Wei Jiang
|
||||
*/
|
||||
static class StaticRegionProvider implements AwsRegionProvider {
|
||||
|
||||
private final Region region;
|
||||
|
||||
@@ -46,6 +46,12 @@ public class BedrockAwsConnectionProperties {
|
||||
*/
|
||||
private String secretKey;
|
||||
|
||||
/**
|
||||
* AWS session token. (optional) When provided the AwsSessionCredentials are used.
|
||||
* Otherwise the AwsBasicCredentials are used.
|
||||
*/
|
||||
private String sessionToken;
|
||||
|
||||
/**
|
||||
* Set model timeout, Defaults 5 min.
|
||||
*/
|
||||
@@ -83,4 +89,12 @@ public class BedrockAwsConnectionProperties {
|
||||
this.timeout = timeout;
|
||||
}
|
||||
|
||||
public String getSessionToken() {
|
||||
return this.sessionToken;
|
||||
}
|
||||
|
||||
public void setSessionToken(String sessionToken) {
|
||||
this.sessionToken = sessionToken;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import io.micrometer.observation.ObservationRegistry;
|
||||
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.providers.AwsRegionProvider;
|
||||
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient;
|
||||
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionConfiguration;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
|
||||
import org.springframework.ai.model.function.FunctionCallback;
|
||||
import org.springframework.ai.model.function.FunctionCallbackContext;
|
||||
import org.springframework.beans.factory.ObjectProvider;
|
||||
import org.springframework.boot.autoconfigure.AutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
|
||||
import org.springframework.boot.context.properties.EnableConfigurationProperties;
|
||||
import org.springframework.context.ApplicationContext;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Import;
|
||||
|
||||
/**
|
||||
* {@link AutoConfiguration Auto-configuration} for Bedrock Converse Proxy Chat Client.
|
||||
*
|
||||
* Leverages the Spring Cloud AWS to resolve the {@link AwsCredentialsProvider}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Wei Jiang
|
||||
*/
|
||||
@AutoConfiguration
|
||||
@EnableConfigurationProperties({ BedrockConverseProxyChatProperties.class, BedrockAwsConnectionConfiguration.class })
|
||||
@ConditionalOnClass({ BedrockRuntimeClient.class, BedrockRuntimeAsyncClient.class })
|
||||
@ConditionalOnProperty(prefix = BedrockConverseProxyChatProperties.CONFIG_PREFIX, name = "enabled",
|
||||
havingValue = "true", matchIfMissing = true)
|
||||
@Import(BedrockAwsConnectionConfiguration.class)
|
||||
public class BedrockConverseProxyChatAutoConfiguration {
|
||||
|
||||
@Bean
|
||||
@ConditionalOnMissingBean
|
||||
@ConditionalOnBean({ AwsCredentialsProvider.class, AwsRegionProvider.class })
|
||||
public BedrockProxyChatModel bedrockProxyChatModel(AwsCredentialsProvider credentialsProvider,
|
||||
AwsRegionProvider regionProvider, BedrockAwsConnectionProperties connectionProperties,
|
||||
BedrockConverseProxyChatProperties chatProperties, FunctionCallbackContext functionCallbackContext,
|
||||
List<FunctionCallback> toolFunctionCallbacks, ObjectProvider<ObservationRegistry> observationRegistry,
|
||||
ObjectProvider<ChatModelObservationConvention> observationConvention,
|
||||
ObjectProvider<BedrockRuntimeClient> bedrockRuntimeClient,
|
||||
ObjectProvider<BedrockRuntimeAsyncClient> bedrockRuntimeAsyncClient) {
|
||||
|
||||
var chatModel = BedrockProxyChatModel.builder()
|
||||
.withCredentialsProvider(credentialsProvider)
|
||||
.withRegion(regionProvider.getRegion())
|
||||
.withTimeout(connectionProperties.getTimeout())
|
||||
.withDefaultOptions(chatProperties.getOptions())
|
||||
.withObservationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
|
||||
.withFunctionCallbackContext(functionCallbackContext)
|
||||
.withToolFunctionCallbacks(toolFunctionCallbacks)
|
||||
.withBedrockRuntimeClient(bedrockRuntimeClient.getIfAvailable())
|
||||
.withBedrockRuntimeAsyncClient(bedrockRuntimeAsyncClient.getIfAvailable())
|
||||
.build();
|
||||
|
||||
observationConvention.ifAvailable(chatModel::setObservationConvention);
|
||||
|
||||
return chatModel;
|
||||
}
|
||||
|
||||
@Bean
|
||||
@ConditionalOnMissingBean
|
||||
public FunctionCallbackContext springAiFunctionManager(ApplicationContext context) {
|
||||
FunctionCallbackContext manager = new FunctionCallbackContext();
|
||||
manager.setApplicationContext(context);
|
||||
return manager;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse;
|
||||
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.NestedConfigurationProperty;
|
||||
import org.springframework.util.Assert;
|
||||
|
||||
/**
|
||||
* Configuration properties for Bedrock Converse.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@ConfigurationProperties(BedrockConverseProxyChatProperties.CONFIG_PREFIX)
|
||||
public class BedrockConverseProxyChatProperties {
|
||||
|
||||
public static final String CONFIG_PREFIX = "spring.ai.bedrock.converse.chat";
|
||||
|
||||
/**
|
||||
* Enable Bedrock Converse chat model.
|
||||
*/
|
||||
private boolean enabled = true;
|
||||
|
||||
/**
|
||||
* The generative id to use. See the {@link BedrockProxyChatModel} for the supported
|
||||
* models.
|
||||
*/
|
||||
private String model = "anthropic.claude-3-5-sonnet-20240620-v1:0";
|
||||
|
||||
@NestedConfigurationProperty
|
||||
private PortableFunctionCallingOptions options = PortableFunctionCallingOptions.builder()
|
||||
.withTemperature(0.7)
|
||||
.withMaxTokens(300)
|
||||
.withTopK(10)
|
||||
.build();
|
||||
|
||||
public boolean isEnabled() {
|
||||
return this.enabled;
|
||||
}
|
||||
|
||||
public void setEnabled(boolean enabled) {
|
||||
this.enabled = enabled;
|
||||
}
|
||||
|
||||
public String getModel() {
|
||||
return this.model;
|
||||
}
|
||||
|
||||
public void setModel(String model) {
|
||||
this.model = model;
|
||||
}
|
||||
|
||||
public PortableFunctionCallingOptions getOptions() {
|
||||
return this.options;
|
||||
}
|
||||
|
||||
public void setOptions(PortableFunctionCallingOptions options) {
|
||||
Assert.notNull(options, "PortableFunctionCallingOptions must not be null");
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -30,6 +30,7 @@ org.springframework.ai.autoconfigure.bedrock.anthropic.BedrockAnthropicChatAutoC
|
||||
org.springframework.ai.autoconfigure.bedrock.anthropic3.BedrockAnthropic3ChatAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.bedrock.titan.BedrockTitanChatAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.bedrock.titan.BedrockTitanEmbeddingAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.bedrock.converse.BedrockConverseProxyChatAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.chat.observation.ChatObservationAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.embedding.observation.EmbeddingObservationAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.image.observation.ImageObservationAutoConfiguration
|
||||
|
||||
@@ -17,7 +17,6 @@
|
||||
package org.springframework.ai.autoconfigure.bedrock;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import software.amazon.awssdk.auth.credentials.AwsCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
@@ -27,7 +26,6 @@ import org.springframework.boot.autoconfigure.AutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
|
||||
import org.springframework.boot.context.properties.EnableConfigurationProperties;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Import;
|
||||
|
||||
@@ -35,18 +33,15 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Wei Jiang
|
||||
* @since 0.8.1
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockAwsConnectionConfigurationIT {
|
||||
|
||||
@Test
|
||||
public void autoConfigureAWSCredentialAndRegionProvider() {
|
||||
new ApplicationContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id())
|
||||
BedrockTestUtils.getContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(TestAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
var awsCredentialsProvider = context.getBean(AwsCredentialsProvider.class);
|
||||
@@ -66,10 +61,7 @@ public class BedrockAwsConnectionConfigurationIT {
|
||||
|
||||
@Test
|
||||
public void autoConfigureWithCustomAWSCredentialAndRegionProvider() {
|
||||
new ApplicationContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id())
|
||||
BedrockTestUtils.getContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(TestAutoConfiguration.class,
|
||||
CustomAwsCredentialsProviderAndAwsRegionProviderAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
|
||||
public final class BedrockTestUtils {
|
||||
|
||||
private BedrockTestUtils() {
|
||||
} // Prevent instantiation
|
||||
|
||||
public static ApplicationContextRunner getContextRunner() {
|
||||
return new ApplicationContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.session-token=" + System.getenv("AWS_SESSION_TOKEN"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id())
|
||||
.withUserConfiguration(Config.class);
|
||||
}
|
||||
|
||||
public static ApplicationContextRunner getContextRunnerWithUserConfiguration() {
|
||||
return new ApplicationContextRunner().withUserConfiguration(Config.class);
|
||||
}
|
||||
|
||||
@Configuration
|
||||
static class Config {
|
||||
|
||||
@Bean
|
||||
public ObjectMapper objectMapper() {
|
||||
return new ObjectMapper();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock;
|
||||
|
||||
import java.lang.annotation.ElementType;
|
||||
import java.lang.annotation.Retention;
|
||||
import java.lang.annotation.RetentionPolicy;
|
||||
import java.lang.annotation.Target;
|
||||
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
|
||||
@Target({ ElementType.TYPE, ElementType.METHOD })
|
||||
@Retention(RetentionPolicy.RUNTIME)
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SESSION_TOKEN", matches = ".*")
|
||||
public @interface RequiresAwsCredentials {
|
||||
|
||||
// You can add custom properties here if needed
|
||||
|
||||
}
|
||||
@@ -21,11 +21,12 @@ import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.anthropic.BedrockAnthropicChatModel;
|
||||
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatModel;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
@@ -39,20 +40,18 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.AssertionsForClassTypes.within;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.0
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockAnthropicChatAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.anthropic.chat.model=" + AnthropicChatModel.CLAUDE_V2.id(),
|
||||
"spring.ai.bedrock.anthropic.chat.options.temperature=0.5")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class));
|
||||
@@ -102,11 +101,11 @@ public class BedrockAnthropicChatAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.anthropic.chat.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.anthropic.chat.options.temperature=0.55")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
@@ -114,9 +113,9 @@ public class BedrockAnthropicChatAutoConfigurationIT {
|
||||
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(anthropicChatProperties.isEnabled()).isTrue();
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
|
||||
assertThat(anthropicChatProperties.getOptions().getTemperature()).isEqualTo(0.55f);
|
||||
assertThat(anthropicChatProperties.getOptions().getTemperature()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(anthropicChatProperties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
|
||||
@@ -128,7 +127,7 @@ public class BedrockAnthropicChatAutoConfigurationIT {
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAnthropicChatProperties.class)).isEmpty();
|
||||
@@ -136,7 +135,8 @@ public class BedrockAnthropicChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAnthropicChatProperties.class)).isNotEmpty();
|
||||
@@ -144,7 +144,8 @@ public class BedrockAnthropicChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAnthropicChatProperties.class)).isEmpty();
|
||||
|
||||
@@ -21,11 +21,12 @@ import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.anthropic3.BedrockAnthropic3ChatModel;
|
||||
import org.springframework.ai.bedrock.anthropic3.api.Anthropic3ChatBedrockApi.AnthropicChatModel;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
@@ -39,20 +40,17 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.AssertionsForClassTypes.within;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockAnthropic3ChatAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic3.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.anthropic3.chat.model=" + AnthropicChatModel.CLAUDE_V3_SONNET.id(),
|
||||
"spring.ai.bedrock.anthropic3.chat.options.temperature=0.5")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropic3ChatAutoConfiguration.class));
|
||||
@@ -102,11 +100,11 @@ public class BedrockAnthropic3ChatAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic3.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.anthropic3.chat.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.anthropic3.chat.options.temperature=0.55")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropic3ChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
@@ -114,9 +112,9 @@ public class BedrockAnthropic3ChatAutoConfigurationIT {
|
||||
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(anthropicChatProperties.isEnabled()).isTrue();
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
|
||||
assertThat(anthropicChatProperties.getOptions().getTemperature()).isEqualTo(0.55f);
|
||||
assertThat(anthropicChatProperties.getOptions().getTemperature()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(anthropicChatProperties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
|
||||
@@ -128,15 +126,17 @@ public class BedrockAnthropic3ChatAutoConfigurationIT {
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropic3ChatAutoConfiguration.class))
|
||||
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAnthropic3ChatProperties.class)).isEmpty();
|
||||
assertThat(context.getBeansOfType(BedrockAnthropic3ChatModel.class)).isEmpty();
|
||||
});
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.anthropic3.chat.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic3.chat.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropic3ChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAnthropic3ChatProperties.class)).isNotEmpty();
|
||||
@@ -144,7 +144,8 @@ public class BedrockAnthropic3ChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.anthropic3.chat.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.anthropic3.chat.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAnthropic3ChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAnthropic3ChatProperties.class)).isEmpty();
|
||||
|
||||
@@ -21,11 +21,12 @@ import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.cohere.BedrockCohereChatModel;
|
||||
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatModel;
|
||||
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.ReturnLikelihoods;
|
||||
@@ -41,20 +42,18 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.AssertionsForClassTypes.within;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.0
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockCohereChatAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.cohere.chat.model=" + CohereChatModel.COHERE_COMMAND_V14.id(),
|
||||
"spring.ai.bedrock.cohere.chat.options.temperature=0.5",
|
||||
"spring.ai.bedrock.cohere.chat.options.maxTokens=500")
|
||||
@@ -105,11 +104,11 @@ public class BedrockCohereChatAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.cohere.chat.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.cohere.chat.options.temperature=0.55",
|
||||
"spring.ai.bedrock.cohere.chat.options.topP=0.55", "spring.ai.bedrock.cohere.chat.options.topK=10",
|
||||
"spring.ai.bedrock.cohere.chat.options.stopSequences=END1,END2",
|
||||
@@ -123,11 +122,11 @@ public class BedrockCohereChatAutoConfigurationIT {
|
||||
var aswProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(chatProperties.isEnabled()).isTrue();
|
||||
assertThat(aswProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(aswProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
assertThat(chatProperties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isEqualTo(0.55f);
|
||||
assertThat(chatProperties.getOptions().getTopP()).isEqualTo(0.55f);
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(chatProperties.getOptions().getTopP()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(chatProperties.getOptions().getTopK()).isEqualTo(10);
|
||||
assertThat(chatProperties.getOptions().getStopSequences()).isEqualTo(List.of("END1", "END2"));
|
||||
assertThat(chatProperties.getOptions().getReturnLikelihoods()).isEqualTo(ReturnLikelihoods.ALL);
|
||||
@@ -144,7 +143,7 @@ public class BedrockCohereChatAutoConfigurationIT {
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockCohereChatProperties.class)).isEmpty();
|
||||
@@ -152,7 +151,8 @@ public class BedrockCohereChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=true")
|
||||
BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockCohereChatProperties.class)).isNotEmpty();
|
||||
@@ -160,7 +160,8 @@ public class BedrockCohereChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=false")
|
||||
BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockCohereChatProperties.class)).isEmpty();
|
||||
|
||||
@@ -19,10 +19,11 @@ package org.springframework.ai.autoconfigure.bedrock.cohere;
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.cohere.BedrockCohereEmbeddingModel;
|
||||
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
|
||||
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
|
||||
@@ -35,17 +36,14 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.0
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockCohereEmbeddingAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.cohere.embedding.model=" + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
|
||||
"spring.ai.bedrock.cohere.embedding.options.inputType=SEARCH_DOCUMENT",
|
||||
"spring.ai.bedrock.cohere.embedding.options.truncate=NONE")
|
||||
@@ -86,10 +84,10 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.cohere.embedding.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.cohere.embedding.options.inputType=CLASSIFICATION",
|
||||
"spring.ai.bedrock.cohere.embedding.options.truncate=START")
|
||||
@@ -99,7 +97,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
|
||||
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(properties.isEnabled()).isTrue();
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
assertThat(properties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
assertThat(properties.getOptions().getInputType()).isEqualTo(InputType.CLASSIFICATION);
|
||||
@@ -114,7 +112,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
|
||||
public void embeddingDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isEmpty();
|
||||
@@ -122,7 +120,8 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly enable the embedding auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isNotEmpty();
|
||||
@@ -130,7 +129,8 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the embedding auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isEmpty();
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.apache.commons.logging.Log;
|
||||
import org.apache.commons.logging.LogFactory;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockConverseProxyChatAutoConfigurationIT {
|
||||
|
||||
private static final Log logger = LogFactory.getLog(BedrockConverseProxyChatAutoConfigurationIT.class);
|
||||
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues(
|
||||
"spring.ai.bedrock.converse.chat.options.model=" + "anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"spring.ai.bedrock.converse.chat.options.temperature=0.5")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class));
|
||||
|
||||
@Test
|
||||
void call() {
|
||||
this.contextRunner.run(context -> {
|
||||
BedrockProxyChatModel chatModel = context.getBean(BedrockProxyChatModel.class);
|
||||
String response = chatModel.call("Hello");
|
||||
assertThat(response).isNotEmpty();
|
||||
logger.info("Response: " + response);
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void stream() {
|
||||
this.contextRunner.run(context -> {
|
||||
BedrockProxyChatModel chatModel = context.getBean(BedrockProxyChatModel.class);
|
||||
Flux<ChatResponse> responseFlux = chatModel.stream(new Prompt(new UserMessage("Hello")));
|
||||
|
||||
String response = responseFlux.collectList()
|
||||
.block()
|
||||
.stream()
|
||||
.map(ChatResponse::getResults)
|
||||
.flatMap(List::stream)
|
||||
.map(Generation::getOutput)
|
||||
.map(AssistantMessage::getContent)
|
||||
.collect(Collectors.joining());
|
||||
|
||||
assertThat(response).isNotEmpty();
|
||||
logger.info("Response: " + response);
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
*
|
||||
* Unit Tests for {@link BedrockConverseProxyChatProperties}.
|
||||
*/
|
||||
public class BedrockConverseProxyChatPropertiesTests {
|
||||
|
||||
@Test
|
||||
public void chatOptionsTest() {
|
||||
|
||||
new ApplicationContextRunner().withPropertyValues(
|
||||
// @formatter:off
|
||||
"spring.ai.bedrock.converse.chat.options.model=MODEL_XYZ",
|
||||
|
||||
"spring.ai.bedrock.converse.chat.options.max-tokens=123",
|
||||
"spring.ai.bedrock.converse.chat.options.metadata.user-id=MyUserId",
|
||||
"spring.ai.bedrock.converse.chat.options.stop_sequences=boza,koza",
|
||||
|
||||
"spring.ai.bedrock.converse.chat.options.temperature=0.55",
|
||||
"spring.ai.bedrock.converse.chat.options.top-p=0.56",
|
||||
"spring.ai.bedrock.converse.chat.options.top-k=100"
|
||||
)
|
||||
// @formatter:on
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
var chatProperties = context.getBean(BedrockConverseProxyChatProperties.class);
|
||||
|
||||
assertThat(chatProperties.isEnabled()).isTrue();
|
||||
|
||||
assertThat(chatProperties.getOptions().getModel()).isEqualTo("MODEL_XYZ");
|
||||
assertThat(chatProperties.getOptions().getMaxTokens()).isEqualTo(123);
|
||||
assertThat(chatProperties.getOptions().getStopSequences()).contains("boza", "koza");
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isEqualTo(0.55);
|
||||
assertThat(chatProperties.getOptions().getTopP()).isEqualTo(0.56);
|
||||
assertThat(chatProperties.getOptions().getTopK()).isEqualTo(100);
|
||||
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is enabled by default
|
||||
new ApplicationContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class))
|
||||
.run(context -> assertThat(context.getBeansOfType(BedrockConverseProxyChatProperties.class)).isNotEmpty());
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.converse.chat..enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class))
|
||||
.run(context -> assertThat(context.getBeansOfType(BedrockConverseProxyChatProperties.class)).isNotEmpty());
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.converse.chat..enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class))
|
||||
.run(context -> assertThat(context.getBeansOfType(BedrockConverseProxyChatProperties.class)).isEmpty());
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse.tool;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.function.Function;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.autoconfigure.bedrock.converse.BedrockConverseProxyChatAutoConfiguration;
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptionsBuilder.PortableFunctionCallingOptions;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.context.annotation.Description;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
@RequiresAwsCredentials
|
||||
class FunctionCallWithFunctionBeanIT {
|
||||
|
||||
private final Logger logger = LoggerFactory.getLogger(FunctionCallWithFunctionBeanIT.class);
|
||||
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class))
|
||||
.withUserConfiguration(Config.class);
|
||||
|
||||
@Test
|
||||
void functionCallTest() {
|
||||
|
||||
this.contextRunner
|
||||
.withPropertyValues(
|
||||
"spring.ai.bedrock.converse.chat.options.model=" + "anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.run(context -> {
|
||||
|
||||
BedrockProxyChatModel chatModel = context.getBean(BedrockProxyChatModel.class);
|
||||
|
||||
var userMessage = new UserMessage(
|
||||
"What's the weather like in San Francisco, in Paris, France and in Tokyo, Japan? Return the temperature in Celsius.");
|
||||
|
||||
ChatResponse response = chatModel.call(new Prompt(List.of(userMessage),
|
||||
PortableFunctionCallingOptions.builder().withFunction("weatherFunction").build()));
|
||||
|
||||
logger.info("Response: {}", response);
|
||||
|
||||
assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
|
||||
|
||||
response = chatModel.call(new Prompt(List.of(userMessage),
|
||||
FunctionCallingOptions.builder().withFunction("weatherFunction3").build()));
|
||||
|
||||
logger.info("Response: {}", response);
|
||||
|
||||
assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void functionStreamTest() {
|
||||
|
||||
this.contextRunner
|
||||
.withPropertyValues(
|
||||
"spring.ai.bedrock.converse.chat.options.model=" + "anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.run(context -> {
|
||||
|
||||
BedrockProxyChatModel chatModel = context.getBean(BedrockProxyChatModel.class);
|
||||
|
||||
var userMessage = new UserMessage(
|
||||
"What's the weather like in San Francisco, in Paris, France and in Tokyo, Japan? Return the temperature in Celsius.");
|
||||
|
||||
Flux<ChatResponse> responses = chatModel.stream(new Prompt(List.of(userMessage),
|
||||
FunctionCallingOptions.builder().withFunction("weatherFunction").build()));
|
||||
|
||||
String content = responses.collectList()
|
||||
.block()
|
||||
.stream()
|
||||
.filter(cr -> cr.getResult() != null)
|
||||
.map(cr -> cr.getResult().getOutput().getContent())
|
||||
.collect(Collectors.joining());
|
||||
|
||||
logger.info("Response: {}", content);
|
||||
assertThat(content).contains("30", "10", "15");
|
||||
|
||||
});
|
||||
}
|
||||
|
||||
@Configuration
|
||||
static class Config {
|
||||
|
||||
@Bean
|
||||
@Description("Get the weather in location. Return temperature in 36°F or 36°C format.")
|
||||
public Function<MockWeatherService.Request, MockWeatherService.Response> weatherFunction() {
|
||||
return new MockWeatherService();
|
||||
}
|
||||
|
||||
// Relies on the Request's JsonClassDescription annotation to provide the
|
||||
// function description.
|
||||
@Bean
|
||||
public Function<MockWeatherService.Request, MockWeatherService.Response> weatherFunction3() {
|
||||
MockWeatherService weatherService = new MockWeatherService();
|
||||
return (weatherService::apply);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse.tool;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.autoconfigure.bedrock.converse.BedrockConverseProxyChatAutoConfiguration;
|
||||
import org.springframework.ai.bedrock.converse.BedrockProxyChatModel;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.model.function.FunctionCallbackWrapper;
|
||||
import org.springframework.ai.model.function.FunctionCallingOptions;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
@RequiresAwsCredentials
|
||||
public class FunctionCallWithPromptFunctionIT {
|
||||
|
||||
private final Logger logger = LoggerFactory.getLogger(FunctionCallWithPromptFunctionIT.class);
|
||||
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockConverseProxyChatAutoConfiguration.class));
|
||||
|
||||
@Test
|
||||
void functionCallTest() {
|
||||
this.contextRunner
|
||||
.withPropertyValues(
|
||||
"spring.ai.bedrock.converse.chat.options.model=" + "anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
.run(context -> {
|
||||
|
||||
BedrockProxyChatModel chatModel = context.getBean(BedrockProxyChatModel.class);
|
||||
|
||||
UserMessage userMessage = new UserMessage(
|
||||
"What's the weather like in San Francisco, in Paris and in Tokyo? Return the temperature in Celsius.");
|
||||
|
||||
var promptOptions = FunctionCallingOptions.builder()
|
||||
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
|
||||
.withName("CurrentWeatherService")
|
||||
.withDescription("Get the weather in location. Return temperature in 36°F or 36°C format.")
|
||||
.build()))
|
||||
.build();
|
||||
|
||||
ChatResponse response = chatModel.call(new Prompt(List.of(userMessage), promptOptions));
|
||||
|
||||
logger.info("Response: {}", response);
|
||||
|
||||
assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
/*
|
||||
* 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.autoconfigure.bedrock.converse.tool;
|
||||
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonClassDescription;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude.Include;
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
|
||||
|
||||
/**
|
||||
* Mock 3rd party weather service.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
public class MockWeatherService implements Function<MockWeatherService.Request, MockWeatherService.Response> {
|
||||
|
||||
@Override
|
||||
public Response apply(Request request) {
|
||||
|
||||
double temperature = 0;
|
||||
if (request.location().contains("Paris")) {
|
||||
temperature = 15;
|
||||
}
|
||||
else if (request.location().contains("Tokyo")) {
|
||||
temperature = 10;
|
||||
}
|
||||
else if (request.location().contains("San Francisco")) {
|
||||
temperature = 30;
|
||||
}
|
||||
|
||||
return new Response(temperature, 15, 20, 2, 53, 45, Unit.C);
|
||||
}
|
||||
|
||||
/**
|
||||
* Temperature units.
|
||||
*/
|
||||
public enum Unit {
|
||||
|
||||
/**
|
||||
* Celsius.
|
||||
*/
|
||||
C("metric"),
|
||||
/**
|
||||
* Fahrenheit.
|
||||
*/
|
||||
F("imperial");
|
||||
|
||||
/**
|
||||
* Human readable unit name.
|
||||
*/
|
||||
public final String unitName;
|
||||
|
||||
Unit(String text) {
|
||||
this.unitName = text;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Weather Function request.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
@JsonClassDescription("Weather API request")
|
||||
public record Request(@JsonProperty(required = true,
|
||||
value = "location") @JsonPropertyDescription("The city and state e.g. San Francisco, CA") String location,
|
||||
@JsonProperty(required = true, value = "unit") @JsonPropertyDescription("Temperature unit") Unit unit) {
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Weather Function response.
|
||||
*/
|
||||
public record Response(double temp, double feels_like, double temp_min, double temp_max, int pressure, int humidity,
|
||||
Unit unit) {
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -20,10 +20,11 @@ import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.autoconfigure.bedrock.jurrasic2.BedrockAi21Jurassic2ChatAutoConfiguration;
|
||||
import org.springframework.ai.autoconfigure.bedrock.jurrasic2.BedrockAi21Jurassic2ChatProperties;
|
||||
import org.springframework.ai.bedrock.jurassic2.BedrockAi21Jurassic2ChatModel;
|
||||
@@ -37,20 +38,17 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.AssertionsForClassTypes.within;
|
||||
|
||||
/**
|
||||
* @author Ahmed Yousri
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockAi21Jurassic2ChatAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.jurassic2.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.jurassic2.chat.model="
|
||||
+ Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatModel.AI21_J2_ULTRA_V1.id(),
|
||||
"spring.ai.bedrock.jurassic2.chat.options.temperature=0.5",
|
||||
@@ -80,7 +78,7 @@ public class BedrockAi21Jurassic2ChatAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.jurassic2.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.jurassic2.chat.model=MODEL_XYZ",
|
||||
@@ -95,7 +93,7 @@ public class BedrockAi21Jurassic2ChatAutoConfigurationIT {
|
||||
assertThat(chatProperties.isEnabled()).isTrue();
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isEqualTo(0.55f);
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(chatProperties.getOptions().getMaxTokens()).isEqualTo(123);
|
||||
assertThat(chatProperties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
@@ -108,7 +106,7 @@ public class BedrockAi21Jurassic2ChatAutoConfigurationIT {
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAi21Jurassic2ChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAi21Jurassic2ChatProperties.class)).isEmpty();
|
||||
@@ -116,7 +114,8 @@ public class BedrockAi21Jurassic2ChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.jurassic2.chat.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.jurassic2.chat.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAi21Jurassic2ChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAi21Jurassic2ChatProperties.class)).isNotEmpty();
|
||||
@@ -124,7 +123,8 @@ public class BedrockAi21Jurassic2ChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.jurassic2.chat.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.jurassic2.chat.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockAi21Jurassic2ChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockAi21Jurassic2ChatProperties.class)).isEmpty();
|
||||
|
||||
@@ -21,11 +21,12 @@ import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.llama.BedrockLlamaChatModel;
|
||||
import org.springframework.ai.bedrock.llama.api.LlamaChatBedrockApi.LlamaChatModel;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
@@ -39,21 +40,19 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.AssertionsForClassTypes.within;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @author Wei Jiang
|
||||
* @since 0.8.0
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockLlamaChatAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.llama.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.llama.chat.model=" + LlamaChatModel.LLAMA3_70B_INSTRUCT_V1.id(),
|
||||
"spring.ai.bedrock.llama.chat.options.temperature=0.5",
|
||||
"spring.ai.bedrock.llama.chat.options.maxGenLen=500")
|
||||
@@ -67,7 +66,7 @@ public class BedrockLlamaChatAutoConfigurationIT {
|
||||
""").createMessage(Map.of("name", "Bob", "voice", "pirate"));
|
||||
|
||||
private final UserMessage userMessage = new UserMessage(
|
||||
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
|
||||
"Describe 3 of the most feared and legendary pirates from the Golden Age of Piracy, particularly those known for their intimidating tactics and whose stories influenced popular culture.");
|
||||
|
||||
@Test
|
||||
public void chatCompletion() {
|
||||
@@ -104,11 +103,11 @@ public class BedrockLlamaChatAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.llama.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.llama.chat.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.llama.chat.options.temperature=0.55",
|
||||
"spring.ai.bedrock.llama.chat.options.maxGenLen=123")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockLlamaChatAutoConfiguration.class))
|
||||
@@ -117,9 +116,9 @@ public class BedrockLlamaChatAutoConfigurationIT {
|
||||
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(llamaChatProperties.isEnabled()).isTrue();
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
|
||||
assertThat(llamaChatProperties.getOptions().getTemperature()).isEqualTo(0.55f);
|
||||
assertThat(llamaChatProperties.getOptions().getTemperature()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(llamaChatProperties.getOptions().getMaxGenLen()).isEqualTo(123);
|
||||
assertThat(llamaChatProperties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
@@ -132,14 +131,16 @@ public class BedrockLlamaChatAutoConfigurationIT {
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner().withConfiguration(AutoConfigurations.of(BedrockLlamaChatAutoConfiguration.class))
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockLlamaChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockLlamaChatProperties.class)).isEmpty();
|
||||
assertThat(context.getBeansOfType(BedrockLlamaChatModel.class)).isEmpty();
|
||||
});
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.llama.chat.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.llama.chat.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockLlamaChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockLlamaChatProperties.class)).isNotEmpty();
|
||||
@@ -147,7 +148,8 @@ public class BedrockLlamaChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.llama.chat.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.llama.chat.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockLlamaChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockLlamaChatProperties.class)).isEmpty();
|
||||
|
||||
@@ -21,11 +21,12 @@ import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.titan.BedrockTitanChatModel;
|
||||
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatModel;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
@@ -39,20 +40,18 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.AssertionsForClassTypes.within;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.0
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockTitanChatAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.titan.chat.model=" + TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
|
||||
"spring.ai.bedrock.titan.chat.options.temperature=0.5",
|
||||
"spring.ai.bedrock.titan.chat.options.maxTokenCount=500")
|
||||
@@ -102,11 +101,11 @@ public class BedrockTitanChatAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.titan.chat.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.titan.chat.options.temperature=0.55",
|
||||
"spring.ai.bedrock.titan.chat.options.topP=0.55",
|
||||
"spring.ai.bedrock.titan.chat.options.stopSequences=END1,END2",
|
||||
@@ -117,11 +116,11 @@ public class BedrockTitanChatAutoConfigurationIT {
|
||||
var aswProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(chatProperties.isEnabled()).isTrue();
|
||||
assertThat(aswProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(aswProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
assertThat(chatProperties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isEqualTo(0.55f);
|
||||
assertThat(chatProperties.getOptions().getTopP()).isEqualTo(0.55f);
|
||||
assertThat(chatProperties.getOptions().getTemperature()).isCloseTo(0.55, within(0.0001));
|
||||
assertThat(chatProperties.getOptions().getTopP()).isCloseTo(0.55, within(0.0001));
|
||||
|
||||
assertThat(chatProperties.getOptions().getStopSequences()).isEqualTo(List.of("END1", "END2"));
|
||||
assertThat(chatProperties.getOptions().getMaxTokenCount()).isEqualTo(123);
|
||||
@@ -135,14 +134,16 @@ public class BedrockTitanChatAutoConfigurationIT {
|
||||
public void chatCompletionDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner().withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockTitanChatProperties.class)).isEmpty();
|
||||
assertThat(context.getBeansOfType(BedrockTitanChatModel.class)).isEmpty();
|
||||
});
|
||||
|
||||
// Explicitly enable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockTitanChatProperties.class)).isNotEmpty();
|
||||
@@ -150,7 +151,8 @@ public class BedrockTitanChatAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the chat auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.chat.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.chat.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockTitanChatProperties.class)).isEmpty();
|
||||
|
||||
@@ -20,10 +20,11 @@ import java.util.Base64;
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
|
||||
import org.springframework.ai.autoconfigure.bedrock.BedrockTestUtils;
|
||||
import org.springframework.ai.autoconfigure.bedrock.RequiresAwsCredentials;
|
||||
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingModel;
|
||||
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingModel.InputType;
|
||||
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel;
|
||||
@@ -36,13 +37,13 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.0
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
|
||||
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
|
||||
@RequiresAwsCredentials
|
||||
public class BedrockTitanEmbeddingAutoConfigurationIT {
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
|
||||
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
|
||||
@@ -84,17 +85,19 @@ public class BedrockTitanEmbeddingAutoConfigurationIT {
|
||||
@Test
|
||||
public void propertiesTest() {
|
||||
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
|
||||
"spring.ai.bedrock.titan.embedding.model=MODEL_XYZ", "spring.ai.bedrock.titan.embedding.inputType=TEXT")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true",
|
||||
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
|
||||
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
|
||||
"spring.ai.bedrock.titan.embedding.model=MODEL_XYZ",
|
||||
"spring.ai.bedrock.titan.embedding.inputType=TEXT")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
var properties = context.getBean(BedrockTitanEmbeddingProperties.class);
|
||||
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
|
||||
|
||||
assertThat(properties.isEnabled()).isTrue();
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
|
||||
assertThat(awsProperties.getRegion()).isEqualTo(Region.US_EAST_1.id());
|
||||
assertThat(properties.getModel()).isEqualTo("MODEL_XYZ");
|
||||
|
||||
assertThat(properties.getInputType()).isEqualTo(InputType.TEXT);
|
||||
@@ -108,7 +111,7 @@ public class BedrockTitanEmbeddingAutoConfigurationIT {
|
||||
public void embeddingDisabled() {
|
||||
|
||||
// It is disabled by default
|
||||
new ApplicationContextRunner()
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isEmpty();
|
||||
@@ -116,7 +119,8 @@ public class BedrockTitanEmbeddingAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly enable the embedding auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isNotEmpty();
|
||||
@@ -124,7 +128,8 @@ public class BedrockTitanEmbeddingAutoConfigurationIT {
|
||||
});
|
||||
|
||||
// Explicitly disable the embedding auto-configuration.
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=false")
|
||||
BedrockTestUtils.getContextRunnerWithUserConfiguration()
|
||||
.withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=false")
|
||||
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isEmpty();
|
||||
|
||||
@@ -20,6 +20,7 @@ import java.util.List;
|
||||
|
||||
import org.apache.commons.logging.Log;
|
||||
import org.apache.commons.logging.LogFactory;
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
|
||||
@@ -33,6 +34,7 @@ import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
// NOTE: works only with US location. Use VPN if you are outside US.
|
||||
@EnabledIfEnvironmentVariable(named = "PALM_API_KEY", matches = ".*")
|
||||
@Disabled("Disabled due to the PALM API being decommissioned by Google.")
|
||||
public class VertexAiPaLm2AutoConfigurationIT {
|
||||
|
||||
private static final Log logger = LogFactory.getLog(VertexAiPaLm2AutoConfigurationIT.class);
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!--
|
||||
~ 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.
|
||||
-->
|
||||
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai</artifactId>
|
||||
<version>1.0.0-SNAPSHOT</version>
|
||||
<relativePath>../../pom.xml</relativePath>
|
||||
</parent>
|
||||
<artifactId>spring-ai-bedrock-converse-spring-boot-starter</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>Spring AI Starter - Bedrock Converse API</name>
|
||||
<description>Spring AI Bedrock Converse API Auto Configuration</description>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
|
||||
<scm>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
<connection>git://github.com/spring-projects/spring-ai.git</connection>
|
||||
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
|
||||
</scm>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bedrock-converse</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
Reference in New Issue
Block a user