Add AWS Bedrock AI support

- Add spring-ai-bedrock project with support for Cohere, Llama2, Ai21 Jurassic 2, Titan and Anthropic LLM models.
 - Add native API clients for CohereChat CohereEmbedding , Llama2Chat, JurassicChat, TitanChat and TitanEmbedding models, supporting both single shot and streaming completions (for the models that allows it)
 - Add ITs tests for the native API clients.
 - Implement Chat (AiClient) and ChatStreaming (AiStreamingClient) and EmbeddingClients (according to the models’ support for those) for Cohere, Llama2, and Anthropica. Titan and Jurassic2 are WIP.
 - Add ITs for the ChatClient, ChatStreamingClient and EmbeddingClient implementations.
 - Add Spring Boot Auto-configurations with flexible properties for the Llama2, Anthropic and Cohere modes + ITs
 - Add Spring Boot Starter configurations for all Bedrock models.
 - Add README documentations for all models.
 - Add BedrockAi APIs AOT hints
 - Add Ai21Jurassic2ChatBedrockApi, TitanEmbeddingBedrockApi, TitanChatBedrockApi
 - Add TitanChatBedrockApi

Resolves #66
This commit is contained in:
Christian Tzolov
2023-12-17 23:31:28 +01:00
committed by Mark Pollack
parent 5b46361373
commit 0e3192a8df
72 changed files with 6115 additions and 43 deletions

View File

@@ -47,6 +47,8 @@
<module>vector-stores/spring-ai-redis</module>
<module>spring-ai-vertex-ai</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-vertex-ai</module>
<module>spring-ai-bedrock</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai</module>
</modules>
@@ -98,6 +100,7 @@
<azure-open-ai-client.version>1.0.0-beta.3</azure-open-ai-client.version>
<jtokkit.version>0.6.1</jtokkit.version>
<victools.version>4.31.1</victools.version>
<bedrockruntime.version>2.22.0</bedrockruntime.version>
<!-- readers/writer/stores dependencies-->
<pdfbox.version>3.0.0</pdfbox.version>

View File

@@ -114,7 +114,7 @@ public class AzureOpenAiClient implements AiClient {
List<ChatMessage> azureMessages = new ArrayList<>();
for (Message message : messages) {
String messageType = message.getMessageTypeValue();
String messageType = message.getMessageType().getValue();
ChatRole chatRole = ChatRole.fromString(messageType);
azureMessages.add(new ChatMessage(chatRole, message.getContent()));
}

View File

@@ -0,0 +1,83 @@
# Bedrock AI Chat and Embedding Clients
[Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) is a managed service that provides foundation models from various AI providers, available through a unified API.
Spring AI implements `API` clients for the [Bedrock models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html) along with implementations for the `AiClient`, `AiStreamingClient` and the `EmbeddingClient`.
The API clients provide structured, type-safe implementation for the Bedrock models, while the `AiClient`, `AiStreamingClient` and the `EmbeddingClient` implementations provide Chat and Embedding clients compliant with the Spring-AI API. Later can be used interchangeably with the other (e.g. OpenAI, Azure OpenAI,
Ollama) model clients.
Also Spring-AI provides Spring Auto-Configurations and Boot Starters for all clients, making it easy to bootstrap and configure for the Bedrocks models.
## Prerequisite
* AWS credentials.
If you dont have AWS account and AWS Cli configured yet then this video guide can help you to configure it: [AWS CLI & SDK Setup in Less Than 4 Minutes!](https://youtu.be/gswVHTrRX8I?si=buaY7aeI0l3-bBVb).
You should be able to obtain your access and security keys.
* Enable Bedrock models to use
Go to [Amazon Bedrock](https://us-east-1.console.aws.amazon.com/bedrock/home) and from the [Model Access](https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess) menu on the left configure the access to the models you are going to use.
## Quick start
Add the `spring-ai-bedrock-ai-spring-boot-starter` dependency to your project POM:
```xml
<dependency>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<groupId>org.springframework.ai</groupId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
### Connect to AWS Bedrock
Use the `BedrockAwsConnectionProperties` to configure the AWS credentials and region:
```shell
spring.ai.bedrock.aws.region=us-east-1
spring.ai.bedrock.aws.access-key=YOUR_ACCESS_KEY
spring.ai.bedrock.aws.secret-key=YOUR_SECRET_KEY
```
The `region` property is compulsory.
The AWS credentials are resolved in the following this order:
* Spring-AI Bedrock `spring.ai.bedrock.aws.access-key` and `spring.ai.bedrock.aws.secret-key` properties.
* Java System Properties - `aws.accessKeyId` and `aws.secretAccessKey`
* Environment Variables - `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`
* Web Identity Token credentials from system properties or environment variables
* Credential profiles file at the default location (`~/.aws/credentials`) shared by all AWS SDKs and the AWS CLI
* Credentials delivered through the Amazon EC2 container service if `AWS_CONTAINER_CREDENTIALS_RELATIVE_URI`" environment variable is set and security manager has permission to access the variable,
* Instance profile credentials delivered through the Amazon EC2 metadata service or set the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables.
### Enable selected Bedrock model
> **NOTE**: By default all models are disabled. You have to enable the chosen Bedrock models explicitly, using the `spring.ai.bedrock.<model>.<chat|embedding>.enabled=true` property.
Here are the supported `<model>` and `<chat|embedding>` combinations:
| Model | Chat | Chat Streaming | Embedding |
| ------------- | ------------- | ------------- | ------------- |
| llama2 | Yes | Yes | No |
| cohere | Yes | Yes | Yes |
| anthropic | Yes | Yes | No |
| jurassic2 | Yes | No | No |
| titan | Yes | Yes | Yes (no batch mode!) |
For example to enable the bedrock Llama2 Chat client you need to set the
`spring.ai.bedrock.llama2.chat.enabled=true`.
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`
* (WIP) [Spring AI Bedrock Titan Embedding](./README_TITAN_EMBEDING.md) - `spring.ai.bedrock.titan.embedding=true`
* (WIP) [Spring AI Bedrock Ai21 Jurassic2 Chat](./README_JURASSIC2_CHAT.md) - `spring.ai.bedrock.jurassic2.chat=true`

View File

@@ -0,0 +1,90 @@
# 1. Bedrock Anthropic
Provides Bedrock Anthropic Chat API and Spring-AI chat clients.
## 1.1 AnthropicChatBedrockApi
[AnthropicChatBedrockApi](./src/main/java/org/springframework/ai/bedrock/anthropic/api/AnthropicChatBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Anthropic Claude models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html).
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
![AnthropicChatBedrockApi Class Diagram](./src/test/resources/doc/Bedrock-Anthropic-Chat-API.jpg)
The AnthropicChatBedrockApi supports the `anthropic.claude-instant-v1` and `anthropic.claude-v2` models.
Also the AnthropicChatBedrockApi supports both synchronous (e.g. `chatCompletion()`) and streaming (e.g. `chatCompletionStream()`) responses.
Here is a simple snippet how to use the api programmatically:
```java
AnthropicChatBedrockApi anthropicChatApi = new AnthropicChatBedrockApi(
AnthropicModel.CLAUDE_V2.id(),
Region.EU_CENTRAL_1.id());
AnthropicChatRequest request = AnthropicChatRequest
.builder(String.format(AnthropicChatBedrockApi.PROMPT_TEMPLATE, "Name 3 famous pirates"))
.withTemperature(0.8f)
.withMaxTokensToSample(300)
.withTopK(10)
// .withStopSequences(List.of("\n\nHuman:"))
.build();
AnthropicChatResponse response = anthropicChatApi.chatCompletion(request);
System.out.println(response.completion());
// Streaming response
Flux<AnthropicChatResponse> responseStream = anthropicChatApi.chatCompletionStream(request);
List<AnthropicChatResponse> responses = responseStream.collectList().block();
System.out.println(responses);
```
Follow the [AnthropicChatBedrockApi.java](./src/main/java/org/springframework/ai/bedrock/anthropic/api/AnthropicChatBedrockApi.java)'s JavaDoc for further information.
## 1.2 BedrockAnthropicChatClient
[BedrockAnthropicChatClient](./src/main/java/org/springframework/ai/bedrock/anthropic/BedrockAnthropicChatClient.java) implements the Spring-Ai `AiClient` and `AiStreamingClient` on top of the `AnthropicChatBedrockApi`.
You can use like this:
```java
@Bean
public AnthropicChatBedrockApi anthropicApi() {
return new AnthropicChatBedrockApi(
AnthropicChatBedrockApi.AnthropicModel.CLAUDE_V2.id(),
EnvironmentVariableCredentialsProvider.create(),
Region.EU_CENTRAL_1.id(),
new ObjectMapper());
}
@Bean
public BedrockAnthropicChatClient anthropicChatClient(AnthropicChatBedrockApi anthropicApi) {
return new BedrockAnthropicChatClient(anthropicApi);
}
```
or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Spring Boot starter:
```xml
<dependency>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<groupId>org.springframework.ai</groupId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
And set `spring.ai.bedrock.anthropic.chat.enabled=true`.
By default the client is disabled.
Use the `BedrockAnthropicChatProperties` to configure the Bedrock Llama2 Chat client:
| Property | Description | Default |
| ------------- | ------------- | ------------- |
| spring.ai.bedrock.anthropic.chat.enable | Enable Bedrock Llama2 chat client. Disabled by default | false |
| spring.ai.bedrock.anthropic.chat.awsRegion | AWS region to use. | us-east-1 |
| spring.ai.bedrock.anthropic.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.8 |
| spring.ai.bedrock.anthropic.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default |
| spring.ai.bedrock.anthropic.chat.maxGenLen | Specify the maximum number of tokens to use in the generated response. | 300 |
| spring.ai.bedrock.anthropic.chat.model | The model id to use. See the `Llama2ChatCompletionModel` for the supported models. | meta.llama2-70b-chat-v1 |

View File

@@ -0,0 +1,178 @@
# 1. Bedrock Cohere Chat
Provides Bedrock Cohere Chat clients.
## 1.1 CohereChatBedrockApi
[CohereChatBedrockApi](./src/main/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Cohere Command models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command.html).
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
![CohereChatBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Cohere%20Chat%20API.jpg)
The CohereChatBedrockApi supports the `cohere.command-light-text-v14` and `cohere.command-text-v14` 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
CohereChatBedrockApi cohereChatApi = new CohereChatBedrockApi(
CohereChatModel.COHERE_COMMAND_V14.id(),
Region.US_EAST_1.id());
var request = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withStream(false)
.withTemperature(0.5f)
.withTopP(0.8f)
.withTopK(15)
.withMaxTokens(100)
.withStopSequences(List.of("END"))
.withReturnLikelihoods(CohereChatRequest.ReturnLikelihoods.ALL)
.withNumGenerations(3)
.withLogitBias(null)
.withTruncate(Truncate.NONE)
.build();
CohereChatResponse response = cohereChatApi.chatCompletion(request);
var request = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withStream(true)
.withTemperature(0.5f)
.withTopP(0.8f)
.withTopK(15)
.withMaxTokens(100)
.withStopSequences(List.of("END"))
.withReturnLikelihoods(CohereChatRequest.ReturnLikelihoods.ALL)
.withNumGenerations(3)
.withLogitBias(null)
.withTruncate(Truncate.NONE)
.build();
Flux<CohereChatResponse.Generation> responseStream = cohereChatApi.chatCompletionStream(request);
List<CohereChatResponse.Generation> responses = responseStream.collectList().block();
```
## 1.2 BedrockCohereChatClient
[BedrockCohereChatClient](./src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClient.java) implements the Spring-Ai `AiClient` and `AiStreamingClient` on top of the `CohereChatBedrockApi`.
You can use like this:
```java
@Bean
public CohereChatBedrockApi cohereApi() {
return new CohereChatBedrockApi(
CohereChatModel.COHERE_COMMAND_V14.id(),
EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(),
new ObjectMapper());
}
@Bean
public BedrockCohereChatClient cohereChatClient(CohereChatBedrockApi cohereApi) {
return new BedrockCohereChatClient(cohereApi);
}
```
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.chat.enabled=true`.
By default the client is disabled.
Use the `BedrockCohereChatProperties` to configure the Bedrock Cohere Chat client:
| Property | Description | Default |
| ------------- | ------------- | ------------- |
| spring.ai.bedrock.cohere.chat.enable | Enable Bedrock Cohere chat client. Disabled by default | false |
| spring.ai.bedrock.cohere.chat.awsRegion | AWS region to use. | us-east-1 |
| spring.ai.bedrock.cohere.chat.model | The model id to use. See the `CohereChatModel` for the supported models. | cohere.command-text-v14 |
| spring.ai.bedrock.cohere.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.7 |
| spring.ai.bedrock.cohere.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.topK | Specify the number of token choices the model uses to generate the next token | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.maxTokens | Specify the maximum number of tokens to use in the generated response. | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.stopSequences | Configure up to four sequences that the model recognizes. | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.returnLikelihoods | The token likelihoods are returned with the response. | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.numGenerations | The maximum number of generations that the model should return. | AWS Bedrock default |
| spring.ai.bedrock.cohere.chat.logitBiasToken | Prevents the model from generating unwanted tokens or incentivize the model to include desired tokens. | AWS Bedrock default |
| 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 |

View File

@@ -0,0 +1,72 @@
# 1. Bedrock Cohere Embedding
## 1.1 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);
```
## 1.2 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 |

View File

@@ -0,0 +1,7 @@
# (WIP) Ai21 Jurassic2
Use the [Ai21Jurassic2ChatBedrockApi.java](src/main/java/org/springframework/ai/bedrock/jurassic/api/Ai21Jurassic2ChatBedrockApi.java) Bedrock Jurassic client to implement `AiClient`.
The Ai21Jurassic2 model doesn't support streaming!
Consult the the existing Llama2 and Cohere implementation.

View File

@@ -0,0 +1,89 @@
# 1. Bedrock Llama2
Provides Bedrock Llama2 Chat API and Spring-AI chat clients.
## 1.1 Llama2ChatBedrockApi
[Llama2ChatBedrockApi](./src/main/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Meta Llama 2 and Llama 2 Chat models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html).
Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
![Llama2ChatBedrockApi Class Diagram](./src/test/resources/doc/Bedrock-Llama2-Chat-API.jpg)
The Llama2ChatBedrockApi supports the `meta.llama2-13b-chat-v1` and `meta.llama2-70b-chat-v1` models.
Also the Llama2ChatBedrockApi supports both synchronous (e.g. `chatCompletion()`) and streaming (e.g. `chatCompletionStream()`) responses.
Here is a simple snippet how to use the api programmatically:
```java
Llama2ChatBedrockApi llama2ChatApi = new Llama2ChatBedrockApi(
Llama2ChatCompletionModel.LLAMA2_70B_CHAT_V1.id(),
Region.US_EAST_1.id());
Llama2ChatRequest request = Llama2ChatRequest.builder("Hello, my name is")
.withTemperature(0.9f)
.withTopP(0.9f)
.withMaxGenLen(20)
.build();
Llama2ChatResponse response = llama2ChatApi.chatCompletion(request);
System.out.println(response.generation());
// Streaming response
Flux<Llama2ChatResponse> responseStream = llama2ChatApi.chatCompletionStream(request);
List<Llama2ChatResponse> responses = responseStream.collectList().block();
System.out.println(responses);
```
Follow the [Llama2ChatBedrockApi.java](./src/main/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApi.java)'s JavaDoc for further information.
## 1.2 BedrockLlama2ChatClient
[BedrockLlama2ChatClient](./src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClient.java) implements the Spring-Ai `AiClient` and `AiStreamingClient` on top of the `Llama2ChatBedrockApi`.
You can use like this:
```java
@Bean
public Llama2ChatBedrockApi llama2Api() {
return new Llama2ChatBedrockApi(
Llama2ChatCompletionModel.LLAMA2_70B_CHAT_V1.id(),
EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(),
new ObjectMapper());
}
@Bean
public BedrockLlama2ChatClient llama2ChatClient(Llama2ChatBedrockApi llama2Api) {
return new BedrockLlama2ChatClient(llama2Api);
}
```
or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Spring Boot starter:
```xml
<dependency>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<groupId>org.springframework.ai</groupId>
<version>0.8.0-SNAPSHOT</version>
</dependency>
```
And set `spring.ai.bedrock.llama2.chat.enabled=true`.
By default the client is disabled.
Use the `BedrockLlama2ChatProperties` to configure the Bedrock Llama2 Chat client:
| Property | Description | Default |
| ------------- | ------------- | ------------- |
| spring.ai.bedrock.llama2.chat.enable | Enable Bedrock Llama2 chat client. Disabled by default | false |
| spring.ai.bedrock.llama2.chat.awsRegion | AWS region to use. | us-east-1 |
| spring.ai.bedrock.llama2.chat.temperature | Controls the randomness of the output. Values can range over [0.0,1.0] | 0.7 |
| spring.ai.bedrock.llama2.chat.topP | The maximum cumulative probability of tokens to consider when sampling. | AWS Bedrock default |
| spring.ai.bedrock.llama2.chat.maxGenLen | Specify the maximum number of tokens to use in the generated response. | 300 |
| spring.ai.bedrock.llama2.chat.model | The model id to use. See the `Llama2ChatCompletionModel` for the supported models. | meta.llama2-70b-chat-v1 |

View File

@@ -0,0 +1,5 @@
# (WIP) 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`.
Consult the the existing Llama2 and Cohere implementation.

View File

@@ -0,0 +1,6 @@
# (WIP) Bedrock Titan Embedding
Use the [TitanEmbeddingBedrockApi.java](src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java) Bedrock Embedding client to implement `EmbeddingClient`.
Consult the the existing Cohere embedding client implementation.
Mind that Titan doesn't support batch embedding. You have to either emulate it (could be very expensive) or throw a not supported exception.

58
spring-ai-bedrock/pom.xml Normal file
View File

@@ -0,0 +1,58 @@
<?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>0.8.0-SNAPSHOT</version>
</parent>
<artifactId>spring-ai-bedrock</artifactId>
<packaging>jar</packaging>
<name>Spring AI Bedrock</name>
<description>Bedrock AI support</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects/spring-ai</url>
<connection>git://github.com/spring-projects/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
</scm>
<dependencies>
<!-- production dependencies -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-core</artifactId>
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
<version>${spring-framework.version}</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>bedrockruntime</artifactId>
<version>${bedrockruntime.version}</version>
</dependency>
<!-- test dependencies -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-test</artifactId>
<version>${project.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -0,0 +1,61 @@
/*
* 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.bedrock;
import org.springframework.ai.bedrock.api.AbstractBedrockApi.AmazonBedrockInvocationMetrics;
import org.springframework.ai.metadata.Usage;
import org.springframework.util.Assert;
/**
* {@link Usage} implementation for Bedrock API.
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class BedrockUsage implements Usage {
public static BedrockUsage from(AmazonBedrockInvocationMetrics usage) {
return new BedrockUsage(usage);
}
private final AmazonBedrockInvocationMetrics usage;
protected BedrockUsage(AmazonBedrockInvocationMetrics usage) {
Assert.notNull(usage, "OpenAI Usage must not be null");
this.usage = usage;
}
protected AmazonBedrockInvocationMetrics getUsage() {
return this.usage;
}
@Override
public Long getPromptTokens() {
return getUsage().inputTokenCount().longValue();
}
@Override
public Long getGenerationTokens() {
return getUsage().outputTokenCount().longValue();
}
@Override
public String toString() {
return getUsage().toString();
}
}

View File

@@ -0,0 +1,91 @@
/*
* 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;
import java.util.List;
import java.util.stream.Collectors;
import org.springframework.ai.prompt.messages.Message;
import org.springframework.ai.prompt.messages.MessageType;
/**
* Converts a list of messages to a prompt for bedrock models.
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class MessageToPromptConverter {
private static final String HUMAN_PROMPT = "Human:";
private static final String ASSISTANT_PROMPT = "Assistant:";
private String humanPrompt = HUMAN_PROMPT;
private String assistantPrompt = ASSISTANT_PROMPT;
private MessageToPromptConverter() {
}
public static MessageToPromptConverter create() {
return new MessageToPromptConverter();
}
public MessageToPromptConverter withHumanPrompt(String humanPrompt) {
this.humanPrompt = humanPrompt;
return this;
}
public MessageToPromptConverter withAssistantPrompt(String assistantPrompt) {
this.assistantPrompt = assistantPrompt;
return this;
}
public String toPrompt(List<Message> messages) {
final String systemMessages = messages.stream()
.filter(message -> message.getMessageType() == MessageType.SYSTEM)
.map(Message::getContent)
.collect(Collectors.joining("\n"));
final String userMessages = messages.stream()
.filter(message -> message.getMessageType() == MessageType.USER
|| message.getMessageType() == MessageType.ASSISTANT)
.map(this::messageToString)
.collect(Collectors.joining("\n"));
final String prompt = String.format("%s%n%n%s%n%s", systemMessages, userMessages, ASSISTANT_PROMPT);
return prompt;
}
protected String messageToString(Message message) {
switch (message.getMessageType()) {
case SYSTEM:
return message.getContent();
case USER:
return humanPrompt + " " + message.getContent();
case ASSISTANT:
return assistantPrompt + " " + message.getContent();
case FUNCTION:
throw new IllegalArgumentException("Tool execution results are not supported for Bedrock models");
}
throw new IllegalArgumentException("Unknown message type: " + message.getMessageType());
}
}

View File

@@ -0,0 +1,135 @@
/*
* 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.anthropic;
import java.util.List;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.MessageToPromptConverter;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatRequest;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatResponse;
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.prompt.Prompt;
/**
* Java {@link AiClient} and {@link AiStreamClient} for the Bedrock Anthropic chat model.
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class BedrockAnthropicChatClient implements AiClient, AiStreamClient {
private final AnthropicChatBedrockApi anthropicChatApi;
private Float temperature = 0.8f;
private Float topP;
private Integer maxTokensToSample = 500;
private Integer topK = 10;
private List<String> stopSequences;
private String anthropicVersion = AnthropicChatBedrockApi.DEFAULT_ANTHROPIC_VERSION;
public BedrockAnthropicChatClient(AnthropicChatBedrockApi chatApi) {
this.anthropicChatApi = chatApi;
}
public BedrockAnthropicChatClient withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public BedrockAnthropicChatClient withMaxTokensToSample(Integer maxTokensToSample) {
this.maxTokensToSample = maxTokensToSample;
return this;
}
public BedrockAnthropicChatClient withTopK(Integer topK) {
this.topK = topK;
return this;
}
public BedrockAnthropicChatClient withTopP(Float tpoP) {
this.topP = tpoP;
return this;
}
public BedrockAnthropicChatClient withStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
return this;
}
public BedrockAnthropicChatClient withAnthropicVersion(String anthropicVersion) {
this.anthropicVersion = anthropicVersion;
return this;
}
@Override
public AiResponse generate(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getMessages());
AnthropicChatRequest request = AnthropicChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withMaxTokensToSample(this.maxTokensToSample)
.withTopK(this.topK)
.withTopP(this.topP)
.withStopSequences(this.stopSequences)
.withAnthropicVersion(this.anthropicVersion)
.build();
AnthropicChatResponse response = this.anthropicChatApi.chatCompletion(request);
return new AiResponse(List.of(new Generation(response.completion())));
}
@Override
public Flux<AiResponse> generateStream(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getMessages());
AnthropicChatRequest request = AnthropicChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withMaxTokensToSample(this.maxTokensToSample)
.withTopK(this.topK)
.withTopP(this.topP)
.withStopSequences(this.stopSequences)
.withAnthropicVersion(this.anthropicVersion)
.build();
Flux<AnthropicChatResponse> fluxResponse = this.anthropicChatApi.chatCompletionStream(request);
return fluxResponse.map(response -> {
String stopReason = response.stopReason() != null ? response.stopReason() : null;
var generation = new Generation(response.completion());
if (response.amazonBedrockInvocationMetrics() != null) {
generation = generation
.withChoiceMetadata(ChoiceMetadata.from(stopReason, response.amazonBedrockInvocationMetrics()));
}
return new AiResponse(List.of(generation));
});
}
}

View File

@@ -0,0 +1,224 @@
/*
* 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.anthropic.api;
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.anthropic.api.AnthropicChatBedrockApi.AnthropicChatRequest;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatResponse;
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
import org.springframework.util.Assert;
/**
* @author Christian Tzolov
*/
// @formatter:off
public class AnthropicChatBedrockApi extends
AbstractBedrockApi<AnthropicChatRequest, AnthropicChatResponse, AnthropicChatResponse> {
public static final String PROMPT_TEMPLATE = "\n\nHuman:%s\n\nAssistant:";
/**
* Default version of the Anthropic chat model.
*/
public static final String DEFAULT_ANTHROPIC_VERSION = "bedrock-2023-05-31";
/**
* Create a new AnthropicChatBedrockApi instance using the default credentials provider chain, the default object.
* @param modelId The model id to use. See the {@link AnthropicChatModel} for the supported models.
* @param region The AWS region to use.
*/
public AnthropicChatBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* Create a new AnthropicChatBedrockApi instance using the provided credentials provider, region and object mapper.
*
* @param modelId The model id to use. See the {@link AnthropicChatModel} 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 AnthropicChatBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
super(modelId, credentialsProvider, region, objectMapper);
}
// https://github.com/build-on-aws/amazon-bedrock-java-examples/blob/main/example_code/bedrock-runtime/src/main/java/aws/community/examples/InvokeBedrockStreamingAsync.java
// https://docs.anthropic.com/claude/reference/complete_post
// https://docs.aws.amazon.com/bedrock/latest/userguide/br-product-ids.html
// Anthropic Claude models: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html
/**
* AnthropicChatRequest encapsulates the request parameters for the Anthropic chat model.
* https://docs.anthropic.com/claude/reference/complete_post
*
* @param prompt The prompt to use for the chat.
* @param temperature (default 0.5) The temperature to use for the chat. You should either alter temperature or
* top_p, but not both.
* @param maxTokensToSample (default 200) Specify the maximum number of tokens to use in the generated response.
* Note that the models may stop before reaching this maximum. This parameter only specifies the absolute maximum
* number of tokens to generate. We recommend a limit of 4,000 tokens for optimal performance.
* @param topK (default 250) Specify the number of token choices the model uses to generate the next token.
* @param topP (default 1) Nucleus sampling to specify the cumulative probability of the next token in range [0,1].
* In nucleus sampling, we compute the cumulative distribution over all the options for each subsequent token in
* decreasing probability order and cut it off once it reaches a particular probability specified by top_p. You
* should either alter temperature or top_p, but not both.
* @param stopSequences (defaults to "\n\nHuman:") 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.
* @param anthropicVersion The version of the model to use. The default value is bedrock-2023-05-31.
*/
@JsonInclude(Include.NON_NULL)
public record AnthropicChatRequest(
@JsonProperty("prompt") String prompt,
@JsonProperty("temperature") Float temperature,
@JsonProperty("max_tokens_to_sample") Integer maxTokensToSample,
@JsonProperty("top_k") Integer topK,
@JsonProperty("top_p") Float topP,
@JsonProperty("stop_sequences") List<String> stopSequences,
@JsonProperty("anthropic_version") String anthropicVersion) {
public static Builder builder(String prompt) {
return new Builder(prompt);
}
public static class Builder {
private final String prompt;
private Float temperature = 0.7f;
private Integer maxTokensToSample = 500;
private Integer topK = 10;
private Float topP;
private List<String> stopSequences;
private String anthropicVersion = DEFAULT_ANTHROPIC_VERSION;
private Builder(String prompt) {
this.prompt = prompt;
}
public Builder withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public Builder withMaxTokensToSample(Integer maxTokensToSample) {
this.maxTokensToSample = maxTokensToSample;
return this;
}
public Builder withTopK(Integer topK) {
this.topK = topK;
return this;
}
public Builder withTopP(Float tpoP) {
this.topP = tpoP;
return this;
}
public Builder withStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
return this;
}
public Builder withAnthropicVersion(String anthropicVersion) {
this.anthropicVersion = anthropicVersion;
return this;
}
public AnthropicChatRequest build() {
return new AnthropicChatRequest(
prompt,
temperature,
maxTokensToSample,
topK,
topP,
stopSequences,
anthropicVersion
);
}
}
}
/**
* AnthropicChatResponse encapsulates the response parameters for the Anthropic chat model.
*
* @param completion The generated text.
* @param stopReason The reason the model stopped generating text.
* @param stop The stop sequence that caused the model to stop generating text.
* @param amazonBedrockInvocationMetrics Metrics about the model invocation.
*/
@JsonInclude(Include.NON_NULL)
public record AnthropicChatResponse(
@JsonProperty("completion") String completion,
@JsonProperty("stop_reason") String stopReason,
@JsonProperty("stop") String stop,
@JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) {
}
/**
* Anthropic models version.
*/
public enum AnthropicChatModel {
/**
* anthropic.claude-instant-v1
*/
CLAUDE_INSTANT_V1("anthropic.claude-instant-v1"),
/**
* anthropic.claude-v2
*/
CLAUDE_V2("anthropic.claude-v2");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
AnthropicChatModel(String value) {
this.id = value;
}
}
@Override
public AnthropicChatResponse chatCompletion(AnthropicChatRequest anthropicRequest) {
Assert.notNull(anthropicRequest, "'anthropicRequest' must not be null");
return this.internalInvocation(anthropicRequest, AnthropicChatResponse.class);
}
@Override
public Flux<AnthropicChatResponse> chatCompletionStream(AnthropicChatRequest anthropicRequest) {
Assert.notNull(anthropicRequest, "'anthropicRequest' must not be null");
return this.internalInvocationStream(anthropicRequest, AnthropicChatResponse.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,266 @@
/*
* 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.
*/
// @formatter:off
package org.springframework.ai.bedrock.api;
import java.io.UncheckedIOException;
import java.nio.charset.StandardCharsets;
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.ObjectMapper;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider;
import software.amazon.awssdk.core.SdkBytes;
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.InvokeModelRequest;
import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelResponse;
import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelWithResponseStreamRequest;
import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelWithResponseStreamResponseHandler;
import software.amazon.awssdk.services.bedrockruntime.model.ResponseStream;
/**
* Abstract class for the Bedrock API. It provides the basic functionality to invoke the chat completion model and
* receive the response for streaming and non-streaming requests.
* <p>
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html
* <p>
* https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess
*
* @param <I> The input request type.
* @param <O> The output response type.
* @param <SO> The streaming response type. For some models this type can be the same as the output response type.
*
* @see <a href="https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html">Model Parameters</a>
* @author Christian Tzolov
* @since 0.8.0
*/
public abstract class AbstractBedrockApi<I, O, SO> {
private static final Log logger = LogFactory.getLog(AbstractBedrockApi.class);
private final String modelId;
private final ObjectMapper objectMapper;
private final AwsCredentialsProvider credentialsProvider;
private final String region;
private final BedrockRuntimeClient client;
private final BedrockRuntimeAsyncClient clientStreaming;
private final Sinks.Many<SO> eventSink;
/**
* Create a new AbstractBedrockApi instance using default credentials provider and object mapper.
*
* @param modelId The model id to use.
* @param region The AWS region to use.
*/
public AbstractBedrockApi(String modelId, String region) {
this(modelId, ProfileCredentialsProvider.builder().build(), region, new ObjectMapper());
}
/**
* Create a new AbstractBedrockApi instance using the provided credentials provider, region and object mapper.
*
* @param modelId The model id to use.
* @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 AbstractBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
this.modelId = modelId;
this.objectMapper = objectMapper;
this.credentialsProvider = credentialsProvider;
this.region = region;
this.eventSink = Sinks.many().unicast().onBackpressureError();
this.client = BedrockRuntimeClient.builder()
.region(Region.of(this.region))
.credentialsProvider(this.credentialsProvider)
.build();
this.clientStreaming = BedrockRuntimeAsyncClient.builder()
.region(Region.of(this.region))
.credentialsProvider(this.credentialsProvider)
.build();
}
/**
* Encapsulates the metrics about the model invocation.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-claude.html
*
* @param inputTokenCount The number of tokens in the input prompt.
* @param firstByteLatency The time in milliseconds between the request being sent and the first byte of the
* response being received.
* @param outputTokenCount The number of tokens in the generated text.
* @param invocationLatency The time in milliseconds between the request being sent and the response being received.
*/
@JsonInclude(Include.NON_NULL)
public record AmazonBedrockInvocationMetrics(
@JsonProperty("inputTokenCount") Long inputTokenCount,
@JsonProperty("firstByteLatency") Long firstByteLatency,
@JsonProperty("outputTokenCount") Long outputTokenCount,
@JsonProperty("invocationLatency") Long invocationLatency) {
}
/**
* Compute the embedding for the given text.
*
* @param request The embedding request.
* @return Returns the embedding response.
*/
protected O embedding(I request) {
throw new UnsupportedOperationException("Embedding is not supported for this model: " + this.modelId);
}
/**
* Chat completion invocation.
*
* @param request The chat completion request.
* @return The chat completion response.
*/
protected O chatCompletion(I request) {
throw new UnsupportedOperationException("Chat completion is not supported for this model: " + this.modelId);
}
/**
* Chat completion invocation with streaming response.
*
* @param request The chat completion request.
* @return The chat completion response stream.
*/
protected Flux<SO> chatCompletionStream(I request) {
throw new UnsupportedOperationException(
"Streaming chat completion is not supported for this model: " + this.modelId);
}
/**
* Internal method to 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_InvokeModel.html
* https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/bedrockruntime/BedrockRuntimeClient.html#invokeModel
*
* @param request Model invocation request.
* @param clazz The response class type
* @return The model invocation response.
*
*/
protected O internalInvocation(I request, Class<O> clazz) {
SdkBytes body;
try {
body = SdkBytes.fromUtf8String(new ObjectMapper().writeValueAsString(request));
}
catch (JsonProcessingException e) {
throw new IllegalArgumentException("Invalid JSON format for the input request: " + request, e);
}
InvokeModelRequest invokeRequest = InvokeModelRequest.builder()
.modelId(this.modelId)
.body(body)
.build();
InvokeModelResponse response = this.client.invokeModel(invokeRequest);
String responseBody = response.body().asString(StandardCharsets.UTF_8);
try {
return this.objectMapper.readValue(responseBody, clazz);
}
catch (JsonProcessingException | UncheckedIOException e) {
throw new IllegalArgumentException("Invalid JSON format for the response: " + responseBody, e);
}
}
/**
* Internal method to invoke the model and return the response stream.
*
* @param request Model invocation request.
* @param clazz Response class type.
* @return The model invocation response stream.
*/
protected Flux<SO> internalInvocationStream(I request, Class<SO> clazz) {
SdkBytes body;
try {
body = SdkBytes.fromUtf8String(this.objectMapper.writeValueAsString(request));
}
catch (JsonProcessingException e) {
this.eventSink.tryEmitError(e);
return this.eventSink.asFlux();
}
InvokeModelWithResponseStreamRequest invokeRequest = InvokeModelWithResponseStreamRequest.builder()
.modelId(this.modelId)
.body(body)
.build();
InvokeModelWithResponseStreamResponseHandler.Visitor visitor = InvokeModelWithResponseStreamResponseHandler.Visitor
.builder()
.onChunk((chunk) -> {
try {
logger.debug("Received chunk: " + chunk.bytes().asString(StandardCharsets.UTF_8));
SO response = this.objectMapper.readValue(chunk.bytes().asByteArray(), clazz);
this.eventSink.tryEmitNext(response);
}
catch (Exception e) {
logger.error("Failed to unmarshall", e);
this.eventSink.tryEmitError(e);
}
})
.onDefault((event) -> {
logger.error("Unknown or unhandled event: " + event.toString());
this.eventSink.tryEmitError(new Throwable("Unknown or unhandled event: " + event.toString()));
})
.build();
InvokeModelWithResponseStreamResponseHandler responseHandler = InvokeModelWithResponseStreamResponseHandler
.builder()
.onComplete(
() -> {
this.eventSink.tryEmitComplete();
logger.debug("\nCompleted streaming response.");
})
.onError((error) -> {
logger.error("\n\nError streaming response: " + error.getMessage());
this.eventSink.tryEmitError(error);
})
.onEventStream((stream) -> {
stream.subscribe(
(ResponseStream e) -> {
e.accept(visitor);
});
})
.build();
this.clientStreaming.invokeModelWithResponseStream(invokeRequest, responseHandler);
return this.eventSink.asFlux();
}
}
// @formatter:on

View File

@@ -0,0 +1,154 @@
/*
* 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.cohere;
import java.util.List;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.BedrockUsage;
import org.springframework.ai.bedrock.MessageToPromptConverter;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatResponse;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.LogitBias;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.ReturnLikelihoods;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.Truncate;
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 BedrockCohereChatClient implements AiClient, AiStreamClient {
private final CohereChatBedrockApi chatApi;
private Float temperature;
private Float topP;
private Integer topK;
private Integer maxTokens;
private List<String> stopSequences;
private ReturnLikelihoods returnLikelihoods;
private Integer numGenerations;
private LogitBias logitBias;
private Truncate truncate;
public BedrockCohereChatClient(CohereChatBedrockApi chatApi) {
this.chatApi = chatApi;
}
public BedrockCohereChatClient withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public BedrockCohereChatClient withTopP(Float topP) {
this.topP = topP;
return this;
}
public BedrockCohereChatClient withTopK(Integer topK) {
this.topK = topK;
return this;
}
public BedrockCohereChatClient withMaxTokens(Integer maxTokens) {
this.maxTokens = maxTokens;
return this;
}
public BedrockCohereChatClient withStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
return this;
}
public BedrockCohereChatClient withReturnLikelihoods(ReturnLikelihoods returnLikelihoods) {
this.returnLikelihoods = returnLikelihoods;
return this;
}
public BedrockCohereChatClient withNumGenerations(Integer numGenerations) {
this.numGenerations = numGenerations;
return this;
}
public BedrockCohereChatClient withLogitBias(LogitBias logitBias) {
this.logitBias = logitBias;
return this;
}
public BedrockCohereChatClient withTruncate(Truncate truncate) {
this.truncate = truncate;
return this;
}
@Override
public AiResponse generate(Prompt prompt) {
CohereChatResponse response = this.chatApi.chatCompletion(this.createRequest(prompt, false));
List<Generation> generations = response.generations().stream().map(g -> {
return new Generation(g.text());
}).toList();
return new AiResponse(generations);
}
@Override
public Flux<AiResponse> generateStream(Prompt prompt) {
return this.chatApi.chatCompletionStream(this.createRequest(prompt, true)).map(g -> {
if (g.isFinished()) {
String finishReason = g.finishReason().name();
Usage usage = BedrockUsage.from(g.amazonBedrockInvocationMetrics());
return new AiResponse(
List.of(new Generation("").withChoiceMetadata(ChoiceMetadata.from(finishReason, usage))));
}
return new AiResponse(List.of(new Generation(g.text())));
});
}
private CohereChatRequest createRequest(Prompt prompt, boolean stream) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getMessages());
return CohereChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withTopP(this.topP)
.withTopK(this.topK)
.withMaxTokens(this.maxTokens)
.withStopSequences(this.stopSequences)
.withReturnLikelihoods(this.returnLikelihoods)
.withStream(stream)
.withNumGenerations(this.numGenerations)
.withLogitBias(this.logitBias)
.withTruncate(this.truncate)
.build();
}
}

View File

@@ -0,0 +1,102 @@
/*
* 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.cohere;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingResponse;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.AbstractEmbeddingClient;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.util.Assert;
/**
* {@link org.springframework.ai.embedding.EmbeddingClient} implementation that uses the
* Bedrock Cohere Embedding API. Note: The invocation metrics are not exposed by AWS for
* this API. If this change in the future we will add it as metadata.
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class BedrockCohereEmbeddingClient extends AbstractEmbeddingClient {
private final CohereEmbeddingBedrockApi embeddingApi;
private CohereEmbeddingRequest.InputType inputType = CohereEmbeddingRequest.InputType.search_document;
private CohereEmbeddingRequest.Truncate truncate = CohereEmbeddingRequest.Truncate.NONE;
public BedrockCohereEmbeddingClient(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi) {
this.embeddingApi = cohereEmbeddingBedrockApi;
}
/**
* Cohere Embedding API input types.
* @param inputType the input type to use.
* @return this client.
*/
public BedrockCohereEmbeddingClient withInputType(CohereEmbeddingRequest.InputType inputType) {
this.inputType = inputType;
return this;
}
/**
* Specifies how the API handles inputs longer than the maximum token length. If you
* specify LEFT or RIGHT, the model discards the input until the remaining input is
* exactly the maximum input token length for the model.
* @param truncate the truncate option to use.
* @return this client.
*/
public BedrockCohereEmbeddingClient withTruncate(CohereEmbeddingRequest.Truncate truncate) {
this.truncate = truncate;
return this;
}
@Override
public List<Double> embed(String text) {
return this.embed(List.of(text)).iterator().next();
}
@Override
public List<Double> embed(Document document) {
return embed(document.getContent());
}
@Override
public List<List<Double>> embed(List<String> texts) {
Assert.notEmpty(texts, "At least one text is required!");
var request = new CohereEmbeddingRequest(texts, this.inputType, this.truncate);
CohereEmbeddingResponse response = this.embeddingApi.embedding(request);
return response.embeddings();
}
@Override
public EmbeddingResponse embedForResponse(List<String> texts) {
var indexCounter = new AtomicInteger(0);
List<Embedding> embeddings = this.embed(texts)
.stream()
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
.toList();
return new EmbeddingResponse(embeddings);
}
}

View File

@@ -0,0 +1,364 @@
/*
* 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.
*/
// @formatter:off
package org.springframework.ai.bedrock.cohere.api;
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.cohere.api.CohereChatBedrockApi.CohereChatRequest;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatResponse;
import org.springframework.util.Assert;
/**
* Java client for the Bedrock Cohere chat model.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere.html
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class CohereChatBedrockApi extends
AbstractBedrockApi<CohereChatRequest, CohereChatResponse, CohereChatResponse.Generation> {
/**
* Create a new CohereChatBedrockApi instance using the default credentials provider chain, the default object
* mapper, default temperature and topP values.
*
* @param modelId The model id to use. See the {@link CohereChatModel} for the supported models.
* @param region The AWS region to use.
*/
public CohereChatBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* Create a new CohereChatBedrockApi instance using the provided credentials provider, region and object mapper.
*
* @param modelId The model id to use. See the {@link CohereChatModel} 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 CohereChatBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
super(modelId, credentialsProvider, region, objectMapper);
}
/**
* CohereChatRequest encapsulates the request parameters for the Cohere command model.
*
* @param prompt The input prompt to generate the response from.
* @param temperature (optional) Use a lower value to decrease randomness in the response.
* @param topP (optional) Use a lower value to ignore less probable options. Set to 0 or 1.0 to disable.
* @param topK (optional) Specify the number of token choices the model uses to generate the next token.
* @param maxTokens (optional) Specify the maximum number of tokens to use in the generated response.
* @param stopSequences (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.
* @param returnLikelihoods (optional) Specify how and if the token likelihoods are returned with the response.
* @param stream (optional) Specify true to return the response piece-by-piece in real-time and false to return the
* complete response after the process finishes.
* @param numGenerations (optional) The maximum number of generations that the model should return.
* @param logitBias (optional) prevents the model from generating unwanted tokens or incentivize the model to
* include desired tokens. The format is {token_id: bias} where bias is a float between -10 and 10. Tokens can be
* obtained from text using any tokenization service, such as Coheres Tokenize endpoint.
* @param truncate (optional) Specifies how the API handles inputs longer than the maximum token length.
*/
@JsonInclude(Include.NON_NULL)
public record CohereChatRequest(
@JsonProperty("prompt") String prompt,
@JsonProperty("temperature") Float temperature,
@JsonProperty("p") Float topP,
@JsonProperty("k") Integer topK,
@JsonProperty("max_tokens") Integer maxTokens,
@JsonProperty("stop_sequences") List<String> stopSequences,
@JsonProperty("return_likelihoods") ReturnLikelihoods returnLikelihoods,
@JsonProperty("stream") boolean stream,
@JsonProperty("num_generations") Integer numGenerations,
@JsonProperty("logit_bias") LogitBias logitBias,
@JsonProperty("truncate") Truncate truncate) {
/**
* Prevents the model from generating unwanted tokens or incentivize the model to include desired tokens.
*
* @param token The token likelihoods.
* @param bias A float between -10 and 10.
*/
@JsonInclude(Include.NON_NULL)
public record LogitBias(
@JsonProperty("token") String token,
@JsonProperty("bias") Float bias) {
}
/**
* (optional) Specify how and if the token likelihoods are returned with the response.
*/
public enum ReturnLikelihoods {
/**
* Only return likelihoods for generated tokens.
*/
GENERATION,
/**
* Return likelihoods for all tokens.
*/
ALL,
/**
* (Default) Don't return any likelihoods.
*/
NONE
}
/**
* Specifies how the API handles inputs longer than the maximum token length. If you specify START or END, the
* model discards the input until the remaining input is exactly the maximum input token length for the model.
*/
public enum Truncate {
/**
* Returns an error when the input exceeds the maximum input token length.
*/
NONE,
/**
* Discard the start of the input.
*/
START,
/**
* (Default) Discards the end of the input.
*/
END
}
/**
* Get CohereChatRequest builder.
* @param prompt compulsory request prompt parameter.
* @return CohereChatRequest builder.
*/
public static Builder builder(String prompt) {
return new Builder(prompt);
}
/**
* Builder for the CohereChatRequest.
*/
public static class Builder {
private final String prompt;
private Float temperature;
private Float topP;
private Integer topK;
private Integer maxTokens;
private List<String> stopSequences;
private ReturnLikelihoods returnLikelihoods;
private boolean stream;
private Integer numGenerations;
private LogitBias logitBias;
private Truncate truncate;
public Builder(String prompt) {
this.prompt = prompt;
}
public Builder withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public Builder withTopP(Float topP) {
this.topP = topP;
return this;
}
public Builder withTopK(Integer topK) {
this.topK = topK;
return this;
}
public Builder withMaxTokens(Integer maxTokens) {
this.maxTokens = maxTokens;
return this;
}
public Builder withStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
return this;
}
public Builder withReturnLikelihoods(ReturnLikelihoods returnLikelihoods) {
this.returnLikelihoods = returnLikelihoods;
return this;
}
public Builder withStream(boolean stream) {
this.stream = stream;
return this;
}
public Builder withNumGenerations(Integer numGenerations) {
this.numGenerations = numGenerations;
return this;
}
public Builder withLogitBias(LogitBias logitBias) {
this.logitBias = logitBias;
return this;
}
public Builder withTruncate(Truncate truncate) {
this.truncate = truncate;
return this;
}
public CohereChatRequest build() {
return new CohereChatRequest(
prompt,
temperature,
topP,
topK,
maxTokens,
stopSequences,
returnLikelihoods,
stream,
numGenerations,
logitBias,
truncate
);
}
}
}
/**
* CohereChatResponse encapsulates the response parameters for the Cohere command model.
*
* @param id An identifier for the request (always returned).
* @param prompt The prompt from the input request. (Always returned).
* @param generations A list of generated results along with the likelihoods for tokens requested. (Always
* returned).
*/
@JsonInclude(Include.NON_NULL)
public record CohereChatResponse(
@JsonProperty("id") String id,
@JsonProperty("prompt") String prompt,
@JsonProperty("generations") List<Generation> generations) {
/**
* Generated result along with the likelihoods for tokens requested.
*
* @param id An identifier for the generation. (Always returned).
* @param likelihood The likelihood of the output. The value is the average of the token likelihoods in
* token_likelihoods. Returned if you specify the return_likelihoods input parameter.
* @param tokenLikelihoods An array of per token likelihoods. Returned if you specify the return_likelihoods
* input parameter.
* @param finishReason states the reason why the model finished generating tokens.
* @param isFinished A boolean field used only when stream is true, signifying whether or not there are
* additional tokens that will be generated as part of the streaming response. (Not always returned).
* @param text The generated text.
* @param index In a streaming response, use to determine which generation a given token belongs to. When only
* one response is streamed, all tokens belong to the same generation and index is not returned. index therefore
* is only returned in a streaming request with a value for num_generations that is larger than one.
* @param amazonBedrockInvocationMetrics Encapsulates the metrics about the model invocation.
*/
@JsonInclude(Include.NON_NULL)
public record Generation(
@JsonProperty("id") String id,
@JsonProperty("likelihood") Float likelihood,
@JsonProperty("token_likelihoods") List<TokenLikelihood> tokenLikelihoods,
@JsonProperty("finish_reason") FinishReason finishReason,
@JsonProperty("is_finished") Boolean isFinished,
@JsonProperty("text") String text,
@JsonProperty("index") Integer index,
@JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) {
/**
* @param token The token.
* @param likelihood The likelihood of the token.
*/
@JsonInclude(Include.NON_NULL)
public record TokenLikelihood(
@JsonProperty("token") String token,
@JsonProperty("likelihood") Float likelihood) {
}
/**
* The reason the response finished being generated.
*/
public enum FinishReason {
/**
* The model sent back a finished reply.
*/
COMPLETE,
/**
* The reply was cut off because the model reached the maximum number of tokens for its context length.
*/
MAX_TOKENS,
/**
* Something went wrong when generating the reply.
*/
ERROR,
/**
* the model generated a reply that was deemed toxic. finish_reason is returned only when
* is_finished=true. (Not always returned).
*/
ERROR_TOXIC
}
}
}
/**
* Cohere models version.
*/
public enum CohereChatModel {
/**
* cohere.command-light-text-v14
*/
COHERE_COMMAND_LIGHT_V14("cohere.command-light-text-v14"),
/**
* cohere.command-text-v14
*/
COHERE_COMMAND_V14("cohere.command-text-v14");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
CohereChatModel(String value) {
this.id = value;
}
}
@Override
public CohereChatResponse chatCompletion(CohereChatRequest request) {
Assert.isTrue(!request.stream(), "The request must be configured to return the complete response!");
return this.internalInvocation(request, CohereChatResponse.class);
}
@Override
public Flux<CohereChatResponse.Generation> chatCompletionStream(CohereChatRequest request) {
Assert.isTrue(request.stream(), "The request must be configured to stream the response!");
return this.internalInvocationStream(request, CohereChatResponse.Generation.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,181 @@
/*
* 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.
*/
// @formatter:off
package org.springframework.ai.bedrock.cohere.api;
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 software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingResponse;
/**
* Cohere Embedding API.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere.html#model-parameters-embed
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class CohereEmbeddingBedrockApi extends
AbstractBedrockApi<CohereEmbeddingRequest, CohereEmbeddingResponse, CohereEmbeddingResponse> {
/**
* Create a new CohereEmbeddingBedrockApi instance using the default credentials provider chain, the default object
* mapper, default temperature and topP values.
*
* @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the supported models.
* @param region The AWS region to use.
*/
public CohereEmbeddingBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* Create a new CohereEmbeddingBedrockApi instance using the provided credentials provider, region and object
* mapper.
*
* @param modelId The model id to use. See the {@link CohereEmbeddingModel} 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 CohereEmbeddingBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
super(modelId, credentialsProvider, region, objectMapper);
}
/**
* The Cohere Embed model request.
*
* @param texts An array of strings for the model to embed. For optimal performance, we recommend reducing the
* length of each text to less than 512 tokens. 1 token is about 4 characters.
* @param 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.
* @param truncate Specifies how the API handles inputs longer than the maximum token length. If you specify LEFT or
* RIGHT, the model discards the input until the remaining input is exactly the maximum input token length for the
* model.
*/
@JsonInclude(Include.NON_NULL)
public record CohereEmbeddingRequest(
@JsonProperty("texts") List<String> texts,
@JsonProperty("input_type") InputType inputType,
@JsonProperty("truncate") Truncate truncate) {
/**
* Cohere Embedding API input types.
*/
public enum InputType {
/**
* In search use-cases, use search_document when you encode documents for embeddings that you store in a
* vector database.
*/
search_document,
/**
* Use search_query when querying your vector DB to find relevant documents.
*/
search_query,
/**
* Use classification when using embeddings as an input to a text classifier.
*/
classification,
/**
* Use clustering to cluster the embeddings.
*/
clustering
}
/**
* Specifies how the API handles inputs longer than the maximum token length. If you specify LEFT or RIGHT, the
* model discards the input until the remaining input is exactly the maximum input token length for the model.
*/
public enum Truncate {
/**
* (Default) Returns an error when the input exceeds the maximum input token length.
*/
NONE,
/**
* Discard the start of the input.
*/
LEFT,
/**
* Discards the end of the input.
*/
RIGHT
}
}
/**
* Cohere Embedding response.
*
* @param id An identifier for the response.
* @param embeddings An array of embeddings, where each embedding is an array of floats with 1024 elements. The
* length of the embeddings array will be the same as the length of the original texts array.
* @param texts An array containing the text entries for which embeddings were returned.
* @param amazonBedrockInvocationMetrics Bedrock invocation metrics. Currently bedrock doesn't return
* invocationMetrics for the cohere embedding model.
*/
@JsonInclude(Include.NON_NULL)
public record CohereEmbeddingResponse(
@JsonProperty("id") String id,
@JsonProperty("embeddings") List<List<Double>> embeddings,
@JsonProperty("texts") List<String> texts,
// For future use: Currently bedrock doesn't return invocationMetrics for the cohere embedding model.
@JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) {
}
/**
* Cohere Embedding model ids. https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html
*/
public enum CohereEmbeddingModel {
/**
* cohere.embed-multilingual-v3
*/
COHERE_EMBED_MULTILINGUAL_V1("cohere.embed-multilingual-v3"),
/**
* cohere.embed-english-v3
*/
COHERE_EMBED_ENGLISH_V3("cohere.embed-english-v3");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
CohereEmbeddingModel(String value) {
this.id = value;
}
}
@Override
public CohereEmbeddingResponse embedding(CohereEmbeddingRequest request) {
return this.internalInvocation(request, CohereEmbeddingResponse.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,277 @@
/*
* 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.
*/
// @formatter:off
package org.springframework.ai.bedrock.jurassic2.api;
import java.util.List;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatRequest;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatResponse;
/**
* Java client for the Bedrock Jurassic2 chat model.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-jurassic2.html
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class Ai21Jurassic2ChatBedrockApi extends
AbstractBedrockApi<Ai21Jurassic2ChatRequest, Ai21Jurassic2ChatResponse, Ai21Jurassic2ChatResponse> {
/**
* Create a new Ai21Jurassic2ChatBedrockApi instance using the default credentials provider chain, the default
* object mapper, default temperature and topP values.
*
* @param modelId The model id to use. See the {@link Ai21Jurassic2ChatCompletionModel} for the supported models.
* @param region The AWS region to use.
*/
public Ai21Jurassic2ChatBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* AI21 Jurassic2 chat request parameters.
*
* @param prompt The prompt to use for the chat.
* @param temperature The temperature value controls the randomness of the generated text.
* @param topP The topP value controls the diversity of the generated text. Use a lower value to ignore less
* probable options.
* @param maxTokens Specify the maximum number of tokens to use in the generated response.
* @param stopSequences Configure stop sequences that the model recognizes and after which it stops generating
* further tokens. Press the Enter key to insert a newline character in a stop sequence. Use the Tab key to finish
* inserting a stop sequence.
* @param countPenalty Control repetition in the generated response. Use a higher value to lower the probability of
* generating new tokens that already appear at least once in the prompt or in the completion. Proportional to the
* number of appearances.
* @param presencePenalty Control repetition in the generated response. Use a higher value to lower the probability
* of generating new tokens that already appear at least once in the prompt or in the completion.
* @param frequencyPenalty Control repetition in the generated response. Use a high value to lower the probability
* of generating new tokens that already appear at least once in the prompt or in the completion. The value is
* proportional to the frequency of the token appearances (normalized to text length).
*/
@JsonInclude(Include.NON_NULL)
public record Ai21Jurassic2ChatRequest(
@JsonProperty("prompt") String prompt,
@JsonProperty("temperature") Float temperature,
@JsonProperty("topP") Float topP,
@JsonProperty("maxTokens") Integer maxTokens,
@JsonProperty("stopSequences") List<String> stopSequences,
@JsonProperty("countPenalty") IntegerScalePenalty countPenalty,
@JsonProperty("presencePenalty") FloatScalePenalty presencePenalty,
@JsonProperty("frequencyPenalty") IntegerScalePenalty frequencyPenalty) {
/**
* Penalty with integer scale value.
*
* @param scale The scale value controls the strength of the penalty. Use a higher value to lower the
* probability of generating new tokens that already appear at least once in the prompt or in the completion.
* @param applyToWhitespaces Reduce the probability of repetition of special characters. A true value applies
* the penalty to whitespaces and new lines.
* @param applyToPunctuations Reduce the probability of repetition of special characters. A true value applies
* the penalty to punctuations.
* @param applyToNumbers Reduce the probability of repetition of special characters. A true value applies the
* penalty to numbers.
* @param applyToStopwords Reduce the probability of repetition of special characters. A true value applies the
* penalty to stopwords.
* @param applyToEmojis Reduce the probability of repetition of special characters. A true value applies the
* penalty to emojis.
*/
@JsonInclude(Include.NON_NULL)
public record IntegerScalePenalty(
@JsonProperty("scale") Integer scale,
@JsonProperty("applyToWhitespaces") boolean applyToWhitespaces,
@JsonProperty("applyToPunctuations") boolean applyToPunctuations,
@JsonProperty("applyToNumbers") boolean applyToNumbers,
@JsonProperty("applyToStopwords") boolean applyToStopwords,
@JsonProperty("applyToEmojis") boolean applyToEmojis) {
}
/**
* Penalty with float scale value.
*
* @param scale The scale value controls the strength of the penalty. Use a higher value to lower the
* probability of generating new tokens that already appear at least once in the prompt or in the completion.
* @param applyToWhitespaces Reduce the probability of repetition of special characters. A true value applies
* the penalty to whitespaces and new lines.
* @param applyToPunctuations Reduce the probability of repetition of special characters. A true value applies
* the penalty to punctuations.
* @param applyToNumbers Reduce the probability of repetition of special characters. A true value applies the
* penalty to numbers.
* @param applyToStopwords Reduce the probability of repetition of special characters. A true value applies the
* penalty to stopwords.
* @param applyToEmojis Reduce the probability of repetition of special characters. A true value applies the
* penalty to emojis.
*/
@JsonInclude(Include.NON_NULL)
public record FloatScalePenalty(@JsonProperty("scale") Float scale,
@JsonProperty("applyToWhitespaces") boolean applyToWhitespaces,
@JsonProperty("applyToPunctuations") boolean applyToPunctuations,
@JsonProperty("applyToNumbers") boolean applyToNumbers,
@JsonProperty("applyToStopwords") boolean applyToStopwords,
@JsonProperty("applyToEmojis") boolean applyToEmojis) {
}
}
/**
* Ai21 Jurassic2 chat response.
* https://docs.ai21.com/reference/j2-complete-api-ref#response
*
* @param id The unique identifier of the response.
* @param prompt The prompt used for the chat.
* @param amazonBedrockInvocationMetrics The metrics about the model invocation.
*/
@JsonInclude(Include.NON_NULL)
public record Ai21Jurassic2ChatResponse(
@JsonProperty("id") String id,
@JsonProperty("prompt") Prompt prompt,
@JsonProperty("completions") List<Completion> completions,
@JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) {
/**
*/
@JsonInclude(Include.NON_NULL)
public record Completion(
@JsonProperty("data") Prompt data,
@JsonProperty("finishReason") FinishReason finishReason) {
}
/**
* Provides detailed information about each token in both the prompt and the completions.
*
* @param generatedToken The generatedToken fields.
* @param topTokens The topTokens field is a list of the top K alternative tokens for this position, sorted by
* probability, according to the topKReturn request parameter. If topKReturn is set to 0, this field will be
* null.
* @param textRange The textRange field indicates the start and end offsets of the token in the decoded text
* string.
*/
@JsonInclude(Include.NON_NULL)
public record Token(
@JsonProperty("generatedToken") GeneratedToken generatedToken,
@JsonProperty("topTokens") List<TopToken> topTokens,
@JsonProperty("textRange") TextRange textRange) {
}
/**
* The generatedToken fields.
*
* @param token TThe string representation of the token.
* @param logprob The predicted log probability of the token after applying the sampling parameters as a float
* value.
* @param rawLogprob The raw predicted log probability of the token as a float value. For the indifferent values
* (namely, temperature=1, topP=1) we get raw_logprob=logprob.
*/
@JsonInclude(Include.NON_NULL)
public record GeneratedToken(
@JsonProperty("token") String token,
@JsonProperty("logprob") Float logprob,
@JsonProperty("raw_logprob") Float rawLogprob) {
}
/**
* The topTokens field is a list of the top K alternative tokens for this position, sorted by probability,
* according to the topKReturn request parameter. If topKReturn is set to 0, this field will be null.
*
* @param token The string representation of the alternative token.
* @param logprob The predicted log probability of the alternative token.
*/
@JsonInclude(Include.NON_NULL)
public record TopToken(
@JsonProperty("token") String token,
@JsonProperty("logprob") Float logprob) {
}
/**
* The textRange field indicates the start and end offsets of the token in the decoded text string.
*
* @param start The starting index of the token in the decoded text string.
* @param end The ending index of the token in the decoded text string.
*/
@JsonInclude(Include.NON_NULL)
public record TextRange(
@JsonProperty("start") Integer start,
@JsonProperty("end") Integer end) {
}
/**
* The prompt includes the raw text, the tokens with their log probabilities, and the top-K alternative tokens
* at each position, if requested.
*
* @param text The raw text of the prompt.
* @param tokens Provides detailed information about each token in both the prompt and the completions.
*/
@JsonInclude(Include.NON_NULL)
public record Prompt(
@JsonProperty("text") String text,
@JsonProperty("tokens") List<Token> tokens) {
}
/**
* Explains why the generation process was halted for a specific completion.
*
* @param reason The reason field indicates the reason for the completion to stop.
*
*/
@JsonInclude(Include.NON_NULL)
public record FinishReason(
@JsonProperty("reason") String reason,
@JsonProperty("length") String length,
@JsonProperty("sequence") String sequence) {
}
}
/**
* Ai21 Jurassic2 models version.
*/
public enum Ai21Jurassic2ChatCompletionModel {
/**
* ai21.j2-mid-v1
*/
AI21_J2_MID_V1("ai21.j2-mid-v1"),
/**
* ai21.j2-ultra-v1
*/
AI21_J2_ULTRA_V1("ai21.j2-ultra-v1");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
Ai21Jurassic2ChatCompletionModel(String value) {
this.id = value;
}
}
@Override
public Ai21Jurassic2ChatResponse chatCompletion(Ai21Jurassic2ChatRequest request) {
return this.internalInvocation(request, Ai21Jurassic2ChatResponse.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,121 @@
/*
* 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.llama2;
import java.util.List;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.MessageToPromptConverter;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatRequest;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatResponse;
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;
/**
* Java {@link AiClient} and {@link AiStreamClient} for the Bedrock Llama2 chat model.
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class BedrockLlama2ChatClient implements AiClient, AiStreamClient {
private final Llama2ChatBedrockApi chatApi;
private Float temperature;
private Float topP;
private Integer maxGenLen;
public BedrockLlama2ChatClient(Llama2ChatBedrockApi chatApi) {
this.chatApi = chatApi;
}
public BedrockLlama2ChatClient withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public BedrockLlama2ChatClient withTopP(Float topP) {
this.topP = topP;
return this;
}
public BedrockLlama2ChatClient withMaxGenLen(Integer maxGenLen) {
this.maxGenLen = maxGenLen;
return this;
}
@Override
public AiResponse generate(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getMessages());
var request = Llama2ChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withTopP(this.topP)
.withMaxGenLen(this.maxGenLen)
.build();
Llama2ChatResponse response = this.chatApi.chatCompletion(request);
return new AiResponse(List.of(new Generation(response.generation())
.withChoiceMetadata(ChoiceMetadata.from(response.stopReason().name(), extractUsage(response)))));
}
@Override
public Flux<AiResponse> generateStream(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getMessages());
var request = Llama2ChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withTopP(this.topP)
.withMaxGenLen(this.maxGenLen)
.build();
Flux<Llama2ChatResponse> fluxResponse = this.chatApi.chatCompletionStream(request);
return fluxResponse.map(response -> {
String stopReason = response.stopReason() != null ? response.stopReason().name() : null;
return new AiResponse(List.of(new Generation(response.generation())
.withChoiceMetadata(ChoiceMetadata.from(stopReason, extractUsage(response)))));
});
}
private Usage extractUsage(Llama2ChatResponse response) {
return new Usage() {
@Override
public Long getPromptTokens() {
return response.promptTokenCount().longValue();
}
@Override
public Long getGenerationTokens() {
return response.generationTokenCount().longValue();
}
};
}
}

View File

@@ -0,0 +1,200 @@
/*
* 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.llama2.api;
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.llama2.api.Llama2ChatBedrockApi.Llama2ChatRequest;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatResponse;
// @formatter:off
/**
* Java client for the Bedrock Llama2 chat model.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html
*
* @author Christian Tzolov
* @since 0.8.0
*/
public class Llama2ChatBedrockApi extends
AbstractBedrockApi<Llama2ChatRequest, Llama2ChatResponse, Llama2ChatResponse> {
/**
* Create a new Llama2ChatBedrockApi instance using the default credentials provider chain, the default object
* mapper, default temperature and topP values.
*
* @param modelId The model id to use. See the {@link Llama2ChatCompletionModel} for the supported models.
* @param region The AWS region to use.
*/
public Llama2ChatBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* Create a new Llama2ChatBedrockApi instance using the provided credentials provider, region and object mapper.
*
* @param modelId The model id to use. See the {@link Llama2ChatCompletionModel} 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 Llama2ChatBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
super(modelId, credentialsProvider, region, objectMapper);
}
/**
* Llama2ChatRequest encapsulates the request parameters for the Meta Llama2 chat model.
*
* @param prompt The prompt to use for the chat.
* @param temperature The temperature value controls the randomness of the generated text. Use a lower value to
* decrease randomness in the response.
* @param topP The topP value controls the diversity of the generated text. Use a lower value to ignore less
* probable options. Set to 0 or 1.0 to disable.
* @param maxGenLen The maximum length of the generated text.
*/
@JsonInclude(Include.NON_NULL)
public record Llama2ChatRequest(
@JsonProperty("prompt") String prompt,
@JsonProperty("temperature") Float temperature,
@JsonProperty("top_p") Float topP,
@JsonProperty("max_gen_len") Integer maxGenLen) {
/**
* Create a new Llama2ChatRequest builder.
* @param prompt compulsory prompt parameter.
* @return a new Llama2ChatRequest builder.
*/
public static Builder builder(String prompt) {
return new Builder(prompt);
}
public static class Builder {
private String prompt;
private Float temperature;
private Float topP;
private Integer maxGenLen;
public Builder(String prompt) {
this.prompt = prompt;
}
public Builder withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public Builder withTopP(Float topP) {
this.topP = topP;
return this;
}
public Builder withMaxGenLen(Integer maxGenLen) {
this.maxGenLen = maxGenLen;
return this;
}
public Llama2ChatRequest build() {
return new Llama2ChatRequest(
prompt,
temperature,
topP,
maxGenLen
);
}
}
}
/**
* Llama2ChatResponse encapsulates the response parameters for the Meta Llama2 chat model.
*
* @param generation The generated text.
* @param promptTokenCount The number of tokens in the prompt.
* @param generationTokenCount The number of tokens in the response.
* @param stopReason The reason why the response stopped generating text. Possible values are: (1) stop The model
* has finished generating text for the input prompt. (2) length The length of the tokens for the generated text
* exceeds the value of max_gen_len in the call. The response is truncated to max_gen_len tokens. Consider
* increasing the value of max_gen_len and trying again.
*/
@JsonInclude(Include.NON_NULL)
public record Llama2ChatResponse(
@JsonProperty("generation") String generation,
@JsonProperty("prompt_token_count") Integer promptTokenCount,
@JsonProperty("generation_token_count") Integer generationTokenCount,
@JsonProperty("stop_reason") StopReason stopReason,
@JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) {
/**
* The reason the response finished being generated.
*/
public enum StopReason {
/**
* The model has finished generating text for the input prompt.
*/
stop,
/**
* The response was truncated because of the response length you set.
*/
length
}
}
/**
* Llama2 models version.
*/
public enum Llama2ChatCompletionModel {
/**
* meta.llama2-13b-chat-v1
*/
LLAMA2_13B_CHAT_V1("meta.llama2-13b-chat-v1"),
/**
* meta.llama2-70b-chat-v1
*/
LLAMA2_70B_CHAT_V1("meta.llama2-70b-chat-v1");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
Llama2ChatCompletionModel(String value) {
this.id = value;
}
}
@Override
public Llama2ChatResponse chatCompletion(Llama2ChatRequest request) {
return this.internalInvocation(request, Llama2ChatResponse.class);
}
@Override
public Flux<Llama2ChatResponse> chatCompletionStream(Llama2ChatRequest request) {
return this.internalInvocationStream(request, Llama2ChatResponse.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,230 @@
/*
* 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.api;
import java.util.List;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
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.TitanChatResponse.CompletionReason;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi.TitanChatResponseChunk;
/**
* Java client for the Bedrock Titan chat model.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html
* <p>
* https://docs.aws.amazon.com/bedrock/latest/userguide/titan-text-models.html
*
* @author Christian Tzolov
* @since 0.8.0
*/
// @formatter:off
public class TitanChatBedrockApi extends
AbstractBedrockApi<TitanChatRequest, TitanChatResponse, TitanChatResponseChunk> {
TitanChatBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* TitanChatRequest encapsulates the request parameters for the Titan chat model.
*
* @param inputText The prompt to use for the chat.
* @param textGenerationConfig The text generation configuration.
*/
@JsonInclude(Include.NON_NULL)
public record TitanChatRequest(
@JsonProperty("inputText") String inputText,
@JsonProperty("textGenerationConfig") TextGenerationConfig textGenerationConfig) {
/**
* Titan request text generation configuration.
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html
*
* @param temperature The temperature value controls the randomness of the generated text.
* @param topP The topP value controls the diversity of the generated text. Use a lower value to ignore less
* probable options.
* @param maxTokenCount The maximum number of tokens to generate.
* @param stopSequences A list of sequences to stop the generation at. Specify character sequences to indicate
* where the model should stop. Use the | (pipe) character to separate different sequences (maximum 20
* characters).
*/
@JsonInclude(Include.NON_NULL)
public record TextGenerationConfig(
@JsonProperty("temperature") Float temperature,
@JsonProperty("topP") Float topP,
@JsonProperty("maxTokenCount") Integer maxTokenCount,
@JsonProperty("stopSequences") List<String> stopSequences) {
}
public static Builder builder(String inputText) {
return new Builder(inputText);
}
public static class Builder {
private final String inputText;
private Float temperature;
private Float topP;
private Integer maxTokenCount;
private List<String> stopSequences;
public Builder(String inputText) {
this.inputText = inputText;
}
public Builder withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public Builder withTopP(Float topP) {
this.topP = topP;
return this;
}
public Builder withMaxTokenCount(Integer maxTokenCount) {
this.maxTokenCount = maxTokenCount;
return this;
}
public Builder withStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
return this;
}
public TitanChatRequest build() {
if (this.temperature == null && this.topP == null && this.maxTokenCount == null
&& this.stopSequences == null) {
return new TitanChatRequest(this.inputText, null);
} else {
return new TitanChatRequest(this.inputText,
new TextGenerationConfig(
this.temperature,
this.topP,
this.maxTokenCount,
this.stopSequences
));
}
}
}
}
/**
* TitanChatResponse encapsulates the response parameters for the Titan chat model.
*
* @param inputTextTokenCount The number of tokens in the input text.
* @param results The list of generated responses.
*/
@JsonInclude(Include.NON_NULL)
public record TitanChatResponse(
@JsonProperty("inputTextTokenCount") Integer inputTextTokenCount,
@JsonProperty("results") List<Result> results) {
/**
* Titan response result.
*
* @param tokenCount The number of tokens in the generated text.
* @param outputText The generated text.
* @param completionReason The reason the response finished being generated.
*/
@JsonInclude(Include.NON_NULL)
public record Result(
@JsonProperty("tokenCount") Integer tokenCount,
@JsonProperty("outputText") String outputText,
@JsonProperty("completionReason") CompletionReason completionReason) {
}
/**
* The reason the response finished being generated.
*/
public enum CompletionReason {
/**
* The response was fully generated.
*/
FINISH,
/**
* The response was truncated because of the response length you set.
*/
LENGTH
}
}
/**
* Titan chat model streaming response.
*
* @param outputText The generated text in this chunk.
* @param index The index of the chunk in the streaming response.
* @param inputTextTokenCount The number of tokens in the prompt.
* @param totalOutputTextTokenCount The number of tokens in the response.
* @param completionReason The reason the response finished being generated.
*/
@JsonInclude(Include.NON_NULL)
public record TitanChatResponseChunk(
@JsonProperty("outputText") String outputText,
@JsonProperty("index") Integer index,
@JsonProperty("inputTextTokenCount") Integer inputTextTokenCount,
@JsonProperty("totalOutputTextTokenCount") Integer totalOutputTextTokenCount,
@JsonProperty("completionReason") CompletionReason completionReason,
@JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) {
}
/**
* Titan models version.
*/
public enum TitanChatCompletionModel {
/**
* amazon.titan-text-lite-v1
*/
TITAN_TEXT_LITE_V1("amazon.titan-text-lite-v1"),
/**
* amazon.titan-text-express-v1
*/
TITAN_TEXT_EXPRESS_V1("amazon.titan-text-express-v1");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
TitanChatCompletionModel(String value) {
this.id = value;
}
}
@Override
public TitanChatResponse chatCompletion(TitanChatRequest request) {
return this.internalInvocation(request, TitanChatResponse.class);
}
@Override
public Flux<TitanChatResponseChunk> chatCompletionStream(TitanChatRequest request) {
return this.internalInvocationStream(request, TitanChatResponseChunk.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,159 @@
/*
* 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.api;
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 software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse;
import org.springframework.util.Assert;
/**
* Java client for the Bedrock Titan Embedding model.
* https://docs.aws.amazon.com/bedrock/latest/userguide/titan-multiemb-models.html
*
* @author Christian Tzolov
* @since 0.8.0
*/
// @formatter:off
public class TitanEmbeddingBedrockApi extends
AbstractBedrockApi<TitanEmbeddingRequest, TitanEmbeddingResponse, TitanEmbeddingResponse> {
/**
* Create a new TitanEmbeddingBedrockApi instance using the default credentials provider and default object
* mapper.
* @param modelId The model id to use. See the {@link TitanEmbeddingModel} for the supported models.
* @param region The AWS region to use.
*/
public TitanEmbeddingBedrockApi(String modelId, String region) {
super(modelId, region);
}
/**
* Create a new TitanEmbeddingBedrockApi instance.
*
* @param modelId The model id to use. See the {@link TitanEmbeddingModel} 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 TitanEmbeddingBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region,
ObjectMapper objectMapper) {
super(modelId, credentialsProvider, region, objectMapper);
}
/**
* Titan Embedding request parameters.
*
* @param inputText The text to compute the embedding for.
* @param inputImage The image to compute the embedding for. Only applicable for the 'Titan Multimodal Embeddings
* G1' model.
*/
@JsonInclude(Include.NON_NULL)
public record TitanEmbeddingRequest(
@JsonProperty("inputText") String inputText,
@JsonProperty("inputImage") String inputImage) {
public static Builder builder() {
return new Builder();
}
/**
* TitanEmbeddingRequest builder.
*/
public static class Builder {
private String inputText;
private String inputImage;
public Builder withInputText(String inputText) {
this.inputText = inputText;
return this;
}
public Builder withInputImage(String inputImage) {
this.inputImage = inputImage;
return this;
}
public TitanEmbeddingRequest build() {
Assert.isTrue(this.inputText != null || this.inputImage != null,
"At least one of the inputText or inputImage parameters must be provided!");
Assert.isTrue(!(this.inputText != null && this.inputImage != null),
"Only one of the inputText or inputImage parameters must be provided!");
return new TitanEmbeddingRequest(this.inputText, this.inputImage);
}
}
}
/**
* Titan Embedding response.
*
* @param embedding The embedding vector.
* @param inputTextTokenCount The number of tokens in the input text.
* @param message No idea what this is.
*/
@JsonInclude(Include.NON_NULL)
public record TitanEmbeddingResponse(
@JsonProperty("embedding") List<Double> embedding,
@JsonProperty("inputTextTokenCount") Integer inputTextTokenCount,
@JsonProperty("message") Object message) {
}
/**
* Titan Embedding model ids.
*/
public enum TitanEmbeddingModel {
/**
* amazon.titan-embed-image-v1
*/
TITAN_EMBED_IMAGE_V1("amazon.titan-embed-image-v1"),
/**
* amazon.titan-embed-text-v1
*/
TITAN_EMBED_TEXT_V1("amazon.titan-embed-text-v1");
private final String id;
/**
* @return The model id.
*/
public String id() {
return id;
}
TitanEmbeddingModel(String value) {
this.id = value;
}
}
@Override
public TitanEmbeddingResponse embedding(TitanEmbeddingRequest request) {
return this.internalInvocation(request, TitanEmbeddingResponse.class);
}
}
// @formatter:on

View File

@@ -0,0 +1,167 @@
package org.springframework.ai.bedrock.anthropic;
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.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.anthropic.api.AnthropicChatBedrockApi;
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 BedrockAnthropicChatClientIT {
@Autowired
private BedrockAnthropicChatClient client;
@Value("classpath:/prompts/system-message.st")
private Resource systemResource;
@Test
void roleTest() {
UserMessage userMessage = new UserMessage(
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate"));
Prompt prompt = new Prompt(List.of(userMessage, systemMessage));
AiResponse response = client.generate(prompt);
assertThat(response.getGeneration().getContent()).contains("Blackbeard");
}
@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);
}
@Test
void mapOutputParser() {
MapOutputParser outputParser = new MapOutputParser();
String format = outputParser.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 = 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) {
}
@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.
Remove non JSON tex blocks from the output.
{format}
Provide your answer in the JSON format with the feature names as the keys.
""";
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);
}
@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 AnthropicChatBedrockApi anthropicApi() {
return new AnthropicChatBedrockApi(AnthropicChatBedrockApi.AnthropicChatModel.CLAUDE_V2.id(),
EnvironmentVariableCredentialsProvider.create(), Region.EU_CENTRAL_1.id(), new ObjectMapper());
}
@Bean
public BedrockAnthropicChatClient anthropicChatClient(AnthropicChatBedrockApi anthropicApi) {
return new BedrockAnthropicChatClient(anthropicApi);
}
}
}

View File

@@ -0,0 +1,86 @@
/*
* 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.anthropic.api;
import java.util.List;
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.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatRequest;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatResponse;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatModel;
import static org.assertj.core.api.Assertions.assertThat;;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class AnthropicChatBedrockApiIT {
private AnthropicChatBedrockApi anthropicChatApi = new AnthropicChatBedrockApi(AnthropicChatModel.CLAUDE_V2.id(),
Region.EU_CENTRAL_1.id());
@Test
public void chatCompletion() {
AnthropicChatRequest request = AnthropicChatRequest
.builder(String.format(AnthropicChatBedrockApi.PROMPT_TEMPLATE, "Name 3 famous pirates"))
.withTemperature(0.8f)
.withMaxTokensToSample(300)
.withTopK(10)
.build();
AnthropicChatResponse response = anthropicChatApi.chatCompletion(request);
System.out.println(response.completion());
assertThat(response).isNotNull();
assertThat(response.completion()).isNotEmpty();
assertThat(response.completion()).contains("Blackbeard");
assertThat(response.stopReason()).isEqualTo("stop_sequence");
assertThat(response.stop()).isEqualTo("\n\nHuman:");
assertThat(response.amazonBedrockInvocationMetrics()).isNull();
System.out.println(response);
}
@Test
public void chatCompletionStream() {
AnthropicChatRequest request = AnthropicChatRequest
.builder(String.format(AnthropicChatBedrockApi.PROMPT_TEMPLATE, "Name 3 famous pirates"))
.withTemperature(0.8f)
.withMaxTokensToSample(300)
.withTopK(10)
.withStopSequences(List.of("\n\nHuman:"))
.build();
Flux<AnthropicChatResponse> responseStream = anthropicChatApi.chatCompletionStream(request);
List<AnthropicChatResponse> responses = responseStream.collectList().block();
assertThat(responses).isNotNull();
assertThat(responses).hasSizeGreaterThan(10);
assertThat(responses.stream().map(AnthropicChatResponse::completion).collect(Collectors.joining()))
.contains("Blackbeard");
}
}

View File

@@ -0,0 +1,167 @@
package org.springframework.ai.bedrock.cohere;
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.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.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatModel;
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 BedrockCohereChatClientIT {
@Autowired
private BedrockCohereChatClient 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");
}
@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);
}
@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) {
}
@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);
}
@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 CohereChatBedrockApi cohereApi() {
return new CohereChatBedrockApi(CohereChatModel.COHERE_COMMAND_V14.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
}
@Bean
public BedrockCohereChatClient cohereChatClient(CohereChatBedrockApi cohereApi) {
return new BedrockCohereChatClient(cohereApi);
}
}
}

View File

@@ -0,0 +1,68 @@
package org.springframework.ai.bedrock.cohere;
import java.util.List;
import com.fasterxml.jackson.databind.ObjectMapper;
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.cohere.api.CohereEmbeddingBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
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;
import static org.assertj.core.api.Assertions.assertThat;
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
class BedrockCohereEmbeddingClientIT {
@Autowired
private BedrockCohereEmbeddingClient embeddingClient;
@Test
void singleEmbedding() {
assertThat(embeddingClient).isNotNull();
EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of("Hello World"));
assertThat(embeddingResponse.getData()).hasSize(1);
assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
}
@Test
void batchEmbedding() {
assertThat(embeddingClient).isNotNull();
EmbeddingResponse embeddingResponse = embeddingClient
.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
assertThat(embeddingResponse.getData()).hasSize(2);
assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getData().get(1).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
}
@SpringBootConfiguration
public static class TestConfiguration {
@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);
}
}
}

View File

@@ -0,0 +1,147 @@
/*
* 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.cohere.api;
import java.util.List;
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.bedrock.cohere.api.CohereChatBedrockApi.CohereChatModel;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.Truncate;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatResponse;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatResponse.Generation.FinishReason;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class CohereChatBedrockApiIT {
private CohereChatBedrockApi cohereChatApi = new CohereChatBedrockApi(CohereChatModel.COHERE_COMMAND_V14.id(),
Region.US_EAST_1.id());
@Test
public void requestBuilder() {
CohereChatRequest request1 = new CohereChatRequest(
"What is the capital of Bulgaria and what is the size? What it the national anthem?", 0.5f, 0.9f, 15,
40, List.of("END"), CohereChatRequest.ReturnLikelihoods.ALL, false, 1, null, Truncate.NONE);
var request2 = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withTemperature(0.5f)
.withTopP(0.9f)
.withTopK(15)
.withMaxTokens(40)
.withStopSequences(List.of("END"))
.withReturnLikelihoods(CohereChatRequest.ReturnLikelihoods.ALL)
.withStream(false)
.withNumGenerations(1)
.withLogitBias(null)
.withTruncate(Truncate.NONE)
.build();
assertThat(request1).isEqualTo(request2);
}
@Test
public void chatCompletion() {
var request = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withStream(false)
.withTemperature(0.5f)
.withTopP(0.8f)
.withTopK(15)
.withMaxTokens(100)
.withStopSequences(List.of("END"))
.withReturnLikelihoods(CohereChatRequest.ReturnLikelihoods.ALL)
.withNumGenerations(3)
.withLogitBias(null)
.withTruncate(Truncate.NONE)
.build();
CohereChatResponse response = cohereChatApi.chatCompletion(request);
assertThat(response).isNotNull();
assertThat(response.prompt()).isEqualTo(request.prompt());
assertThat(response.generations()).hasSize(request.numGenerations());
assertThat(response.generations().get(0).text()).isNotEmpty();
}
@Test
public void chatCompletionStream() {
var request = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withStream(true)
.withTemperature(0.5f)
.withTopP(0.8f)
.withTopK(15)
.withMaxTokens(100)
.withStopSequences(List.of("END"))
.withReturnLikelihoods(CohereChatRequest.ReturnLikelihoods.ALL)
.withNumGenerations(3)
.withLogitBias(null)
.withTruncate(Truncate.NONE)
.build();
Flux<CohereChatResponse.Generation> responseStream = cohereChatApi.chatCompletionStream(request);
List<CohereChatResponse.Generation> responses = responseStream.collectList().block();
assertThat(responses).isNotNull();
assertThat(responses).hasSizeGreaterThan(10);
assertThat(responses.get(0).text()).isNotEmpty();
CohereChatResponse.Generation lastResponse = responses.get(responses.size() - 1);
assertThat(lastResponse.text()).isNull();
assertThat(lastResponse.isFinished()).isTrue();
assertThat(lastResponse.finishReason()).isEqualTo(FinishReason.MAX_TOKENS);
assertThat(lastResponse.amazonBedrockInvocationMetrics()).isNotNull();
}
@Test
public void testStreamConfigurations() {
var streamRequest = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withStream(true)
.build();
assertThatThrownBy(() -> cohereChatApi.chatCompletion(streamRequest))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("The request must be configured to return the complete response!");
var notStreamRequest = CohereChatRequest
.builder("What is the capital of Bulgaria and what is the size? What it the national anthem?")
.withStream(false)
.build();
assertThatThrownBy(() -> cohereChatApi.chatCompletionStream(notStreamRequest))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("The request must be configured to stream the response!");
}
}

View File

@@ -0,0 +1,59 @@
/*
* 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.cohere.api;
import java.util.List;
import com.fasterxml.jackson.databind.ObjectMapper;
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.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingResponse;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class CohereEmbeddingBedrockApiIT {
CohereEmbeddingBedrockApi api = new CohereEmbeddingBedrockApi(
CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(), EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(), new ObjectMapper());
@Test
public void embedText() {
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).isNotNull();
assertThat(response.texts()).isEqualTo(request.texts());
assertThat(response.embeddings()).hasSize(2);
assertThat(response.embeddings().get(0)).hasSize(1024);
}
}

View File

@@ -0,0 +1,64 @@
/*
* 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.jurassic2.api;
import java.util.stream.Collectors;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatCompletionModel;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatRequest;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatResponse;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class Ai21Jurassic2ChatBedrockApiIT {
Ai21Jurassic2ChatBedrockApi api = new Ai21Jurassic2ChatBedrockApi(
Ai21Jurassic2ChatCompletionModel.AI21_J2_ULTRA_V1.id(), Region.US_EAST_1.id());
@Test
public void chatCompletion() {
Ai21Jurassic2ChatRequest request = new Ai21Jurassic2ChatRequest("Give me the names of 3 famous pirates?", 0.9f,
0.9f, 100, null, // List.of("END"),
new Ai21Jurassic2ChatRequest.IntegerScalePenalty(1, true, true, true, true, true),
new Ai21Jurassic2ChatRequest.FloatScalePenalty(0.5f, true, true, true, true, true),
new Ai21Jurassic2ChatRequest.IntegerScalePenalty(1, true, true, true, true, true));
Ai21Jurassic2ChatResponse response = api.chatCompletion(request);
assertThat(response).isNotNull();
assertThat(response.completions()).isNotEmpty();
assertThat(response.amazonBedrockInvocationMetrics()).isNull();
String responseContent = response.completions()
.stream()
.map(c -> c.data().text())
.collect(Collectors.joining("\n"));
assertThat(responseContent).contains("Blackbeard");
}
// Note: Ai21Jurassic2 doesn't support streaming yet!
}

View File

@@ -0,0 +1,172 @@
package org.springframework.ai.bedrock.llama2;
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.llama2.api.Llama2ChatBedrockApi;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatCompletionModel;
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 BedrockLlama2ChatClientIT {
@Autowired
private BedrockLlama2ChatClient client;
@Value("classpath:/prompts/system-message.st")
private Resource systemResource;
@Test
void roleTest() {
UserMessage userMessage = new UserMessage(
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate"));
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);
}
@Test
void mapOutputParser() {
MapOutputParser outputParser = new MapOutputParser();
String format = outputParser.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 = 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 non JSON tex blocks from the output.
Provide your answer in the JSON format with the feature names as the keys.
""";
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 Llama2ChatBedrockApi llama2Api() {
return new Llama2ChatBedrockApi(Llama2ChatCompletionModel.LLAMA2_70B_CHAT_V1.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
}
@Bean
public BedrockLlama2ChatClient llama2ChatClient(Llama2ChatBedrockApi llama2Api) {
return new BedrockLlama2ChatClient(llama2Api);
}
}
}

View File

@@ -0,0 +1,79 @@
/*
* 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.llama2.api;
import java.util.List;
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.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatCompletionModel;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatRequest;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatResponse;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class Llama2ChatBedrockApiIT {
private Llama2ChatBedrockApi llama2ChatApi = new Llama2ChatBedrockApi(
Llama2ChatCompletionModel.LLAMA2_70B_CHAT_V1.id(), Region.US_EAST_1.id());
@Test
public void chatCompletion() {
Llama2ChatRequest request = Llama2ChatRequest.builder("Hello, my name is")
.withTemperature(0.9f)
.withTopP(0.9f)
.withMaxGenLen(20)
.build();
Llama2ChatResponse response = llama2ChatApi.chatCompletion(request);
System.out.println(response.generation());
assertThat(response).isNotNull();
assertThat(response.generation()).isNotEmpty();
assertThat(response.promptTokenCount()).isEqualTo(6);
assertThat(response.generationTokenCount()).isGreaterThan(10);
assertThat(response.generationTokenCount()).isLessThanOrEqualTo(20);
assertThat(response.stopReason()).isNotNull();
assertThat(response.amazonBedrockInvocationMetrics()).isNull();
}
@Test
public void chatCompletionStream() {
Llama2ChatRequest request = new Llama2ChatRequest("Hello, my name is", 0.9f, 0.9f, 20);
Flux<Llama2ChatResponse> responseStream = llama2ChatApi.chatCompletionStream(request);
List<Llama2ChatResponse> responses = responseStream.collectList().block();
assertThat(responses).isNotNull();
assertThat(responses).hasSizeGreaterThan(10);
assertThat(responses.get(0).generation()).isNotEmpty();
Llama2ChatResponse lastResponse = responses.get(responses.size() - 1);
assertThat(lastResponse.stopReason()).isNotNull();
assertThat(lastResponse.amazonBedrockInvocationMetrics()).isNotNull();
}
}

View File

@@ -0,0 +1,67 @@
/*
* 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.api;
import java.util.List;
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.bedrock.titan.api.TitanChatBedrockApi.TitanChatCompletionModel;
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 static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class TitanChatBedrockApiIT {
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();
@Test
public void chatCompletion() {
TitanChatResponse response = titanBedrockApi.chatCompletion(titanChatRequest);
assertThat(response.results()).hasSize(1);
assertThat(response.results().get(0).outputText()).contains("Blackbeard");
}
@Test
public void chatCompletionStream() {
Flux<TitanChatResponseChunk> response = titanBedrockApi.chatCompletionStream(titanChatRequest);
List<TitanChatResponseChunk> results = response.collectList().block();
assertThat(results.stream().map(TitanChatResponseChunk::outputText).collect(Collectors.joining("\n")))
.contains("Blackbeard");
}
}

View File

@@ -0,0 +1,76 @@
/*
* 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.api;
import java.io.IOException;
import java.util.Base64;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse;
import org.springframework.core.io.DefaultResourceLoader;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class TitanEmbeddingBedrockApiIT {
@Test
public void embedText() {
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
TitanEmbeddingModel.TITAN_EMBED_TEXT_V1.id(), Region.US_EAST_1.id());
TitanEmbeddingRequest request = TitanEmbeddingRequest.builder().withInputText("I like to eat apples.").build();
TitanEmbeddingResponse response = titanEmbedApi.embedding(request);
assertThat(response).isNotNull();
assertThat(response.inputTextTokenCount()).isEqualTo(6);
assertThat(response.embedding()).hasSize(1536);
}
@Test
public void embedImage() throws IOException {
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id());
byte[] image = new DefaultResourceLoader().getResource("classpath:/spring_framework.png")
.getContentAsByteArray();
String imageBase64 = Base64.getEncoder().encodeToString(image);
System.out.println(imageBase64.length());
TitanEmbeddingRequest request = TitanEmbeddingRequest.builder().withInputImage(imageBase64).build();
TitanEmbeddingResponse response = titanEmbedApi.embedding(request);
assertThat(response).isNotNull();
assertThat(response.inputTextTokenCount()).isEqualTo(0); // e.g. image input
assertThat(response.embedding()).hasSize(1024);
}
}

Binary file not shown.

After

Width:  |  Height:  |  Size: 429 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 379 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 307 KiB

View File

@@ -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}.

Binary file not shown.

After

Width:  |  Height:  |  Size: 43 KiB

View File

@@ -38,23 +38,16 @@
<version>${antlr.version}</version>
</dependency>
<dependency>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-core</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-messaging</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-json</artifactId>
</dependency>
<dependency>
<groupId>com.knuddels</groupId>
<artifactId>jtokkit</artifactId>

View File

@@ -3,21 +3,35 @@ package org.springframework.ai.embedding;
import java.util.List;
import java.util.Objects;
/**
* Represents a single embedding vector.
*/
public class Embedding {
private List<Double> embedding;
private Integer index;
/**
* Creates a new {@link Embedding} instance.
* @param embedding the embedding vector values.
* @param index the embedding index in a list of embeddings.
*/
public Embedding(List<Double> embedding, Integer index) {
this.embedding = embedding;
this.index = index;
}
/**
* @return Get the embedding vector values.
*/
public List<Double> getEmbedding() {
return embedding;
}
/**
* @return Get the embedding index in a list of embeddings.
*/
public Integer getIndex() {
return index;
}

View File

@@ -5,21 +5,49 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
/**
* Embedding response object.
*/
public class EmbeddingResponse {
/**
* Embedding data.
*/
private List<Embedding> data;
/**
* Embedding metadata.
*/
private Map<String, Object> metadata = new HashMap<>();
/**
* Creates a new {@link EmbeddingResponse} instance with empty metadata.
* @param data the embedding data.
*/
public EmbeddingResponse(List<Embedding> data) {
this(data, new HashMap<>());
}
/**
* Creates a new {@link EmbeddingResponse} instance.
* @param data the embedding data.
* @param metadata the embedding metadata.
*/
public EmbeddingResponse(List<Embedding> data, Map<String, Object> metadata) {
this.data = data;
this.metadata = metadata;
}
/**
* @return Get the embedding data.
*/
public List<Embedding> getData() {
return data;
}
/**
* @return Get the embedding metadata.
*/
public Map<String, Object> getMetadata() {
return metadata;
}

View File

@@ -48,7 +48,7 @@ public class BeanOutputParser<T> implements OutputParser<T> {
/** The Java class representing the target type. */
@SuppressWarnings({ "FieldMayBeFinal", "rawtypes" })
private Class clazz;
private Class<T> clazz;
/** The object mapper used for deserialization and other JSON operations. */
@SuppressWarnings("FieldMayBeFinal")
@@ -95,7 +95,6 @@ public class BeanOutputParser<T> implements OutputParser<T> {
*/
public T parse(String text) {
try {
// noinspection unchecked
return (T) this.objectMapper.readValue(text, this.clazz);
}
catch (JsonProcessingException e) {

View File

@@ -81,11 +81,6 @@ public abstract class AbstractMessage implements Message {
return this.messageType;
}
@Override
public String getMessageTypeValue() {
return this.messageType.getValue();
}
@Override
public int hashCode() {
final int prime = 31;

View File

@@ -26,6 +26,4 @@ public interface Message {
MessageType getMessageType();
String getMessageTypeValue();
}

View File

@@ -42,6 +42,11 @@
<!-- Spring Framework -->
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context-support</artifactId>

View File

@@ -55,7 +55,11 @@
<version>${victools.version}</version>
</dependency>
<!-- Spring Framework -->
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context-support</artifactId>
@@ -64,12 +68,6 @@
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-logging</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
<version>${spring-framework.version}</version>
</dependency>
<!-- test dependencies -->
<dependency>

View File

@@ -36,14 +36,14 @@ class OpenAiClientIT extends AbstractIT {
@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);
UserMessage userMessage = new UserMessage(
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", name, "voice", voice));
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate"));
Prompt prompt = new Prompt(List.of(userMessage, systemMessage));
AiResponse response = openAiClient.generate(prompt);
assertThat(response.getGenerations()).hasSize(1);
assertThat(response.getGenerations().get(0).getContent()).contains("Blackbeard");
// needs fine tuning... evaluateQuestionAndAnswer(request, response, false);
}

View File

@@ -156,6 +156,13 @@
<optional>true</optional>
</dependency>
<!-- Bedrock -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-bedrock</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>

View File

@@ -3,6 +3,14 @@ package org.springframework.ai.autoconfigure;
import com.fasterxml.jackson.annotation.JsonInclude;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vertex.api.VertexAiApi;
import org.springframework.aot.hint.MemberCategory;
@@ -31,7 +39,8 @@ public class NativeHints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
for (var h : Set.of(new VertexAiHints(), new OpenAiHints(), new PdfReaderHints(), new KnuddelsHints()))
for (var h : Set.of(new BedrockAiHints(), new VertexAiHints(), new OpenAiHints(), new PdfReaderHints(),
new KnuddelsHints()))
h.registerHints(hints, classLoader);
hints.resources().registerResource(new ClassPathResource("embedding/embedding-model-dimensions.properties"));
@@ -62,6 +71,29 @@ public class NativeHints implements RuntimeHintsRegistrar {
}
static class BedrockAiHints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
var mcs = MemberCategory.values();
for (var tr : findJsonAnnotatedClasses(Ai21Jurassic2ChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClasses(CohereChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClasses(CohereEmbeddingBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClasses(Llama2ChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClasses(TitanChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClasses(TitanEmbeddingBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClasses(AnthropicChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
}
}
static class OpenAiHints implements RuntimeHintsRegistrar {
@Override

View File

@@ -0,0 +1,49 @@
/*
* 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;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.StringUtils;
/**
* @author Christian Tzolov
*/
@Configuration
@EnableConfigurationProperties({ BedrockAwsConnectionProperties.class })
public class BedrockAwsConnectionConfiguration {
@Bean
@ConditionalOnMissingBean
public AwsCredentialsProvider credentialsProvider(BedrockAwsConnectionProperties properties) {
if (StringUtils.hasText(properties.getAccessKey()) && StringUtils.hasText(properties.getSecretKey())) {
return StaticCredentialsProvider
.create(AwsBasicCredentials.create(properties.getAccessKey(), properties.getSecretKey()));
}
return DefaultCredentialsProvider.create();
}
}

View File

@@ -0,0 +1,71 @@
/*
* 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;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* Configuration properties for Bedrock AWS connection.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@ConfigurationProperties(BedrockAwsConnectionProperties.CONFIG_PREFIX)
public class BedrockAwsConnectionProperties {
public static final String CONFIG_PREFIX = "spring.ai.bedrock.aws";
/**
* AWS region to use. Defaults to us-east-1.
*/
private String region = "us-east-1";
/**
* AWS access key.
*/
private String accessKey;
/**
* AWS secret key.
*/
private String secretKey;
public String getRegion() {
return region;
}
public void setRegion(String awsRegion) {
this.region = awsRegion;
}
public String getAccessKey() {
return accessKey;
}
public void setAccessKey(String accessKey) {
this.accessKey = accessKey;
}
public String getSecretKey() {
return secretKey;
}
public void setSecretKey(String secretKey) {
this.secretKey = secretKey;
}
}

View File

@@ -0,0 +1,73 @@
/*
* 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.anthropic;
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.anthropic.BedrockAnthropicChatClient;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi;
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 Anthropic Chat Client.
*
* Leverages the Spring Cloud AWS to resolve the {@link AwsCredentialsProvider}.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@AutoConfiguration
@ConditionalOnClass(AnthropicChatBedrockApi.class)
@EnableConfigurationProperties({ BedrockAnthropicChatProperties.class, BedrockAwsConnectionProperties.class })
@ConditionalOnProperty(prefix = BedrockAnthropicChatProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@Import(BedrockAwsConnectionConfiguration.class)
@ImportRuntimeHints(NativeHints.class)
public class BedrockAnthropicChatAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public AnthropicChatBedrockApi anthropicApi(AwsCredentialsProvider credentialsProvider,
BedrockAnthropicChatProperties properties, BedrockAwsConnectionProperties awsProperties) {
return new AnthropicChatBedrockApi(properties.getModel(), credentialsProvider, awsProperties.getRegion(),
new ObjectMapper());
}
@Bean
public BedrockAnthropicChatClient anthropicChatClient(AnthropicChatBedrockApi anthropicApi,
BedrockAnthropicChatProperties properties) {
return new BedrockAnthropicChatClient(anthropicApi).withTemperature(properties.getTemperature())
.withTopP(properties.getTopP())
.withMaxTokensToSample(properties.getMaxTokensToSample())
.withTopK(properties.getTopK())
.withStopSequences(properties.getStopSequences())
.withAnthropicVersion(properties.getAnthropicVersion());
}
}

View File

@@ -0,0 +1,151 @@
/*
* 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.anthropic;
import java.util.List;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatModel;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* Configuration properties for Bedrock Anthropic.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@ConfigurationProperties(BedrockAnthropicChatProperties.CONFIG_PREFIX)
public class BedrockAnthropicChatProperties {
public static final String CONFIG_PREFIX = "spring.ai.bedrock.anthropic.chat";
/**
* Enable Bedrock Anthropic chat client. Disabled by default.
*/
private boolean enabled = false;
/**
* The model id to use. See the {@link AnthropicChatModel} for the supported models.
*/
private String model = AnthropicChatModel.CLAUDE_V2.id();
/**
* Controls the randomness of the output. Values can range over [0.0,1.0], inclusive.
* A value closer to 1.0 will produce responses that are more varied, while a value
* closer to 0.0 will typically result in less surprising responses from the model.
* This value specifies default to be used by the backend while making the call to the
* model.
*/
private Float temperature = 0.7f;
/**
* 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 = null;
/**
* Specify the maximum number of tokens to use in the generated response. Note that
* the models may stop before reaching this maximum. This parameter only specifies the
* absolute maximum number of tokens to generate. We recommend a limit of 4,000 tokens
* for optimal performance.
*/
private Integer maxTokensToSample = 300;
/**
* Specify the number of token choices the model uses to generate the next token.
*/
private Integer topK = 10;
/**
* 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 = List.of("\n\nHuman:");
/**
* The version of the model to use. The default value is bedrock-2023-05-31.
*/
private String anthropicVersion = AnthropicChatBedrockApi.DEFAULT_ANTHROPIC_VERSION;
public boolean isEnabled() {
return this.enabled;
}
public void setEnabled(boolean enabled) {
this.enabled = enabled;
}
public String getModel() {
return this.model;
}
public void setModel(String model) {
this.model = model;
}
public Float getTemperature() {
return this.temperature;
}
public void setTemperature(Float temperature) {
this.temperature = temperature;
}
public Float getTopP() {
return this.topP;
}
public void setTopP(Float topP) {
this.topP = topP;
}
public Integer getMaxTokensToSample() {
return maxTokensToSample;
}
public void setMaxTokensToSample(Integer maxTokensToSample) {
this.maxTokensToSample = maxTokensToSample;
}
public Integer getTopK() {
return topK;
}
public void setTopK(Integer topK) {
this.topK = topK;
}
public List<String> getStopSequences() {
return stopSequences;
}
public void setStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
}
public String getAnthropicVersion() {
return anthropicVersion;
}
public void setAnthropicVersion(String anthropicVersion) {
this.anthropicVersion = anthropicVersion;
}
}

View File

@@ -0,0 +1,74 @@
/*
* 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.cohere;
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.cohere.BedrockCohereChatClient;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.LogitBias;
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 Cohere Chat Client.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@AutoConfiguration
@ConditionalOnClass(CohereChatBedrockApi.class)
@EnableConfigurationProperties({ BedrockCohereChatProperties.class, BedrockAwsConnectionProperties.class })
@ConditionalOnProperty(prefix = BedrockCohereChatProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@Import(BedrockAwsConnectionConfiguration.class)
@ImportRuntimeHints(NativeHints.class)
public class BedrockCohereChatAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public CohereChatBedrockApi cohereApi(AwsCredentialsProvider credentialsProvider,
BedrockCohereChatProperties properties, BedrockAwsConnectionProperties awsProperties) {
return new CohereChatBedrockApi(properties.getModel(), credentialsProvider, awsProperties.getRegion(),
new ObjectMapper());
}
@Bean
public BedrockCohereChatClient cohereChatClient(CohereChatBedrockApi cohereChatApi,
BedrockCohereChatProperties properties) {
LogitBias logitBiasBias = (properties.getLogitBiasBias() != null && properties.getLogitBiasToken() != null)
? new LogitBias(properties.getLogitBiasToken(), properties.getLogitBiasBias()) : null;
return new BedrockCohereChatClient(cohereChatApi).withTemperature(properties.getTemperature())
.withTopP(properties.getTopP())
.withTopK(properties.getTopK())
.withMaxTokens(properties.getMaxTokens())
.withStopSequences(properties.getStopSequences())
.withLogitBias(logitBiasBias)
.withTruncate(properties.getTruncate());
}
}

View File

@@ -0,0 +1,202 @@
/*
* 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.cohere;
import java.util.List;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.ReturnLikelihoods;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.Truncate;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* Bedrock Cohere Chat autoconfiguration properties.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@ConfigurationProperties(BedrockCohereChatProperties.CONFIG_PREFIX)
public class BedrockCohereChatProperties {
public static final String CONFIG_PREFIX = "spring.ai.bedrock.cohere.chat";
/**
* Enable Bedrock Cohere Chat Client. False by default.
*/
private boolean enabled = false;
/**
* Bedrock Cohere Chat model name. Defaults to 'cohere-command-v14'.
*/
private String model = CohereChatBedrockApi.CohereChatModel.COHERE_COMMAND_V14.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 number of token choices the model uses to generate the next
* token.
*/
private Integer topK;
/**
* (optional) Specify the maximum number of tokens to use in the generated response.
*/
private Integer maxTokens;
/**
* (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;
/**
* (optional) Specify how and if the token likelihoods are returned with the response.
*/
private ReturnLikelihoods returnLikelihoods;
/**
* (optional) The maximum number of generations that the model should return.
*/
private Integer numGenerations;
/**
* LogitBias prevents the model from generating unwanted tokens or incentivize the
* model to include desired tokens. The token likelihoods.
*/
private String logitBiasToken;
/**
* LogitBias prevents the model from generating unwanted tokens or incentivize the
* model to include desired tokens. A float between -10 and 10.
*/
private Float logitBiasBias;
/**
* (optional) Specifies how the API handles inputs longer than the maximum token
* length.
*/
private Truncate truncate;
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 getTopK() {
return topK;
}
public void setTopK(Integer topK) {
this.topK = topK;
}
public Integer getMaxTokens() {
return maxTokens;
}
public void setMaxTokens(Integer maxTokens) {
this.maxTokens = maxTokens;
}
public List<String> getStopSequences() {
return stopSequences;
}
public void setStopSequences(List<String> stopSequences) {
this.stopSequences = stopSequences;
}
public ReturnLikelihoods getReturnLikelihoods() {
return returnLikelihoods;
}
public void setReturnLikelihoods(ReturnLikelihoods returnLikelihoods) {
this.returnLikelihoods = returnLikelihoods;
}
public Integer getNumGenerations() {
return numGenerations;
}
public void setNumGenerations(Integer numGenerations) {
this.numGenerations = numGenerations;
}
public String getLogitBiasToken() {
return logitBiasToken;
}
public void setLogitBiasToken(String logitBiasToken) {
this.logitBiasToken = logitBiasToken;
}
public Float getLogitBiasBias() {
return logitBiasBias;
}
public void setLogitBiasBias(Float logitBiasBias) {
this.logitBiasBias = logitBiasBias;
}
public Truncate getTruncate() {
return truncate;
}
public void setTruncate(Truncate truncate) {
this.truncate = truncate;
}
}

View File

@@ -0,0 +1,67 @@
/*
* 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.cohere;
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.cohere.BedrockCohereEmbeddingClient;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
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 Cohere Embedding Client.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@AutoConfiguration
@ConditionalOnClass(CohereEmbeddingBedrockApi.class)
@EnableConfigurationProperties({ BedrockCohereEmbeddingProperties.class, BedrockAwsConnectionProperties.class })
@ConditionalOnProperty(prefix = BedrockCohereEmbeddingProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@Import(BedrockAwsConnectionConfiguration.class)
@ImportRuntimeHints(NativeHints.class)
public class BedrockCohereEmbeddingAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public CohereEmbeddingBedrockApi cohereApi(AwsCredentialsProvider credentialsProvider,
BedrockCohereEmbeddingProperties properties, BedrockAwsConnectionProperties awsProperties) {
return new CohereEmbeddingBedrockApi(properties.getModel(), credentialsProvider, awsProperties.getRegion(),
new ObjectMapper());
}
@Bean
@ConditionalOnMissingBean
public BedrockCohereEmbeddingClient cohereEmbeddingClient(CohereEmbeddingBedrockApi cohereEmbeddingApi,
BedrockCohereEmbeddingProperties properties) {
return new BedrockCohereEmbeddingClient(cohereEmbeddingApi).withInputType(properties.getInputType())
.withTruncate(properties.getTruncate());
}
}

View File

@@ -0,0 +1,94 @@
/*
* 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.cohere;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* Bedrock Cohere Embedding autoconfiguration properties.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@ConfigurationProperties(BedrockCohereEmbeddingProperties.CONFIG_PREFIX)
public class BedrockCohereEmbeddingProperties {
public static final String CONFIG_PREFIX = "spring.ai.bedrock.cohere.embedding";
/**
* Enable Bedrock Cohere Embedding Client. False by default.
*/
private boolean enabled = false;
/**
* Bedrock Cohere Embedding model name. Defaults to 'cohere.embed-multilingual-v3'.
*/
private String model = CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id();
/**
* 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.
*/
private InputType inputType = InputType.search_document;
/**
* Specifies how the API handles inputs longer than the maximum token length.
*/
private CohereEmbeddingRequest.Truncate truncate = CohereEmbeddingRequest.Truncate.NONE;
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 static String getConfigPrefix() {
return CONFIG_PREFIX;
}
public void setInputType(InputType inputType) {
this.inputType = inputType;
}
public InputType getInputType() {
return inputType;
}
public CohereEmbeddingRequest.Truncate getTruncate() {
return truncate;
}
public void setTruncate(CohereEmbeddingRequest.Truncate truncate) {
this.truncate = truncate;
}
}

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.llama2;
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.llama2.BedrockLlama2ChatClient;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi;
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 Llama2 Chat Client.
*
* Leverages the Spring Cloud AWS to resolve the {@link AwsCredentialsProvider}.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@AutoConfiguration
@ConditionalOnClass(Llama2ChatBedrockApi.class)
@EnableConfigurationProperties({ BedrockLlama2ChatProperties.class, BedrockAwsConnectionProperties.class })
@ConditionalOnProperty(prefix = BedrockLlama2ChatProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@Import(BedrockAwsConnectionConfiguration.class)
@ImportRuntimeHints(NativeHints.class)
public class BedrockLlama2ChatAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public Llama2ChatBedrockApi llama2Api(AwsCredentialsProvider credentialsProvider,
BedrockLlama2ChatProperties properties, BedrockAwsConnectionProperties awsProperties) {
return new Llama2ChatBedrockApi(properties.getModel(), credentialsProvider, awsProperties.getRegion(),
new ObjectMapper());
}
@Bean
public BedrockLlama2ChatClient llama2ChatClient(Llama2ChatBedrockApi llama2Api,
BedrockLlama2ChatProperties properties) {
return new BedrockLlama2ChatClient(llama2Api).withTemperature(properties.getTemperature())
.withTopP(properties.getTopP())
.withMaxGenLen(properties.getMaxGenLen());
}
}

View File

@@ -0,0 +1,106 @@
/*
* 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.llama2;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatCompletionModel;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* Configuration properties for Bedrock Llama2.
*
* @author Christian Tzolov
* @since 0.8.0
*/
@ConfigurationProperties(BedrockLlama2ChatProperties.CONFIG_PREFIX)
public class BedrockLlama2ChatProperties {
public static final String CONFIG_PREFIX = "spring.ai.bedrock.llama2.chat";
/**
* Enable Bedrock Llama2 chat client. Disabled by default.
*/
private boolean enabled = false;
/**
* Controls the randomness of the output. Values can range over [0.0,1.0], inclusive.
* A value closer to 1.0 will produce responses that are more varied, while a value
* closer to 0.0 will typically result in less surprising responses from the model.
* This value specifies default to be used by the backend while making the call to the
* model.
*/
private Float temperature = 0.7f;
/**
* 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 = null;
/**
* Specify the maximum number of tokens to use in the generated response. The model
* truncates the response once the generated text exceeds maxGenLen.
*/
private Integer maxGenLen = 300;
/**
* The model id to use. See the {@link Llama2ChatCompletionModel} for the supported
* models.
*/
private String model = Llama2ChatCompletionModel.LLAMA2_70B_CHAT_V1.id();
public boolean isEnabled() {
return this.enabled;
}
public void setEnabled(boolean enabled) {
this.enabled = enabled;
}
public String getModel() {
return this.model;
}
public void setModel(String model) {
this.model = model;
}
public Float getTemperature() {
return this.temperature;
}
public void setTemperature(Float temperature) {
this.temperature = temperature;
}
public Float getTopP() {
return this.topP;
}
public void setTopP(Float topP) {
this.topP = topP;
}
public Integer getMaxGenLen() {
return this.maxGenLen;
}
public void setMaxGenLen(Integer topK) {
this.maxGenLen = topK;
}
}

View File

@@ -39,12 +39,10 @@ public class VertexAiAutoConfiguration {
@ConditionalOnMissingBean
public VertexAiChatClient vertexAiClient(VertexAiApi vertexAiApi, VertexAiChatProperties chatProperties) {
VertexAiChatClient client = new VertexAiChatClient(vertexAiApi);
client.setTemperature(chatProperties.getTemperature());
client.setTopP(chatProperties.getTopP());
client.setTopK(chatProperties.getTopK());
client.setCandidateCount(chatProperties.getCandidateCount());
VertexAiChatClient client = new VertexAiChatClient(vertexAiApi).withTemperature(chatProperties.getTemperature())
.withTopP(chatProperties.getTopP())
.withTopK(chatProperties.getTopK())
.withCandidateCount(chatProperties.getCandidateCount());
return client;
}

View File

@@ -11,3 +11,8 @@ org.springframework.ai.autoconfigure.vectorstore.chroma.ChromaVectorStoreAutoCon
org.springframework.ai.autoconfigure.vectorstore.azure.AzureVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.weaviate.WeaviateVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.vertexai.VertexAiAutoConfiguration
org.springframework.ai.autoconfigure.bedrock.llama2.BedrockLlama2AutoConfiguration
org.springframework.ai.autoconfigure.bedrock.cohere.BedrockCohereChatAutoConfiguration
org.springframework.ai.autoconfigure.bedrock.cohere.BedrockCohereEmbeddingAutoConfiguration

View File

@@ -0,0 +1,153 @@
/*
* 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.anthropic;
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.anthropic.BedrockAnthropicChatClient;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatModel;
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 BedrockAnthropicChatAutoConfigurationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
"spring.ai.bedrock.anthropic.chat.model=" + AnthropicChatModel.CLAUDE_V2.id(),
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.anthropic.chat.temperature=0.5", "spring.ai.bedrock.anthropic.chat.maxGenLen=500")
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.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 -> {
BedrockAnthropicChatClient anthropicChatClient = context.getBean(BedrockAnthropicChatClient.class);
AiResponse response = anthropicChatClient.generate(new Prompt(List.of(userMessage, systemMessage)));
assertThat(response.getGeneration().getContent()).contains("Blackbeard");
});
}
@Test
public void chatCompletionStreaming() {
contextRunner.run(context -> {
BedrockAnthropicChatClient anthropicChatClient = context.getBean(BedrockAnthropicChatClient.class);
Flux<AiResponse> response = anthropicChatClient
.generateStream(new Prompt(List.of(userMessage, systemMessage)));
List<AiResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(2);
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.anthropic.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
"spring.ai.bedrock.anthropic.chat.model=MODEL_XYZ",
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.anthropic.chat.temperature=0.55")
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
.run(context -> {
var anthropicChatProperties = context.getBean(BedrockAnthropicChatProperties.class);
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
assertThat(anthropicChatProperties.isEnabled()).isTrue();
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
assertThat(anthropicChatProperties.getTemperature()).isEqualTo(0.55f);
assertThat(anthropicChatProperties.getModel()).isEqualTo("MODEL_XYZ");
assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
assertThat(awsProperties.getSecretKey()).isEqualTo("SECRET_KEY");
});
}
@Test
public void chatCompletionDisabled() {
// It is disabled by default
new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockAnthropicChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockAnthropicChatClient.class)).isEmpty();
});
// Explicitly enable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=true")
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockAnthropicChatProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(BedrockAnthropicChatClient.class)).isNotEmpty();
});
// Explicitly disable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.anthropic.chat.enabled=false")
.withConfiguration(AutoConfigurations.of(BedrockAnthropicChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockAnthropicChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockAnthropicChatClient.class)).isEmpty();
});
}
}

View File

@@ -0,0 +1,164 @@
/*
* 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.cohere;
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.cohere.BedrockCohereChatClient;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatModel;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.ReturnLikelihoods;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi.CohereChatRequest.Truncate;
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 BedrockCohereChatAutoConfigurationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
"spring.ai.bedrock.cohere.chat.model=" + CohereChatModel.COHERE_COMMAND_V14.id(),
"spring.ai.bedrock.cohere.chat.temperature=0.5", "spring.ai.bedrock.cohere.chat.maxTokens=500")
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.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 -> {
BedrockCohereChatClient cohereChatClient = context.getBean(BedrockCohereChatClient.class);
AiResponse response = cohereChatClient.generate(new Prompt(List.of(userMessage, systemMessage)));
assertThat(response.getGeneration().getContent()).contains("Blackbeard");
});
}
@Test
public void chatCompletionStreaming() {
contextRunner.run(context -> {
BedrockCohereChatClient cohereChatClient = context.getBean(BedrockCohereChatClient.class);
Flux<AiResponse> response = cohereChatClient
.generateStream(new Prompt(List.of(userMessage, systemMessage)));
List<AiResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(2);
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.cohere.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
"spring.ai.bedrock.cohere.chat.model=MODEL_XYZ",
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.cohere.chat.temperature=0.55", "spring.ai.bedrock.cohere.chat.topP=0.55",
"spring.ai.bedrock.cohere.chat.topK=10", "spring.ai.bedrock.cohere.chat.stopSequences=END1,END2",
"spring.ai.bedrock.cohere.chat.returnLikelihoods=ALL", "spring.ai.bedrock.cohere.chat.numGenerations=3",
"spring.ai.bedrock.cohere.chat.truncate=START", "spring.ai.bedrock.cohere.chat.maxTokens=123")
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.class))
.run(context -> {
var chatProperties = context.getBean(BedrockCohereChatProperties.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.getTopK()).isEqualTo(10);
assertThat(chatProperties.getStopSequences()).isEqualTo(List.of("END1", "END2"));
assertThat(chatProperties.getReturnLikelihoods()).isEqualTo(ReturnLikelihoods.ALL);
assertThat(chatProperties.getNumGenerations()).isEqualTo(3);
assertThat(chatProperties.getTruncate()).isEqualTo(Truncate.START);
assertThat(chatProperties.getMaxTokens()).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(BedrockCohereChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockCohereChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockCohereChatClient.class)).isEmpty();
});
// Explicitly enable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=true")
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockCohereChatProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(BedrockCohereChatClient.class)).isNotEmpty();
});
// Explicitly disable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.chat.enabled=false")
.withConfiguration(AutoConfigurations.of(BedrockCohereChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockCohereChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockCohereChatClient.class)).isEmpty();
});
}
}

View File

@@ -0,0 +1,141 @@
/*
* 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.cohere;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
import org.springframework.ai.bedrock.cohere.BedrockCohereEmbeddingClient;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType;
import org.springframework.ai.embedding.EmbeddingResponse;
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 BedrockCohereEmbeddingAutoConfigurationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true",
"spring.ai.bedrock.aws.access-key=" + System.getenv("AWS_ACCESS_KEY_ID"),
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
"spring.ai.bedrock.cohere.embedding.model=" + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
"spring.ai.bedrock.cohere.chat.inputType=search_document",
"spring.ai.bedrock.cohere.chat.truncate=NONE")
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class));
@Test
public void singleEmbedding() {
contextRunner.run(context -> {
BedrockCohereEmbeddingClient embeddingClient = context.getBean(BedrockCohereEmbeddingClient.class);
assertThat(embeddingClient).isNotNull();
EmbeddingResponse embeddingResponse = embeddingClient.embedForResponse(List.of("Hello World"));
assertThat(embeddingResponse.getData()).hasSize(1);
assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
});
}
@Test
public void batchEmbedding() {
contextRunner.run(context -> {
BedrockCohereEmbeddingClient embeddingClient = context.getBean(BedrockCohereEmbeddingClient.class);
assertThat(embeddingClient).isNotNull();
EmbeddingResponse embeddingResponse = embeddingClient
.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
assertThat(embeddingResponse.getData()).hasSize(2);
assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getData().get(1).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
});
}
@Test
public void propertiesTest() {
new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true",
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.cohere.embedding.model=MODEL_XYZ",
"spring.ai.bedrock.cohere.embedding.inputType=classification",
"spring.ai.bedrock.cohere.embedding.truncate=RIGHT")
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
.run(context -> {
var chatProperties = context.getBean(BedrockCohereEmbeddingProperties.class);
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
assertThat(chatProperties.isEnabled()).isTrue();
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
assertThat(chatProperties.getModel()).isEqualTo("MODEL_XYZ");
assertThat(chatProperties.getInputType()).isEqualTo(InputType.classification);
assertThat(chatProperties.getTruncate()).isEqualTo(CohereEmbeddingRequest.Truncate.RIGHT);
assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
assertThat(awsProperties.getSecretKey()).isEqualTo("SECRET_KEY");
});
}
@Test
public void chatCompletionDisabled() {
// It is disabled by default
new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockCohereEmbeddingClient.class)).isEmpty();
});
// Explicitly enable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true")
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(BedrockCohereEmbeddingClient.class)).isNotEmpty();
});
// Explicitly disable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=false")
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockCohereEmbeddingClient.class)).isEmpty();
});
}
}

View File

@@ -0,0 +1,154 @@
/*
* 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.llama2;
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.llama2.BedrockLlama2ChatClient;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi.Llama2ChatCompletionModel;
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 BedrockLlama2ChatAutoConfigurationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withPropertyValues("spring.ai.bedrock.llama2.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.llama2.chat.model=" + Llama2ChatCompletionModel.LLAMA2_70B_CHAT_V1.id(),
"spring.ai.bedrock.llama2.chat.temperature=0.5", "spring.ai.bedrock.llama2.chat.maxGenLen=500")
.withConfiguration(AutoConfigurations.of(BedrockLlama2ChatAutoConfiguration.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 -> {
BedrockLlama2ChatClient llama2ChatClient = context.getBean(BedrockLlama2ChatClient.class);
AiResponse response = llama2ChatClient.generate(new Prompt(List.of(userMessage, systemMessage)));
assertThat(response.getGeneration().getContent()).contains("Blackbeard");
});
}
@Test
public void chatCompletionStreaming() {
contextRunner.run(context -> {
BedrockLlama2ChatClient llama2ChatClient = context.getBean(BedrockLlama2ChatClient.class);
Flux<AiResponse> response = llama2ChatClient
.generateStream(new Prompt(List.of(userMessage, systemMessage)));
List<AiResponse> responses = response.collectList().block();
assertThat(responses.size()).isGreaterThan(2);
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.llama2.chat.enabled=true",
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
"spring.ai.bedrock.llama2.chat.model=MODEL_XYZ",
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.llama2.chat.temperature=0.55", "spring.ai.bedrock.llama2.chat.maxGenLen=123")
.withConfiguration(AutoConfigurations.of(BedrockLlama2ChatAutoConfiguration.class))
.run(context -> {
var llama2ChatProperties = context.getBean(BedrockLlama2ChatProperties.class);
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
assertThat(llama2ChatProperties.isEnabled()).isTrue();
assertThat(awsProperties.getRegion()).isEqualTo(Region.EU_CENTRAL_1.id());
assertThat(llama2ChatProperties.getTemperature()).isEqualTo(0.55f);
assertThat(llama2ChatProperties.getMaxGenLen()).isEqualTo(123);
assertThat(llama2ChatProperties.getModel()).isEqualTo("MODEL_XYZ");
assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
assertThat(awsProperties.getSecretKey()).isEqualTo("SECRET_KEY");
});
}
@Test
public void chatCompletionDisabled() {
// It is disabled by default
new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(BedrockLlama2ChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockLlama2ChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockLlama2ChatClient.class)).isEmpty();
});
// Explicitly enable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.llama2.chat.enabled=true")
.withConfiguration(AutoConfigurations.of(BedrockLlama2ChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockLlama2ChatProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(BedrockLlama2ChatClient.class)).isNotEmpty();
});
// Explicitly disable the chat auto-configuration.
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.llama2.chat.enabled=false")
.withConfiguration(AutoConfigurations.of(BedrockLlama2ChatAutoConfiguration.class))
.run(context -> {
assertThat(context.getBeansOfType(BedrockLlama2ChatProperties.class)).isEmpty();
assertThat(context.getBeansOfType(BedrockLlama2ChatClient.class)).isEmpty();
});
}
}

View File

@@ -0,0 +1,70 @@
/*
* 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.openai;
import java.util.List;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.openai.client.OpenAiClient;
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import static org.assertj.core.api.Assertions.assertThat;
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
public class OpenAiAutoConfigurationIT {
private static final Log logger = LogFactory.getLog(OpenAiAutoConfigurationIT.class);
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"))
.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class));
@Test
void generate() {
contextRunner.run(context -> {
OpenAiClient client = context.getBean(OpenAiClient.class);
String response = client.generate("Hello");
assertThat(response).isNotEmpty();
logger.info("Response: " + response);
});
}
@Test
void embedding() {
contextRunner.run(context -> {
OpenAiEmbeddingClient embeddingClient = context.getBean(OpenAiEmbeddingClient.class);
EmbeddingResponse embeddingResponse = embeddingClient
.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
assertThat(embeddingResponse.getData()).hasSize(2);
assertThat(embeddingResponse.getData().get(0).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
assertThat(embeddingResponse.getData().get(1).getEmbedding()).isNotEmpty();
assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
assertThat(embeddingClient.dimensions()).isEqualTo(1536);
});
}
}

View File

@@ -0,0 +1,42 @@
<?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>0.8.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
<packaging>jar</packaging>
<name>Spring AI Starter - Bedrock AI</name>
<description>Spring AI Bedrock AI Auto Configuration</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects/spring-ai</url>
<connection>git://github.com/spring-projects/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-bedrock</artifactId>
<version>${project.parent.version}</version>
</dependency>
</dependencies>
</project>

View File

@@ -50,20 +50,24 @@ public class VertexAiChatClient implements AiClient {
this.vertexAiApi = vertexAiApi;
}
public void setTemperature(Float temperature) {
public VertexAiChatClient withTemperature(Float temperature) {
this.temperature = temperature;
return this;
}
public void setTopK(Integer candidateCount) {
this.topK = candidateCount;
}
public void setTopP(Float topP) {
public VertexAiChatClient withTopP(Float topP) {
this.topP = topP;
return this;
}
public void setCandidateCount(Integer maxTokens) {
public VertexAiChatClient withTopK(Integer topK) {
this.topK = topK;
return this;
}
public VertexAiChatClient withCandidateCount(Integer maxTokens) {
this.candidateCount = maxTokens;
return this;
}
@Override

View File

@@ -27,6 +27,11 @@
<version>${parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>