Skip to content

Commit 1e67d32

Browse files
committed
Add tests for Mistral moderation API
1 parent ca283da commit 1e67d32

File tree

3 files changed

+84
-0
lines changed

3 files changed

+84
-0
lines changed
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package org.springframework.ai.mistralai;
2+
3+
import org.junit.jupiter.api.Test;
4+
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
5+
import org.slf4j.Logger;
6+
import org.slf4j.LoggerFactory;
7+
import org.springframework.ai.mistralai.moderation.MistralAiModerationModel;
8+
import org.springframework.ai.moderation.Moderation;
9+
import org.springframework.ai.moderation.ModerationPrompt;
10+
import org.springframework.ai.moderation.ModerationResult;
11+
import org.springframework.beans.factory.annotation.Autowired;
12+
import org.springframework.boot.test.context.SpringBootTest;
13+
14+
import static org.assertj.core.api.Assertions.assertThat;
15+
16+
/**
17+
* @author Ricken Bazolo
18+
*/
19+
@SpringBootTest(classes = MistralAiTestConfiguration.class)
20+
@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
21+
public class MistralAiModerationModelIT {
22+
23+
private static final Logger logger = LoggerFactory.getLogger(MistralAiModerationModelIT.class);
24+
25+
@Autowired
26+
private MistralAiModerationModel mistralAiModerationModel;
27+
28+
@Test
29+
void moderationAsPositiveTest() {
30+
var instructions = """
31+
I want to kill them.!".""";
32+
33+
var moderationPrompt = new ModerationPrompt(instructions);
34+
35+
var moderationResponse = this.mistralAiModerationModel.call(moderationPrompt);
36+
37+
assertThat(moderationResponse.getResults()).hasSize(1);
38+
39+
var generation = moderationResponse.getResult();
40+
Moderation moderation = generation.getOutput();
41+
assertThat(moderation.getId()).isNotEmpty();
42+
assertThat(moderation.getResults()).isNotNull();
43+
assertThat(moderation.getResults().size()).isNotZero();
44+
logger.info(moderation.getResults().toString());
45+
46+
assertThat(moderation.getId()).isNotNull();
47+
assertThat(moderation.getModel()).isNotNull();
48+
49+
ModerationResult result = moderation.getResults().get(0);
50+
assertThat(result.isFlagged()).isTrue();
51+
assertThat(result.getCategories().isViolence()).isTrue();
52+
}
53+
54+
}

models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiTestConfiguration.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
import org.springframework.ai.embedding.EmbeddingModel;
2020
import org.springframework.ai.mistralai.api.MistralAiApi;
21+
import org.springframework.ai.mistralai.api.MistralAiModerationApi;
22+
import org.springframework.ai.mistralai.moderation.MistralAiModerationModel;
2123
import org.springframework.boot.SpringBootConfiguration;
2224
import org.springframework.context.annotation.Bean;
2325
import org.springframework.util.StringUtils;
@@ -35,6 +37,16 @@ public MistralAiApi mistralAiApi() {
3537
return new MistralAiApi(apiKey);
3638
}
3739

40+
@Bean
41+
public MistralAiModerationApi mistralAiModerationApi() {
42+
var apiKey = System.getenv("MISTRAL_AI_API_KEY");
43+
if (!StringUtils.hasText(apiKey)) {
44+
throw new IllegalArgumentException(
45+
"Missing MISTRAL_AI_API_KEY environment variable. Please set it to your Mistral AI API key.");
46+
}
47+
return new MistralAiModerationApi(apiKey);
48+
}
49+
3850
@Bean
3951
public EmbeddingModel mistralAiEmbeddingModel(MistralAiApi api) {
4052
return new MistralAiEmbeddingModel(api,
@@ -47,4 +59,9 @@ public MistralAiChatModel mistralAiChatModel(MistralAiApi mistralAiApi) {
4759
MistralAiChatOptions.builder().model(MistralAiApi.ChatModel.OPEN_MIXTRAL_7B.getValue()).build());
4860
}
4961

62+
@Bean
63+
public MistralAiModerationModel mistralAiModerationModel(MistralAiModerationApi mistralAiModerationApi) {
64+
return new MistralAiModerationModel(mistralAiModerationApi);
65+
}
66+
5067
}

spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/mistralai/MistralAiPropertiesTests.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,4 +145,17 @@ public void embeddingOptionsTest() {
145145
});
146146
}
147147

148+
@Test
149+
public void moderationOptionsTest() {
150+
new ApplicationContextRunner()
151+
.withPropertyValues("spring.ai.mistralai.base-url=TEST_BASE_URL", "spring.ai.mistralai.api-key=abc123",
152+
"spring.ai.mistralai.moderation.options.model=MODERATION_MODEL")
153+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
154+
RestClientAutoConfiguration.class, MistralAiAutoConfiguration.class))
155+
.run(context -> {
156+
var moderationProperties = context.getBean(MistralAiModerationProperties.class);
157+
assertThat(moderationProperties.getOptions().getModel()).isEqualTo("MODERATION_MODEL");
158+
});
159+
}
160+
148161
}

0 commit comments

Comments
 (0)