Add observability to Bedrock Titan Embedding model

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

Author: PSriVarshan

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

@@ -41,9 +41,12 @@
 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 SriVarshan P
  */
 @AutoConfiguration
 @ConditionalOnClass(RetryUtils.class)
@@ -63,9 +66,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 +88,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

@@ -79,6 +79,12 @@
 			<optional>true</optional>
 		</dependency>
 
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-test</artifactId>
+			<scope>test</scope>
+    	</dependency>
+
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
 			<artifactId>spring-boot-configuration-processor</artifactId>
@@ -110,6 +116,11 @@
 			<artifactId>mockito-core</artifactId>
 			<scope>test</scope>
 		</dependency>
-	</dependencies>
+	
+   <dependency>
+     <groupId>io.micrometer</groupId>
+     <artifactId>micrometer-observation</artifactId>
+   </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.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -40,6 +43,7 @@
  *
  * @author Christian Tzolov
  * @author Wei Jiang
+ * @author SriVarshan P
  * @since 0.8.0
  */
 @AutoConfiguration
@@ -56,6 +60,12 @@ public class BedrockTitanEmbeddingAutoConfiguration {
 	public TitanEmbeddingBedrockApi titanEmbeddingBedrockApi(AwsCredentialsProvider credentialsProvider,
 			AwsRegionProvider regionProvider, BedrockTitanEmbeddingProperties properties,
 			BedrockAwsConnectionProperties awsProperties, ObjectMapper objectMapper) {
+
+		// Validate required properties
+		if (properties.getModel() == null || awsProperties.getTimeout() == null) {
+			throw new IllegalArgumentException("Required properties for TitanEmbeddingBedrockApi are missing.");
+		}
+
 		return new TitanEmbeddingBedrockApi(properties.getModel(), credentialsProvider, regionProvider.getRegion(),
 				objectMapper, awsProperties.getTimeout());
 	}
@@ -64,8 +74,16 @@ public TitanEmbeddingBedrockApi titanEmbeddingBedrockApi(AwsCredentialsProvider
 	@ConditionalOnMissingBean
 	@ConditionalOnBean(TitanEmbeddingBedrockApi.class)
 	public BedrockTitanEmbeddingModel titanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingApi,
-			BedrockTitanEmbeddingProperties properties) {
-		return new BedrockTitanEmbeddingModel(titanEmbeddingApi).withInputType(properties.getInputType());
+			BedrockTitanEmbeddingProperties properties, ObjectProvider<ObservationRegistry> observationRegistry) {
+
+		// Validate required properties
+		if (properties.getInputType() == null) {
+			throw new IllegalArgumentException("InputType property for BedrockTitanEmbeddingModel is missing.");
+		}
+
+		return new BedrockTitanEmbeddingModel(titanEmbeddingApi,
+				observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
+			.withInputType(properties.getInputType());
 	}
 
 }

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

@@ -34,6 +34,9 @@
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.util.Assert;
 
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.Observation;
+
 /**
  * {@link org.springframework.ai.embedding.EmbeddingModel} implementation that uses the
  * Bedrock Titan Embedding API. Titan Embedding supports text and image (encoded in
@@ -51,13 +54,17 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
 
 	private final TitanEmbeddingBedrockApi embeddingApi;
 
+	private final ObservationRegistry observationRegistry;
+
 	/**
 	 * Titan Embedding API input types. Could be either text or image (encoded in base64).
 	 */
 	private InputType inputType = InputType.TEXT;
 
-	public BedrockTitanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingBedrockApi) {
+	public BedrockTitanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingBedrockApi,
+			ObservationRegistry observationRegistry) {
 		this.embeddingApi = titanEmbeddingBedrockApi;
+		this.observationRegistry = observationRegistry;
 	}
 
 	/**
@@ -78,17 +85,42 @@ public float[] embed(Document document) {
 	public EmbeddingResponse call(EmbeddingRequest request) {
 		Assert.notEmpty(request.getInstructions(), "At least one text is required!");
 		if (request.getInstructions().size() != 1) {
-			logger.warn(
-					"Titan Embedding does not support batch embedding. Will make multiple API calls to embed(Document)");
+			logger.warn("Titan Embedding does not support batch embedding. Multiple API calls will be made.");
 		}
 
 		List<Embedding> embeddings = new ArrayList<>();
 		var indexCounter = new AtomicInteger(0);
+
 		for (String inputContent : request.getInstructions()) {
 			var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
-			TitanEmbeddingResponse response = this.embeddingApi.embedding(apiRequest);
-			embeddings.add(new Embedding(response.embedding(), indexCounter.getAndIncrement()));
+
+			try {
+				TitanEmbeddingResponse response = Observation
+					.createNotStarted("bedrock.embedding", this.observationRegistry)
+					.lowCardinalityKeyValue("model", "titan")
+					.lowCardinalityKeyValue("input_type", this.inputType.name().toLowerCase())
+					.highCardinalityKeyValue("input_length", String.valueOf(inputContent.length()))
+					.observe(() -> {
+						TitanEmbeddingResponse r = this.embeddingApi.embedding(apiRequest);
+						Assert.notNull(r, "Embedding API returned null response");
+						return r;
+					});
+
+				if (response.embedding() == null || response.embedding().length == 0) {
+					logger.warn("Empty embedding vector returned for input at index {}. Skipping.", indexCounter.get());
+					continue;
+				}
+
+				embeddings.add(new Embedding(response.embedding(), indexCounter.getAndIncrement()));
+			}
+			catch (Exception ex) {
+				logger.error("Titan API embedding failed for input at index {}: {}", indexCounter.get(),
+						summarizeInput(inputContent), ex);
+				throw ex; // Optional: Continue instead of throwing if you want partial
+							// success
+			}
 		}
+
 		return new EmbeddingResponse(embeddings);
 	}
 
@@ -117,6 +149,13 @@ public int dimensions() {
 
 	}
 
+	private String summarizeInput(String input) {
+		if (this.inputType == InputType.IMAGE) {
+			return "[image content omitted, length=" + input.length() + "]";
+		}
+		return input.length() > 100 ? input.substring(0, 100) + "..." : input;
+	}
+
 	public enum InputType {
 
 		TEXT, IMAGE

models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModelIT.java

@@ -40,13 +40,18 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import io.micrometer.observation.tck.TestObservationRegistry;
+
 @SpringBootTest
 @RequiresAwsCredentials
 class BedrockTitanEmbeddingModelIT {
 
 	@Autowired
 	private BedrockTitanEmbeddingModel embeddingModel;
 
+	@Autowired
+	TestObservationRegistry observationRegistry;
+
 	@Test
 	void singleEmbedding() {
 		assertThat(this.embeddingModel).isNotNull();
@@ -82,8 +87,9 @@ public TitanEmbeddingBedrockApi titanEmbeddingApi() {
 		}
 
 		@Bean
-		public BedrockTitanEmbeddingModel titanEmbedding(TitanEmbeddingBedrockApi titanEmbeddingApi) {
-			return new BedrockTitanEmbeddingModel(titanEmbeddingApi);
+		public BedrockTitanEmbeddingModel titanEmbedding(TitanEmbeddingBedrockApi titanEmbeddingApi,
+				TestObservationRegistry observationRegistry) {
+			return new BedrockTitanEmbeddingModel(titanEmbeddingApi, observationRegistry);
 		}
 
 	}