Skip to content

Commit c036931

Browse files
hemeda3tzolov
authored andcommitted
Implemented Bedrock Jurassic 2 ChatClient
- Implemented Bedrock Jurassic ChatClient - added documentation reference - implement auto-configuration and boot starter - Disable the BedrockAi21Jurassic2ChatClientIT.emojiPenaltyWhenTrueByDefaultApplyPenaltyTest() test as it fails when run in combination with the other tests.
1 parent db43fe3 commit c036931

File tree

14 files changed

+1290
-1
lines changed

14 files changed

+1290
-1
lines changed

models/spring-ai-bedrock/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,6 @@
66
- [Llama2 Chat Documentation](https://docs.spring.io/spring-ai/reference/api/chat/bedrock/bedrock-llama2.html)
77
- [Titan Chat Documentation](https://docs.spring.io/spring-ai/reference/api/chat/bedrock/bedrock-titan.html)
88
- [Titan Embedding Documentation](https://docs.spring.io/spring-ai/reference/api/embeddings/bedrock-titan-embedding.html)
9+
- [Jurassic2 Chat Documentation](https://docs.spring.io/spring-ai/reference/api/chat/bedrock/bedrock-jurassic2.html)
910

1011
NOTE: There is not yet an implementation for Jurassic, but you can use the lower level client [Ai21Jurassic2ChatBedrockApi.java](https://github.com/spring-projects/spring-ai/blob/main/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApi.java) in the meantime. See [Issue 343](https://github.com/spring-projects/spring-ai/issues/343)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
/*
2+
* Copyright 2023 - 2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.bedrock.jurassic2;
18+
19+
import org.springframework.ai.bedrock.MessageToPromptConverter;
20+
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi;
21+
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi.Ai21Jurassic2ChatRequest;
22+
import org.springframework.ai.chat.ChatClient;
23+
import org.springframework.ai.chat.ChatResponse;
24+
import org.springframework.ai.chat.Generation;
25+
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
26+
import org.springframework.ai.chat.prompt.ChatOptions;
27+
import org.springframework.ai.chat.prompt.Prompt;
28+
import org.springframework.ai.model.ModelOptionsUtils;
29+
import org.springframework.util.Assert;
30+
31+
/**
32+
* Java {@link ChatClient} for the Bedrock Jurassic2 chat generative model.
33+
*
34+
* @author Ahmed Yousri
35+
* @since 1.0.0
36+
*/
37+
public class BedrockAi21Jurassic2ChatClient implements ChatClient {
38+
39+
private final Ai21Jurassic2ChatBedrockApi chatApi;
40+
41+
private final BedrockAi21Jurassic2ChatOptions defaultOptions;
42+
43+
public BedrockAi21Jurassic2ChatClient(Ai21Jurassic2ChatBedrockApi chatApi,
44+
BedrockAi21Jurassic2ChatOptions options) {
45+
Assert.notNull(chatApi, "Ai21Jurassic2ChatBedrockApi must not be null");
46+
Assert.notNull(options, "BedrockAi21Jurassic2ChatOptions must not be null");
47+
48+
this.chatApi = chatApi;
49+
this.defaultOptions = options;
50+
}
51+
52+
public BedrockAi21Jurassic2ChatClient(Ai21Jurassic2ChatBedrockApi chatApi) {
53+
this(chatApi,
54+
BedrockAi21Jurassic2ChatOptions.builder()
55+
.withTemperature(0.8f)
56+
.withTopP(0.9f)
57+
.withMaxTokens(100)
58+
.build());
59+
}
60+
61+
@Override
62+
public ChatResponse call(Prompt prompt) {
63+
var request = createRequest(prompt);
64+
var response = this.chatApi.chatCompletion(request);
65+
66+
return new ChatResponse(response.completions()
67+
.stream()
68+
.map(completion -> new Generation(completion.data().text())
69+
.withGenerationMetadata(ChatGenerationMetadata.from(completion.finishReason().reason(), null)))
70+
.toList());
71+
}
72+
73+
private Ai21Jurassic2ChatRequest createRequest(Prompt prompt) {
74+
75+
final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());
76+
77+
Ai21Jurassic2ChatRequest request = Ai21Jurassic2ChatRequest.builder(promptValue).build();
78+
79+
if (prompt.getOptions() != null) {
80+
if (prompt.getOptions() instanceof ChatOptions runtimeOptions) {
81+
BedrockAi21Jurassic2ChatOptions updatedRuntimeOptions = ModelOptionsUtils.copyToTarget(runtimeOptions,
82+
ChatOptions.class, BedrockAi21Jurassic2ChatOptions.class);
83+
request = ModelOptionsUtils.merge(updatedRuntimeOptions, request, Ai21Jurassic2ChatRequest.class);
84+
}
85+
else {
86+
throw new IllegalArgumentException("Prompt options are not of type ChatOptions: "
87+
+ prompt.getOptions().getClass().getSimpleName());
88+
}
89+
}
90+
91+
if (this.defaultOptions != null) {
92+
request = ModelOptionsUtils.merge(request, this.defaultOptions, Ai21Jurassic2ChatRequest.class);
93+
}
94+
95+
return request;
96+
}
97+
98+
public static Builder builder(Ai21Jurassic2ChatBedrockApi chatApi) {
99+
return new Builder(chatApi);
100+
}
101+
102+
public static class Builder {
103+
104+
private final Ai21Jurassic2ChatBedrockApi chatApi;
105+
106+
private BedrockAi21Jurassic2ChatOptions options;
107+
108+
public Builder(Ai21Jurassic2ChatBedrockApi chatApi) {
109+
this.chatApi = chatApi;
110+
}
111+
112+
public Builder withOptions(BedrockAi21Jurassic2ChatOptions options) {
113+
this.options = options;
114+
return this;
115+
}
116+
117+
public BedrockAi21Jurassic2ChatClient build() {
118+
return new BedrockAi21Jurassic2ChatClient(chatApi,
119+
options != null ? options : BedrockAi21Jurassic2ChatOptions.builder().build());
120+
}
121+
122+
}
123+
124+
}

0 commit comments

Comments
 (0)