Add Bedrock Titan Chat Options + Docs

Author: tzolov

README.md

@@ -6,6 +6,8 @@ The Spring AI project provides a Spring-friendly API and abstractions for develo
 
 Let's make your `@Beans` intelligent!
 
+For further information go to our [Spring AI documentation](https://docs.spring.io/spring-ai/reference/).
+
 ## Project Update
 
 :partying_face: The Spring AI project has graduated out of the repository!

models/spring-ai-bedrock/README_TITAN_CHAT.md

@@ -1,83 +1,3 @@
-# 1. Bedrock Titan Chat
+# Bedrock Titan Chat
 
-## 1.1 TitanChatBedrockApi
-
-[TitanChatBedrockApi](./src/main/java/org/springframework/ai/bedrock/titan/api/TitanChatBedrockApi.java) provides is lightweight Java client on top of AWS Bedrock [Titan text models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html).
-
-Following class diagram illustrates the Llama2ChatBedrockApi interface and building blocks:
-
-![TitanChatBedrockApi Class Diagram](./src/test/resources/doc/Bedrock%20Titan%20Chat%20API.jpg)
-
-The TitanChatBedrockApi supports the `amazon.titan-text-lite-v1` and `amazon.titan-text-express-v1` models for bot synchronous (e.g. `chatCompletion()`) and streaming (e.g. `chatCompletionStream()`) responses.
-
-Here is a simple snippet how to use the api programmatically:
-
-```java
-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();
-
-TitanChatResponse response = titanBedrockApi.chatCompletion(titanChatRequest);
-
-assertThat(response.results()).hasSize(1);
-assertThat(response.results().get(0).outputText()).contains("Blackbeard");
-
-Flux<TitanChatResponseChunk> response = titanBedrockApi.chatCompletionStream(titanChatRequest);
-
-List<TitanChatResponseChunk> results = response.collectList().block();
-assertThat(results.stream().map(TitanChatResponseChunk::outputText).collect(Collectors.joining("\n")))
-	.contains("Blackbeard");
-```
-
-## 1.2 BedrockTitanChatClient
-
-[BedrockTitanChatClient](./src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClient.java) implements the Spring-Ai `ChatClient` and `StreamingChatClient` on top of the `TitanChatBedrockApi`.
-
-You can use like this:
-
-```java
-@Bean
-public TitanChatBedrockApi titanApi() {
-	return new TitanChatBedrockApi(TitanChatModel.TITAN_TEXT_EXPRESS_V1.id(),
-			EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper());
-}
-
-@Bean
-public BedrockTitanChatClient titanChatClient(TitanChatBedrockApi titanApi) {
-	return new BedrockTitanChatClient(titanApi);
-}
-
-```
-
-or you can leverage the `spring-ai-bedrock-ai-spring-boot-starter` Boot starter. For this add the following dependency:
-
-```xml
-<dependency>
-	<artifactId>spring-ai-bedrock-ai-spring-boot-starter</artifactId>
-	<groupId>org.springframework.ai</groupId>
-    <version>0.8.0-SNAPSHOT</version>
-</dependency>
-```
-
-**NOTE:** You have to enable the Bedrock Titan chat client with `spring.ai.bedrock.titan.chat.enabled=true`.
-By default the client is disabled.
-
-Use the `BedrockTitanChatProperties` to configure the Bedrock Titan Chat client:
-
-| Property  | Description | Default |
-| ------------- | ------------- | ------------- |
-| spring.ai.bedrock.aws.region  | AWS region to use.  | us-east-1 |
-| spring.ai.bedrock.aws.accessKey  | AWS credentials access key.  |  |
-| spring.ai.bedrock.aws.secretKey  | AWS credentials secret key.  |  |
-| spring.ai.bedrock.titan.chat.enable | Enable Bedrock Titan chat client. Disabled by default | false |
-| spring.ai.bedrock.titan.chat.model  | The model id to use. See the `TitanChatModel` for the supported models.  | amazon.titan-text-express-v1 |
-| spring.ai.bedrock.titan.chat.temperature  | Controls the randomness of the output. Values can range over [0.0,1.0]  | 0.7 |
-| spring.ai.bedrock.titan.chat.topP  | The maximum cumulative probability of tokens to consider when sampling.  | AWS Bedrock default |
-| spring.ai.bedrock.titan.chat.maxTokenCount  | Specify the maximum number of tokens to use in the generated response. | AWS Bedrock default |
-| spring.ai.bedrock.titan.chat.stopSequences  | Configure up to four sequences that the model recognizes. | AWS Bedrock default |
+Visit the Spring AI [Bedrock Titan Chat Documentation](https://docs.spring.io/spring-ai/reference/api/clients/bedrock/bedrock-titan.html).

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatClient.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2023 the original author or authors.
+ * 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.
@@ -18,20 +18,23 @@
 
 import java.util.List;
 
-import org.springframework.ai.chat.ChatClient;
-import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
 import reactor.core.publisher.Flux;
 
 import org.springframework.ai.bedrock.MessageToPromptConverter;
 import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
 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 org.springframework.ai.chat.ChatClient;
+import org.springframework.ai.chat.ChatOptions;
 import org.springframework.ai.chat.ChatResponse;
-import org.springframework.ai.chat.StreamingChatClient;
 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.Prompt;
+import org.springframework.ai.model.ModelOptionsUtils;
+import org.springframework.util.Assert;
 
 /**
  * @author Christian Tzolov
@@ -41,41 +44,22 @@ public class BedrockTitanChatClient implements ChatClient, StreamingChatClient {
 
 	private final TitanChatBedrockApi chatApi;
 
-	private Float temperature;
-
-	private Float topP;
-
-	private Integer maxTokenCount;
-
-	private List<String> stopSequences;
+	private final BedrockTitanChatOptions defaultOptions;
 
 	public BedrockTitanChatClient(TitanChatBedrockApi chatApi) {
-		this.chatApi = chatApi;
+		this(chatApi, BedrockTitanChatOptions.builder().withTemperature(0.8f).build());
 	}
 
-	public BedrockTitanChatClient withTemperature(Float temperature) {
-		this.temperature = temperature;
-		return this;
-	}
-
-	public BedrockTitanChatClient withTopP(Float topP) {
-		this.topP = topP;
-		return this;
-	}
-
-	public BedrockTitanChatClient withMaxTokenCount(Integer maxTokens) {
-		this.maxTokenCount = maxTokens;
-		return this;
-	}
-
-	public BedrockTitanChatClient withStopSequences(List<String> stopSequences) {
-		this.stopSequences = stopSequences;
-		return this;
+	public BedrockTitanChatClient(TitanChatBedrockApi chatApi, BedrockTitanChatOptions defaultOptions) {
+		Assert.notNull(chatApi, "ChatApi must not be null");
+		Assert.notNull(defaultOptions, "DefaultOptions must not be null");
+		this.chatApi = chatApi;
+		this.defaultOptions = defaultOptions;
 	}
 
 	@Override
 	public ChatResponse call(Prompt prompt) {
-		TitanChatResponse response = this.chatApi.chatCompletion(this.createRequest(prompt, false));
+		TitanChatResponse response = this.chatApi.chatCompletion(this.createRequest(prompt));
 		List<Generation> generations = response.results().stream().map(result -> {
 			return new Generation(result.outputText());
 		}).toList();
@@ -85,7 +69,7 @@ public ChatResponse call(Prompt prompt) {
 
 	@Override
 	public Flux<ChatResponse> stream(Prompt prompt) {
-		return this.chatApi.chatCompletionStream(this.createRequest(prompt, true)).map(chunk -> {
+		return this.chatApi.chatCompletionStream(this.createRequest(prompt)).map(chunk -> {
 
 			Generation generation = new Generation(chunk.outputText());
 
@@ -104,15 +88,48 @@ else if (chunk.inputTextTokenCount() != null && chunk.totalOutputTextTokenCount(
 		});
 	}
 
-	private TitanChatRequest createRequest(Prompt prompt, boolean stream) {
+	/**
+	 * Test access.
+	 */
+	TitanChatRequest createRequest(Prompt prompt) {
 		final String promptValue = MessageToPromptConverter.create().toPrompt(prompt.getInstructions());
 
-		return TitanChatRequest.builder(promptValue)
-			.withTemperature(this.temperature)
-			.withTopP(this.topP)
-			.withMaxTokenCount(this.maxTokenCount)
-			.withStopSequences(this.stopSequences)
-			.build();
+		var requestBuilder = TitanChatRequest.builder(promptValue);
+
+		if (this.defaultOptions != null) {
+			requestBuilder = update(requestBuilder, this.defaultOptions);
+		}
+
+		if (prompt.getOptions() != null) {
+			if (prompt.getOptions() instanceof ChatOptions runtimeOptions) {
+				BedrockTitanChatOptions updatedRuntimeOptions = ModelOptionsUtils.copyToTarget(runtimeOptions,
+						ChatOptions.class, BedrockTitanChatOptions.class);
+
+				requestBuilder = update(requestBuilder, updatedRuntimeOptions);
+			}
+			else {
+				throw new IllegalArgumentException("Prompt options are not of type ChatOptions: "
+						+ prompt.getOptions().getClass().getSimpleName());
+			}
+		}
+
+		return requestBuilder.build();
+	}
+
+	private TitanChatRequest.Builder update(TitanChatRequest.Builder builder, BedrockTitanChatOptions options) {
+		if (options.getTemperature() != null) {
+			builder.withTemperature(options.getTemperature());
+		}
+		if (options.getTopP() != null) {
+			builder.withTopP(options.getTopP());
+		}
+		if (options.getMaxTokenCount() != null) {
+			builder.withMaxTokenCount(options.getMaxTokenCount());
+		}
+		if (options.getStopSequences() != null) {
+			builder.withStopSequences(options.getStopSequences());
+		}
+		return builder;
 	}
 
 	private Usage extractUsage(TitanChatResponseChunk response) {

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanChatOptions.java

@@ -0,0 +1,133 @@
+/*
+ * Copyright 2024-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.bedrock.titan;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonInclude.Include;
+
+import org.springframework.ai.chat.ChatOptions;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * @author Christian Tzolov
+ * @since 0.8.0
+ */
+@JsonInclude(Include.NON_NULL)
+public class BedrockTitanChatOptions implements ChatOptions {
+
+	// @formatter:off
+	/**
+	 * The temperature value controls the randomness of the generated text.
+	 */
+	private @JsonProperty("temperature") Float temperature;
+
+	/**
+	 * The topP value controls the diversity of the generated text. Use a lower value to ignore less probable options.
+	 */
+	private @JsonProperty("topP") Float topP;
+
+	/**
+	 * Maximum number of tokens to generate.
+	 */
+	private @JsonProperty("maxTokenCount") Integer maxTokenCount;
+
+	/**
+	 * A list of tokens that the model should stop generating after.
+	 */
+	private @JsonProperty("stopSequences") List<String> stopSequences;
+	// @formatter:on
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public static class Builder {
+
+		private BedrockTitanChatOptions options = new BedrockTitanChatOptions();
+
+		public Builder withTemperature(Float temperature) {
+			this.options.temperature = temperature;
+			return this;
+		}
+
+		public Builder withTopP(Float topP) {
+			this.options.topP = topP;
+			return this;
+		}
+
+		public Builder withMaxTokenCount(Integer maxTokenCount) {
+			this.options.maxTokenCount = maxTokenCount;
+			return this;
+		}
+
+		public Builder withStopSequences(List<String> stopSequences) {
+			this.options.stopSequences = stopSequences;
+			return this;
+		}
+
+		public BedrockTitanChatOptions build() {
+			return this.options;
+		}
+
+	}
+
+	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 getMaxTokenCount() {
+		return maxTokenCount;
+	}
+
+	public void setMaxTokenCount(Integer maxTokenCount) {
+		this.maxTokenCount = maxTokenCount;
+	}
+
+	public List<String> getStopSequences() {
+		return stopSequences;
+	}
+
+	public void setStopSequences(List<String> stopSequences) {
+		this.stopSequences = stopSequences;
+	}
+
+	@Override
+	public Integer getTopK() {
+		throw new UnsupportedOperationException("Bedrock Titian Chat does not support the 'TopK' option.");
+	}
+
+	@Override
+	public void setTopK(Integer topK) {
+		throw new UnsupportedOperationException("Bedrock Titian Chat does not support the 'TopK' option.'");
+	}
+
+}