From 73df85d5eb3b08616f2a019fdd105ee2117d5c75 Mon Sep 17 00:00:00 2001
From: GR <gr@fastball.dev>
Date: Tue, 8 Oct 2024 19:58:45 +0800
Subject: [PATCH]  MiniMax embedding model autoconfiguration

- Add context isolation in embedding unit test
---
 .../ai/minimax/MiniMaxEmbeddingModel.java          |  4 ++++
 .../ai/minimax/embedding/EmbeddingIT.java          |  3 ++-
 .../minimax/MiniMaxAutoConfiguration.java          | 14 +++++++++++---
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java
index 06c219a0e..0ba752c38 100644
--- a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java
+++ b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxEmbeddingModel.java
@@ -205,4 +205,8 @@ public class MiniMaxEmbeddingModel extends AbstractEmbeddingModel {
 		return optionBuilder.build();
 	}
 
+	public void setObservationConvention(EmbeddingModelObservationConvention observationConvention) {
+		this.observationConvention = observationConvention;
+	}
+
 }
diff --git a/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/embedding/EmbeddingIT.java b/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/embedding/EmbeddingIT.java
index 6e2968673..551c9fef6 100644
--- a/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/embedding/EmbeddingIT.java
+++ b/models/spring-ai-minimax/src/test/java/org/springframework/ai/minimax/embedding/EmbeddingIT.java
@@ -19,6 +19,7 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
 import org.springframework.ai.embedding.EmbeddingResponse;
 import org.springframework.ai.minimax.MiniMaxEmbeddingModel;
+import org.springframework.ai.minimax.MiniMaxTestConfiguration;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 
@@ -29,7 +30,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 /**
  * @author Geng Rong
  */
-@SpringBootTest
+@SpringBootTest(classes = MiniMaxTestConfiguration.class)
 @EnabledIfEnvironmentVariable(named = "MINIMAX_API_KEY", matches = ".+")
 class EmbeddingIT {
 
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/minimax/MiniMaxAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/minimax/MiniMaxAutoConfiguration.java
index 03ccbdea7..43ea7928b 100644
--- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/minimax/MiniMaxAutoConfiguration.java
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/minimax/MiniMaxAutoConfiguration.java
@@ -18,6 +18,7 @@ package org.springframework.ai.autoconfigure.minimax;
 import io.micrometer.observation.ObservationRegistry;
 import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
 import org.springframework.ai.chat.observation.ChatModelObservationConvention;
+import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
 import org.springframework.ai.minimax.MiniMaxChatModel;
 import org.springframework.ai.minimax.MiniMaxEmbeddingModel;
 import org.springframework.ai.minimax.api.MiniMaxApi;
@@ -80,13 +81,20 @@ public class MiniMaxAutoConfiguration {
 			matchIfMissing = true)
 	public MiniMaxEmbeddingModel miniMaxEmbeddingModel(MiniMaxConnectionProperties commonProperties,
 			MiniMaxEmbeddingProperties embeddingProperties, RestClient.Builder restClientBuilder,
-			RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) {
+			RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler,
+			ObjectProvider<ObservationRegistry> observationRegistry,
+			ObjectProvider<EmbeddingModelObservationConvention> observationConvention) {
 
 		var miniMaxApi = miniMaxApi(embeddingProperties.getBaseUrl(), commonProperties.getBaseUrl(),
 				embeddingProperties.getApiKey(), commonProperties.getApiKey(), restClientBuilder, responseErrorHandler);
 
-		return new MiniMaxEmbeddingModel(miniMaxApi, embeddingProperties.getMetadataMode(),
-				embeddingProperties.getOptions(), retryTemplate);
+		var embeddingModel = new MiniMaxEmbeddingModel(miniMaxApi, embeddingProperties.getMetadataMode(),
+				embeddingProperties.getOptions(), retryTemplate,
+				observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
+
+		observationConvention.ifAvailable(embeddingModel::setObservationConvention);
+
+		return embeddingModel;
 	}
 
 	private MiniMaxApi miniMaxApi(String baseUrl, String commonBaseUrl, String apiKey, String commonApiKey,