Add observability to Bedrock Titan Embedding model

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

Author: PSriVarshan

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

@@ -34,9 +34,14 @@
 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
+ * {@link org.springframework.ai.embedding.EmbeddingModel} implementation that
+ * uses the
+ * Bedrock Titan Embedding API. Titan Embedding supports text and image (encoded
+ * in
  * base64) inputs.
  *
  * Note: Titan Embedding does not support batch embedding.
@@ -51,17 +56,24 @@ 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).
+	 * 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;
 	}
 
 	/**
-	 * Titan Embedding API input types. Could be either text or image (encoded in base64).
+	 * Titan Embedding API input types. Could be either text or image (encoded in
+	 * base64).
+	 * 
 	 * @param inputType the input type to use.
 	 */
 	public BedrockTitanEmbeddingModel withInputType(InputType inputType) {
@@ -78,17 +90,40 @@ 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 +152,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