Fixing issues in BedrockTitanEmbeddingModel.java

Signed-off-by: PSriVarshan <[email protected]>

Author: PSriVarshan

auto-configurations/common/spring-ai-autoconfigure-retry/src/main/java/org/springframework/ai/retry/autoconfigure/SpringAiRetryAutoConfiguration.java

@@ -41,9 +41,11 @@
 import org.springframework.web.client.ResponseErrorHandler;
 
 /**
- * {@link AutoConfiguration Auto-configuration} for AI Retry.
+ * {@link AutoConfiguration Auto-configuration} for AI Retry. Provides beans for retry
+ * template and response error handling. Handles transient and non-transient exceptions
+ * based on HTTP status codes.
  *
- * @author Christian Tzolov
+ * Author: Christian Tzolov
  */
 @AutoConfiguration
 @ConditionalOnClass(RetryUtils.class)
@@ -63,9 +65,10 @@ public RetryTemplate retryTemplate(SpringAiRetryProperties properties) {
 			.withListener(new RetryListener() {
 
 				@Override
-				public <T extends Object, E extends Throwable> void onError(RetryContext context,
-						RetryCallback<T, E> callback, Throwable throwable) {
-					logger.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
+				public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback,
+						Throwable throwable) {
+					logger.warn("Retry error. Retry count: {}, Exception: {}", context.getRetryCount(),
+							throwable.getMessage(), throwable);
 				}
 			})
 			.build();
@@ -84,29 +87,35 @@ public boolean hasError(@NonNull ClientHttpResponse response) throws IOException
 
 			@Override
 			public void handleError(@NonNull ClientHttpResponse response) throws IOException {
-				if (response.getStatusCode().isError()) {
-					String error = StreamUtils.copyToString(response.getBody(), StandardCharsets.UTF_8);
-					String message = String.format("%s - %s", response.getStatusCode().value(), error);
-
-					// Explicitly configured transient codes
-					if (properties.getOnHttpCodes().contains(response.getStatusCode().value())) {
-						throw new TransientAiException(message);
-					}
-
-					// onClientErrors - If true, do not throw a NonTransientAiException,
-					// and do not attempt retry for 4xx client error codes, false by
-					// default.
-					if (!properties.isOnClientErrors() && response.getStatusCode().is4xxClientError()) {
-						throw new NonTransientAiException(message);
-					}
-
-					// Explicitly configured non-transient codes
-					if (!CollectionUtils.isEmpty(properties.getExcludeOnHttpCodes())
-							&& properties.getExcludeOnHttpCodes().contains(response.getStatusCode().value())) {
-						throw new NonTransientAiException(message);
-					}
+				if (!response.getStatusCode().isError()) {
+					return;
+				}
+
+				String error = StreamUtils.copyToString(response.getBody(), StandardCharsets.UTF_8);
+				if (error == null || error.isEmpty()) {
+					error = "No response body available";
+				}
+
+				String message = String.format("HTTP %s - %s", response.getStatusCode().value(), error);
+
+				// Explicitly configured transient codes
+				if (properties.getOnHttpCodes().contains(response.getStatusCode().value())) {
 					throw new TransientAiException(message);
 				}
+
+				// Handle client errors (4xx)
+				if (!properties.isOnClientErrors() && response.getStatusCode().is4xxClientError()) {
+					throw new NonTransientAiException(message);
+				}
+
+				// Explicitly configured non-transient codes
+				if (!CollectionUtils.isEmpty(properties.getExcludeOnHttpCodes())
+						&& properties.getExcludeOnHttpCodes().contains(response.getStatusCode().value())) {
+					throw new NonTransientAiException(message);
+				}
+
+				// Default to transient exception
+				throw new TransientAiException(message);
 			}
 		};
 	}

auto-configurations/models/spring-ai-autoconfigure-model-bedrock-ai/pom.xml

@@ -110,6 +110,12 @@
 			<artifactId>mockito-core</artifactId>
 			<scope>test</scope>
 		</dependency>
-	</dependencies>
+	
+   <dependency>
+     <groupId>io.micrometer</groupId>
+     <artifactId>micrometer-observation</artifactId>
+     <version>1.15.0-RC1</version>
+   </dependency>
+ </dependencies>
 
 </project>

auto-configurations/models/spring-ai-autoconfigure-model-bedrock-ai/src/main/java/org/springframework/ai/model/bedrock/titan/autoconfigure/BedrockTitanEmbeddingAutoConfiguration.java

@@ -17,6 +17,8 @@
 package org.springframework.ai.model.bedrock.titan.autoconfigure;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+
+import io.micrometer.observation.ObservationRegistry;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.regions.providers.AwsRegionProvider;
 
@@ -26,6 +28,7 @@
 import org.springframework.ai.model.SpringAIModels;
 import org.springframework.ai.model.bedrock.autoconfigure.BedrockAwsConnectionConfiguration;
 import org.springframework.ai.model.bedrock.autoconfigure.BedrockAwsConnectionProperties;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -50,6 +53,9 @@
 @Import(BedrockAwsConnectionConfiguration.class)
 public class BedrockTitanEmbeddingAutoConfiguration {
 
+	@Autowired
+	private ObservationRegistry observationRegistry;
+
 	@Bean
 	@ConditionalOnMissingBean
 	@ConditionalOnBean({ AwsCredentialsProvider.class, AwsRegionProvider.class })
@@ -65,7 +71,8 @@ public TitanEmbeddingBedrockApi titanEmbeddingBedrockApi(AwsCredentialsProvider
 	@ConditionalOnBean(TitanEmbeddingBedrockApi.class)
 	public BedrockTitanEmbeddingModel titanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingApi,
 			BedrockTitanEmbeddingProperties properties) {
-		return new BedrockTitanEmbeddingModel(titanEmbeddingApi).withInputType(properties.getInputType());
+		return new BedrockTitanEmbeddingModel(titanEmbeddingApi, observationRegistry)
+			.withInputType(properties.getInputType());
 	}
 
 }