Add Bedrock Cohere Command R model support.

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereCommandRChatClient.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2023 - 2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.ai.bedrock.cohere;
+
+import java.util.List;
+
+import org.springframework.ai.bedrock.BedrockUsage;
+import org.springframework.ai.bedrock.MessageToPromptConverter;
+import org.springframework.ai.bedrock.cohere.api.CohereCommandRChatBedrockApi;
+import org.springframework.ai.bedrock.cohere.api.CohereCommandRChatBedrockApi.CohereCommandRChatRequest;
+import org.springframework.ai.bedrock.cohere.api.CohereCommandRChatBedrockApi.CohereCommandRChatResponse;
+import org.springframework.ai.chat.ChatClient;
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.Generation;
+import org.springframework.ai.chat.StreamingChatClient;
+import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.model.ModelOptionsUtils;
+import org.springframework.util.Assert;
+
+import reactor.core.publisher.Flux;
+
+/**
+ * @author Wei Jiang
+ * @since 1.0.0
+ */
+public class BedrockCohereCommandRChatClient implements ChatClient, StreamingChatClient {
+
+	private final CohereCommandRChatBedrockApi chatApi;
+
+	private final BedrockCohereCommandRChatOptions defaultOptions;
+
+	public BedrockCohereCommandRChatClient(CohereCommandRChatBedrockApi chatApi) {
+		this(chatApi, BedrockCohereCommandRChatOptions.builder().build());
+	}
+
+	public BedrockCohereCommandRChatClient(CohereCommandRChatBedrockApi chatApi,
+			BedrockCohereCommandRChatOptions options) {
+		Assert.notNull(chatApi, "CohereCommandRChatBedrockApi must not be null");
+		Assert.notNull(options, "BedrockCohereCommandRChatOptions must not be null");
+
+		this.chatApi = chatApi;
+		this.defaultOptions = options;
+	}
+
+	@Override
+	public ChatResponse call(Prompt prompt) {
+		CohereCommandRChatResponse response = this.chatApi.chatCompletion(this.createRequest(prompt));
+
+		Generation generation = new Generation(response.text());
+
+		return new ChatResponse(List.of(generation));
+	}
+
+	@Override
+	public Flux<ChatResponse> stream(Prompt prompt) {
+		return this.chatApi.chatCompletionStream(this.createRequest(prompt)).map(g -> {
+			if (g.isFinished()) {
+				String finishReason = g.finishReason().name();
+				Usage usage = BedrockUsage.from(g.amazonBedrockInvocationMetrics());
+				return new ChatResponse(List
+					.of(new Generation("").withGenerationMetadata(ChatGenerationMetadata.from(finishReason, usage))));
+			}
+			return new ChatResponse(List.of(new Generation(g.text())));
+		});
+	}
+
+	CohereCommandRChatRequest createRequest(Prompt prompt) {
+		final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());
+
+		var request = CohereCommandRChatRequest.builder(promptValue)
+			.withSearchQueriesOnly(this.defaultOptions.getSearchQueriesOnly())
+			.withPreamble(this.defaultOptions.getPreamble())
+			.withMaxTokens(this.defaultOptions.getMaxTokens())
+			.withTemperature(this.defaultOptions.getTemperature())
+			.withTopP(this.defaultOptions.getTopP())
+			.withTopK(this.defaultOptions.getTopK())
+			.withPromptTruncation(this.defaultOptions.getPromptTruncation())
+			.withFrequencyPenalty(this.defaultOptions.getFrequencyPenalty())
+			.withPresencePenalty(this.defaultOptions.getPresencePenalty())
+			.withSeed(this.defaultOptions.getSeed())
+			.withReturnPrompt(this.defaultOptions.getReturnPrompt())
+			.withStopSequences(this.defaultOptions.getStopSequences())
+			.withRawPrompting(this.defaultOptions.getRawPrompting())
+			.build();
+
+		if (prompt.getOptions() != null) {
+			if (prompt.getOptions() instanceof ChatOptions runtimeOptions) {
+				BedrockCohereCommandRChatOptions updatedRuntimeOptions = ModelOptionsUtils.copyToTarget(runtimeOptions,
+						ChatOptions.class, BedrockCohereCommandRChatOptions.class);
+				request = ModelOptionsUtils.merge(updatedRuntimeOptions, request, CohereCommandRChatRequest.class);
+			}
+			else {
+				throw new IllegalArgumentException("Prompt options are not of type ChatOptions: "
+						+ prompt.getOptions().getClass().getSimpleName());
+			}
+		}
+
+		return request;
+	}
+
+}

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereCommandRChatOptions.java

@@ -0,0 +1,276 @@
+/*
+ * Copyright 2023 - 2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.ai.bedrock.cohere;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonInclude.Include;
+
+import org.springframework.ai.bedrock.cohere.api.CohereCommandRChatBedrockApi.CohereCommandRChatRequest.PromptTruncation;
+import org.springframework.ai.chat.prompt.ChatOptions;
+
+/**
+ * @author Wei Jiang
+ * @since 1.0.0
+ */
+@JsonInclude(Include.NON_NULL)
+public class BedrockCohereCommandRChatOptions implements ChatOptions {
+
+	// @formatter:off
+	/**
+	 * (optional) When enabled, it will only generate potential search queries without performing
+	 * searches or providing a response.
+	 */
+	@JsonProperty("search_queries_only") Boolean searchQueriesOnly;
+	/**
+	 * (optional) Overrides the default preamble for search query generation.
+	 */
+	@JsonProperty("preamble") String preamble;
+	/**
+	 * (optional) Specify the maximum number of tokens to use in the generated response.
+	 */
+	@JsonProperty("max_tokens") Integer maxTokens;
+	/**
+	 * (optional) Use a lower value to decrease randomness in the response.
+	 */
+	@JsonProperty("temperature") Float temperature;
+	/**
+	 * Top P. Use a lower value to ignore less probable options. Set to 0 or 1.0 to disable.
+	 */
+	@JsonProperty("p") Float topP;
+	/**
+	 * Top K. Specify the number of token choices the model uses to generate the next token.
+	 */
+	@JsonProperty("k") Integer topK;
+	/**
+	 * (optional) Dictates how the prompt is constructed.
+	 */
+	@JsonProperty("prompt_truncation") PromptTruncation promptTruncation;
+	/**
+	 * (optional) Used to reduce repetitiveness of generated tokens.
+	 */
+	@JsonProperty("frequency_penalty") Float frequencyPenalty;
+	/**
+	 * (optional) Used to reduce repetitiveness of generated tokens.
+	 */
+	@JsonProperty("presence_penalty") Float presencePenalty;
+	/**
+	 * (optional) Specify the best effort to sample tokens deterministically.
+	 */
+	@JsonProperty("seed") Integer seed;
+	/**
+	 * (optional) Specify true to return the full prompt that was sent to the model.
+	 */
+	@JsonProperty("return_prompt") Boolean returnPrompt;
+	/**
+	 * (optional) A list of stop sequences.
+	 */
+	@JsonProperty("stop_sequences") List<String> stopSequences;
+	/**
+	 * (optional) Specify true, to send the user’s message to the model without any preprocessing.
+	 */
+	@JsonProperty("raw_prompting") Boolean rawPrompting;
+	// @formatter:on
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public static class Builder {
+
+		private final BedrockCohereCommandRChatOptions options = new BedrockCohereCommandRChatOptions();
+
+		public Builder withSearchQueriesOnly(Boolean searchQueriesOnly) {
+			options.setSearchQueriesOnly(searchQueriesOnly);
+			return this;
+		}
+
+		public Builder withPreamble(String preamble) {
+			options.setPreamble(preamble);
+			return this;
+		}
+
+		public Builder withMaxTokens(Integer maxTokens) {
+			options.setMaxTokens(maxTokens);
+			return this;
+		}
+
+		public Builder withTemperature(Float temperature) {
+			options.setTemperature(temperature);
+			return this;
+		}
+
+		public Builder withTopP(Float topP) {
+			options.setTopP(topP);
+			return this;
+		}
+
+		public Builder withTopK(Integer topK) {
+			options.setTopK(topK);
+			return this;
+		}
+
+		public Builder withPromptTruncation(PromptTruncation promptTruncation) {
+			options.setPromptTruncation(promptTruncation);
+			return this;
+		}
+
+		public Builder withFrequencyPenalty(Float frequencyPenalty) {
+			options.setFrequencyPenalty(frequencyPenalty);
+			return this;
+		}
+
+		public Builder withPresencePenalty(Float presencePenalty) {
+			options.setPresencePenalty(presencePenalty);
+			return this;
+		}
+
+		public Builder withSeed(Integer seed) {
+			options.setSeed(seed);
+			return this;
+		}
+
+		public Builder withReturnPrompt(Boolean returnPrompt) {
+			options.setReturnPrompt(returnPrompt);
+			return this;
+		}
+
+		public Builder withStopSequences(List<String> stopSequences) {
+			options.setStopSequences(stopSequences);
+			return this;
+		}
+
+		public Builder withRawPrompting(Boolean rawPrompting) {
+			options.setRawPrompting(rawPrompting);
+			return this;
+		}
+
+		public BedrockCohereCommandRChatOptions build() {
+			return this.options;
+		}
+
+	}
+
+	public Boolean getSearchQueriesOnly() {
+		return searchQueriesOnly;
+	}
+
+	public void setSearchQueriesOnly(Boolean searchQueriesOnly) {
+		this.searchQueriesOnly = searchQueriesOnly;
+	}
+
+	public String getPreamble() {
+		return preamble;
+	}
+
+	public void setPreamble(String preamble) {
+		this.preamble = preamble;
+	}
+
+	public Integer getMaxTokens() {
+		return maxTokens;
+	}
+
+	public void setMaxTokens(Integer maxTokens) {
+		this.maxTokens = maxTokens;
+	}
+
+	@Override
+	public Float getTemperature() {
+		return temperature;
+	}
+
+	public void setTemperature(Float temperature) {
+		this.temperature = temperature;
+	}
+
+	@Override
+	public Float getTopP() {
+		return topP;
+	}
+
+	public void setTopP(Float topP) {
+		this.topP = topP;
+	}
+
+	@Override
+	public Integer getTopK() {
+		return topK;
+	}
+
+	public void setTopK(Integer topK) {
+		this.topK = topK;
+	}
+
+	public PromptTruncation getPromptTruncation() {
+		return promptTruncation;
+	}
+
+	public void setPromptTruncation(PromptTruncation promptTruncation) {
+		this.promptTruncation = promptTruncation;
+	}
+
+	public Float getFrequencyPenalty() {
+		return frequencyPenalty;
+	}
+
+	public void setFrequencyPenalty(Float frequencyPenalty) {
+		this.frequencyPenalty = frequencyPenalty;
+	}
+
+	public Float getPresencePenalty() {
+		return presencePenalty;
+	}
+
+	public void setPresencePenalty(Float presencePenalty) {
+		this.presencePenalty = presencePenalty;
+	}
+
+	public Integer getSeed() {
+		return seed;
+	}
+
+	public void setSeed(Integer seed) {
+		this.seed = seed;
+	}
+
+	public Boolean getReturnPrompt() {
+		return returnPrompt;
+	}
+
+	public void setReturnPrompt(Boolean returnPrompt) {
+		this.returnPrompt = returnPrompt;
+	}
+
+	public List<String> getStopSequences() {
+		return stopSequences;
+	}
+
+	public void setStopSequences(List<String> stopSequences) {
+		this.stopSequences = stopSequences;
+	}
+
+	public Boolean getRawPrompting() {
+		return rawPrompting;
+	}
+
+	public void setRawPrompting(Boolean rawPrompting) {
+		this.rawPrompting = rawPrompting;
+	}
+
+}