Add integration for Mistral AI Chat and Embedding models
- Implement MistrealAiApi as a REST client for the Mistral REST API. - Add MistralAiChatClient implementing ChatClient and StreamingChatClient. Add MistralAiChatOptions implementing ChatOptions. - Add MistralAiEmbeddingClient impelementing EmbeddingClient. Add MistralAiEmbeddingOptions implementing EmbeddingOptions. - Add Unit and IT tests for api and clients. - Add mistral auto-configuration and boot strarter - update auto-config property for Mistral AI embedding client - update auto-config property for Mistral AI embedding client Additional, review changes: - Add missing license headers. - Fix the intialization of defaul options. - Code re-formatting to improve the readabitliy. - Many improvements - Add missing sine annotations - Remove stream field from MistralAiChatOptions. This handled internally. - Add MistralAiApi ChatModel and EmbeddingModel enums. - Add MistralChatClientIT for testeing chat, stream , parsers... - Add MistralAiApiIT tests - Refactor and streamline the MistralAiAutoConfiguration. - Add MistralAiAutoConfigurationIT
This commit is contained in:
ricken07
committed by
Christian Tzolov
Christian Tzolov
parent
0d0b858d60
commit
30c3530561
70
models/spring-ai-mistral-ai/pom.xml
Normal file
70
models/spring-ai-mistral-ai/pom.xml
Normal file
@@ -0,0 +1,70 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai</artifactId>
|
||||
<version>0.8.1-SNAPSHOT</version>
|
||||
<relativePath>../../pom.xml</relativePath>
|
||||
</parent>
|
||||
<artifactId>spring-ai-mistral-ai</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>Spring AI Mistral AI</name>
|
||||
<description>Mistral AI support</description>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
|
||||
<scm>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
<connection>git://github.com/spring-projects/spring-ai.git</connection>
|
||||
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
|
||||
</scm>
|
||||
|
||||
<dependencies>
|
||||
|
||||
<!-- production dependencies -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-core</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework</groupId>
|
||||
<artifactId>spring-web</artifactId>
|
||||
<version>${spring-framework.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.retry</groupId>
|
||||
<artifactId>spring-retry</artifactId>
|
||||
<version>2.0.4</version>
|
||||
</dependency>
|
||||
|
||||
|
||||
<!-- Spring Framework -->
|
||||
<dependency>
|
||||
<groupId>org.springframework</groupId>
|
||||
<artifactId>spring-webflux</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework</groupId>
|
||||
<artifactId>spring-context-support</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-logging</artifactId>
|
||||
</dependency>
|
||||
|
||||
<!-- test dependencies -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-test</artifactId>
|
||||
<version>${project.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
@@ -0,0 +1,174 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import org.springframework.ai.chat.ChatClient;
|
||||
import org.springframework.ai.chat.ChatResponse;
|
||||
import org.springframework.ai.chat.Generation;
|
||||
import org.springframework.ai.chat.StreamingChatClient;
|
||||
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
|
||||
import org.springframework.ai.chat.prompt.ChatOptions;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi;
|
||||
import org.springframework.ai.model.ModelOptionsUtils;
|
||||
import org.springframework.retry.RetryCallback;
|
||||
import org.springframework.retry.RetryContext;
|
||||
import org.springframework.retry.RetryListener;
|
||||
import org.springframework.retry.support.RetryTemplate;
|
||||
import org.springframework.util.Assert;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @since 0.8.1
|
||||
*/
|
||||
public class MistralAiChatClient implements ChatClient, StreamingChatClient {
|
||||
|
||||
private final Logger log = LoggerFactory.getLogger(getClass());
|
||||
|
||||
/**
|
||||
* The default options used for the chat completion requests.
|
||||
*/
|
||||
private MistralAiChatOptions defaultOptions;
|
||||
|
||||
/**
|
||||
* Low-level access to the OpenAI API.
|
||||
*/
|
||||
private final MistralAiApi mistralAiApi;
|
||||
|
||||
private final RetryTemplate retryTemplate = RetryTemplate.builder()
|
||||
.maxAttempts(10)
|
||||
.retryOn(MistralAiApi.MistralAiApiException.class)
|
||||
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
|
||||
.withListener(new RetryListener() {
|
||||
public <T extends Object, E extends Throwable> void onError(RetryContext context,
|
||||
RetryCallback<T, E> callback, Throwable throwable) {
|
||||
log.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
|
||||
};
|
||||
})
|
||||
.build();
|
||||
|
||||
public MistralAiChatClient(MistralAiApi mistralAiApi, MistralAiChatOptions options) {
|
||||
Assert.notNull(mistralAiApi, "MistralAiApi must not be null");
|
||||
Assert.notNull(options, "Options must not be null");
|
||||
this.mistralAiApi = mistralAiApi;
|
||||
this.defaultOptions = options;
|
||||
}
|
||||
|
||||
public MistralAiChatClient(MistralAiApi mistralAiApi) {
|
||||
this(mistralAiApi,
|
||||
MistralAiChatOptions.builder()
|
||||
.withTemperature(0.7f)
|
||||
.withTopP(1f)
|
||||
.withSafePrompt(false)
|
||||
.withModel(MistralAiApi.ChatModel.TINY.getValue())
|
||||
.build());
|
||||
}
|
||||
|
||||
/**
|
||||
* Accessible for testing.
|
||||
*/
|
||||
public MistralAiApi.ChatCompletionRequest createRequest(Prompt prompt, boolean stream) {
|
||||
var chatCompletionMessages = prompt.getInstructions()
|
||||
.stream()
|
||||
.map(m -> new MistralAiApi.ChatCompletionMessage(m.getContent(),
|
||||
MistralAiApi.ChatCompletionMessage.Role.valueOf(m.getMessageType().name())))
|
||||
.toList();
|
||||
|
||||
var request = new MistralAiApi.ChatCompletionRequest(chatCompletionMessages, stream);
|
||||
|
||||
if (this.defaultOptions != null) {
|
||||
request = ModelOptionsUtils.merge(request, this.defaultOptions, MistralAiApi.ChatCompletionRequest.class);
|
||||
}
|
||||
|
||||
if (prompt.getOptions() != null) {
|
||||
if (prompt.getOptions() instanceof ChatOptions runtimeOptions) {
|
||||
var updatedRuntimeOptions = ModelOptionsUtils.copyToTarget(runtimeOptions, ChatOptions.class,
|
||||
MistralAiChatOptions.class);
|
||||
request = ModelOptionsUtils.merge(updatedRuntimeOptions, request,
|
||||
MistralAiApi.ChatCompletionRequest.class);
|
||||
}
|
||||
else {
|
||||
throw new IllegalArgumentException("Prompt options are not of type ChatOptions: "
|
||||
+ prompt.getOptions().getClass().getSimpleName());
|
||||
}
|
||||
}
|
||||
|
||||
return request;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatResponse call(Prompt prompt) {
|
||||
return retryTemplate.execute(ctx -> {
|
||||
var request = createRequest(prompt, false);
|
||||
|
||||
var completionEntity = this.mistralAiApi.chatCompletionEntity(request);
|
||||
|
||||
var chatCompletion = completionEntity.getBody();
|
||||
if (chatCompletion == null) {
|
||||
log.warn("No chat completion returned for prompt: {}", prompt);
|
||||
return new ChatResponse(List.of());
|
||||
}
|
||||
|
||||
List<Generation> generations = chatCompletion.choices()
|
||||
.stream()
|
||||
.map(choice -> new Generation(choice.message().content(),
|
||||
Map.of("role", choice.message().role().name()))
|
||||
.withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null)))
|
||||
.toList();
|
||||
|
||||
return new ChatResponse(generations);
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public Flux<ChatResponse> stream(Prompt prompt) {
|
||||
return retryTemplate.execute(ctx -> {
|
||||
var request = createRequest(prompt, true);
|
||||
|
||||
var completionChunks = 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<String, String> roleMap = new ConcurrentHashMap<>();
|
||||
|
||||
return completionChunks.map(chunk -> {
|
||||
String chunkId = chunk.id();
|
||||
List<Generation> generations = chunk.choices().stream().map(choice -> {
|
||||
if (choice.delta().role() != null) {
|
||||
roleMap.putIfAbsent(chunkId, choice.delta().role().name());
|
||||
}
|
||||
var generation = new Generation(choice.delta().content(), Map.of("role", roleMap.get(chunkId)));
|
||||
if (choice.finishReason() != null) {
|
||||
generation = generation
|
||||
.withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null));
|
||||
}
|
||||
return generation;
|
||||
}).toList();
|
||||
return new ChatResponse(generations);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonIgnore;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import org.springframework.ai.chat.prompt.ChatOptions;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@JsonInclude(JsonInclude.Include.NON_NULL)
|
||||
public class MistralAiChatOptions implements ChatOptions {
|
||||
|
||||
/**
|
||||
* ID of the model to use
|
||||
*/
|
||||
private @JsonProperty("model") String model;
|
||||
|
||||
/**
|
||||
* What sampling temperature to use, between 0.0 and 1.0. Higher values like 0.8 will
|
||||
* make the output more random, while lower values like 0.2 will make it more focused
|
||||
* and deterministic. We generally recommend altering this or top_p but not both.
|
||||
*/
|
||||
private @JsonProperty("temperature") Float temperature;
|
||||
|
||||
/**
|
||||
* Nucleus sampling, where the model considers the results of the tokens with top_p
|
||||
* probability mass. So 0.1 means only the tokens comprising the top 10% probability
|
||||
* mass are considered. We generally recommend altering this or temperature but not
|
||||
* both.
|
||||
*/
|
||||
private @JsonProperty("top_p") Float topP;
|
||||
|
||||
/**
|
||||
* The maximum number of tokens to generate in the completion. The token count of your
|
||||
* prompt plus max_tokens cannot exceed the model's context length.
|
||||
*/
|
||||
private @JsonProperty("max_tokens") Integer maxTokens;
|
||||
|
||||
/**
|
||||
* Whether to inject a safety prompt before all conversations.
|
||||
*/
|
||||
private @JsonProperty("safe_prompt") Boolean safePrompt;
|
||||
|
||||
/**
|
||||
* The seed to use for random sampling. If set, different calls will generate
|
||||
* deterministic results.
|
||||
*/
|
||||
private @JsonProperty("random_seed") Integer randomSeed;
|
||||
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
public static class Builder {
|
||||
|
||||
private final MistralAiChatOptions options = new MistralAiChatOptions();
|
||||
|
||||
public Builder withModel(String model) {
|
||||
this.options.setModel(model);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withMaxToken(Integer maxTokens) {
|
||||
this.options.setMaxTokens(maxTokens);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withSafePrompt(Boolean safePrompt) {
|
||||
this.options.setSafePrompt(safePrompt);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withRandomSeed(Integer randomSeed) {
|
||||
this.options.setRandomSeed(randomSeed);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withTemperature(Float temperature) {
|
||||
this.options.setTemperature(temperature);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withTopP(Float topP) {
|
||||
this.options.setTopP(topP);
|
||||
return this;
|
||||
}
|
||||
|
||||
public MistralAiChatOptions build() {
|
||||
return this.options;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public String getModel() {
|
||||
return this.model;
|
||||
}
|
||||
|
||||
public void setModel(String model) {
|
||||
this.model = model;
|
||||
}
|
||||
|
||||
public Integer getMaxTokens() {
|
||||
return this.maxTokens;
|
||||
}
|
||||
|
||||
public void setMaxTokens(Integer maxTokens) {
|
||||
this.maxTokens = maxTokens;
|
||||
}
|
||||
|
||||
public Boolean getSafePrompt() {
|
||||
return this.safePrompt;
|
||||
}
|
||||
|
||||
public void setSafePrompt(Boolean safePrompt) {
|
||||
this.safePrompt = safePrompt;
|
||||
}
|
||||
|
||||
public Integer getRandomSeed() {
|
||||
return this.randomSeed;
|
||||
}
|
||||
|
||||
public void setRandomSeed(Integer randomSeed) {
|
||||
this.randomSeed = randomSeed;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Float getTemperature() {
|
||||
return this.temperature;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void setTemperature(Float temperature) {
|
||||
this.temperature = temperature;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Float getTopP() {
|
||||
return this.topP;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void setTopP(Float topP) {
|
||||
this.topP = topP;
|
||||
}
|
||||
|
||||
@Override
|
||||
@JsonIgnore
|
||||
public Integer getTopK() {
|
||||
throw new UnsupportedOperationException("Unsupported option: 'TopK'");
|
||||
}
|
||||
|
||||
@Override
|
||||
@JsonIgnore
|
||||
public void setTopK(Integer topK) {
|
||||
throw new UnsupportedOperationException("Unsupported option: 'TopK'");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral;
|
||||
|
||||
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.*;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.MistralAiApiException;
|
||||
import org.springframework.ai.model.ModelOptionsUtils;
|
||||
import org.springframework.retry.RetryCallback;
|
||||
import org.springframework.retry.RetryContext;
|
||||
import org.springframework.retry.RetryListener;
|
||||
import org.springframework.retry.support.RetryTemplate;
|
||||
import org.springframework.util.Assert;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @since 0.8.1
|
||||
*/
|
||||
public class MistralAiEmbeddingClient extends AbstractEmbeddingClient {
|
||||
|
||||
private final Logger log = LoggerFactory.getLogger(getClass());
|
||||
|
||||
private final MistralAiEmbeddingOptions defaultOptions;
|
||||
|
||||
private final MetadataMode metadataMode;
|
||||
|
||||
private final MistralAiApi mistralAiApi;
|
||||
|
||||
private final RetryTemplate retryTemplate = RetryTemplate.builder()
|
||||
.maxAttempts(10)
|
||||
.retryOn(MistralAiApiException.class)
|
||||
.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
|
||||
.withListener(new RetryListener() {
|
||||
public <T extends Object, E extends Throwable> void onError(RetryContext context,
|
||||
RetryCallback<T, E> callback, Throwable throwable) {
|
||||
log.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
|
||||
};
|
||||
})
|
||||
.build();
|
||||
|
||||
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi) {
|
||||
this(mistralAiApi, MetadataMode.EMBED);
|
||||
}
|
||||
|
||||
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi, MetadataMode metadataMode) {
|
||||
this(mistralAiApi, metadataMode,
|
||||
MistralAiEmbeddingOptions.builder().withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()).build());
|
||||
}
|
||||
|
||||
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi, MistralAiEmbeddingOptions options) {
|
||||
this(mistralAiApi, MetadataMode.EMBED, options);
|
||||
}
|
||||
|
||||
public MistralAiEmbeddingClient(MistralAiApi mistralAiApi, MetadataMode metadataMode,
|
||||
MistralAiEmbeddingOptions options) {
|
||||
Assert.notNull(mistralAiApi, "MistralAiApi must not be null");
|
||||
Assert.notNull(metadataMode, "metadataMode must not be null");
|
||||
Assert.notNull(options, "options must not be null");
|
||||
|
||||
this.mistralAiApi = mistralAiApi;
|
||||
this.metadataMode = metadataMode;
|
||||
this.defaultOptions = options;
|
||||
}
|
||||
|
||||
@Override
|
||||
@SuppressWarnings("unchecked")
|
||||
public EmbeddingResponse call(EmbeddingRequest request) {
|
||||
return this.retryTemplate.execute(ctx -> {
|
||||
|
||||
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());
|
||||
|
||||
if (request.getOptions() != null && !EmbeddingOptions.EMPTY.equals(request.getOptions())) {
|
||||
apiRequest = ModelOptionsUtils.merge(request.getOptions(), apiRequest,
|
||||
MistralAiApi.EmbeddingRequest.class);
|
||||
}
|
||||
|
||||
var apiEmbeddingResponse = this.mistralAiApi.embeddings(apiRequest).getBody();
|
||||
|
||||
if (apiEmbeddingResponse == null) {
|
||||
log.warn("No embeddings returned for request: {}", request);
|
||||
return new EmbeddingResponse(List.of());
|
||||
}
|
||||
|
||||
var metadata = generateResponseMetadata(apiEmbeddingResponse.model(), apiEmbeddingResponse.usage());
|
||||
|
||||
var embeddings = apiEmbeddingResponse.data()
|
||||
.stream()
|
||||
.map(e -> new Embedding(e.embedding(), e.index()))
|
||||
.toList();
|
||||
|
||||
return new EmbeddingResponse(embeddings, metadata);
|
||||
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Double> embed(Document document) {
|
||||
Assert.notNull(document, "Document must not be null");
|
||||
return this.embed(document.getFormattedContent(this.metadataMode));
|
||||
}
|
||||
|
||||
private EmbeddingResponseMetadata generateResponseMetadata(String model, MistralAiApi.Usage usage) {
|
||||
var metadata = new EmbeddingResponseMetadata();
|
||||
metadata.put("model", model);
|
||||
metadata.put("prompt-tokens", usage.promptTokens());
|
||||
metadata.put("total-tokens", usage.totalTokens());
|
||||
return metadata;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
import com.fasterxml.jackson.annotation.JsonInclude.Include;
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import org.springframework.ai.embedding.EmbeddingOptions;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public class MistralAiEmbeddingOptions implements EmbeddingOptions {
|
||||
|
||||
/**
|
||||
* ID of the model to use.
|
||||
*/
|
||||
private @JsonProperty("model") String model;
|
||||
|
||||
/**
|
||||
* The format to return the embeddings in. Can be either float or base64.
|
||||
*/
|
||||
private @JsonProperty("encoding_format") String encodingFormat;
|
||||
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
public String getModel() {
|
||||
return this.model;
|
||||
}
|
||||
|
||||
public void setModel(String model) {
|
||||
this.model = model;
|
||||
}
|
||||
|
||||
public String getEncodingFormat() {
|
||||
return this.encodingFormat;
|
||||
}
|
||||
|
||||
public void setEncodingFormat(String encodingFormat) {
|
||||
this.encodingFormat = encodingFormat;
|
||||
}
|
||||
|
||||
public static class Builder {
|
||||
|
||||
protected MistralAiEmbeddingOptions options;
|
||||
|
||||
public Builder() {
|
||||
this.options = new MistralAiEmbeddingOptions();
|
||||
}
|
||||
|
||||
public Builder withModel(String model) {
|
||||
this.options.setModel(model);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withEncodingFormat(String encodingFormat) {
|
||||
this.options.setEncodingFormat(encodingFormat);
|
||||
return this;
|
||||
}
|
||||
|
||||
public MistralAiEmbeddingOptions build() {
|
||||
return this.options;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,646 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral.api;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.function.Consumer;
|
||||
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 com.fasterxml.jackson.core.type.TypeReference;
|
||||
import com.fasterxml.jackson.databind.DeserializationFeature;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
import org.springframework.core.ParameterizedTypeReference;
|
||||
import org.springframework.http.HttpHeaders;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.http.client.ClientHttpResponse;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.CollectionUtils;
|
||||
import org.springframework.web.client.ResponseErrorHandler;
|
||||
import org.springframework.web.client.RestClient;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Implementation of the MistralAI Embedding API:
|
||||
* <a href="https://docs.mistral.ai/api/#operation/createEmbedding">...</a> and Chat
|
||||
* Completion API:
|
||||
* <a href="https://docs.mistral.ai/api/#operation/createChatCompletion">...</a>
|
||||
*
|
||||
* @author Ricken Bazolo
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
public class MistralAiApi {
|
||||
|
||||
private static final String DEFAULT_BASE_URL = "https://api.mistral.ai";
|
||||
|
||||
private static final Predicate<String> SSE_DONE_PREDICATE = "[DONE]"::equals;
|
||||
|
||||
private final RestClient restClient;
|
||||
|
||||
private WebClient webClient;
|
||||
|
||||
private final ObjectMapper objectMapper;
|
||||
|
||||
/**
|
||||
* Create a new client api with DEFAULT_BASE_URL
|
||||
* @param mistralAiApiKey Mistral api Key.
|
||||
*/
|
||||
public MistralAiApi(String mistralAiApiKey) {
|
||||
this(DEFAULT_BASE_URL, mistralAiApiKey);
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a new client api.
|
||||
* @param baseUrl api base URL.
|
||||
* @param mistralAiApiKey Mistral api Key.
|
||||
*/
|
||||
public MistralAiApi(String baseUrl, String mistralAiApiKey) {
|
||||
this(baseUrl, mistralAiApiKey, RestClient.builder());
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a new client api.
|
||||
* @param baseUrl api base URL.
|
||||
* @param mistralAiApiKey Mistral api Key.
|
||||
* @param restClientBuilder RestClient builder.
|
||||
*/
|
||||
public MistralAiApi(String baseUrl, String mistralAiApiKey, RestClient.Builder restClientBuilder) {
|
||||
|
||||
this.objectMapper = new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
|
||||
|
||||
Consumer<HttpHeaders> jsonContentHeaders = headers -> {
|
||||
headers.setBearerAuth(mistralAiApiKey);
|
||||
headers.setContentType(MediaType.APPLICATION_JSON);
|
||||
};
|
||||
|
||||
var responseErrorHandler = new ResponseErrorHandler() {
|
||||
|
||||
@Override
|
||||
public boolean hasError(ClientHttpResponse response) throws IOException {
|
||||
return response.getStatusCode().isError();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void handleError(ClientHttpResponse response) throws IOException {
|
||||
if (response.getStatusCode().isError()) {
|
||||
if (response.getStatusCode().is4xxClientError()) {
|
||||
throw new MistralAiApiClientErrorException(String.format("%s - %s",
|
||||
response.getStatusCode().value(),
|
||||
MistralAiApi.this.objectMapper.readValue(response.getBody(), ResponseError.class)));
|
||||
}
|
||||
throw new MistralAiApiException(String.format("%s - %s", response.getStatusCode().value(),
|
||||
MistralAiApi.this.objectMapper.readValue(response.getBody(), ResponseError.class)));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
this.restClient = restClientBuilder.baseUrl(baseUrl)
|
||||
.defaultHeaders(jsonContentHeaders)
|
||||
.defaultStatusHandler(responseErrorHandler)
|
||||
.build();
|
||||
this.webClient = WebClient.builder().baseUrl(baseUrl).defaultHeaders(jsonContentHeaders).build();
|
||||
}
|
||||
|
||||
public static class MistralAiApiException extends RuntimeException {
|
||||
|
||||
public MistralAiApiException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
public MistralAiApiException(String message, Throwable t) {
|
||||
super(message, t);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Thrown on 4xx client errors, such as 401 - Incorrect API key provided, 401 - You
|
||||
* must be a member of an organization to use the API, 429 - Rate limit reached for
|
||||
* requests, 429 - You exceeded your current quota , please check your plan and
|
||||
* billing details.
|
||||
*/
|
||||
public static class MistralAiApiClientErrorException extends RuntimeException {
|
||||
|
||||
public MistralAiApiClientErrorException(String message) {
|
||||
super(message);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* API error response.
|
||||
*
|
||||
* @param error Error details.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record ResponseError(@JsonProperty("error") Error error) {
|
||||
|
||||
/**
|
||||
* Error details.
|
||||
*
|
||||
* @param message Error message.
|
||||
* @param type Error type.
|
||||
* @param param Error parameter.
|
||||
* @param code Error code.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record Error(
|
||||
// @formatter:off
|
||||
@JsonProperty("message") String message,
|
||||
@JsonProperty("type") String type,
|
||||
@JsonProperty("param") String param,
|
||||
@JsonProperty("code") String code) {
|
||||
// @formatter:on
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Usage statistics.
|
||||
*
|
||||
* @param promptTokens Number of tokens in the prompt.
|
||||
* @param totalTokens Total number of tokens used in the request (prompt +
|
||||
* completion).
|
||||
* @param completionTokens Number of tokens in the generated completion. Only
|
||||
* applicable for completion requests.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record Usage(
|
||||
// @formatter:off
|
||||
@JsonProperty("prompt_tokens") Integer promptTokens,
|
||||
@JsonProperty("total_tokens") Integer totalTokens,
|
||||
@JsonProperty("completion_tokens") Integer completionTokens) {
|
||||
// @formatter:on
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents an embedding vector returned by embedding endpoint.
|
||||
*
|
||||
* @param index The index of the embedding in the list of embeddings.
|
||||
* @param embedding The embedding vector, which is a list of floats. The length of
|
||||
* vector depends on the model.
|
||||
* @param object The object type, which is always 'embedding'.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record Embedding(
|
||||
// @formatter:off
|
||||
@JsonProperty("index") Integer index,
|
||||
@JsonProperty("embedding") List<Double> embedding,
|
||||
@JsonProperty("object") String object) {
|
||||
// @formatter:on
|
||||
|
||||
/**
|
||||
* Create an embedding with the given index, embedding and object type set to
|
||||
* 'embedding'.
|
||||
* @param index The index of the embedding in the list of embeddings.
|
||||
* @param embedding The embedding vector, which is a list of floats. The length of
|
||||
* vector depends on the model.
|
||||
*/
|
||||
public Embedding(Integer index, List<Double> embedding) {
|
||||
this(index, embedding, "embedding");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates an embedding vector representing the input text.
|
||||
*
|
||||
* @param input Input text to embed, encoded as a string or array of tokens
|
||||
* @param model ID of the model to use.
|
||||
* @param encodingFormat The format to return the embeddings in. Can be either float
|
||||
* or base64.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record EmbeddingRequest<T>(
|
||||
// @formatter:off
|
||||
@JsonProperty("input") T input,
|
||||
@JsonProperty("model") String model,
|
||||
@JsonProperty("encoding_format") String encodingFormat) {
|
||||
// @formatter:on
|
||||
|
||||
/**
|
||||
* Create an embedding request with the given input, model and encoding format set
|
||||
* to float.
|
||||
* @param input Input text to embed.
|
||||
* @param model ID of the model to use.
|
||||
*/
|
||||
public EmbeddingRequest(T input, String model) {
|
||||
this(input, model, "float");
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an embedding request with the given input. Encoding format is set to
|
||||
* float and user is null and the model is set to 'mistral-embed'.
|
||||
* @param input Input text to embed.
|
||||
*/
|
||||
public EmbeddingRequest(T input) {
|
||||
this(input, EmbeddingModel.EMBED.getValue());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* List of multiple embedding responses.
|
||||
*
|
||||
* @param <T> Type of the entities in the data list.
|
||||
* @param object Must have value "list".
|
||||
* @param data List of entities.
|
||||
* @param model ID of the model to use.
|
||||
* @param usage Usage statistics for the completion request.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record EmbeddingList<T>(
|
||||
// @formatter:off
|
||||
@JsonProperty("object") String object,
|
||||
@JsonProperty("data") List<T> data,
|
||||
@JsonProperty("model") String model,
|
||||
@JsonProperty("usage") Usage usage) {
|
||||
// @formatter:on
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates an embedding vector representing the input text or token array.
|
||||
* @param embeddingRequest The embedding request.
|
||||
* @return Returns list of {@link Embedding} wrapped in {@link EmbeddingList}.
|
||||
* @param <T> Type of the entity in the data list. Can be a {@link String} or
|
||||
* {@link List} of tokens (e.g. Integers). For embedding multiple inputs in a single
|
||||
* request, You can pass a {@link List} of {@link String} or {@link List} of
|
||||
* {@link List} of tokens. For example:
|
||||
*
|
||||
* <pre>{@code List.of("text1", "text2", "text3") or List.of(List.of(1, 2, 3), List.of(3, 4, 5))} </pre>
|
||||
*/
|
||||
public <T> ResponseEntity<EmbeddingList<Embedding>> embeddings(EmbeddingRequest<T> embeddingRequest) {
|
||||
|
||||
Assert.notNull(embeddingRequest, "The request body can not be null.");
|
||||
|
||||
// Input text to embed, encoded as a string or array of tokens. To embed multiple
|
||||
// inputs in a single
|
||||
// request, pass an array of strings or array of token arrays.
|
||||
Assert.notNull(embeddingRequest.input(), "The input can not be null.");
|
||||
Assert.isTrue(embeddingRequest.input() instanceof String || embeddingRequest.input() instanceof List,
|
||||
"The input must be either a String, or a List of Strings or List of List of integers.");
|
||||
|
||||
// The input must not an empty string, and any array must be 1024 dimensions or
|
||||
// less.
|
||||
if (embeddingRequest.input() instanceof List list) {
|
||||
Assert.isTrue(!CollectionUtils.isEmpty(list), "The input list can not be empty.");
|
||||
Assert.isTrue(list.size() <= 1024, "The list must be 1024 dimensions or less");
|
||||
Assert.isTrue(
|
||||
list.get(0) instanceof String || list.get(0) instanceof Integer || list.get(0) instanceof List,
|
||||
"The input must be either a String, or a List of Strings or list of list of integers.");
|
||||
}
|
||||
|
||||
return this.restClient.post()
|
||||
.uri("/v1/embeddings")
|
||||
.body(embeddingRequest)
|
||||
.retrieve()
|
||||
.toEntity(new ParameterizedTypeReference<>() {
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a model request for chat conversation.
|
||||
*
|
||||
* @param model ID of the model to use.
|
||||
* @param messages The prompt(s) to generate completions for, encoded as a list of
|
||||
* dict with role and content. The first prompt role should be user or system.
|
||||
* @param temperature What sampling temperature to use, between 0.0 and 1.0. Higher
|
||||
* values like 0.8 will make the output more random, while lower values like 0.2 will
|
||||
* make it more focused and deterministic. We generally recommend altering this or
|
||||
* top_p but not both.
|
||||
* @param topP Nucleus sampling, where the model considers the results of the tokens
|
||||
* with top_p probability mass. So 0.1 means only the tokens comprising the top 10%
|
||||
* probability mass are considered. We generally recommend altering this or
|
||||
* temperature but not both.
|
||||
* @param maxTokens The maximum number of tokens to generate in the completion. The
|
||||
* token count of your prompt plus max_tokens cannot exceed the model's context
|
||||
* length.
|
||||
* @param stream Whether to stream back partial progress. If set, tokens will be sent
|
||||
* as data-only server-sent events as they become available, with the stream
|
||||
* terminated by a data: [DONE] message. Otherwise, the server will hold the request
|
||||
* open until the timeout or until completion, with the response containing the full
|
||||
* result as JSON.
|
||||
* @param safePrompt Whether to inject a safety prompt before all conversations.
|
||||
* @param randomSeed The seed to use for random sampling. If set, different calls will
|
||||
* generate deterministic results.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record ChatCompletionRequest(
|
||||
// @formatter:off
|
||||
@JsonProperty("model") String model,
|
||||
@JsonProperty("messages") List<ChatCompletionMessage> messages,
|
||||
@JsonProperty("temperature") Float temperature,
|
||||
@JsonProperty("top_p") Float topP,
|
||||
@JsonProperty("max_tokens") Integer maxTokens,
|
||||
@JsonProperty("stream") Boolean stream,
|
||||
@JsonProperty("safe_prompt") Boolean safePrompt,
|
||||
@JsonProperty("random_seed") Integer randomSeed) {
|
||||
// @formatter:on
|
||||
|
||||
/**
|
||||
* Shortcut constructor for a chat completion request with the given messages and
|
||||
* model.
|
||||
* @param messages The prompt(s) to generate completions for, encoded as a list of
|
||||
* dict with role and content. The first prompt role should be user or system.
|
||||
* @param model ID of the model to use.
|
||||
*/
|
||||
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model) {
|
||||
this(model, messages, 0.7f, 1f, null, false, false, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Shortcut constructor for a chat completion request with the given messages,
|
||||
* model and temperature.
|
||||
* @param messages The prompt(s) to generate completions for, encoded as a list of
|
||||
* dict with role and content. The first prompt role should be user or system.
|
||||
* @param model ID of the model to use.
|
||||
* @param temperature What sampling temperature to use, between 0.0 and 1.0.
|
||||
* @param stream Whether to stream back partial progress. If set, tokens will be
|
||||
* sent
|
||||
*/
|
||||
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, Float temperature,
|
||||
boolean stream) {
|
||||
this(model, messages, temperature, 1f, null, stream, false, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Shortcut constructor for a chat completion request with the given messages,
|
||||
* model and temperature.
|
||||
* @param messages The prompt(s) to generate completions for, encoded as a list of
|
||||
* dict with role and content. The first prompt role should be user or system.
|
||||
* @param model ID of the model to use.
|
||||
* @param temperature What sampling temperature to use, between 0.0 and 1.0.
|
||||
*
|
||||
*/
|
||||
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, Float temperature) {
|
||||
this(model, messages, temperature, 1f, null, false, false, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Shortcut constructor for a chat completion request with the given messages and
|
||||
* stream.
|
||||
*/
|
||||
public ChatCompletionRequest(List<ChatCompletionMessage> messages, Boolean stream) {
|
||||
this(null, messages, 0.7f, 1f, null, stream, false, null);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Message comprising the conversation.
|
||||
*
|
||||
* @param content The contents of the message.
|
||||
* @param role The role of the messages author. Could be one of the {@link Role}
|
||||
* types.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record ChatCompletionMessage(
|
||||
// @formatter:off
|
||||
@JsonProperty("content") String content,
|
||||
@JsonProperty("role") Role role) {
|
||||
// @formatter:on
|
||||
|
||||
/**
|
||||
* The role of the author of this message.
|
||||
*
|
||||
* NOTE: Mistral expects the system message to be before the user message or will
|
||||
* fail with 400 error.
|
||||
*/
|
||||
public enum Role {
|
||||
|
||||
// @formatter:off
|
||||
@JsonProperty("system") SYSTEM,
|
||||
@JsonProperty("user") USER,
|
||||
@JsonProperty("assistant") ASSISTANT
|
||||
// @formatter:on
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The reason the model stopped generating tokens.
|
||||
*/
|
||||
public enum ChatCompletionFinishReason {
|
||||
|
||||
// @formatter:off
|
||||
/**
|
||||
* The model hit a natural stop point or a provided stop sequence.
|
||||
*/
|
||||
@JsonProperty("stop") STOP,
|
||||
/**
|
||||
* The maximum number of tokens specified in the request was reached.
|
||||
*/
|
||||
@JsonProperty("length") LENGTH,
|
||||
/**
|
||||
* The content was omitted due to a flag from our content filters.
|
||||
*/
|
||||
@JsonProperty("model_length") MODEL_LENGTH
|
||||
// @formatter:on
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents a chat completion response returned by model, based on the provided
|
||||
* input.
|
||||
*
|
||||
* @param id A unique identifier for the chat completion.
|
||||
* @param object The object type, which is always chat.completion.
|
||||
* @param created The Unix timestamp (in seconds) of when the chat completion was
|
||||
* created.
|
||||
* @param model The model used for the chat completion.
|
||||
* @param choices A list of chat completion choices.
|
||||
* @param usage Usage statistics for the completion request.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record ChatCompletion(
|
||||
// @formatter:off
|
||||
@JsonProperty("id") String id,
|
||||
@JsonProperty("object") String object,
|
||||
@JsonProperty("created") Long created,
|
||||
@JsonProperty("model") String model,
|
||||
@JsonProperty("choices") List<Choice> choices,
|
||||
@JsonProperty("usage") Usage usage) {
|
||||
// @formatter:on
|
||||
|
||||
/**
|
||||
* Chat completion choice.
|
||||
*
|
||||
* @param index The index of the choice in the list of choices.
|
||||
* @param message A chat completion message generated by the model.
|
||||
* @param finishReason The reason the model stopped generating tokens.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record Choice(
|
||||
// @formatter:off
|
||||
@JsonProperty("index") Integer index,
|
||||
@JsonProperty("message") ChatCompletionMessage message,
|
||||
@JsonProperty("finish_reason") ChatCompletionFinishReason finishReason) {
|
||||
// @formatter:on
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents a streamed chunk of a chat completion response returned by model, based
|
||||
* on the provided input.
|
||||
*
|
||||
* @param id A unique identifier for the chat completion. Each chunk has the same ID.
|
||||
* @param object The object type, which is always 'chat.completion.chunk'.
|
||||
* @param created The Unix timestamp (in seconds) of when the chat completion was
|
||||
* created. Each chunk has the same timestamp.
|
||||
* @param model The model used for the chat completion.
|
||||
* @param choices A list of chat completion choices. Can be more than one if n is
|
||||
* greater than 1.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record ChatCompletionChunk(
|
||||
// @formatter:off
|
||||
@JsonProperty("id") String id,
|
||||
@JsonProperty("object") String object,
|
||||
@JsonProperty("created") Long created,
|
||||
@JsonProperty("model") String model,
|
||||
@JsonProperty("choices") List<ChunkChoice> choices) {
|
||||
// @formatter:on
|
||||
|
||||
/**
|
||||
* Chat completion choice.
|
||||
*
|
||||
* @param index The index of the choice in the list of choices.
|
||||
* @param delta A chat completion delta generated by streamed model responses.
|
||||
* @param finishReason The reason the model stopped generating tokens.
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public record ChunkChoice(
|
||||
// @formatter:off
|
||||
@JsonProperty("index") Integer index,
|
||||
@JsonProperty("delta") ChatCompletionMessage delta,
|
||||
@JsonProperty("finish_reason") ChatCompletionFinishReason finishReason) {
|
||||
// @formatter:on
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* List of well-known Mistral chat models.
|
||||
* https://docs.mistral.ai/platform/endpoints/#mistral-ai-generative-models
|
||||
*/
|
||||
public enum ChatModel {
|
||||
|
||||
// @formatter:off
|
||||
@JsonProperty("mistral-tiny") TINY("mistral-tiny"),
|
||||
@JsonProperty("mistral-small") SMALL("mistral-small"),
|
||||
@JsonProperty("mistral-medium") MEDIUM("mistral-medium"),
|
||||
@JsonProperty("mistral-large") LARGE("mistral-large"),
|
||||
@JsonProperty("mistral-xlarge") XLARGE("mistral-xlarge");
|
||||
// @formatter:on
|
||||
|
||||
private final String value;
|
||||
|
||||
ChatModel(String value) {
|
||||
this.value = value;
|
||||
}
|
||||
|
||||
public String getValue() {
|
||||
return this.value;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* List of well-known Mistral embedding models.
|
||||
* https://docs.mistral.ai/platform/endpoints/#mistral-ai-embedding-model
|
||||
*/
|
||||
public enum EmbeddingModel {
|
||||
|
||||
// @formatter:off
|
||||
@JsonProperty("mistral-embed") EMBED("mistral-embed");
|
||||
// @formatter:on
|
||||
|
||||
private final String value;
|
||||
|
||||
EmbeddingModel(String value) {
|
||||
this.value = value;
|
||||
}
|
||||
|
||||
public String getValue() {
|
||||
return this.value;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a model response for the given chat conversation.
|
||||
* @param chatRequest The chat completion request.
|
||||
* @return Entity response with {@link ChatCompletion} as a body and HTTP status code
|
||||
* and headers.
|
||||
*/
|
||||
public ResponseEntity<ChatCompletion> chatCompletionEntity(ChatCompletionRequest chatRequest) {
|
||||
|
||||
Assert.notNull(chatRequest, "The request body can not be null.");
|
||||
Assert.isTrue(!chatRequest.stream(), "Request must set the steam property to false.");
|
||||
|
||||
return this.restClient.post()
|
||||
.uri("/v1/chat/completions")
|
||||
.body(chatRequest)
|
||||
.retrieve()
|
||||
.toEntity(ChatCompletion.class);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a streaming chat response for the given chat conversation.
|
||||
* @param chatRequest The chat completion request. Must have the stream property set
|
||||
* to true.
|
||||
* @return Returns a {@link Flux} stream from chat completion chunks.
|
||||
*/
|
||||
public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chatRequest) {
|
||||
|
||||
Assert.notNull(chatRequest, "The request body can not be null.");
|
||||
Assert.isTrue(chatRequest.stream(), "Request must set the steam property to true.");
|
||||
|
||||
return this.webClient.post()
|
||||
.uri("/v1/chat/completions")
|
||||
.body(Mono.just(chatRequest), ChatCompletionRequest.class)
|
||||
.retrieve()
|
||||
.bodyToFlux(String.class)
|
||||
.takeUntil(SSE_DONE_PREDICATE)
|
||||
.filter(SSE_DONE_PREDICATE.negate())
|
||||
.map(content -> parseJson(content, ChatCompletionChunk.class));
|
||||
}
|
||||
|
||||
public static Map<String, Object> parseJson(String jsonSchema) {
|
||||
try {
|
||||
return new ObjectMapper().readValue(jsonSchema, new TypeReference<Map<String, Object>>() {
|
||||
});
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new MistralAiApiException("Failed to parse schema: " + jsonSchema, e);
|
||||
}
|
||||
}
|
||||
|
||||
private <T> T parseJson(String json, Class<T> type) {
|
||||
try {
|
||||
return this.objectMapper.readValue(json, type);
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new MistralAiApiException("Failed to parse schema: " + json, e);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral;
|
||||
|
||||
import org.springframework.ai.embedding.EmbeddingClient;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi;
|
||||
import org.springframework.boot.SpringBootConfiguration;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
@SpringBootConfiguration
|
||||
public class MistralAiTestConfiguration {
|
||||
|
||||
@Bean
|
||||
public MistralAiApi mistralAiApi() {
|
||||
var apiKey = System.getenv("MISTRAL_AI_API_KEY");
|
||||
if (!StringUtils.hasText(apiKey)) {
|
||||
throw new IllegalArgumentException(
|
||||
"Missing MISTRAL_AI_API_KEY environment variable. Please set it to your Mistral AI API key.");
|
||||
}
|
||||
return new MistralAiApi(apiKey);
|
||||
}
|
||||
|
||||
@Bean
|
||||
public EmbeddingClient mistralAiEmbeddingClient(MistralAiApi api) {
|
||||
return new MistralAiEmbeddingClient(api,
|
||||
MistralAiEmbeddingOptions.builder().withModel(MistralAiApi.EmbeddingModel.EMBED.getValue()).build());
|
||||
}
|
||||
|
||||
@Bean
|
||||
public MistralAiChatClient mistralAiChatClient(MistralAiApi mistralAiApi) {
|
||||
return new MistralAiChatClient(mistralAiApi,
|
||||
MistralAiChatOptions.builder().withModel(MistralAiApi.ChatModel.SMALL.getValue()).build());
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral.chat;
|
||||
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import org.springframework.ai.chat.ChatClient;
|
||||
import org.springframework.ai.chat.ChatResponse;
|
||||
import org.springframework.ai.chat.Generation;
|
||||
import org.springframework.ai.chat.StreamingChatClient;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.Message;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
|
||||
import org.springframework.ai.mistral.MistralAiTestConfiguration;
|
||||
import org.springframework.ai.parser.BeanOutputParser;
|
||||
import org.springframework.ai.parser.ListOutputParser;
|
||||
import org.springframework.ai.parser.MapOutputParser;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.core.convert.support.DefaultConversionService;
|
||||
import org.springframework.core.io.Resource;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@SpringBootTest(classes = MistralAiTestConfiguration.class)
|
||||
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
|
||||
class MistralAiChatClientIT {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(MistralAiChatClientIT.class);
|
||||
|
||||
@Autowired
|
||||
protected ChatClient chatClient;
|
||||
|
||||
@Autowired
|
||||
protected StreamingChatClient streamingChatClient;
|
||||
|
||||
@Value("classpath:/prompts/system-message.st")
|
||||
private Resource systemResource;
|
||||
|
||||
@Value("classpath:/prompts/eval/qa-evaluator-accurate-answer.st")
|
||||
protected Resource qaEvaluatorAccurateAnswerResource;
|
||||
|
||||
@Value("classpath:/prompts/eval/qa-evaluator-not-related-message.st")
|
||||
protected Resource qaEvaluatorNotRelatedResource;
|
||||
|
||||
@Value("classpath:/prompts/eval/qa-evaluator-fact-based-answer.st")
|
||||
protected Resource qaEvalutaorFactBasedAnswerResource;
|
||||
|
||||
@Value("classpath:/prompts/eval/user-evaluator-message.st")
|
||||
protected Resource userEvaluatorResource;
|
||||
|
||||
@Test
|
||||
void roleTest() {
|
||||
UserMessage userMessage = new UserMessage(
|
||||
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.");
|
||||
SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
|
||||
Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate"));
|
||||
// NOTE: Mistral expects the system message to be before the user message or will
|
||||
// fail with 400 error.
|
||||
Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
|
||||
ChatResponse response = chatClient.call(prompt);
|
||||
assertThat(response.getResults()).hasSize(1);
|
||||
assertThat(response.getResults().get(0).getOutput().getContent()).contains("Blackbeard");
|
||||
}
|
||||
|
||||
@Test
|
||||
void outputParser() {
|
||||
DefaultConversionService conversionService = new DefaultConversionService();
|
||||
ListOutputParser outputParser = new ListOutputParser(conversionService);
|
||||
|
||||
String format = outputParser.getFormat();
|
||||
String template = """
|
||||
List five {subject}
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template,
|
||||
Map.of("subject", "ice cream flavors", "format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
Generation generation = this.chatClient.call(prompt).getResult();
|
||||
|
||||
List<String> list = outputParser.parse(generation.getOutput().getContent());
|
||||
assertThat(list).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void mapOutputParser() {
|
||||
MapOutputParser outputParser = new MapOutputParser();
|
||||
|
||||
String format = outputParser.getFormat();
|
||||
String template = """
|
||||
Provide me a List of {subject}
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template,
|
||||
Map.of("subject", "an array of numbers from 1 to 9 under they key name 'numbers'", "format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
Generation generation = chatClient.call(prompt).getResult();
|
||||
|
||||
Map<String, Object> result = outputParser.parse(generation.getOutput().getContent());
|
||||
assertThat(result.get("numbers")).isEqualTo(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
|
||||
|
||||
}
|
||||
|
||||
record ActorsFilmsRecord(String actor, List<String> movies) {
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanOutputParserRecords() {
|
||||
|
||||
BeanOutputParser<ActorsFilmsRecord> outputParser = new BeanOutputParser<>(ActorsFilmsRecord.class);
|
||||
|
||||
String format = outputParser.getFormat();
|
||||
String template = """
|
||||
Generate the filmography of 5 movies for Tom Hanks.
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
Generation generation = chatClient.call(prompt).getResult();
|
||||
|
||||
ActorsFilmsRecord actorsFilms = outputParser.parse(generation.getOutput().getContent());
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
|
||||
assertThat(actorsFilms.movies()).hasSize(5);
|
||||
}
|
||||
|
||||
@Test
|
||||
void beanStreamOutputParserRecords() {
|
||||
|
||||
BeanOutputParser<ActorsFilmsRecord> outputParser = new BeanOutputParser<>(ActorsFilmsRecord.class);
|
||||
|
||||
String format = outputParser.getFormat();
|
||||
String template = """
|
||||
Generate the filmography of 5 movies for Tom Hanks.
|
||||
{format}
|
||||
""";
|
||||
PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format));
|
||||
Prompt prompt = new Prompt(promptTemplate.createMessage());
|
||||
|
||||
String generationTextFromStream = streamingChatClient.stream(prompt)
|
||||
.collectList()
|
||||
.block()
|
||||
.stream()
|
||||
.map(ChatResponse::getResults)
|
||||
.flatMap(List::stream)
|
||||
.map(Generation::getOutput)
|
||||
.map(AssistantMessage::getContent)
|
||||
.collect(Collectors.joining());
|
||||
|
||||
ActorsFilmsRecord actorsFilms = outputParser.parse(generationTextFromStream);
|
||||
logger.info("" + actorsFilms);
|
||||
assertThat(actorsFilms.actor()).isEqualTo("Tom Hanks");
|
||||
assertThat(actorsFilms.movies()).hasSize(5);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral.chat;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.mistral.MistralAiChatClient;
|
||||
import org.springframework.ai.mistral.MistralAiChatOptions;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@SpringBootTest
|
||||
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
|
||||
public class MistralChatCompletionRequestTest {
|
||||
|
||||
MistralAiChatClient chatClient = new MistralAiChatClient(new MistralAiApi("test"));
|
||||
|
||||
@Test
|
||||
void chatCompletionDefaultRequestTest() {
|
||||
|
||||
var request = chatClient.createRequest(new Prompt("test content"), false);
|
||||
|
||||
assertThat(request.messages()).hasSize(1);
|
||||
assertThat(request.topP()).isEqualTo(1);
|
||||
assertThat(request.temperature()).isEqualTo(0.7f);
|
||||
assertThat(request.safePrompt()).isFalse();
|
||||
assertThat(request.maxTokens()).isNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void chatCompletionRequestWithOptionsTest() {
|
||||
|
||||
var options = MistralAiChatOptions.builder().withTemperature(0.5f).withTopP(0.8f).build();
|
||||
|
||||
var request = chatClient.createRequest(new Prompt("test content", options), false);
|
||||
|
||||
assertThat(request.messages().size()).isEqualTo(1);
|
||||
assertThat(request.topP()).isEqualTo(0.8f);
|
||||
assertThat(request.temperature()).isEqualTo(0.5f);
|
||||
assertThat(request.stream()).isTrue();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral.chat.api;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import org.springframework.ai.mistral.api.MistralAiApi;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.ChatCompletionMessage.Role;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.ChatCompletionRequest;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.Embedding;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.EmbeddingList;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.ChatCompletion;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.ChatCompletionChunk;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi.ChatCompletionMessage;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
|
||||
public class MistralAiApiIT {
|
||||
|
||||
MistralAiApi mistralAiApi = new MistralAiApi(System.getenv("MISTRAL_AI_API_KEY"));
|
||||
|
||||
@Test
|
||||
void chatCompletionEntity() {
|
||||
ChatCompletionMessage chatCompletionMessage = new ChatCompletionMessage("Hello world", Role.USER);
|
||||
ResponseEntity<ChatCompletion> response = mistralAiApi.chatCompletionEntity(new ChatCompletionRequest(
|
||||
List.of(chatCompletionMessage), MistralAiApi.ChatModel.TINY.getValue(), 0.8f, false));
|
||||
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.getBody()).isNotNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void chatCompletionEntityWithSystemMessage() {
|
||||
ChatCompletionMessage userMessage = new ChatCompletionMessage(
|
||||
"Tell me about 3 famous pirates from the Golden Age of Piracy and why they did?", Role.USER);
|
||||
ChatCompletionMessage systemMessage = new ChatCompletionMessage("""
|
||||
You are an AI assistant that helps people find information.
|
||||
Your name is Bob.
|
||||
You should reply to the user's request with your name and also in the style of a pirate.
|
||||
""", Role.SYSTEM);
|
||||
|
||||
ResponseEntity<ChatCompletion> response = mistralAiApi.chatCompletionEntity(new ChatCompletionRequest(
|
||||
List.of(systemMessage, userMessage), MistralAiApi.ChatModel.TINY.getValue(), 0.8f, false));
|
||||
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.getBody()).isNotNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void chatCompletionStream() {
|
||||
ChatCompletionMessage chatCompletionMessage = new ChatCompletionMessage("Hello world", Role.USER);
|
||||
Flux<ChatCompletionChunk> response = mistralAiApi.chatCompletionStream(new ChatCompletionRequest(
|
||||
List.of(chatCompletionMessage), MistralAiApi.ChatModel.TINY.getValue(), 0.8f, true));
|
||||
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.collectList().block()).isNotNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void embeddings() {
|
||||
ResponseEntity<EmbeddingList<Embedding>> response = mistralAiApi
|
||||
.embeddings(new MistralAiApi.EmbeddingRequest<String>("Hello world"));
|
||||
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.getBody().data()).hasSize(1);
|
||||
assertThat(response.getBody().data().get(0).embedding()).hasSize(1024);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
/*
|
||||
* Copyright 2024-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.mistral.embedding;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
|
||||
import org.springframework.ai.embedding.EmbeddingRequest;
|
||||
import org.springframework.ai.mistral.MistralAiEmbeddingClient;
|
||||
import org.springframework.ai.mistral.MistralAiEmbeddingOptions;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
@SpringBootTest
|
||||
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
|
||||
class MistralEmbeddingIT {
|
||||
|
||||
@Autowired
|
||||
private MistralAiEmbeddingClient mistralAiEmbeddingClient;
|
||||
|
||||
@Test
|
||||
void defaultEmbedding() {
|
||||
assertThat(mistralAiEmbeddingClient).isNotNull();
|
||||
var embeddingResponse = mistralAiEmbeddingClient.embedForResponse(List.of("Hello World"));
|
||||
assertThat(embeddingResponse.getResults()).hasSize(1);
|
||||
assertThat(embeddingResponse.getResults().get(0)).isNotNull();
|
||||
assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024);
|
||||
assertThat(embeddingResponse.getMetadata()).containsEntry("model", "mistral-embed");
|
||||
assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 4);
|
||||
assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 4);
|
||||
assertThat(mistralAiEmbeddingClient.dimensions()).isEqualTo(1024);
|
||||
}
|
||||
|
||||
@Test
|
||||
void embeddingTest() {
|
||||
assertThat(mistralAiEmbeddingClient).isNotNull();
|
||||
var embeddingResponse = mistralAiEmbeddingClient.call(new EmbeddingRequest(
|
||||
List.of("Hello World", "World is big"),
|
||||
MistralAiEmbeddingOptions.builder().withModel("mistral-embed").withEncodingFormat("float").build()));
|
||||
assertThat(embeddingResponse.getResults()).hasSize(2);
|
||||
assertThat(embeddingResponse.getResults().get(0)).isNotNull();
|
||||
assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024);
|
||||
assertThat(embeddingResponse.getMetadata()).containsEntry("model", "mistral-embed");
|
||||
assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 9);
|
||||
assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 9);
|
||||
assertThat(mistralAiEmbeddingClient.dimensions()).isEqualTo(1024);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
You're assisting with questions about products in a bicycle catalog.
|
||||
Use the information from the DOCUMENTS section to provide accurate answers.
|
||||
The the answer involves referring to the price or the dimension of the bicycle, include the bicycle name in the response.
|
||||
If unsure, simply state that you don't know.
|
||||
|
||||
DOCUMENTS:
|
||||
{documents}
|
||||
@@ -0,0 +1,3 @@
|
||||
You are an AI assistant who helps users to evaluate if the answers to questions are accurate.
|
||||
You will be provided with a QUESTION and an ANSWER.
|
||||
Your goal is to evaluate the QUESTION and ANSWER and reply with a YES or NO answer.
|
||||
@@ -0,0 +1,7 @@
|
||||
You are an AI evaluator. Your task is to verify if the provided ANSWER is a direct and accurate response to the given QUESTION. If the ANSWER is correct and directly answers the QUESTION, reply with "YES". If the ANSWER is not a direct response or is inaccurate, reply with "NO".
|
||||
|
||||
For example:
|
||||
|
||||
If the QUESTION is "What is the capital of France?" and the ANSWER is "Paris.", you should respond with "YES".
|
||||
If the QUESTION is "What is the capital of France?" and the ANSWER is "France is in Europe.", respond with "NO".
|
||||
Now, evaluate the following:
|
||||
@@ -0,0 +1,4 @@
|
||||
You are an AI assistant who helps users to evaluate if the answers to questions are accurate.
|
||||
You will be provided with a QUESTION and an ANSWER.
|
||||
A previous evaluation has determined that QUESTION and ANSWER are not related.
|
||||
Give an explanation as to why they are not related.
|
||||
@@ -0,0 +1,6 @@
|
||||
The question and answer to evaluate are:
|
||||
|
||||
QUESTION: ```{question}```
|
||||
|
||||
ANSWER: ```{answer}```
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
You are an AI assistant that helps people find information.
|
||||
Your name is {name}.
|
||||
You should reply to the user's request with your name and also in the style of a {voice}.
|
||||
2
pom.xml
2
pom.xml
@@ -24,6 +24,7 @@
|
||||
<module>models/spring-ai-openai</module>
|
||||
<module>models/spring-ai-vertex-ai</module>
|
||||
<module>models/spring-ai-stabilityai</module>
|
||||
<module>models/spring-ai-mistral-ai</module>
|
||||
<module>spring-ai-test</module>
|
||||
<module>spring-ai-spring-boot-autoconfigure</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-openai</module>
|
||||
@@ -53,6 +54,7 @@
|
||||
<module>vector-stores/spring-ai-redis</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-vertex-ai</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai</module>
|
||||
<module>spring-ai-spring-boot-starters/spring-ai-starter-mistral-ai</module>
|
||||
</modules>
|
||||
|
||||
<organization>
|
||||
|
||||
@@ -101,6 +101,12 @@
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-mistral-ai</artifactId>
|
||||
<version>${project.version}</version>
|
||||
</dependency>
|
||||
|
||||
<!-- Vector Databses -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
|
||||
@@ -179,6 +179,14 @@
|
||||
<optional>true</optional>
|
||||
</dependency>
|
||||
|
||||
<!-- Mistral AI LLM -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-mistral-ai</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
<optional>true</optional>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-configuration-processor</artifactId>
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
/*
|
||||
* Copyright 2024-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.autoconfigure.mistralai;
|
||||
|
||||
import org.springframework.ai.embedding.EmbeddingClient;
|
||||
import org.springframework.ai.mistral.MistralAiChatClient;
|
||||
import org.springframework.ai.mistral.MistralAiEmbeddingClient;
|
||||
import org.springframework.ai.mistral.api.MistralAiApi;
|
||||
import org.springframework.boot.autoconfigure.AutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
|
||||
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
|
||||
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
|
||||
import org.springframework.boot.context.properties.EnableConfigurationProperties;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.StringUtils;
|
||||
import org.springframework.web.client.RestClient;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@AutoConfiguration(after = { RestClientAutoConfiguration.class })
|
||||
@EnableConfigurationProperties({ MistralAiEmbeddingProperties.class, MistralAiCommonProperties.class,
|
||||
MistralAiChatProperties.class })
|
||||
@ConditionalOnClass(MistralAiApi.class)
|
||||
public class MistralAiAutoConfiguration {
|
||||
|
||||
@Bean
|
||||
@ConditionalOnMissingBean
|
||||
@ConditionalOnProperty(prefix = MistralAiEmbeddingProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true",
|
||||
matchIfMissing = true)
|
||||
public EmbeddingClient mistralAiEmbeddingClient(MistralAiCommonProperties commonProperties,
|
||||
MistralAiEmbeddingProperties embeddingProperties, RestClient.Builder restClientBuilder) {
|
||||
|
||||
var mistralAiApi = mistralAiApi(embeddingProperties.getApiKey(), commonProperties.getApiKey(),
|
||||
embeddingProperties.getBaseUrl(), commonProperties.getBaseUrl(), restClientBuilder);
|
||||
|
||||
return new MistralAiEmbeddingClient(mistralAiApi, embeddingProperties.getMetadataMode(),
|
||||
embeddingProperties.getOptions());
|
||||
}
|
||||
|
||||
@Bean
|
||||
@ConditionalOnMissingBean
|
||||
@ConditionalOnProperty(prefix = MistralAiChatProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true",
|
||||
matchIfMissing = true)
|
||||
public MistralAiChatClient mistralAiChatClient(MistralAiCommonProperties commonProperties,
|
||||
MistralAiChatProperties chatProperties, RestClient.Builder restClientBuilder) {
|
||||
|
||||
var mistralAiApi = mistralAiApi(chatProperties.getApiKey(), commonProperties.getApiKey(),
|
||||
chatProperties.getBaseUrl(), commonProperties.getBaseUrl(), restClientBuilder);
|
||||
|
||||
return new MistralAiChatClient(mistralAiApi, chatProperties.getOptions());
|
||||
}
|
||||
|
||||
private MistralAiApi mistralAiApi(String apiKey, String commonApiKey, String baseUrl, String commonBaseUrl,
|
||||
RestClient.Builder restClientBuilder) {
|
||||
|
||||
var resolvedApiKey = StringUtils.hasText(apiKey) ? apiKey : commonApiKey;
|
||||
var resoledBaseUrl = StringUtils.hasText(baseUrl) ? baseUrl : commonBaseUrl;
|
||||
|
||||
Assert.hasText(resolvedApiKey, "Mistral API key must be set");
|
||||
Assert.hasText(resoledBaseUrl, "Mistral base URL must be set");
|
||||
|
||||
return new MistralAiApi(resoledBaseUrl, resolvedApiKey, restClientBuilder);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
/*
|
||||
* Copyright 2024-204 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.autoconfigure.mistralai;
|
||||
|
||||
import org.springframework.ai.mistral.MistralAiChatOptions;
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.NestedConfigurationProperty;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@ConfigurationProperties(MistralAiChatProperties.CONFIG_PREFIX)
|
||||
public class MistralAiChatProperties extends MistralAiParentProperties {
|
||||
|
||||
public static final String CONFIG_PREFIX = "spring.ai.mistral.chat";
|
||||
|
||||
public static final String DEFAULT_CHAT_MODEL = "mistral-tiny";
|
||||
|
||||
private static final Double DEFAULT_TEMPERATURE = 0.7;
|
||||
|
||||
private static final Float DEFAULT_TOP_P = 1.0f;
|
||||
|
||||
private static final Boolean IS_ENABLED = false;
|
||||
|
||||
public MistralAiChatProperties() {
|
||||
super.setBaseUrl(MistralAiCommonProperties.DEFAULT_BASE_URL);
|
||||
}
|
||||
|
||||
/**
|
||||
* Enable OpenAI chat client.
|
||||
*/
|
||||
private boolean enabled = true;
|
||||
|
||||
@NestedConfigurationProperty
|
||||
private MistralAiChatOptions options = MistralAiChatOptions.builder()
|
||||
.withModel(DEFAULT_CHAT_MODEL)
|
||||
.withTemperature(DEFAULT_TEMPERATURE.floatValue())
|
||||
.withSafePrompt(!IS_ENABLED)
|
||||
.withTopP(DEFAULT_TOP_P)
|
||||
.build();
|
||||
|
||||
public MistralAiChatOptions getOptions() {
|
||||
return this.options;
|
||||
}
|
||||
|
||||
public void setOptions(MistralAiChatOptions options) {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
public boolean isEnabled() {
|
||||
return this.enabled;
|
||||
}
|
||||
|
||||
public void setEnabled(boolean enabled) {
|
||||
this.enabled = enabled;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
/*
|
||||
* Copyright 2024-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.autoconfigure.mistralai;
|
||||
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@ConfigurationProperties(MistralAiCommonProperties.CONFIG_PREFIX)
|
||||
public class MistralAiCommonProperties extends MistralAiParentProperties {
|
||||
|
||||
public static final String CONFIG_PREFIX = "spring.ai.mistralai";
|
||||
|
||||
public static final String DEFAULT_BASE_URL = "https://api.mistral.ai";
|
||||
|
||||
public MistralAiCommonProperties() {
|
||||
super.setBaseUrl(DEFAULT_BASE_URL);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
/*
|
||||
* Copyright 2024-204 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.autoconfigure.mistralai;
|
||||
|
||||
import org.springframework.ai.document.MetadataMode;
|
||||
import org.springframework.ai.mistral.MistralAiEmbeddingOptions;
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.context.properties.NestedConfigurationProperty;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@ConfigurationProperties(MistralAiEmbeddingProperties.CONFIG_PREFIX)
|
||||
public class MistralAiEmbeddingProperties extends MistralAiParentProperties {
|
||||
|
||||
public static final String CONFIG_PREFIX = "spring.ai.mistralai.embedding";
|
||||
|
||||
public static final String DEFAULT_EMBEDDING_MODEL = "mistral-embed";
|
||||
|
||||
public static final String DEFAULT_ENCODING_FORMAT = "float";
|
||||
|
||||
/**
|
||||
* Enable MistralAI embedding client.
|
||||
*/
|
||||
private boolean enabled = true;
|
||||
|
||||
public MetadataMode metadataMode = MetadataMode.EMBED;
|
||||
|
||||
@NestedConfigurationProperty
|
||||
private MistralAiEmbeddingOptions options = MistralAiEmbeddingOptions.builder()
|
||||
.withModel(DEFAULT_EMBEDDING_MODEL)
|
||||
.withEncodingFormat(DEFAULT_ENCODING_FORMAT)
|
||||
.build();
|
||||
|
||||
public MistralAiEmbeddingProperties() {
|
||||
super.setBaseUrl(MistralAiCommonProperties.DEFAULT_BASE_URL);
|
||||
}
|
||||
|
||||
public MistralAiEmbeddingOptions getOptions() {
|
||||
return this.options;
|
||||
}
|
||||
|
||||
public void setOptions(MistralAiEmbeddingOptions options) {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
public MetadataMode getMetadataMode() {
|
||||
return this.metadataMode;
|
||||
}
|
||||
|
||||
public void setMetadataMode(MetadataMode metadataMode) {
|
||||
this.metadataMode = metadataMode;
|
||||
}
|
||||
|
||||
public boolean isEnabled() {
|
||||
return this.enabled;
|
||||
}
|
||||
|
||||
public void setEnabled(boolean enabled) {
|
||||
this.enabled = enabled;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
/*
|
||||
* Copyright 2024-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.autoconfigure.mistralai;
|
||||
|
||||
/**
|
||||
* @author Ricken Bazolo
|
||||
* @since 0.8.1
|
||||
*/
|
||||
public class MistralAiParentProperties {
|
||||
|
||||
private String apiKey;
|
||||
|
||||
private String baseUrl;
|
||||
|
||||
public String getApiKey() {
|
||||
return this.apiKey;
|
||||
}
|
||||
|
||||
public void setApiKey(String apiKey) {
|
||||
this.apiKey = apiKey;
|
||||
}
|
||||
|
||||
public String getBaseUrl() {
|
||||
return this.baseUrl;
|
||||
}
|
||||
|
||||
public void setBaseUrl(String baseUrl) {
|
||||
this.baseUrl = baseUrl;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -11,6 +11,7 @@ org.springframework.ai.autoconfigure.bedrock.anthropic.BedrockAnthropicChatAutoC
|
||||
org.springframework.ai.autoconfigure.bedrock.titan.BedrockTitanChatAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.bedrock.titan.BedrockTitanEmbeddingAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.ollama.OllamaAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.mistralai.MistralAiAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.vectorstore.pgvector.PgVectorStoreAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.vectorstore.pinecone.PineconeVectorStoreAutoConfiguration
|
||||
org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreAutoConfiguration
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
/*
|
||||
* Copyright 2023-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.autoconfigure.mistralai;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.apache.commons.logging.Log;
|
||||
import org.apache.commons.logging.LogFactory;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import org.springframework.ai.chat.ChatResponse;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.embedding.EmbeddingResponse;
|
||||
import org.springframework.ai.mistral.MistralAiChatClient;
|
||||
import org.springframework.ai.mistral.MistralAiEmbeddingClient;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 0.8.1
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".*")
|
||||
public class MistralAiAutoConfigurationIT {
|
||||
|
||||
private static final Log logger = LogFactory.getLog(MistralAiAutoConfigurationIT.class);
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
.withPropertyValues("spring.ai.mistralai.apiKey=" + System.getenv("MISTRAL_AI_API_KEY"))
|
||||
.withConfiguration(AutoConfigurations.of(RestClientAutoConfiguration.class, MistralAiAutoConfiguration.class));
|
||||
|
||||
@Test
|
||||
void generate() {
|
||||
contextRunner.run(context -> {
|
||||
MistralAiChatClient client = context.getBean(MistralAiChatClient.class);
|
||||
String response = client.call("Hello");
|
||||
assertThat(response).isNotEmpty();
|
||||
logger.info("Response: " + response);
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void generateStreaming() {
|
||||
contextRunner.run(context -> {
|
||||
MistralAiChatClient client = context.getBean(MistralAiChatClient.class);
|
||||
Flux<ChatResponse> responseFlux = client.stream(new Prompt(new UserMessage("Hello")));
|
||||
String response = responseFlux.collectList().block().stream().map(chatResponse -> {
|
||||
return chatResponse.getResults().get(0).getOutput().getContent();
|
||||
}).collect(Collectors.joining());
|
||||
|
||||
assertThat(response).isNotEmpty();
|
||||
logger.info("Response: " + response);
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void embedding() {
|
||||
contextRunner.run(context -> {
|
||||
MistralAiEmbeddingClient embeddingClient = context.getBean(MistralAiEmbeddingClient.class);
|
||||
|
||||
EmbeddingResponse embeddingResponse = embeddingClient
|
||||
.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
|
||||
assertThat(embeddingResponse.getResults()).hasSize(2);
|
||||
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
|
||||
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
|
||||
assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty();
|
||||
assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1);
|
||||
|
||||
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
/*
|
||||
* Copyright 2024-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.autoconfigure.mistralai;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.boot.autoconfigure.AutoConfigurations;
|
||||
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Unit Tests for {@link MistralAiCommonProperties}, {@link MistralAiEmbeddingProperties}.
|
||||
*/
|
||||
public class MistralAiPropertiesTests {
|
||||
|
||||
@Test
|
||||
public void embeddingProperties() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
.withPropertyValues("spring.ai.mistralai.base-url=TEST_BASE_URL", "spring.ai.mistralai.api-key=abc123",
|
||||
"spring.ai.mistralai.embedding.options.model=MODEL_XYZ")
|
||||
.withConfiguration(
|
||||
AutoConfigurations.of(RestClientAutoConfiguration.class, MistralAiAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
var embeddingProperties = context.getBean(MistralAiEmbeddingProperties.class);
|
||||
var connectionProperties = context.getBean(MistralAiCommonProperties.class);
|
||||
|
||||
assertThat(connectionProperties.getApiKey()).isEqualTo("abc123");
|
||||
assertThat(connectionProperties.getBaseUrl()).isEqualTo("TEST_BASE_URL");
|
||||
|
||||
assertThat(embeddingProperties.getApiKey()).isNull();
|
||||
assertThat(embeddingProperties.getBaseUrl()).isNull();
|
||||
|
||||
assertThat(embeddingProperties.getOptions().getModel()).isEqualTo("MODEL_XYZ");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
public void embeddingOverrideConnectionProperties() {
|
||||
|
||||
new ApplicationContextRunner().withPropertyValues("spring.ai.mistralai.base-url=TEST_BASE_URL",
|
||||
"spring.ai.mistralai.api-key=abc123", "spring.ai.mistralai.embedding.base-url=TEST_BASE_URL2",
|
||||
"spring.ai.mistralai.embedding.api-key=456", "spring.ai.mistralai.embedding.options.model=MODEL_XYZ")
|
||||
.withConfiguration(
|
||||
AutoConfigurations.of(RestClientAutoConfiguration.class, MistralAiAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
var embeddingProperties = context.getBean(MistralAiEmbeddingProperties.class);
|
||||
var connectionProperties = context.getBean(MistralAiCommonProperties.class);
|
||||
|
||||
assertThat(connectionProperties.getApiKey()).isEqualTo("abc123");
|
||||
assertThat(connectionProperties.getBaseUrl()).isEqualTo("TEST_BASE_URL");
|
||||
|
||||
assertThat(embeddingProperties.getApiKey()).isEqualTo("456");
|
||||
assertThat(embeddingProperties.getBaseUrl()).isEqualTo("TEST_BASE_URL2");
|
||||
|
||||
assertThat(embeddingProperties.getOptions().getModel()).isEqualTo("MODEL_XYZ");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
public void embeddingOptionsTest() {
|
||||
|
||||
new ApplicationContextRunner()
|
||||
.withPropertyValues("spring.ai.mistralai.api-key=API_KEY", "spring.ai.mistralai.base-url=TEST_BASE_URL",
|
||||
|
||||
"spring.ai.mistralai.embedding.options.model=MODEL_XYZ",
|
||||
"spring.ai.mistralai.embedding.options.encodingFormat=MyEncodingFormat")
|
||||
.withConfiguration(
|
||||
AutoConfigurations.of(RestClientAutoConfiguration.class, MistralAiAutoConfiguration.class))
|
||||
.run(context -> {
|
||||
var connectionProperties = context.getBean(MistralAiCommonProperties.class);
|
||||
var embeddingProperties = context.getBean(MistralAiEmbeddingProperties.class);
|
||||
|
||||
assertThat(connectionProperties.getBaseUrl()).isEqualTo("TEST_BASE_URL");
|
||||
assertThat(connectionProperties.getApiKey()).isEqualTo("API_KEY");
|
||||
|
||||
assertThat(embeddingProperties.getOptions().getModel()).isEqualTo("MODEL_XYZ");
|
||||
assertThat(embeddingProperties.getOptions().getEncodingFormat()).isEqualTo("MyEncodingFormat");
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai</artifactId>
|
||||
<version>0.8.1-SNAPSHOT</version>
|
||||
<relativePath>../../pom.xml</relativePath>
|
||||
</parent>
|
||||
<artifactId>spring-ai-mistral-ai-spring-boot-starter</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>Spring AI Starter - MistralAI</name>
|
||||
<description>Spring AI MistralAI Auto Configuration</description>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
|
||||
<scm>
|
||||
<url>https://github.com/spring-projects/spring-ai</url>
|
||||
<connection>git://github.com/spring-projects/spring-ai.git</connection>
|
||||
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
|
||||
</scm>
|
||||
|
||||
<dependencies>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter</artifactId>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-mistral-ai</artifactId>
|
||||
<version>${project.parent.version}</version>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
Reference in New Issue
Block a user