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:
Christian Tzolov
committed by
Mark Pollack
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 |
Reference in New Issue
Block a user