chatCompletionStream(AnthropicChatRequest anthropicRequest) {
+ Assert.notNull(anthropicRequest, "'anthropicRequest' must not be null");
+ return this.internalInvocationStream(anthropicRequest, AnthropicChatResponse.class);
+ }
+
+}
+// @formatter:on
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java
new file mode 100644
index 000000000..a48baab8a
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/api/AbstractBedrockApi.java
@@ -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.
+ *
+ * https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html
+ *
+ * https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess
+ *
+ * @param The input request type.
+ * @param The output response type.
+ * @param The streaming response type. For some models this type can be the same as the output response type.
+ *
+ * @see Model Parameters
+
+ * @author Christian Tzolov
+ * @since 0.8.0
+ */
+public abstract class AbstractBedrockApi {
+
+ 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 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 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 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 internalInvocationStream(I request, Class 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
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClient.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClient.java
new file mode 100644
index 000000000..5befb9dce
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClient.java
@@ -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 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 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 generations = response.generations().stream().map(g -> {
+ return new Generation(g.text());
+ }).toList();
+
+ return new AiResponse(generations);
+ }
+
+ @Override
+ public Flux 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();
+ }
+
+}
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java
new file mode 100644
index 000000000..201a5c79e
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClient.java
@@ -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 embed(String text) {
+ return this.embed(List.of(text)).iterator().next();
+ }
+
+ @Override
+ public List embed(Document document) {
+ return embed(document.getContent());
+ }
+
+ @Override
+ public List> embed(List 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 texts) {
+ var indexCounter = new AtomicInteger(0);
+ List embeddings = this.embed(texts)
+ .stream()
+ .map(e -> new Embedding(e, indexCounter.getAndIncrement()))
+ .toList();
+ return new EmbeddingResponse(embeddings);
+ }
+
+}
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApi.java
new file mode 100644
index 000000000..6d2cd9fb8
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApi.java
@@ -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 {
+
+ /**
+ * 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 Cohere’s 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 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 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 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 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 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 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
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java
new file mode 100644
index 000000000..af8ffe0aa
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java
@@ -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 {
+
+ /**
+ * 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 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> embeddings,
+ @JsonProperty("texts") List 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
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApi.java
new file mode 100644
index 000000000..9e50b380f
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApi.java
@@ -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 {
+
+ /**
+ * 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 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 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 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 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
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClient.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClient.java
new file mode 100644
index 000000000..6255ff5d0
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClient.java
@@ -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 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 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();
+ }
+ };
+ }
+
+}
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApi.java
new file mode 100644
index 000000000..e7d7bd18a
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApi.java
@@ -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 {
+
+ /**
+ * 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 chatCompletionStream(Llama2ChatRequest request) {
+ return this.internalInvocationStream(request, Llama2ChatResponse.class);
+ }
+}
+// @formatter:on
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java
new file mode 100644
index 000000000..50009cf7f
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java
@@ -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
+ *
+ * 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 {
+
+ 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 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 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 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 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 chatCompletionStream(TitanChatRequest request) {
+ return this.internalInvocationStream(request, TitanChatResponseChunk.class);
+ }
+}
+// @formatter:on
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java
new file mode 100644
index 000000000..63adadf1b
--- /dev/null
+++ b/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java
@@ -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 {
+
+ /**
+ * 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 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
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/BedrockAnthropicChatClientIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/BedrockAnthropicChatClientIT.java
new file mode 100644
index 000000000..3615d83e9
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/BedrockAnthropicChatClientIT.java
@@ -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 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 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 movies) {
+ }
+
+ @Test
+ void beanOutputParserRecords() {
+
+ BeanOutputParser 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 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);
+ }
+
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/api/AnthropicChatBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/api/AnthropicChatBedrockApiIT.java
new file mode 100644
index 000000000..334acfc75
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/anthropic/api/AnthropicChatBedrockApiIT.java
@@ -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 responseStream = anthropicChatApi.chatCompletionStream(request);
+
+ List responses = responseStream.collectList().block();
+ assertThat(responses).isNotNull();
+ assertThat(responses).hasSizeGreaterThan(10);
+ assertThat(responses.stream().map(AnthropicChatResponse::completion).collect(Collectors.joining()))
+ .contains("Blackbeard");
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClientIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClientIT.java
new file mode 100644
index 000000000..5900ec9eb
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereChatClientIT.java
@@ -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 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 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 movies) {
+ }
+
+ @Test
+ void beanOutputParserRecords() {
+
+ BeanOutputParser 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 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);
+ }
+
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java
new file mode 100644
index 000000000..d3e7f3b5d
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java
@@ -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);
+ }
+
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApiIT.java
new file mode 100644
index 000000000..a54ef1481
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereChatBedrockApiIT.java
@@ -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 responseStream = cohereChatApi.chatCompletionStream(request);
+ List 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!");
+
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java
new file mode 100644
index 000000000..5c38aabb8
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java
@@ -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);
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApiIT.java
new file mode 100644
index 000000000..7ec33b5bd
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/jurassic2/api/Ai21Jurassic2ChatBedrockApiIT.java
@@ -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!
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClientIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClientIT.java
new file mode 100644
index 000000000..56fcd3c39
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/BedrockLlama2ChatClientIT.java
@@ -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 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 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 movies) {
+ }
+
+ @Disabled("TODO: Fix the parser instructions to return the correct format")
+ @Test
+ void beanOutputParserRecords() {
+
+ BeanOutputParser 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 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);
+ }
+
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApiIT.java
new file mode 100644
index 000000000..003462998
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/llama2/api/Llama2ChatBedrockApiIT.java
@@ -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 responseStream = llama2ChatApi.chatCompletionStream(request);
+ List 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();
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApiIT.java
new file mode 100644
index 000000000..fcec41ed4
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApiIT.java
@@ -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 response = titanBedrockApi.chatCompletionStream(titanChatRequest);
+ List results = response.collectList().block();
+
+ assertThat(results.stream().map(TitanChatResponseChunk::outputText).collect(Collectors.joining("\n")))
+ .contains("Blackbeard");
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApiIT.java b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApiIT.java
new file mode 100644
index 000000000..7e6515626
--- /dev/null
+++ b/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApiIT.java
@@ -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);
+ }
+
+}
diff --git a/spring-ai-bedrock/src/test/resources/doc/Bedrock Cohere Chat API.jpg b/spring-ai-bedrock/src/test/resources/doc/Bedrock Cohere Chat API.jpg
new file mode 100644
index 000000000..cecfae8e7
Binary files /dev/null and b/spring-ai-bedrock/src/test/resources/doc/Bedrock Cohere Chat API.jpg differ
diff --git a/spring-ai-bedrock/src/test/resources/doc/Bedrock Cohere Embedding API.jpg b/spring-ai-bedrock/src/test/resources/doc/Bedrock Cohere Embedding API.jpg
new file mode 100644
index 000000000..1516817e2
Binary files /dev/null and b/spring-ai-bedrock/src/test/resources/doc/Bedrock Cohere Embedding API.jpg differ
diff --git a/spring-ai-bedrock/src/test/resources/doc/Bedrock-Llama2-Chat-API.jpg b/spring-ai-bedrock/src/test/resources/doc/Bedrock-Llama2-Chat-API.jpg
new file mode 100644
index 000000000..2c0d3c4f6
Binary files /dev/null and b/spring-ai-bedrock/src/test/resources/doc/Bedrock-Llama2-Chat-API.jpg differ
diff --git a/spring-ai-bedrock/src/test/resources/prompts/system-message.st b/spring-ai-bedrock/src/test/resources/prompts/system-message.st
new file mode 100644
index 000000000..dc2cf2dcd
--- /dev/null
+++ b/spring-ai-bedrock/src/test/resources/prompts/system-message.st
@@ -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}.
\ No newline at end of file
diff --git a/spring-ai-bedrock/src/test/resources/spring_framework.png b/spring-ai-bedrock/src/test/resources/spring_framework.png
new file mode 100644
index 000000000..735057a6b
Binary files /dev/null and b/spring-ai-bedrock/src/test/resources/spring_framework.png differ
diff --git a/spring-ai-core/pom.xml b/spring-ai-core/pom.xml
index bf7da5411..6f7090d86 100644
--- a/spring-ai-core/pom.xml
+++ b/spring-ai-core/pom.xml
@@ -38,23 +38,16 @@
${antlr.version}
-
io.projectreactor
reactor-core
-
org.springframework
spring-messaging
-
- org.springframework.boot
- spring-boot-starter-json
-
-
com.knuddels
jtokkit
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java
index add3533b0..01721d6bf 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java
@@ -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 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 embedding, Integer index) {
this.embedding = embedding;
this.index = index;
}
+ /**
+ * @return Get the embedding vector values.
+ */
public List getEmbedding() {
return embedding;
}
+ /**
+ * @return Get the embedding index in a list of embeddings.
+ */
public Integer getIndex() {
return index;
}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java
index 6178357f4..a9fb69c28 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java
@@ -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 data;
+ /**
+ * Embedding metadata.
+ */
private Map metadata = new HashMap<>();
+ /**
+ * Creates a new {@link EmbeddingResponse} instance with empty metadata.
+ * @param data the embedding data.
+ */
+ public EmbeddingResponse(List data) {
+ this(data, new HashMap<>());
+ }
+
+ /**
+ * Creates a new {@link EmbeddingResponse} instance.
+ * @param data the embedding data.
+ * @param metadata the embedding metadata.
+ */
public EmbeddingResponse(List data, Map metadata) {
this.data = data;
this.metadata = metadata;
}
+ /**
+ * @return Get the embedding data.
+ */
public List getData() {
return data;
}
+ /**
+ * @return Get the embedding metadata.
+ */
public Map getMetadata() {
return metadata;
}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java b/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java
index 5364aa0d2..718153a54 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java
@@ -48,7 +48,7 @@ public class BeanOutputParser implements OutputParser {
/** The Java class representing the target type. */
@SuppressWarnings({ "FieldMayBeFinal", "rawtypes" })
- private Class clazz;
+ private Class clazz;
/** The object mapper used for deserialization and other JSON operations. */
@SuppressWarnings("FieldMayBeFinal")
@@ -95,7 +95,6 @@ public class BeanOutputParser implements OutputParser {
*/
public T parse(String text) {
try {
- // noinspection unchecked
return (T) this.objectMapper.readValue(text, this.clazz);
}
catch (JsonProcessingException e) {
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AbstractMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AbstractMessage.java
index 68ec99df0..9f4c123ec 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AbstractMessage.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AbstractMessage.java
@@ -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;
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/Message.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/Message.java
index 190cd9642..a405eb10b 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/Message.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/Message.java
@@ -26,6 +26,4 @@ public interface Message {
MessageType getMessageType();
- String getMessageTypeValue();
-
}
diff --git a/spring-ai-huggingface/pom.xml b/spring-ai-huggingface/pom.xml
index 582524846..b8c13239e 100644
--- a/spring-ai-huggingface/pom.xml
+++ b/spring-ai-huggingface/pom.xml
@@ -42,6 +42,11 @@
+
+ org.springframework
+ spring-web
+
+
org.springframework
spring-context-support
diff --git a/spring-ai-openai/pom.xml b/spring-ai-openai/pom.xml
index 7aeaa0c52..edf232627 100644
--- a/spring-ai-openai/pom.xml
+++ b/spring-ai-openai/pom.xml
@@ -55,7 +55,11 @@
${victools.version}
-
+
+ org.springframework
+ spring-webflux
+
+
org.springframework
spring-context-support
@@ -64,12 +68,6 @@
org.springframework.boot
spring-boot-starter-logging
-
- org.springframework
- spring-webflux
- ${spring-framework.version}
-
-
diff --git a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java
index 5baf05229..277ce7c17 100644
--- a/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java
+++ b/spring-ai-openai/src/test/java/org/springframework/ai/openai/client/OpenAiClientIT.java
@@ -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);
}
diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml
index 3828d25cd..b7569b1c4 100644
--- a/spring-ai-spring-boot-autoconfigure/pom.xml
+++ b/spring-ai-spring-boot-autoconfigure/pom.xml
@@ -156,6 +156,13 @@
true
+
+
+ org.springframework.ai
+ spring-ai-bedrock
+ ${project.parent.version}
+ true
+
org.springframework.boot
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java
index e4cff74d3..0dcbb4311 100644
--- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java
@@ -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
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/BedrockAwsConnectionConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/BedrockAwsConnectionConfiguration.java
new file mode 100644
index 000000000..a1363c3f1
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/BedrockAwsConnectionConfiguration.java
@@ -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();
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/BedrockAwsConnectionProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/BedrockAwsConnectionProperties.java
new file mode 100644
index 000000000..76725bc71
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/BedrockAwsConnectionProperties.java
@@ -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;
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfiguration.java
new file mode 100644
index 000000000..93a4013be
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfiguration.java
@@ -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());
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatProperties.java
new file mode 100644
index 000000000..5b22fff7b
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatProperties.java
@@ -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 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 getStopSequences() {
+ return stopSequences;
+ }
+
+ public void setStopSequences(List stopSequences) {
+ this.stopSequences = stopSequences;
+ }
+
+ public String getAnthropicVersion() {
+ return anthropicVersion;
+ }
+
+ public void setAnthropicVersion(String anthropicVersion) {
+ this.anthropicVersion = anthropicVersion;
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfiguration.java
new file mode 100644
index 000000000..1a3e78a82
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfiguration.java
@@ -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());
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatProperties.java
new file mode 100644
index 000000000..97d58bbf3
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatProperties.java
@@ -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 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 getStopSequences() {
+ return stopSequences;
+ }
+
+ public void setStopSequences(List 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;
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfiguration.java
new file mode 100644
index 000000000..7d2e16b7b
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfiguration.java
@@ -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());
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java
new file mode 100644
index 000000000..5a5762a10
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java
@@ -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;
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfiguration.java
new file mode 100644
index 000000000..af657e887
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfiguration.java
@@ -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());
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatProperties.java
new file mode 100644
index 000000000..b224dacd8
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatProperties.java
@@ -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;
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java
index 8cab8410d..0f33fb599 100644
--- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java
@@ -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;
}
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
index b3a5b3711..6046d8113 100644
--- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
+++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
@@ -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
+
+
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfigurationIT.java
new file mode 100644
index 000000000..2895ae6c7
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/anthropic/BedrockAnthropicChatAutoConfigurationIT.java
@@ -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 response = anthropicChatClient
+ .generateStream(new Prompt(List.of(userMessage, systemMessage)));
+
+ List 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();
+ });
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfigurationIT.java
new file mode 100644
index 000000000..b7041d1f1
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereChatAutoConfigurationIT.java
@@ -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 response = cohereChatClient
+ .generateStream(new Prompt(List.of(userMessage, systemMessage)));
+
+ List 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();
+ });
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java
new file mode 100644
index 000000000..afd747ddb
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java
@@ -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();
+ });
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfigurationIT.java
new file mode 100644
index 000000000..d85b6bf5a
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/llama2/BedrockLlama2ChatAutoConfigurationIT.java
@@ -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 response = llama2ChatClient
+ .generateStream(new Prompt(List.of(userMessage, systemMessage)));
+
+ List 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();
+ });
+ }
+
+}
diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java
new file mode 100644
index 000000000..0f132edd9
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfigurationIT.java
@@ -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);
+ });
+ }
+
+}
diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai/pom.xml
new file mode 100644
index 000000000..9af97333c
--- /dev/null
+++ b/spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai/pom.xml
@@ -0,0 +1,42 @@
+
+
+ 4.0.0
+
+ org.springframework.ai
+ spring-ai
+ 0.8.0-SNAPSHOT
+ ../../pom.xml
+
+ spring-ai-bedrock-ai-spring-boot-starter
+ jar
+ Spring AI Starter - Bedrock AI
+ Spring AI Bedrock AI Auto Configuration
+ https://github.com/spring-projects/spring-ai
+
+
+ https://github.com/spring-projects/spring-ai
+ git://github.com/spring-projects/spring-ai.git
+ git@github.com:spring-projects/spring-ai.git
+
+
+
+
+ org.springframework.boot
+ spring-boot-starter
+
+
+
+ org.springframework.ai
+ spring-ai-spring-boot-autoconfigure
+ ${project.parent.version}
+
+
+
+ org.springframework.ai
+ spring-ai-bedrock
+ ${project.parent.version}
+
+
+
+
diff --git a/spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/generation/VertexAiChatClient.java b/spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/generation/VertexAiChatClient.java
index a231e9ec4..5a2ac9432 100644
--- a/spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/generation/VertexAiChatClient.java
+++ b/spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/generation/VertexAiChatClient.java
@@ -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
diff --git a/vector-stores/spring-ai-chroma/pom.xml b/vector-stores/spring-ai-chroma/pom.xml
index 671a0809d..8aaa88d0d 100644
--- a/vector-stores/spring-ai-chroma/pom.xml
+++ b/vector-stores/spring-ai-chroma/pom.xml
@@ -27,6 +27,11 @@
${parent.version}
+
+ org.springframework
+ spring-web
+
+
org.springframework.ai