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

Add Bedrock Llama2 Options Support

Browse Source
This commit is contained in:
Christian Tzolov
2024-02-08 19:41:21 +01:00
parent ee81558e10
commit 577a605cdc
13 changed files with 450 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClient.java
View File

@@ -22,15 +22,20 @@ import org.springframework.ai.chat.ChatResponse;
import reactor.core.publisher.Flux;
import org.springframework.ai.bedrock.MessageToPromptConverter;
import org.springframework.ai.bedrock.anthropic.AnthropicChatOptions;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi.AnthropicChatRequest;
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.chat.ChatClient;
import org.springframework.ai.chat.ChatOptions;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.Generation;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.util.Assert;
/**
* Java {@link ChatClient} and {@link StreamingChatClient} for the Bedrock Llama2 chat
@@ -43,40 +48,25 @@ public class BedrockLlama2ChatClient implements ChatClient, StreamingChatClient
private final Llama2ChatBedrockApi chatApi;
private Float temperature;
private Float topP;
private Integer maxGenLen;
private final BedrockLlama2ChatOptions defaultOptions;
public BedrockLlama2ChatClient(Llama2ChatBedrockApi chatApi) {
this(chatApi,
BedrockLlama2ChatOptions.builder().withTemperature(0.8f).withTopP(0.9f).withMaxGenLen(100).build());
}
public BedrockLlama2ChatClient(Llama2ChatBedrockApi chatApi, BedrockLlama2ChatOptions options) {
Assert.notNull(chatApi, "Llama2ChatBedrockApi must not be null");
Assert.notNull(options, "BedrockLlama2ChatOptions must not be null");
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;
this.defaultOptions = options;
}
@Override
public ChatResponse call(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());
var request = Llama2ChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withTopP(this.topP)
.withMaxGenLen(this.maxGenLen)
.build();
var request = createRequest(prompt);
Llama2ChatResponse response = this.chatApi.chatCompletion(request);
@@ -87,13 +77,7 @@ public class BedrockLlama2ChatClient implements ChatClient, StreamingChatClient
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());
var request = Llama2ChatRequest.builder(promptValue)
.withTemperature(this.temperature)
.withTopP(this.topP)
.withMaxGenLen(this.maxGenLen)
.build();
var request = createRequest(prompt);
Flux<Llama2ChatResponse> fluxResponse = this.chatApi.chatCompletionStream(request);
@@ -119,4 +103,33 @@ public class BedrockLlama2ChatClient implements ChatClient, StreamingChatClient
};
}
/**
* Accessible for testing.
*/
Llama2ChatRequest createRequest(Prompt prompt) {
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());
Llama2ChatRequest request = Llama2ChatRequest.builder(promptValue).build();
if (this.defaultOptions != null) {
request = ModelOptionsUtils.merge(request, this.defaultOptions, Llama2ChatRequest.class);
}
if (prompt.getOptions() != null) {
if (prompt.getOptions() instanceof ChatOptions runtimeOptions) {
BedrockLlama2ChatOptions updatedRuntimeOptions = ModelOptionsUtils.copyToTarget(runtimeOptions,
ChatOptions.class, BedrockLlama2ChatOptions.class);
request = ModelOptionsUtils.merge(updatedRuntimeOptions, request, Llama2ChatRequest.class);
}
else {
throw new IllegalArgumentException("Prompt options are not of type ChatOptions: "
+ prompt.getOptions().getClass().getSimpleName());
}
}
return request;
}
}

114
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatOptions.java Normal file
View File

@@ -0,0 +1,114 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.bedrock.llama2;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.chat.ChatOptions;
/**
* @author Christian Tzolov
*/
@JsonInclude(Include.NON_NULL)
public class BedrockLlama2ChatOptions implements ChatOptions {
/**
* The temperature value controls the randomness of the generated text. Use a lower
* value to decrease randomness in the response.
*/
private @JsonProperty("temperature") Float temperature;
/**
* 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.
*/
private @JsonProperty("top_p") Float topP;
/**
* The maximum length of the generated text.
*/
private @JsonProperty("max_gen_len") Integer maxGenLen;
public static Builder builder() {
return new Builder();
}
public static class Builder {
private BedrockLlama2ChatOptions options = new BedrockLlama2ChatOptions();
public Builder withTemperature(Float temperature) {
this.options.setTemperature(temperature);
return this;
}
public Builder withTopP(Float topP) {
this.options.setTopP(topP);
return this;
}
public Builder withMaxGenLen(Integer maxGenLen) {
this.options.setMaxGenLen(maxGenLen);
return this;
}
public BedrockLlama2ChatOptions build() {
return this.options;
}
}
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 maxGenLen) {
this.maxGenLen = maxGenLen;
}
@Override
@JsonIgnore
public Integer getTopK() {
throw new UnsupportedOperationException("Unsupported option: 'TopK'");
}
@Override
@JsonIgnore
public void setTopK(Integer topK) {
throw new UnsupportedOperationException("Unsupported option: 'TopK'");
}
}

3
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClientIT.java
View File

@@ -166,7 +166,8 @@ class BedrockLlama2ChatClientIT {
@Bean
public BedrockLlama2ChatClient llama2ChatClient(Llama2ChatBedrockApi llama2Api) {
return new BedrockLlama2ChatClient(llama2Api);
return new BedrockLlama2ChatClient(llama2Api,
BedrockLlama2ChatOptions.builder().withTemperature(0.5f).withMaxGenLen(100).withTopP(0.9f).build());
}
}

60
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/BedrockLlama2CreateRequestTests.java Normal file
View File

@@ -0,0 +1,60 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.bedrock.llama2;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Test;
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.Llama2ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
public class BedrockLlama2CreateRequestTests {
private Llama2ChatBedrockApi api = new Llama2ChatBedrockApi(Llama2ChatModel.LLAMA2_70B_CHAT_V1.id(),
EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
@Test
public void createRequestWithChatOptions() {
var client = new BedrockLlama2ChatClient(api,
BedrockLlama2ChatOptions.builder().withTemperature(66.6f).withMaxGenLen(666).withTopP(0.66f).build());
var request = client.createRequest(new Prompt("Test message content"));
assertThat(request.prompt()).isNotEmpty();
assertThat(request.temperature()).isEqualTo(66.6f);
assertThat(request.topP()).isEqualTo(0.66f);
assertThat(request.maxGenLen()).isEqualTo(666);
request = client.createRequest(new Prompt("Test message content",
BedrockLlama2ChatOptions.builder().withTemperature(99.9f).withMaxGenLen(999).withTopP(0.99f).build()));
assertThat(request.prompt()).isNotEmpty();
assertThat(request.temperature()).isEqualTo(99.9f);
assertThat(request.topP()).isEqualTo(0.99f);
assertThat(request.maxGenLen()).isEqualTo(999);
}
}