[FEATURE] #734 Add missing foundational Bedrock models (#773)

This commit is contained in:
LangChain4j 2024-03-25 11:54:59 +01:00
parent 82c60b98e8
commit 6224530ac7
2 changed files with 14 additions and 14 deletions

View File

@ -17,8 +17,8 @@ import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelResponse;
@Getter
@SuperBuilder
public class BedrockMistralAiChatModel extends AbstractBedrockChatModel<BedrockMistralAIChatModelResponse> {
public class BedrockMistralAiChatModel extends AbstractBedrockChatModel<BedrockMistralAiChatModelResponse> {
@Builder.Default
private final int topK = 200;
@Builder.Default
@ -42,27 +42,27 @@ public class BedrockMistralAiChatModel extends AbstractBedrockChatModel<BedrockM
return parameters;
}
@Override
public Response<AiMessage> generate(List<ChatMessage> messages) {
String prompt = buildPrompt(messages);
final Map<String, Object> requestParameters = getRequestParameters(prompt);
final String body = Json.toJson(requestParameters);
InvokeModelResponse invokeModelResponse = withRetry(() -> invoke(body), getMaxRetries());
final String response = invokeModelResponse.body().asUtf8String().trim();
final BedrockMistralAIChatModelResponse result = Json.fromJson(response, getResponseClassType());
final BedrockMistralAiChatModelResponse result = Json.fromJson(response, getResponseClassType());
return new Response<>(new AiMessage(result.getOutputText()),
result.getTokenUsage(),
result.getFinishReason());
}
private String buildPrompt(List<ChatMessage> messages) {
StringBuilder promptBuilder = new StringBuilder();
promptBuilder.append("<s>");
for (ChatMessage message : messages) {
switch (message.type()) {
case USER:
@ -75,14 +75,14 @@ public class BedrockMistralAiChatModel extends AbstractBedrockChatModel<BedrockM
throw new IllegalArgumentException("Bedrock Mistral AI does not support the message type: " + message.type());
}
}
promptBuilder.append("</s>");
return promptBuilder.toString();
}
@Override
public Class<BedrockMistralAIChatModelResponse> getResponseClassType() {
return BedrockMistralAIChatModelResponse.class;
public Class<BedrockMistralAiChatModelResponse> getResponseClassType() {
return BedrockMistralAiChatModelResponse.class;
}
/**

View File

@ -12,7 +12,7 @@ import lombok.Setter;
*/
@Getter
@Setter
public class BedrockMistralAiChatModelResponse implements BedrockChatModelResponse {
class BedrockMistralAiChatModelResponse implements BedrockChatModelResponse {
private List<Output> outputs;