diff --git a/models/spring-ai-mistral-ai/pom.xml b/models/spring-ai-mistral-ai/pom.xml index a7a6266cf..ea6ba2097 100644 --- a/models/spring-ai-mistral-ai/pom.xml +++ b/models/spring-ai-mistral-ai/pom.xml @@ -54,6 +54,12 @@ test + + io.micrometer + micrometer-observation-test + test + + diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java index fc4d37817..afc6627fd 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java @@ -21,6 +21,9 @@ import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import io.micrometer.observation.Observation; +import io.micrometer.observation.ObservationRegistry; +import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.chat.messages.AssistantMessage; @@ -29,11 +32,13 @@ import org.springframework.ai.chat.messages.ToolResponseMessage; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.metadata.ChatGenerationMetadata; import org.springframework.ai.chat.metadata.ChatResponseMetadata; -import org.springframework.ai.chat.model.AbstractToolCallSupport; -import org.springframework.ai.chat.model.ChatModel; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.model.*; +import org.springframework.ai.chat.observation.ChatModelObservationContext; +import org.springframework.ai.chat.observation.ChatModelObservationConvention; +import org.springframework.ai.chat.observation.ChatModelObservationDocumentation; +import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention; import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.ChatOptionsBuilder; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.mistralai.api.MistralAiApi; import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletion; @@ -57,17 +62,21 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; /** + * Represents a Mistral AI Chat Model. + * * @author Ricken Bazolo * @author Christian Tzolov * @author Grogdunn * @author Thomas Vitale * @author luocongqiu - * @since 0.8.1 + * @since 1.0.0 */ public class MistralAiChatModel extends AbstractToolCallSupport implements ChatModel { private final Logger logger = LoggerFactory.getLogger(getClass()); + private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention(); + /** * The default options used for the chat completion requests. */ @@ -80,6 +89,16 @@ public class MistralAiChatModel extends AbstractToolCallSupport implements ChatM private final RetryTemplate retryTemplate; + /** + * Observation registry used for instrumentation. + */ + private final ObservationRegistry observationRegistry; + + /** + * Conventions to use for generating observations. + */ + private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION; + public MistralAiChatModel(MistralAiApi mistralAiApi) { this(mistralAiApi, MistralAiChatOptions.builder() @@ -102,118 +121,160 @@ public class MistralAiChatModel extends AbstractToolCallSupport implements ChatM public MistralAiChatModel(MistralAiApi mistralAiApi, MistralAiChatOptions options, FunctionCallbackContext functionCallbackContext, List toolFunctionCallbacks, RetryTemplate retryTemplate) { + this(mistralAiApi, options, functionCallbackContext, toolFunctionCallbacks, retryTemplate, + ObservationRegistry.NOOP); + } + + public MistralAiChatModel(MistralAiApi mistralAiApi, MistralAiChatOptions options, + FunctionCallbackContext functionCallbackContext, List toolFunctionCallbacks, + RetryTemplate retryTemplate, ObservationRegistry observationRegistry) { super(functionCallbackContext, options, toolFunctionCallbacks); - Assert.notNull(mistralAiApi, "MistralAiApi must not be null"); - Assert.notNull(options, "Options must not be null"); - Assert.notNull(retryTemplate, "RetryTemplate must not be null"); + Assert.notNull(mistralAiApi, "mistralAiApi must not be null"); + Assert.notNull(options, "options must not be null"); + Assert.notNull(retryTemplate, "retryTemplate must not be null"); + Assert.notNull(observationRegistry, "observationRegistry must not be null"); this.mistralAiApi = mistralAiApi; this.defaultOptions = options; this.retryTemplate = retryTemplate; + this.observationRegistry = observationRegistry; } @Override public ChatResponse call(Prompt prompt) { - var request = createRequest(prompt, false); + MistralAiApi.ChatCompletionRequest request = createRequest(prompt, false); - ResponseEntity completionEntity = retryTemplate - .execute(ctx -> this.mistralAiApi.chatCompletionEntity(request)); + ChatModelObservationContext observationContext = ChatModelObservationContext.builder() + .prompt(prompt) + .provider(MistralAiApi.PROVIDER_NAME) + .requestOptions(buildRequestOptions(request)) + .build(); - ChatCompletion chatCompletion = completionEntity.getBody(); + ChatResponse response = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION + .observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, + this.observationRegistry) + .observe(() -> { - if (chatCompletion == null) { - logger.warn("No chat completion returned for prompt: {}", prompt); - return new ChatResponse(List.of()); - } + ResponseEntity completionEntity = retryTemplate + .execute(ctx -> this.mistralAiApi.chatCompletionEntity(request)); - List generations = chatCompletion.choices().stream().map(choice -> { + ChatCompletion chatCompletion = completionEntity.getBody(); + + if (chatCompletion == null) { + logger.warn("No chat completion returned for prompt: {}", prompt); + return new ChatResponse(List.of()); + } + + List generations = chatCompletion.choices().stream().map(choice -> { // @formatter:off - Map metadata = Map.of( - "id", chatCompletion.id() != null ? chatCompletion.id() : "", - "index", choice.index(), - "role", choice.message().role() != null ? choice.message().role().name() : "", - "finishReason", choice.finishReason() != null ? choice.finishReason().name() : ""); - // @formatter:on - return buildGeneration(choice, metadata); - }).toList(); + Map metadata = Map.of( + "id", chatCompletion.id() != null ? chatCompletion.id() : "", + "index", choice.index(), + "role", choice.message().role() != null ? choice.message().role().name() : "", + "finishReason", choice.finishReason() != null ? choice.finishReason().name() : ""); + // @formatter:on + return buildGeneration(choice, metadata); + }).toList(); - // // Non function calling. - // RateLimit rateLimit = - // OpenAiResponseHeaderExtractor.extractAiResponseHeaders(completionEntity); + ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody())); - ChatResponse chatResponse = new ChatResponse(generations, from(completionEntity.getBody())); + observationContext.setResponse(chatResponse); - if (isToolCall(chatResponse, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(), + return chatResponse; + }); + + if (response != null && isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(), MistralAiApi.ChatCompletionFinishReason.STOP.name()))) { - var toolCallConversation = handleToolCalls(prompt, chatResponse); + var toolCallConversation = handleToolCalls(prompt, response); // Recursively call the call method with the tool call message // conversation that contains the call responses. return this.call(new Prompt(toolCallConversation, prompt.getOptions())); } - return chatResponse; + return response; } @Override public Flux stream(Prompt prompt) { - var request = createRequest(prompt, true); + return Flux.deferContextual(contextView -> { + var request = createRequest(prompt, true); - Flux completionChunks = retryTemplate - .execute(ctx -> this.mistralAiApi.chatCompletionStream(request)); + ChatModelObservationContext observationContext = ChatModelObservationContext.builder() + .prompt(prompt) + .provider(MistralAiApi.PROVIDER_NAME) + .requestOptions(buildRequestOptions(request)) + .build(); - // For chunked responses, only the first chunk contains the choice role. - // The rest of the chunks with same ID share the same role. - ConcurrentHashMap roleMap = new ConcurrentHashMap<>(); + Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation( + this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, + this.observationRegistry); - // Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse - // the function call handling logic. - Flux chatResponse = completionChunks.map(this::toChatCompletion) - .switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> { - try { - @SuppressWarnings("null") - String id = chatCompletion2.id(); + observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start(); + + Flux completionChunks = retryTemplate + .execute(ctx -> this.mistralAiApi.chatCompletionStream(request)); + + // For chunked responses, only the first chunk contains the choice role. + // The rest of the chunks with same ID share the same role. + ConcurrentHashMap roleMap = new ConcurrentHashMap<>(); + + // Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse + // the function call handling logic. + Flux chatResponse = completionChunks.map(this::toChatCompletion) + .switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> { + try { + @SuppressWarnings("null") + String id = chatCompletion2.id(); + + // @formatter:off + List generations = chatCompletion2.choices().stream().map(choice -> { + if (choice.message().role() != null) { + roleMap.putIfAbsent(id, choice.message().role().name()); + } + Map metadata = Map.of( + "id", chatCompletion2.id(), + "role", roleMap.getOrDefault(id, ""), + "index", choice.index(), + "finishReason", choice.finishReason() != null ? choice.finishReason().name() : ""); + return buildGeneration(choice, metadata); + }).toList(); + // @formatter:on + + if (chatCompletion2.usage() != null) { + return new ChatResponse(generations, from(chatCompletion2)); + } + else { + return new ChatResponse(generations); + } + } + catch (Exception e) { + logger.error("Error processing chat completion", e); + return new ChatResponse(List.of()); + } + })); // @formatter:off - List generations = chatCompletion2.choices().stream().map(choice -> { - if (choice.message().role() != null) { - roleMap.putIfAbsent(id, choice.message().role().name()); - } - Map metadata = Map.of( - "id", chatCompletion2.id(), - "role", roleMap.getOrDefault(id, ""), - "index", choice.index(), - "finishReason", choice.finishReason() != null ? choice.finishReason().name() : ""); - return buildGeneration(choice, metadata); - }).toList(); - // @formatter:on - - if (chatCompletion2.usage() != null) { - return new ChatResponse(generations, from(chatCompletion2)); - } - else { - return new ChatResponse(generations); - } + Flux chatResponseFlux = chatResponse.flatMap(response -> { + if (isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name()))) { + var toolCallConversation = handleToolCalls(prompt, response); + // Recursively call the stream method with the tool call message + // conversation that contains the call responses. + return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); } - catch (Exception e) { - logger.error("Error processing chat completion", e); - return new ChatResponse(List.of()); + else { + return Flux.just(response); } + }) + .doOnError(observation::error) + .doFinally(s -> { + observation.stop(); + }) + .contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation)); + // @formatter:on; - })); - - return chatResponse.flatMap(response -> { - - if (isToolCall(response, Set.of(MistralAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(), - MistralAiApi.ChatCompletionFinishReason.STOP.name()))) { - var toolCallConversation = handleToolCalls(prompt, response); - // Recursively call the stream method with the tool call message - // conversation that contains the call responses. - return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); - } - else { - return Flux.just(response); - } + return new MessageAggregator().aggregate(chatResponseFlux, observationContext::setResponse); }); + } private Generation buildGeneration(Choice choice, Map metadata) { @@ -333,9 +394,28 @@ public class MistralAiChatModel extends AbstractToolCallSupport implements ChatM }).toList(); } + private ChatOptions buildRequestOptions(MistralAiApi.ChatCompletionRequest request) { + return ChatOptionsBuilder.builder() + .withModel(request.model()) + .withMaxTokens(request.maxTokens()) + .withStopSequences(request.stop()) + .withTemperature(request.temperature()) + .withTopP(request.topP()) + .build(); + } + @Override public ChatOptions getDefaultOptions() { return MistralAiChatOptions.fromOptions(this.defaultOptions); } + /** + * Use the provided convention for reporting observation data + * @param observationConvention The provided convention + */ + public void setObservationConvention(ChatModelObservationConvention observationConvention) { + Assert.notNull(observationConvention, "observationConvention cannot be null"); + this.observationConvention = observationConvention; + } + } diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java index baf343222..197882130 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiEmbeddingModel.java @@ -17,17 +17,17 @@ package org.springframework.ai.mistralai; import java.util.List; +import io.micrometer.observation.ObservationRegistry; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.document.MetadataMode; -import org.springframework.ai.embedding.AbstractEmbeddingModel; -import org.springframework.ai.embedding.Embedding; -import org.springframework.ai.embedding.EmbeddingOptions; -import org.springframework.ai.embedding.EmbeddingRequest; -import org.springframework.ai.embedding.EmbeddingResponse; -import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import org.springframework.ai.embedding.*; +import org.springframework.ai.embedding.observation.DefaultEmbeddingModelObservationConvention; +import org.springframework.ai.embedding.observation.EmbeddingModelObservationContext; +import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention; +import org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation; import org.springframework.ai.mistralai.api.MistralAiApi; import org.springframework.ai.mistralai.metadata.MistralAiUsage; import org.springframework.ai.model.ModelOptionsUtils; @@ -36,13 +36,18 @@ import org.springframework.retry.support.RetryTemplate; import org.springframework.util.Assert; /** + * Provides the Mistral AI Embedding Model. + * + * @see AbstractEmbeddingModel * @author Ricken Bazolo * @author Thomas Vitale - * @since 0.8.1 + * @since 1.0.0 */ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel { - private final Logger log = LoggerFactory.getLogger(getClass()); + private static final Logger logger = LoggerFactory.getLogger(MistralAiEmbeddingModel.class); + + private static final EmbeddingModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultEmbeddingModelObservationConvention(); private final MistralAiEmbeddingOptions defaultOptions; @@ -52,6 +57,16 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel { private final RetryTemplate retryTemplate; + /** + * Observation registry used for instrumentation. + */ + private final ObservationRegistry observationRegistry; + + /** + * Conventions to use for generating observations. + */ + private EmbeddingModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION; + public MistralAiEmbeddingModel(MistralAiApi mistralAiApi) { this(mistralAiApi, MetadataMode.EMBED); } @@ -68,51 +83,72 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel { public MistralAiEmbeddingModel(MistralAiApi mistralAiApi, MetadataMode metadataMode, MistralAiEmbeddingOptions options, RetryTemplate retryTemplate) { - Assert.notNull(mistralAiApi, "MistralAiApi must not be null"); + this(mistralAiApi, metadataMode, options, retryTemplate, ObservationRegistry.NOOP); + } + + public MistralAiEmbeddingModel(MistralAiApi mistralAiApi, MetadataMode metadataMode, + MistralAiEmbeddingOptions options, RetryTemplate retryTemplate, ObservationRegistry observationRegistry) { + Assert.notNull(mistralAiApi, "mistralAiApi must not be null"); Assert.notNull(metadataMode, "metadataMode must not be null"); Assert.notNull(options, "options must not be null"); Assert.notNull(retryTemplate, "retryTemplate must not be null"); + Assert.notNull(observationRegistry, "observationRegistry must not be null"); this.mistralAiApi = mistralAiApi; this.metadataMode = metadataMode; this.defaultOptions = options; this.retryTemplate = retryTemplate; + this.observationRegistry = observationRegistry; } @Override - @SuppressWarnings("unchecked") public EmbeddingResponse call(EmbeddingRequest request) { - return this.retryTemplate.execute(ctx -> { + var apiRequest = createRequest(request); - var apiRequest = (this.defaultOptions != null) - ? new MistralAiApi.EmbeddingRequest<>(request.getInstructions(), this.defaultOptions.getModel(), - this.defaultOptions.getEncodingFormat()) - : new MistralAiApi.EmbeddingRequest<>(request.getInstructions(), - MistralAiApi.EmbeddingModel.EMBED.getValue()); + var observationContext = EmbeddingModelObservationContext.builder() + .embeddingRequest(request) + .provider(MistralAiApi.PROVIDER_NAME) + .requestOptions(buildRequestOptions(apiRequest)) + .build(); - if (request.getOptions() != null && !EmbeddingOptions.EMPTY.equals(request.getOptions())) { - apiRequest = ModelOptionsUtils.merge(request.getOptions(), apiRequest, - MistralAiApi.EmbeddingRequest.class); - } + return EmbeddingModelObservationDocumentation.EMBEDDING_MODEL_OPERATION + .observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, + this.observationRegistry) + .observe(() -> { + var apiEmbeddingResponse = this.retryTemplate + .execute(ctx -> this.mistralAiApi.embeddings(apiRequest).getBody()); - var apiEmbeddingResponse = this.mistralAiApi.embeddings(apiRequest).getBody(); + if (apiEmbeddingResponse == null) { + logger.warn("No embeddings returned for request: {}", request); + return new EmbeddingResponse(List.of()); + } - if (apiEmbeddingResponse == null) { - log.warn("No embeddings returned for request: {}", request); - return new EmbeddingResponse(List.of()); - } + var metadata = new EmbeddingResponseMetadata(apiEmbeddingResponse.model(), + MistralAiUsage.from(apiEmbeddingResponse.usage())); - var metadata = new EmbeddingResponseMetadata(apiEmbeddingResponse.model(), - MistralAiUsage.from(apiEmbeddingResponse.usage())); + var embeddings = apiEmbeddingResponse.data() + .stream() + .map(e -> new Embedding(e.embedding(), e.index())) + .toList(); - var embeddings = apiEmbeddingResponse.data() - .stream() - .map(e -> new Embedding(e.embedding(), e.index())) - .toList(); + var embeddingResponse = new EmbeddingResponse(embeddings, metadata); - return new EmbeddingResponse(embeddings, metadata); + observationContext.setResponse(embeddingResponse); - }); + return embeddingResponse; + }); + } + + @SuppressWarnings("unchecked") + private MistralAiApi.EmbeddingRequest> createRequest(EmbeddingRequest request) { + var embeddingRequest = new MistralAiApi.EmbeddingRequest<>(request.getInstructions(), + this.defaultOptions.getModel(), this.defaultOptions.getEncodingFormat()); + + if (request.getOptions() != null) { + embeddingRequest = ModelOptionsUtils.merge(request.getOptions(), embeddingRequest, + MistralAiApi.EmbeddingRequest.class); + } + return embeddingRequest; } @Override @@ -121,4 +157,17 @@ public class MistralAiEmbeddingModel extends AbstractEmbeddingModel { return this.embed(document.getFormattedContent(this.metadataMode)); } + private EmbeddingOptions buildRequestOptions(MistralAiApi.EmbeddingRequest> request) { + return EmbeddingOptionsBuilder.builder().withModel(request.model()).build(); + } + + /** + * Use the provided convention for reporting observation data + * @param observationConvention The provided convention + */ + public void setObservationConvention(EmbeddingModelObservationConvention observationConvention) { + Assert.notNull(observationConvention, "observationConvention cannot be null"); + this.observationConvention = observationConvention; + } + } diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java index cd31077d8..3b034da97 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java @@ -15,8 +15,10 @@ */ package org.springframework.ai.mistralai.api; +import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import java.util.function.Predicate; @@ -24,6 +26,7 @@ import java.util.function.Predicate; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonInclude.Include; import com.fasterxml.jackson.annotation.JsonProperty; +import org.springframework.ai.observation.conventions.AiProvider; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -55,12 +58,14 @@ import org.springframework.web.reactive.function.client.WebClient; * @author Ricken Bazolo * @author Christian Tzolov * @author Thomas Vitale - * @since 0.8.1 + * @since 1.0.0 */ public class MistralAiApi { private static final String DEFAULT_BASE_URL = "https://api.mistral.ai"; + public static final String PROVIDER_NAME = AiProvider.MISTRAL_AI.value(); + private static final Predicate SSE_DONE_PREDICATE = "[DONE]"::equals; private final RestClient restClient; @@ -210,6 +215,29 @@ public class MistralAiApi { public Embedding(Integer index, float[] embedding) { this(index, embedding, "embedding"); } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (!(o instanceof Embedding embedding1)) + return false; + return Objects.equals(index, embedding1.index) && Arrays.equals(embedding, embedding1.embedding) + && Objects.equals(object, embedding1.object); + } + + @Override + public int hashCode() { + int result = Objects.hash(index, object); + result = 31 * result + Arrays.hashCode(embedding); + return result; + } + + @Override + public String toString() { + return "Embedding{" + "index=" + index + ", embedding=" + Arrays.toString(embedding) + ", object='" + object + + '\'' + '}'; + } } /** @@ -355,6 +383,7 @@ public class MistralAiApi { @JsonProperty("max_tokens") Integer maxTokens, @JsonProperty("stream") Boolean stream, @JsonProperty("safe_prompt") Boolean safePrompt, + @JsonProperty("stop") List stop, @JsonProperty("random_seed") Integer randomSeed, @JsonProperty("response_format") ResponseFormat responseFormat) { // @formatter:on @@ -367,7 +396,7 @@ public class MistralAiApi { * @param model ID of the model to use. */ public ChatCompletionRequest(List messages, String model) { - this(model, messages, null, null, 0.7f, 1f, null, false, false, null, null); + this(model, messages, null, null, 0.7f, 1f, null, false, false, null, null, null); } /** @@ -382,7 +411,7 @@ public class MistralAiApi { */ public ChatCompletionRequest(List messages, String model, Float temperature, boolean stream) { - this(model, messages, null, null, temperature, 1f, null, stream, false, null, null); + this(model, messages, null, null, temperature, 1f, null, stream, false, null, null, null); } /** @@ -395,7 +424,7 @@ public class MistralAiApi { * */ public ChatCompletionRequest(List messages, String model, Float temperature) { - this(model, messages, null, null, temperature, 1f, null, false, false, null, null); + this(model, messages, null, null, temperature, 1f, null, false, false, null, null, null); } /** @@ -410,7 +439,7 @@ public class MistralAiApi { */ public ChatCompletionRequest(List messages, String model, List tools, ToolChoice toolChoice) { - this(model, messages, tools, toolChoice, null, 1f, null, false, false, null, null); + this(model, messages, tools, toolChoice, null, 1f, null, false, false, null, null, null); } /** @@ -418,7 +447,7 @@ public class MistralAiApi { * stream. */ public ChatCompletionRequest(List messages, Boolean stream) { - this(null, messages, null, null, 0.7f, 1f, null, stream, false, null, null); + this(null, messages, null, null, 0.7f, 1f, null, stream, false, null, null, null); } /** diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatCompletionRequestTest.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatCompletionRequestTest.java index b83fdcd42..893f4d43a 100644 --- a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatCompletionRequestTest.java +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatCompletionRequestTest.java @@ -28,7 +28,7 @@ import static org.assertj.core.api.Assertions.assertThat; * @author Ricken Bazolo * @since 0.8.1 */ -@SpringBootTest +@SpringBootTest(classes = MistralAiTestConfiguration.class) @EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+") public class MistralAiChatCompletionRequestTest { diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatModelObservationIT.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatModelObservationIT.java new file mode 100644 index 000000000..b6a8ea41e --- /dev/null +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatModelObservationIT.java @@ -0,0 +1,176 @@ +/* + * Copyright 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.mistralai; + +import io.micrometer.common.KeyValue; +import io.micrometer.observation.tck.TestObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistryAssert; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.mistralai.api.MistralAiApi; +import org.springframework.ai.model.function.FunctionCallbackContext; +import org.springframework.ai.observation.conventions.AiOperationType; +import org.springframework.ai.observation.conventions.AiProvider; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.context.annotation.Bean; +import org.springframework.retry.support.RetryTemplate; +import org.springframework.util.StringUtils; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.HighCardinalityKeyNames; +import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.LowCardinalityKeyNames; + +/** + * Integration tests for observation instrumentation in {@link MistralAiChatModel}. + * + * @author Thomas Vitale + */ +@SpringBootTest(classes = MistralAiChatModelObservationIT.Config.class) +@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+") +public class MistralAiChatModelObservationIT { + + @Autowired + TestObservationRegistry observationRegistry; + + @Autowired + MistralAiChatModel chatModel; + + @BeforeEach + void beforeEach() { + observationRegistry.clear(); + } + + @Test + void observationForChatOperation() { + var options = MistralAiChatOptions.builder() + .withModel(MistralAiApi.ChatModel.OPEN_MISTRAL_7B.getValue()) + .withMaxTokens(2048) + .withStop(List.of("this-is-the-end")) + .withTemperature(0.7f) + .withTopP(1f) + .build(); + + Prompt prompt = new Prompt("Why does a raven look like a desk?", options); + + ChatResponse chatResponse = chatModel.call(prompt); + assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty(); + + ChatResponseMetadata responseMetadata = chatResponse.getMetadata(); + assertThat(responseMetadata).isNotNull(); + + validate(responseMetadata); + } + + @Test + void observationForStreamingChatOperation() { + var options = MistralAiChatOptions.builder() + .withModel(MistralAiApi.ChatModel.OPEN_MISTRAL_7B.getValue()) + .withMaxTokens(2048) + .withStop(List.of("this-is-the-end")) + .withTemperature(0.7f) + .withTopP(1f) + .build(); + + Prompt prompt = new Prompt("Why does a raven look like a desk?", options); + + Flux chatResponseFlux = chatModel.stream(prompt); + + List responses = chatResponseFlux.collectList().block(); + assertThat(responses).isNotEmpty(); + assertThat(responses).hasSizeGreaterThan(10); + + String aggregatedResponse = responses.subList(0, responses.size() - 1) + .stream() + .map(r -> r.getResult().getOutput().getContent()) + .collect(Collectors.joining()); + assertThat(aggregatedResponse).isNotEmpty(); + + ChatResponse lastChatResponse = responses.get(responses.size() - 1); + + ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata(); + assertThat(responseMetadata).isNotNull(); + + validate(responseMetadata); + } + + private void validate(ChatResponseMetadata responseMetadata) { + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("chat " + MistralAiApi.ChatModel.OPEN_MISTRAL_7B.getValue()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(), + AiOperationType.CHAT.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.MISTRAL_AI.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(), + MistralAiApi.ChatModel.OPEN_MISTRAL_7B.getValue()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_FREQUENCY_PENALTY.asString(), + KeyValue.NONE_VALUE) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_PRESENCE_PENALTY.asString(), + KeyValue.NONE_VALUE) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(), + "[\"this-is-the-end\"]") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(), "0.7") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_K.asString(), KeyValue.NONE_VALUE) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_ID.asString(), responseMetadata.getId()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(), "[\"STOP\"]") + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getPromptTokens())) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getGenerationTokens())) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getTotalTokens())) + .hasBeenStarted() + .hasBeenStopped(); + } + + @SpringBootConfiguration + static class Config { + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + @Bean + public MistralAiApi mistralAiApi() { + return new MistralAiApi(System.getenv("MISTRAL_AI_API_KEY")); + } + + @Bean + public MistralAiChatModel openAiChatModel(MistralAiApi mistralAiApi, + TestObservationRegistry observationRegistry) { + return new MistralAiChatModel(mistralAiApi, MistralAiChatOptions.builder().build(), + new FunctionCallbackContext(), List.of(), RetryTemplate.defaultInstance(), observationRegistry); + } + + } + +} diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java index 99d2aa1a6..de378c57a 100644 --- a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java @@ -26,7 +26,7 @@ import org.springframework.boot.test.context.SpringBootTest; import static org.assertj.core.api.Assertions.assertThat; -@SpringBootTest +@SpringBootTest(classes = MistralAiTestConfiguration.class) @EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+") class MistralAiEmbeddingIT { diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingModelObservationIT.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingModelObservationIT.java new file mode 100644 index 000000000..a49370ce2 --- /dev/null +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingModelObservationIT.java @@ -0,0 +1,116 @@ +/* + * Copyright 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.mistralai; + +import io.micrometer.common.KeyValue; +import io.micrometer.observation.tck.TestObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistryAssert; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.embedding.EmbeddingRequest; +import org.springframework.ai.embedding.EmbeddingResponse; +import org.springframework.ai.embedding.EmbeddingResponseMetadata; +import org.springframework.ai.embedding.observation.DefaultEmbeddingModelObservationConvention; +import org.springframework.ai.mistralai.api.MistralAiApi; +import org.springframework.ai.observation.conventions.AiOperationType; +import org.springframework.ai.observation.conventions.AiProvider; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.context.annotation.Bean; +import org.springframework.retry.support.RetryTemplate; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation.HighCardinalityKeyNames; +import static org.springframework.ai.embedding.observation.EmbeddingModelObservationDocumentation.LowCardinalityKeyNames; + +/** + * Integration tests for observation instrumentation in {@link MistralAiEmbeddingModel}. + * + * @author Thomas Vitale + */ +@SpringBootTest(classes = MistralAiEmbeddingModelObservationIT.Config.class) +@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+") +public class MistralAiEmbeddingModelObservationIT { + + @Autowired + TestObservationRegistry observationRegistry; + + @Autowired + MistralAiEmbeddingModel embeddingModel; + + @Test + void observationForEmbeddingOperation() { + var options = MistralAiEmbeddingOptions.builder() + .withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()) + .withEncodingFormat("float") + .build(); + + EmbeddingRequest embeddingRequest = new EmbeddingRequest(List.of("Here comes the sun"), options); + + EmbeddingResponse embeddingResponse = embeddingModel.call(embeddingRequest); + assertThat(embeddingResponse.getResults()).isNotEmpty(); + + EmbeddingResponseMetadata responseMetadata = embeddingResponse.getMetadata(); + assertThat(responseMetadata).isNotNull(); + + TestObservationRegistryAssert.assertThat(observationRegistry) + .doesNotHaveAnyRemainingCurrentObservation() + .hasObservationWithNameEqualTo(DefaultEmbeddingModelObservationConvention.DEFAULT_NAME) + .that() + .hasContextualNameEqualTo("embedding " + MistralAiApi.EmbeddingModel.EMBED.getValue()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(), + AiOperationType.EMBEDDING.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.MISTRAL_AI.value()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(), + MistralAiApi.EmbeddingModel.EMBED.getValue()) + .hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel()) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_EMBEDDING_DIMENSIONS.asString(), + KeyValue.NONE_VALUE) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getPromptTokens())) + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(), + String.valueOf(responseMetadata.getUsage().getTotalTokens())) + .hasBeenStarted() + .hasBeenStopped(); + } + + @SpringBootConfiguration + static class Config { + + @Bean + public TestObservationRegistry observationRegistry() { + return TestObservationRegistry.create(); + } + + @Bean + public MistralAiApi mistralAiApi() { + return new MistralAiApi(System.getenv("MISTRAL_AI_API_KEY")); + } + + @Bean + public MistralAiEmbeddingModel openAiEmbeddingModel(MistralAiApi mistralAiApi, + TestObservationRegistry observationRegistry) { + return new MistralAiEmbeddingModel(mistralAiApi, MetadataMode.EMBED, + MistralAiEmbeddingOptions.builder().build(), RetryTemplate.defaultInstance(), observationRegistry); + } + + } + +} diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java index d0523397e..e2f765ec2 100644 --- a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiRetryTests.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Optional; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -135,6 +136,7 @@ public class MistralAiRetryTests { } @Test + @Disabled("Currently stream() does not implement retry") public void mistralAiChatStreamTransientError() { var choice = new ChatCompletionChunk.ChunkChoice(0, new ChatCompletionMessage("Response", Role.ASSISTANT), @@ -156,6 +158,7 @@ public class MistralAiRetryTests { } @Test + @Disabled("Currently stream() does not implement retry") public void mistralAiChatStreamNonTransientError() { when(mistralAiApi.chatCompletionStream(isA(ChatCompletionRequest.class))) .thenThrow(new RuntimeException("Non Transient Error")); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/mistralai/MistralAiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/mistralai/MistralAiAutoConfiguration.java index 0d5735f88..6e40c5eaf 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/mistralai/MistralAiAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/mistralai/MistralAiAutoConfiguration.java @@ -17,12 +17,16 @@ package org.springframework.ai.autoconfigure.mistralai; import java.util.List; +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.mistralai.MistralAiChatModel; import org.springframework.ai.mistralai.MistralAiEmbeddingModel; import org.springframework.ai.mistralai.api.MistralAiApi; import org.springframework.ai.model.function.FunctionCallback; import org.springframework.ai.model.function.FunctionCallbackContext; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.ImportAutoConfiguration; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; @@ -42,6 +46,7 @@ import org.springframework.web.client.RestClient; /** * @author Ricken Bazolo * @author Christian Tzolov + * @author Thomas Vitale * @since 0.8.1 */ @AutoConfiguration(after = { RestClientAutoConfiguration.class, SpringAiRetryAutoConfiguration.class }) @@ -58,14 +63,21 @@ public class MistralAiAutoConfiguration { matchIfMissing = true) public MistralAiEmbeddingModel mistralAiEmbeddingModel(MistralAiCommonProperties commonProperties, MistralAiEmbeddingProperties embeddingProperties, RestClient.Builder restClientBuilder, - RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) { + RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler, + ObjectProvider observationRegistry, + ObjectProvider observationConvention) { var mistralAiApi = mistralAiApi(embeddingProperties.getApiKey(), commonProperties.getApiKey(), embeddingProperties.getBaseUrl(), commonProperties.getBaseUrl(), restClientBuilder, responseErrorHandler); - return new MistralAiEmbeddingModel(mistralAiApi, embeddingProperties.getMetadataMode(), - embeddingProperties.getOptions(), retryTemplate); + var embeddingModel = new MistralAiEmbeddingModel(mistralAiApi, embeddingProperties.getMetadataMode(), + embeddingProperties.getOptions(), retryTemplate, + observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)); + + observationConvention.ifAvailable(embeddingModel::setObservationConvention); + + return embeddingModel; } @Bean @@ -75,13 +87,19 @@ public class MistralAiAutoConfiguration { public MistralAiChatModel mistralAiChatModel(MistralAiCommonProperties commonProperties, MistralAiChatProperties chatProperties, RestClient.Builder restClientBuilder, List toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext, - RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) { + RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler, + ObjectProvider observationRegistry, + ObjectProvider observationConvention) { var mistralAiApi = mistralAiApi(chatProperties.getApiKey(), commonProperties.getApiKey(), chatProperties.getBaseUrl(), commonProperties.getBaseUrl(), restClientBuilder, responseErrorHandler); - return new MistralAiChatModel(mistralAiApi, chatProperties.getOptions(), functionCallbackContext, - toolFunctionCallbacks, retryTemplate); + var chatModel = new MistralAiChatModel(mistralAiApi, chatProperties.getOptions(), functionCallbackContext, + toolFunctionCallbacks, retryTemplate, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)); + + observationConvention.ifAvailable(chatModel::setObservationConvention); + + return chatModel; } private MistralAiApi mistralAiApi(String apiKey, String commonApiKey, String baseUrl, String commonBaseUrl,