Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Bedrock: Add Titan Chat support

Browse Source
This commit is contained in:
Christian Tzolov
2023-12-19 12:58:37 +01:00
parent 3b91c51295
commit 26b1ccdc57
12 changed files with 769 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
spring-ai-bedrock/README.md
View File

@@ -74,10 +74,10 @@ For example to enable the bedrock Llama2 Chat client you need to set the
Next you can use the `spring.ai.bedrock.<model>.<chat|embedding>.*` properties to configure each model as provided in its documentation:
* [Spring AI Bedrock Llama2 Chat](./README_LLAMA2_CHAT.md) - `spring.ai.bedrock.llama2.chat=true`
* [Spring AI Bedrock Cohere Chat](./README_COHERE_CHAT.md) - `spring.ai.bedrock.cohere.chat=true`
* [Spring AI Bedrock Cohere Embedding](./README_COHERE_EMBEDDING.md) - `spring.ai.bedrock.cohere.embedding=true`
* [Spring AI Bedrock Anthropic Chat](./README_ANTHROPIC_CHAT.md) - `spring.ai.bedrock.anthropic.chat=true`
* (WIP) [Spring AI Bedrock Titan Chat](./README_TITAN_CHAT.md) - `spring.ai.bedrock.titan.chat=true`
* [Spring AI Bedrock Titan Embedding](./README_TITAN_EMBEDDING.md) - `spring.ai.bedrock.titan.embedding=true`
* (WIP) [Spring AI Bedrock Ai21 Jurassic2 Chat](./README_JURASSIC2_CHAT.md) - `spring.ai.bedrock.jurassic2.chat=true`
* [Spring AI Bedrock Llama2 Chat](./README_LLAMA2_CHAT.md) - `spring.ai.bedrock.llama2.chat.enabled=true`
* [Spring AI Bedrock Cohere Chat](./README_COHERE_CHAT.md) - `spring.ai.bedrock.cohere.chat.enabled=true`
* [Spring AI Bedrock Cohere Embedding](./README_COHERE_EMBEDDING.md) - `spring.ai.bedrock.cohere.embedding.enabled=true`
* [Spring AI Bedrock Anthropic Chat](./README_ANTHROPIC_CHAT.md) - `spring.ai.bedrock.anthropic.chat.enabled=true`
* [Spring AI Bedrock Titan Chat](./README_TITAN_CHAT.md) - `spring.ai.bedrock.titan.chat.enabled=true`
* [Spring AI Bedrock Titan Embedding](./README_TITAN_EMBEDDING.md) - `spring.ai.bedrock.titan.embedding.enabled=true`
* (WIP) [Spring AI Bedrock Ai21 Jurassic2 Chat](./README_JURASSIC2_CHAT.md) - `spring.ai.bedrock.jurassic2.chat.enabled=true`

70
spring-ai-bedrock/README_COHERE_CHAT.md
View File

@@ -108,73 +108,3 @@ Use the `BedrockCohereChatProperties` to configure the Bedrock Cohere Chat clien
| spring.ai.bedrock.cohere.chat.logitBiasBias | Prevents the model from generating unwanted tokens or incentivize the model to include desired tokens. | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.truncate | Specifies how the API handles inputs longer than the maximum token length | AWS Bedrock default |
### 2.3 CohereEmbeddingBedrockApi
[CohereEmbeddingBedrockApi](./src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Cohere Embed models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-embed.html).
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
![CohereEmbeddingBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Cohere%20Embedding%20API.jpg)
The CohereEmbeddingBedrockApi supports the `cohere.embed-english-v3` and `cohere.embed-multilingual-v3` models for single and batch embedding computation.
Here is a simple snippet how to use the api programmatically:
```java
CohereEmbeddingBedrockApi api = new CohereEmbeddingBedrockApi(
CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(), new ObjectMapper());
CohereEmbeddingRequest request = new CohereEmbeddingRequest(
List.of("I like to eat apples", "I like to eat oranges"),
CohereEmbeddingRequest.InputType.search_document,
CohereEmbeddingRequest.Truncate.NONE);
CohereEmbeddingResponse response = api.embedding(request);
assertThat(response.embeddings()).hasSize(2);
assertThat(response.embeddings().get(0)).hasSize(1024);
```
### 2.4 BedrockCohereEmbeddingClient
[BedrockCohereEmbeddingClient](./src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java) implements the Spring-Ai `EmbeddingClient` on top of the `CohereEmbeddingBedrockApi`.
You can use like this:
```java
@Bean
public CohereEmbeddingBedrockApi cohereEmbeddingApi() {
return new CohereEmbeddingBedrockApi(CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
}
@Bean
public BedrockCohereEmbeddingClient cohereAiEmbedding(CohereEmbeddingBedrockApi cohereEmbeddingApi) {
return new BedrockCohereEmbeddingClient(cohereEmbeddingApi);
}
```
or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. For this add the following dependency:
```xml
<dependency>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<groupId>org.springframework.ai</groupId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
**NOTE:** You have to enable the Bedrock Cohere chat client with `spring.ai.bedrock.cohere.embedding.enabled=true`.
By default the client is disabled.
Use the `BedrockCohereEmbeddingProperties` to configure the Bedrock Cohere Chat client:
| Property | Description | Default |
| ------------- | ------------- | ------------- |
| spring.ai.bedrock.cohere.embedding.enable | Enable Bedrock Cohere chat client. Disabled by default | false |
| spring.ai.bedrock.cohere.embedding.awsRegion | AWS region to use. | us-east-1 |
| spring.ai.bedrock.cohere.embedding.model | The model id to use. See the `CohereEmbeddingModel` for the supported models. | cohere.embed-multilingual-v3 |
| spring.ai.bedrock.cohere.embedding.inputType | Prepends special tokens to differentiate each type from one another. You should not mix different types together, except when mixing types for for search and retrieval. In this case, embed your corpus with the search_document type and embedded queries with type search_query type. | search_document |
| spring.ai.bedrock.cohere.embedding.truncate | Specifies how the API handles inputs longer than the maximum token length. | NONE |

84
spring-ai-bedrock/README_TITAN_CHAT.md
View File

@@ -1,5 +1,83 @@
# (WIP) Bedrock Titan Chat
# 1. Bedrock Titan Chat
Use the [TitanChatBedrockApi.java](src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java) Bedrock Titan chat client to implement `AiClient` and `AiStreamClient`.
## 1.1 TitanChatBedrockApi
Consult the the existing Llama2 and Cohere implementation.
[TitanChatBedrockApi](./src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Titan text models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html).
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
![TitanChatBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Titan%20Chat%20API.jpg)
The TitanChatBedrockApi supports the `amazon.titan-text-lite-v1` and `amazon.titan-text-express-v1` models for bot synchronous (e.g. `chatCompletion()`) and streaming (e.g. `chatCompletionStream()`) responses.
Here is a simple snippet how to use the api programmatically:
```java
TitanChatBedrockApi titanBedrockApi = new TitanChatBedrockApi(TitanChatCompletionModel.TITAN_TEXT_EXPRESS_V1.id(),
Region.EU_CENTRAL_1.id());
TitanChatRequest titanChatRequest = TitanChatRequest.builder("Give me the names of 3 famous pirates?")
.withTemperature(0.5f)
.withTopP(0.9f)
.withMaxTokenCount(100)
.withStopSequences(List.of("|"))
.build();
TitanChatResponse response = titanBedrockApi.chatCompletion(titanChatRequest);
assertThat(response.results()).hasSize(1);
assertThat(response.results().get(0).outputText()).contains("Blackbeard");
Flux<TitanChatResponseChunk> response = titanBedrockApi.chatCompletionStream(titanChatRequest);
List<TitanChatResponseChunk> results = response.collectList().block();
assertThat(results.stream().map(TitanChatResponseChunk::outputText).collect(Collectors.joining("\n")))
.contains("Blackbeard");
```
## 1.2 BedrockTitanChatClient
[BedrockTitanChatClient](./src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClient.java) implements the Spring-Ai `AiClient` and `AiStreamingClient` on top of the `TitanChatBedrockApi`.
You can use like this:
```java
@Bean
public TitanChatBedrockApi titanApi() {
return new TitanChatBedrockApi(TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
}
@Bean
public BedrockTitanChatClient titanChatClient(TitanChatBedrockApi titanApi) {
return new BedrockTitanChatClient(titanApi);
}
```
or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. For this add the following dependency:
```xml
<dependency>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<groupId>org.springframework.ai</groupId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
**NOTE:** You have to enable the Bedrock Titan chat client with `spring.ai.bedrock.titan.chat.enabled=true`.
By default the client is disabled.
Use the `BedrockTitanChatProperties` to configure the Bedrock Titan Chat client:
| Property | Description | Default |
| ------------- | ------------- | ------------- |
| spring.ai.bedrock.embedding.aws.region | AWS region to use. | us-east-1 |
| spring.ai.bedrock.embedding.aws.accessKey | AWS credentials access key. | |
| spring.ai.bedrock.embedding.aws.secretKey | AWS credentials secret key. | |
| spring.ai.bedrock.titan.chat.enable | Enable Bedrock Titan chat client. Disabled by default | false |
| spring.ai.bedrock.titan.chat.model | The model id to use. See the `TitanChatModel` for the supported models. | amazon.titan-text-express-v1 |
| spring.ai.bedrock.titan.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.7 |
| spring.ai.bedrock.titan.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default |
| spring.ai.bedrock.titan.chat.maxTokenCount | Specify the maximum number of tokens to use in the generated response. | AWS Bedrock default |
| spring.ai.bedrock.titan.chat.stopSequences | Configure up to four sequences that the model recognizes. | AWS Bedrock default |

3
spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java
View File

@@ -24,6 +24,7 @@ import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
@@ -92,7 +93,7 @@ public abstract class AbstractBedrockApi<I, O, SO> {
ObjectMapper objectMapper) {
this.modelId = modelId;
this.objectMapper = objectMapper;
this.objectMapper = objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
this.credentialsProvider = credentialsProvider;
this.region = region;

132
spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClient.java Normal file
View File

@@ -0,0 +1,132 @@
/*
* Copyright 2023-2023 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.titan;
import java.util.List;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.MessageToPromptConverter;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatRequest;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponse;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponseChunk;
import org.springframework.ai.client.AiClient;
import org.springframework.ai.client.AiResponse;
import org.springframework.ai.client.AiStreamClient;
import org.springframework.ai.client.Generation;
import org.springframework.ai.metadata.ChoiceMetadata;
import org.springframework.ai.metadata.Usage;
import org.springframework.ai.prompt.Prompt;
/**
* @author Christian Tzolov
* @since 0.8.0
*/
public class BedrockTitanChatClient implements AiClient, AiStreamClient {
private final TitanChatBedrockApi chatApi;
private Float temperature;
private Float topP;
private Integer maxTokenCount;
private List<String> stopSequences;
public BedrockTitanChatClient(TitanChatBedrockApi chatApi) {
this.chatApi = chatApi;
}
public BedrockTitanChatClient withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public BedrockTitanChatClient withTopP(Float topP) {
this.topP = topP;
return this;
}
public BedrockTitanChatClient withMaxTokenCount(Integer maxTokens) {
this.maxTokenCount = maxTokens;
return this;
}
public BedrockTitanChatClient withStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
return this;
}
@Override
public AiResponse generate(Prompt prompt) {
TitanChatResponse response = this.chatApi.chatCompletion(this.createRequest(prompt, false));
List<Generation> generations = response.results().stream().map(result -> {
return new Generation(result.outputText());
}).toList();
return new AiResponse(generations);
}
@Override
public Flux<AiResponse> generateStream(Prompt prompt) {
return this.chatApi.chatCompletionStream(this.createRequest(prompt, true)).map(chunk -> {
Generation generation = new Generation(chunk.outputText());
if (chunk.amazonBedrockInvocationMetrics() != null) {
String completionReason = chunk.completionReason().name();
generation = generation
.withChoiceMetadata(ChoiceMetadata.from(completionReason, chunk.amazonBedrockInvocationMetrics()));
}
else if (chunk.inputTextTokenCount() != null && chunk.totalOutputTextTokenCount() != null) {
String completionReason = chunk.completionReason().name();
generation = generation.withChoiceMetadata(ChoiceMetadata.from(completionReason, extractUsage(chunk)));
}
return new AiResponse(List.of(generation));
});
}
private TitanChatRequest createRequest(Prompt prompt, boolean stream) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getMessages());
return TitanChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withTopP(this.topP)
.withMaxTokenCount(this.maxTokenCount)
.withStopSequences(this.stopSequences)
.build();
}
private Usage extractUsage(TitanChatResponseChunk response) {
return new Usage() {
@Override
public Long getPromptTokens() {
return response.inputTextTokenCount().longValue();
}
@Override
public Long getGenerationTokens() {
return response.totalOutputTextTokenCount().longValue();
}
};
}
}

32
spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java
View File

@@ -20,7 +20,9 @@ import java.util.List;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Flux;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatRequest;
@@ -45,6 +47,19 @@ public class TitanChatBedrockApi extends
super(modelId, region);
}
/**
* Create a new TitanChatBedrockApi instance using the provided credentials provider, region and object mapper.
*
* @param modelId The model id to use. See the {@link TitanChatModel} for the supported models.
* @param credentialsProvider The credentials provider to connect to AWS.
* @param region The AWS region to use.
* @param objectMapper The object mapper to use for JSON serialization and deserialization.
*/
public TitanChatBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
super(modelId, credentialsProvider, region, objectMapper);
}
/**
* TitanChatRequest encapsulates the request parameters for the Titan chat model.
*
@@ -76,6 +91,11 @@ public class TitanChatBedrockApi extends
@JsonProperty("stopSequences") List<String> stopSequences) {
}
/**
* Create a new TitanChatRequest builder.
* @param inputText The prompt to use for the chat.
* @return A new TitanChatRequest builder.
*/
public static Builder builder(String inputText) {
return new Builder(inputText);
}
@@ -162,10 +182,16 @@ public class TitanChatBedrockApi extends
* The response was fully generated.
*/
FINISH,
/**
* The response was truncated because of the response length you set.
*/
LENGTH
LENGTH,
/**
* The response was truncated because of restrictions.
*/
CONTENT_FILTERED
}
}
@@ -191,7 +217,7 @@ public class TitanChatBedrockApi extends
/**
* Titan models version.
*/
public enum TitanChatCompletionModel {
public enum TitanChatModel {
/**
* amazon.titan-text-lite-v1
@@ -212,7 +238,7 @@ public class TitanChatBedrockApi extends
return id;
}
TitanChatCompletionModel(String value) {
TitanChatModel(String value) {
this.id = value;
}
}

173
spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClientIT.java Normal file
View File

@@ -0,0 +1,173 @@
package org.springframework.ai.bedrock.titan;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatModel;
import org.springframework.ai.client.AiResponse;
import org.springframework.ai.client.Generation;
import org.springframework.ai.parser.BeanOutputParser;
import org.springframework.ai.parser.ListOutputParser;
import org.springframework.ai.parser.MapOutputParser;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.PromptTemplate;
import org.springframework.ai.prompt.SystemPromptTemplate;
import org.springframework.ai.prompt.messages.Message;
import org.springframework.ai.prompt.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.convert.support.DefaultConversionService;
import org.springframework.core.io.Resource;
import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
class BedrockTitanChatClientIT {
@Autowired
private BedrockTitanChatClient client;
@Value("classpath:/prompts/system-message.st")
private Resource systemResource;
@Test
void roleTest() {
String request = "Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.";
String name = "Bob";
String voice = "pirate";
UserMessage userMessage = new UserMessage(request);
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", name, "voice", voice));
Prompt prompt = new Prompt(List.of(userMessage, systemMessage));
AiResponse response = client.generate(prompt);
assertThat(response.getGeneration().getContent()).contains("Blackbeard");
}
@Disabled("TODO: Fix the parser instructions to return the correct format")
@Test
void outputParser() {
DefaultConversionService conversionService = new DefaultConversionService();
ListOutputParser outputParser = new ListOutputParser(conversionService);
String format = outputParser.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.client.generate(prompt).getGeneration();
List<String> list = outputParser.parse(generation.getContent());
assertThat(list).hasSize(5);
}
@Disabled("TODO: Fix the parser instructions to return the correct format")
@Test
void mapOutputParser() {
MapOutputParser outputParser = new MapOutputParser();
String format = outputParser.getFormat();
String template = """
Remove Markdown code blocks from the output.
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 = client.generate(prompt).getGeneration();
Map<String, Object> result = outputParser.parse(generation.getContent());
assertThat(result.get("numbers")).isEqualTo(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
}
record ActorsFilmsRecord(String actor, List<String> movies) {
}
@Disabled("TODO: Fix the parser instructions to return the correct format")
@Test
void beanOutputParserRecords() {
BeanOutputParser<ActorsFilmsRecord> outputParser = new BeanOutputParser<>(ActorsFilmsRecord.class);
String format = outputParser.getFormat();
String template = """
Generate the filmography of 5 movies for Tom Hanks.
{format}
Remove Markdown code blocks from the output.
""";
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
Generation generation = client.generate(prompt).getGeneration();
ActorsFilmsRecord actorsFilms = outputParser.parse(generation.getContent());
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
assertThat(actorsFilms.movies()).hasSize(5);
}
@Disabled("TODO: Fix the parser instructions to return the correct format")
@Test
void beanStreamOutputParserRecords() {
BeanOutputParser<ActorsFilmsRecord> outputParser = new BeanOutputParser<>(ActorsFilmsRecord.class);
String format = outputParser.getFormat();
String template = """
Generate the filmography of 5 movies for Tom Hanks.
{format}
Remove Markdown code blocks from the output.
""";
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
Prompt prompt = new Prompt(promptTemplate.createMessage());
String generationTextFromStream = client.generateStream(prompt)
.collectList()
.block()
.stream()
.map(AiResponse::getGenerations)
.flatMap(List::stream)
.map(Generation::getContent)
.collect(Collectors.joining());
ActorsFilmsRecord actorsFilms = outputParser.parse(generationTextFromStream);
System.out.println(actorsFilms);
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
assertThat(actorsFilms.movies()).hasSize(5);
}
@SpringBootConfiguration
public static class TestConfiguration {
@Bean
public TitanChatBedrockApi titanApi() {
return new TitanChatBedrockApi(TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
}
@Bean
public BedrockTitanChatClient titanChatClient(TitanChatBedrockApi titanApi) {
return new BedrockTitanChatClient(titanApi);
}
}
}

4
spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApiIT.java
View File

@@ -24,7 +24,7 @@ import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import reactor.core.publisher.Flux;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatCompletionModel;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatModel;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatRequest;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponse;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponseChunk;
@@ -38,7 +38,7 @@ import static org.assertj.core.api.Assertions.assertThat;
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class TitanChatBedrockApiIT {
TitanChatBedrockApi titanBedrockApi = new TitanChatBedrockApi(TitanChatCompletionModel.TITAN_TEXT_EXPRESS_V1.id(),
TitanChatBedrockApi titanBedrockApi = new TitanChatBedrockApi(TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
Region.EU_CENTRAL_1.id());
TitanChatRequest titanChatRequest = TitanChatRequest.builder("Give me the names of 3 famous pirates?")

BIN
spring-ai-bedrock/src/test/resources/doc/Bedrock Titan Chat API.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 416 KiB

68
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanChatAutoConfiguration.java Normal file
View File

@@ -0,0 +1,68 @@
/*
* Copyright 2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.bedrock.titan;
import com.fasterxml.jackson.databind.ObjectMapper;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import org.springframework.ai.autoconfigure.NativeHints;
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionConfiguration;
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
import org.springframework.ai.bedrock.titan.BedrockTitanChatClient;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Import;
import org.springframework.context.annotation.ImportRuntimeHints;
/**
* {@link AutoConfiguration Auto-configuration} for Bedrock Titan Chat Client.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@AutoConfiguration
@ConditionalOnClass(TitanChatBedrockApi.class)
@EnableConfigurationProperties({ BedrockTitanChatProperties.class, BedrockAwsConnectionProperties.class })
@ConditionalOnProperty(prefix = BedrockTitanChatProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@Import(BedrockAwsConnectionConfiguration.class)
@ImportRuntimeHints(NativeHints.class)
public class BedrockTitanChatAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public TitanChatBedrockApi cohereApi(AwsCredentialsProvider credentialsProvider,
BedrockTitanChatProperties properties, BedrockAwsConnectionProperties awsProperties) {
return new TitanChatBedrockApi(properties.getModel(), credentialsProvider, awsProperties.getRegion(),
new ObjectMapper());
}
@Bean
public BedrockTitanChatClient cohereChatClient(TitanChatBedrockApi titanChatApi,
BedrockTitanChatProperties properties) {
return new BedrockTitanChatClient(titanChatApi).withTemperature(properties.getTemperature())
.withTopP(properties.getTopP())
.withMaxTokenCount(properties.getMaxTokenCount())
.withStopSequences(properties.getStopSequences());
}
}

118
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanChatProperties.java Normal file
View File

@@ -0,0 +1,118 @@
/*
* Copyright 2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.bedrock.titan;
import java.util.List;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatModel;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* Bedrock Titan Chat autoconfiguration properties.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@ConfigurationProperties(BedrockTitanChatProperties.CONFIG_PREFIX)
public class BedrockTitanChatProperties {
public static final String CONFIG_PREFIX = "spring.ai.bedrock.titan.chat";
/**
* Enable Bedrock Titan Chat Client. False by default.
*/
private boolean enabled = false;
/**
* Bedrock Titan Chat model name. Defaults to 'amazon.titan-text-express-v1'.
*/
private String model = TitanChatModel.TITAN_TEXT_EXPRESS_V1.id();
/**
* (optional) Use a lower value to decrease randomness in the response. Defaults to
* 0.7.
*/
private Float temperature = 0.7f;
/**
* (optional) The maximum cumulative probability of tokens to consider when sampling.
* The model uses combined Top-k and nucleus sampling. Nucleus sampling considers the
* smallest set of tokens whose probability sum is at least topP.
*/
private Float topP;
/**
* (optional) Specify the maximum number of tokens to use in the generated response.
*/
private Integer maxTokenCount;
/**
* (optional) Configure up to four sequences that the model recognizes. After a stop
* sequence, the model stops generating further tokens. The returned text doesn't
* contain the stop sequence.
*/
private List<String> stopSequences;
public boolean isEnabled() {
return enabled;
}
public void setEnabled(boolean enabled) {
this.enabled = enabled;
}
public String getModel() {
return model;
}
public void setModel(String model) {
this.model = model;
}
public Float getTemperature() {
return temperature;
}
public void setTemperature(Float temperature) {
this.temperature = temperature;
}
public Float getTopP() {
return topP;
}
public void setTopP(Float topP) {
this.topP = topP;
}
public Integer getMaxTokenCount() {
return maxTokenCount;
}
public void setMaxTokenCount(Integer maxTokens) {
this.maxTokenCount = maxTokens;
}
public List<String> getStopSequences() {
return stopSequences;
}
public void setStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
}
}

157
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanChatAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,157 @@
/*
* Copyright 2023-2023 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.bedrock.titan;
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 reactor.core.publisher.Flux;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
import org.springframework.ai.bedrock.titan.BedrockTitanChatClient;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatModel;
import org.springframework.ai.client.AiResponse;
import org.springframework.ai.client.Generation;
import org.springframework.ai.prompt.Prompt;
import org.springframework.ai.prompt.SystemPromptTemplate;
import org.springframework.ai.prompt.messages.Message;
import org.springframework.ai.prompt.messages.UserMessage;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @since 0.8.0
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class BedrockTitanChatAutoConfigurationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
"spring.ai.bedrock.titan.chat.model=" + TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
"spring.ai.bedrock.titan.chat.temperature=0.5", "spring.ai.bedrock.titan.chat.maxTokens=500")
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class));
private final Message systemMessage = new SystemPromptTemplate("""
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}.
""").createMessage(Map.of("name", "Bob", "voice", "pirate"));
private final UserMessage userMessage = new UserMessage(
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
@Test
public void chatCompletion() {
contextRunner.run(context -> {
BedrockTitanChatClient chatClient = context.getBean(BedrockTitanChatClient.class);
AiResponse response = chatClient.generate(new Prompt(List.of(userMessage, systemMessage)));
assertThat(response.getGeneration().getContent()).contains("Blackbeard");
});
}
@Test
public void chatCompletionStreaming() {
contextRunner.run(context -> {
BedrockTitanChatClient chatClient = context.getBean(BedrockTitanChatClient.class);
Flux<AiResponse> response = chatClient.generateStream(new Prompt(List.of(userMessage, systemMessage)));
List<AiResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(1);
String stitchedResponseContent = responses.stream()
.map(AiResponse::getGenerations)
.flatMap(List::stream)
.map(Generation::getContent)
.collect(Collectors.joining());
assertThat(stitchedResponseContent).contains("Blackbeard");
});
}
@Test
public void propertiesTest() {
new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
"spring.ai.bedrock.titan.chat.model=MODEL_XYZ",
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.titan.chat.temperature=0.55", "spring.ai.bedrock.titan.chat.topP=0.55",
"spring.ai.bedrock.titan.chat.stopSequences=END1,END2",
"spring.ai.bedrock.titan.chat.maxTokenCount=123")
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
.run(context -> {
var chatProperties = context.getBean(BedrockTitanChatProperties.class);
var aswProperties = context.getBean(BedrockAwsConnectionProperties.class);
assertThat(chatProperties.isEnabled()).isTrue();
assertThat(aswProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
assertThat(chatProperties.getModel()).isEqualTo("MODEL_XYZ");
assertThat(chatProperties.getTemperature()).isEqualTo(0.55f);
assertThat(chatProperties.getTopP()).isEqualTo(0.55f);
assertThat(chatProperties.getStopSequences()).isEqualTo(List.of("END1", "END2"));
assertThat(chatProperties.getMaxTokenCount()).isEqualTo(123);
assertThat(aswProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
assertThat(aswProperties.getSecretKey()).isEqualTo("SECRET_KEY");
});
}
@Test
public void chatCompletionDisabled() {
// It is disabled by default
new ApplicationContextRunner().withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockTitanChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockTitanChatClient.class)).isEmpty();
});
// Explicitly enable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.chat.enabled=true")
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockTitanChatProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(BedrockTitanChatClient.class)).isNotEmpty();
});
// Explicitly disable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.chat.enabled=false")
.withConfiguration(AutoConfigurations.of(BedrockTitanChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockTitanChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockTitanChatClient.class)).isEmpty();
});
}
}