feat(mistral-ai): Add moderation model support (#2201)

Implement MistralAI moderation capabilities to detect potentially harmful content.
This allows Spring AI applications to use Mistral's content moderation services
to identify and filter inappropriate content before processing

- Add MistralAiModerationApi for interacting with Mistral's moderation endpoints
- Create MistralAiModerationModel implementing the ModerationModel interface
- Add configuration properties and auto-configuration for the moderation model
- Extend Categories and CategoryScores with additional moderation categories
- Add integration tests to verify moderation functionality

Signed-off-by: Ricken Bazolo <ricken.bazolo@gmail.com>

auto-configurations/models/spring-ai-autoconfigure-model-mistral-ai/src/main/java/org/springframework/ai/model/mistralai/autoconfigure/MistralAiChatAutoConfiguration.java

@@ -21,6 +21,8 @@ import io.micrometer.observation.ObservationRegistry;
 import org.springframework.ai.chat.observation.ChatModelObservationConvention;
 import org.springframework.ai.mistralai.MistralAiChatModel;
 import org.springframework.ai.mistralai.api.MistralAiApi;
+import org.springframework.ai.mistralai.api.MistralAiModerationApi;
+import org.springframework.ai.mistralai.moderation.MistralAiModerationModel;
 import org.springframework.ai.model.SpringAIModelProperties;
 import org.springframework.ai.model.SpringAIModels;
 import org.springframework.ai.model.function.DefaultFunctionCallbackResolver;
@@ -57,7 +59,8 @@ import org.springframework.web.client.RestClient;
  */
 @AutoConfiguration(after = { RestClientAutoConfiguration.class, SpringAiRetryAutoConfiguration.class,
 		ToolCallingAutoConfiguration.class })
-@EnableConfigurationProperties({ MistralAiCommonProperties.class, MistralAiChatProperties.class })
+@EnableConfigurationProperties({ MistralAiCommonProperties.class, MistralAiChatProperties.class,
+		MistralAiModerationProperties.class })
 @ConditionalOnProperty(name = SpringAIModelProperties.CHAT_MODEL, havingValue = SpringAIModels.MISTRAL,
 		matchIfMissing = true)
 @ConditionalOnClass(MistralAiApi.class)
@@ -93,6 +96,27 @@ public class MistralAiChatAutoConfiguration {
 		return chatModel;
 	}
 
+	@Bean
+	@ConditionalOnMissingBean
+	public MistralAiModerationModel mistralAiModerationModel(MistralAiCommonProperties commonProperties,
+			MistralAiModerationProperties moderationProperties, RetryTemplate retryTemplate,
+			ObjectProvider<RestClient.Builder> restClientBuilderProvider, ResponseErrorHandler responseErrorHandler) {
+
+		var apiKey = moderationProperties.getApiKey();
+		var baseUrl = moderationProperties.getBaseUrl();
+
+		var resolvedApiKey = StringUtils.hasText(apiKey) ? apiKey : commonProperties.getApiKey();
+		var resoledBaseUrl = StringUtils.hasText(baseUrl) ? baseUrl : commonProperties.getBaseUrl();
+
+		Assert.hasText(resolvedApiKey, "Mistral API key must be set");
+		Assert.hasText(resoledBaseUrl, "Mistral base URL must be set");
+
+		var mistralAiModerationAi = new MistralAiModerationApi(resoledBaseUrl, resolvedApiKey,
+				restClientBuilderProvider.getIfAvailable(RestClient::builder), responseErrorHandler);
+
+		return new MistralAiModerationModel(mistralAiModerationAi, retryTemplate, moderationProperties.getOptions());
+	}
+
 	private MistralAiApi mistralAiApi(String apiKey, String commonApiKey, String baseUrl, String commonBaseUrl,
 			RestClient.Builder restClientBuilder, ResponseErrorHandler responseErrorHandler) {
 

auto-configurations/models/spring-ai-autoconfigure-model-mistral-ai/src/main/java/org/springframework/ai/model/mistralai/autoconfigure/MistralAiModerationProperties.java

@@ -0,0 +1,35 @@
+package org.springframework.ai.model.mistralai.autoconfigure;
+
+import org.springframework.ai.mistralai.moderation.MistralAiModerationOptions;
+import org.springframework.ai.mistralai.api.MistralAiModerationApi;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.context.properties.NestedConfigurationProperty;
+
+/**
+ * @author Ricken Bazolo
+ */
+@ConfigurationProperties(MistralAiModerationProperties.CONFIG_PREFIX)
+public class MistralAiModerationProperties extends MistralAiParentProperties {
+
+	public static final String CONFIG_PREFIX = "spring.ai.mistralai.moderation";
+
+	private static final String DEFAULT_MODERATION_MODEL = MistralAiModerationApi.Model.MISTRAL_MODERATION.getValue();
+
+	@NestedConfigurationProperty
+	private MistralAiModerationOptions options = MistralAiModerationOptions.builder()
+		.model(DEFAULT_MODERATION_MODEL)
+		.build();
+
+	public MistralAiModerationProperties() {
+		super.setBaseUrl(MistralAiCommonProperties.DEFAULT_BASE_URL);
+	}
+
+	public MistralAiModerationOptions getOptions() {
+		return this.options;
+	}
+
+	public void setOptions(MistralAiModerationOptions options) {
+		this.options = options;
+	}
+
+}

auto-configurations/models/spring-ai-autoconfigure-model-mistral-ai/src/test/java/org/springframework/ai/model/mistralai/autoconfigure/MistralAiPropertiesTests.java

@@ -145,4 +145,17 @@ public class MistralAiPropertiesTests {
 			});
 	}
 
+	@Test
+	public void moderationOptionsTest() {
+		new ApplicationContextRunner()
+			.withPropertyValues("spring.ai.mistralai.base-url=TEST_BASE_URL", "spring.ai.mistralai.api-key=abc123",
+					"spring.ai.mistralai.moderation.options.model=MODERATION_MODEL")
+			.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
+					RestClientAutoConfiguration.class, MistralAiChatAutoConfiguration.class))
+			.run(context -> {
+				var moderationProperties = context.getBean(MistralAiModerationProperties.class);
+				assertThat(moderationProperties.getOptions().getModel()).isEqualTo("MODERATION_MODEL");
+			});
+	}
+
 }

models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiModerationApi.java

@@ -0,0 +1,137 @@
+package org.springframework.ai.mistralai.api;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.springframework.ai.retry.RetryUtils;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
+import org.springframework.util.Assert;
+import org.springframework.web.client.ResponseErrorHandler;
+import org.springframework.web.client.RestClient;
+
+import java.util.function.Consumer;
+
+/**
+ * MistralAI Moderation API.
+ *
+ * @author Ricken Bazolo
+ * @see <a href= "https://docs.mistral.ai/capabilities/guardrailing/</a>
+ */
+public class MistralAiModerationApi {
+
+	private static final String DEFAULT_BASE_URL = "https://api.mistral.ai";
+
+	private final RestClient restClient;
+
+	public MistralAiModerationApi(String mistralAiApiKey) {
+		this(DEFAULT_BASE_URL, mistralAiApiKey, RestClient.builder(), RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
+	}
+
+	public MistralAiModerationApi(String baseUrl, String mistralAiApiKey, RestClient.Builder restClientBuilder,
+			ResponseErrorHandler responseErrorHandler) {
+
+		Consumer<HttpHeaders> jsonContentHeaders = headers -> {
+			headers.setBearerAuth(mistralAiApiKey);
+			headers.setContentType(MediaType.APPLICATION_JSON);
+		};
+
+		this.restClient = restClientBuilder.baseUrl(baseUrl)
+			.defaultHeaders(jsonContentHeaders)
+			.defaultStatusHandler(responseErrorHandler)
+			.build();
+	}
+
+	public ResponseEntity<MistralAiModerationResponse> moderate(MistralAiModerationRequest mistralAiModerationRequest) {
+		Assert.notNull(mistralAiModerationRequest, "Moderation request cannot be null.");
+		Assert.hasLength(mistralAiModerationRequest.prompt(), "Prompt cannot be empty.");
+		Assert.notNull(mistralAiModerationRequest.model(), "Model cannot be null.");
+
+		return this.restClient.post()
+			.uri("v1/moderations")
+			.body(mistralAiModerationRequest)
+			.retrieve()
+			.toEntity(MistralAiModerationResponse.class);
+	}
+
+	public enum Model {
+
+		// @formatter:off
+		MISTRAL_MODERATION("mistral-moderation-latest");
+		// @formatter:on
+
+		private final String value;
+
+		Model(String value) {
+			this.value = value;
+		}
+
+		public String getValue() {
+			return this.value;
+		}
+
+	}
+
+	// @formatter:off
+	@JsonInclude(JsonInclude.Include.NON_NULL)
+	public record MistralAiModerationRequest(
+		@JsonProperty("input") String prompt,
+		@JsonProperty("model") String model
+	) {
+
+		public MistralAiModerationRequest(String prompt) {
+			this(prompt, null);
+		}
+	}
+
+	@JsonInclude(JsonInclude.Include.NON_NULL)
+	public record MistralAiModerationResponse(
+			@JsonProperty("id") String id,
+			@JsonProperty("model") String model,
+			@JsonProperty("results") MistralAiModerationResult[] results) {
+
+	}
+
+	@JsonInclude(JsonInclude.Include.NON_NULL)
+	public record MistralAiModerationResult(
+			@JsonProperty("categories") Categories categories,
+			@JsonProperty("category_scores") CategoryScores categoryScores) {
+
+		public boolean flagged() {
+			return categories != null && (categories.sexual() || categories.hateAndDiscrimination() || categories.violenceAndThreats()
+					|| categories.selfHarm() || categories.dangerousAndCriminalContent() || categories.health()
+					|| categories.financial() || categories.law() || categories.pii());
+		}
+
+	}
+
+	@JsonInclude(JsonInclude.Include.NON_NULL)
+	public record Categories(
+			@JsonProperty("sexual") boolean sexual,
+			@JsonProperty("hate_and_discrimination") boolean hateAndDiscrimination,
+			@JsonProperty("violence_and_threats") boolean violenceAndThreats,
+			@JsonProperty("selfharm") boolean selfHarm,
+			@JsonProperty("dangerous_and_criminal_content") boolean dangerousAndCriminalContent,
+			@JsonProperty("health") boolean health,
+			@JsonProperty("financial") boolean financial,
+			@JsonProperty("law") boolean law,
+			@JsonProperty("pii") boolean pii)  {
+
+	}
+
+	@JsonInclude(JsonInclude.Include.NON_NULL)
+	public record CategoryScores(
+			@JsonProperty("sexual") double sexual,
+			@JsonProperty("hate_and_discrimination") double hateAndDiscrimination,
+			@JsonProperty("violence_and_threats") double violenceAndThreats,
+			@JsonProperty("selfharm") double selfHarm,
+			@JsonProperty("dangerous_and_criminal_content") double dangerousAndCriminalContent,
+			@JsonProperty("health") double health,
+			@JsonProperty("financial") double financial,
+			@JsonProperty("law") double law,
+			@JsonProperty("pii") double pii)  {
+
+	}
+	// @formatter:onn
+
+}

models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/moderation/MistralAiModerationModel.java

@@ -0,0 +1,146 @@
+package org.springframework.ai.mistralai.moderation;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.mistralai.api.MistralAiModerationApi;
+import org.springframework.ai.model.ModelOptionsUtils;
+import org.springframework.ai.moderation.*;
+import org.springframework.ai.retry.RetryUtils;
+import org.springframework.http.ResponseEntity;
+import org.springframework.retry.support.RetryTemplate;
+import org.springframework.util.Assert;
+
+import static org.springframework.ai.mistralai.api.MistralAiModerationApi.MistralAiModerationRequest;
+import static org.springframework.ai.mistralai.api.MistralAiModerationApi.MistralAiModerationResponse;
+import static org.springframework.ai.mistralai.api.MistralAiModerationApi.MistralAiModerationResult;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * @author Ricken Bazolo
+ */
+public class MistralAiModerationModel implements ModerationModel {
+
+	private final Logger logger = LoggerFactory.getLogger(getClass());
+
+	private final MistralAiModerationApi mistralAiModerationApi;
+
+	private final RetryTemplate retryTemplate;
+
+	private final MistralAiModerationOptions defaultOptions;
+
+	public MistralAiModerationModel(MistralAiModerationApi mistralAiModerationApi) {
+		this(mistralAiModerationApi, RetryUtils.DEFAULT_RETRY_TEMPLATE,
+				MistralAiModerationOptions.builder()
+					.model(MistralAiModerationApi.Model.MISTRAL_MODERATION.getValue())
+					.build());
+	}
+
+	public MistralAiModerationModel(MistralAiModerationApi mistralAiModerationApi, MistralAiModerationOptions options) {
+		this(mistralAiModerationApi, RetryUtils.DEFAULT_RETRY_TEMPLATE, options);
+	}
+
+	public MistralAiModerationModel(MistralAiModerationApi mistralAiModerationApi, RetryTemplate retryTemplate,
+			MistralAiModerationOptions options) {
+		Assert.notNull(mistralAiModerationApi, "mistralAiModerationApi must not be null");
+		Assert.notNull(retryTemplate, "retryTemplate must not be null");
+		Assert.notNull(options, "options must not be null");
+		this.mistralAiModerationApi = mistralAiModerationApi;
+		this.retryTemplate = retryTemplate;
+		this.defaultOptions = options;
+	}
+
+	@Override
+	public ModerationResponse call(ModerationPrompt moderationPrompt) {
+		return this.retryTemplate.execute(ctx -> {
+
+			var instructions = moderationPrompt.getInstructions().getText();
+
+			var moderationRequest = new MistralAiModerationRequest(instructions);
+
+			if (this.defaultOptions != null) {
+				moderationRequest = ModelOptionsUtils.merge(this.defaultOptions, moderationRequest,
+						MistralAiModerationRequest.class);
+			}
+			else {
+				// moderationPrompt.getOptions() never null but model can be empty, cause
+				// by ModerationPrompt constructor
+				moderationRequest = ModelOptionsUtils.merge(toMistralAiModerationOptions(moderationPrompt.getOptions()),
+						moderationRequest, MistralAiModerationRequest.class);
+			}
+
+			var moderationResponseEntity = this.mistralAiModerationApi.moderate(moderationRequest);
+
+			return convertResponse(moderationResponseEntity, moderationRequest);
+		});
+	}
+
+	private ModerationResponse convertResponse(ResponseEntity<MistralAiModerationResponse> moderationResponseEntity,
+			MistralAiModerationRequest openAiModerationRequest) {
+		var moderationApiResponse = moderationResponseEntity.getBody();
+		if (moderationApiResponse == null) {
+			logger.warn("No moderation response returned for request: {}", openAiModerationRequest);
+			return new ModerationResponse(new Generation());
+		}
+
+		List<ModerationResult> moderationResults = new ArrayList<>();
+		if (moderationApiResponse.results() != null) {
+
+			for (MistralAiModerationResult result : moderationApiResponse.results()) {
+				Categories categories = null;
+				CategoryScores categoryScores = null;
+				if (result.categories() != null) {
+					categories = Categories.builder()
+						.sexual(result.categories().sexual())
+						.pii(result.categories().pii())
+						.law(result.categories().law())
+						.financial(result.categories().financial())
+						.health(result.categories().health())
+						.dangerousAndCriminalContent(result.categories().dangerousAndCriminalContent())
+						.violence(result.categories().violenceAndThreats())
+						.hate(result.categories().hateAndDiscrimination())
+						.selfHarm(result.categories().selfHarm())
+						.build();
+				}
+				if (result.categoryScores() != null) {
+					categoryScores = CategoryScores.builder()
+						.sexual(result.categoryScores().sexual())
+						.pii(result.categoryScores().pii())
+						.law(result.categoryScores().law())
+						.financial(result.categoryScores().financial())
+						.health(result.categoryScores().health())
+						.dangerousAndCriminalContent(result.categoryScores().dangerousAndCriminalContent())
+						.violence(result.categoryScores().violenceAndThreats())
+						.hate(result.categoryScores().hateAndDiscrimination())
+						.selfHarm(result.categoryScores().selfHarm())
+						.build();
+				}
+				var moderationResult = ModerationResult.builder()
+					.categories(categories)
+					.categoryScores(categoryScores)
+					.flagged(result.flagged())
+					.build();
+				moderationResults.add(moderationResult);
+			}
+
+		}
+
+		var moderation = Moderation.builder()
+			.id(moderationApiResponse.id())
+			.model(moderationApiResponse.model())
+			.results(moderationResults)
+			.build();
+
+		return new ModerationResponse(new Generation(moderation));
+	}
+
+	private MistralAiModerationOptions toMistralAiModerationOptions(ModerationOptions runtimeModerationOptions) {
+		var mistralAiModerationOptionsBuilder = MistralAiModerationOptions.builder();
+		if (runtimeModerationOptions != null && runtimeModerationOptions.getModel() != null) {
+			mistralAiModerationOptionsBuilder.model(runtimeModerationOptions.getModel());
+		}
+		return mistralAiModerationOptionsBuilder.build();
+	}
+
+}

models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/moderation/MistralAiModerationOptions.java

@@ -0,0 +1,54 @@
+package org.springframework.ai.mistralai.moderation;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.springframework.ai.mistralai.api.MistralAiModerationApi;
+import org.springframework.ai.moderation.ModerationOptions;
+
+/**
+ * @author Ricken Bazolo
+ */
+@JsonInclude(JsonInclude.Include.NON_NULL)
+public class MistralAiModerationOptions implements ModerationOptions {
+
+	private static final String DEFAULT_MODEL = MistralAiModerationApi.Model.MISTRAL_MODERATION.getValue();
+
+	/**
+	 * The model to use for moderation generation.
+	 */
+	@JsonProperty("model")
+	private String model = DEFAULT_MODEL;
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	@Override
+	public String getModel() {
+		return this.model;
+	}
+
+	public void setModel(String model) {
+		this.model = model;
+	}
+
+	public static final class Builder {
+
+		private final MistralAiModerationOptions options;
+
+		private Builder() {
+			this.options = new MistralAiModerationOptions();
+		}
+
+		public Builder model(String model) {
+			this.options.setModel(model);
+			return this;
+		}
+
+		public MistralAiModerationOptions build() {
+			return this.options;
+		}
+
+	}
+
+}

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

@@ -0,0 +1,54 @@
+package org.springframework.ai.mistralai;
+
+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.mistralai.moderation.MistralAiModerationModel;
+import org.springframework.ai.moderation.Moderation;
+import org.springframework.ai.moderation.ModerationPrompt;
+import org.springframework.ai.moderation.ModerationResult;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Ricken Bazolo
+ */
+@SpringBootTest(classes = MistralAiTestConfiguration.class)
+@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+")
+public class MistralAiModerationModelIT {
+
+	private static final Logger logger = LoggerFactory.getLogger(MistralAiModerationModelIT.class);
+
+	@Autowired
+	private MistralAiModerationModel mistralAiModerationModel;
+
+	@Test
+	void moderationAsPositiveTest() {
+		var instructions = """
+				I want to kill them.!".""";
+
+		var moderationPrompt = new ModerationPrompt(instructions);
+
+		var moderationResponse = this.mistralAiModerationModel.call(moderationPrompt);
+
+		assertThat(moderationResponse.getResults()).hasSize(1);
+
+		var generation = moderationResponse.getResult();
+		Moderation moderation = generation.getOutput();
+		assertThat(moderation.getId()).isNotEmpty();
+		assertThat(moderation.getResults()).isNotNull();
+		assertThat(moderation.getResults().size()).isNotZero();
+		logger.info(moderation.getResults().toString());
+
+		assertThat(moderation.getId()).isNotNull();
+		assertThat(moderation.getModel()).isNotNull();
+
+		ModerationResult result = moderation.getResults().get(0);
+		assertThat(result.isFlagged()).isTrue();
+		assertThat(result.getCategories().isViolence()).isTrue();
+	}
+
+}

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

@@ -18,6 +18,8 @@ package org.springframework.ai.mistralai;
 
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.mistralai.api.MistralAiApi;
+import org.springframework.ai.mistralai.api.MistralAiModerationApi;
+import org.springframework.ai.mistralai.moderation.MistralAiModerationModel;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.util.StringUtils;
@@ -35,6 +37,16 @@ public class MistralAiTestConfiguration {
 		return new MistralAiApi(apiKey);
 	}
 
+	@Bean
+	public MistralAiModerationApi mistralAiModerationApi() {
+		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 MistralAiModerationApi(apiKey);
+	}
+
 	@Bean
 	public EmbeddingModel mistralAiEmbeddingModel(MistralAiApi api) {
 		return new MistralAiEmbeddingModel(api,
@@ -49,4 +61,9 @@ public class MistralAiTestConfiguration {
 			.build();
 	}
 
+	@Bean
+	public MistralAiModerationModel mistralAiModerationModel(MistralAiModerationApi mistralAiModerationApi) {
+		return new MistralAiModerationModel(mistralAiModerationApi);
+	}
+
 }

spring-ai-model/src/main/java/org/springframework/ai/moderation/Categories.java

@@ -25,6 +25,7 @@ import java.util.Objects;
  *
  * @author Ahmed Yousri
  * @author Ilayaperumal Gopinathan
+ * @author Ricken Bazolo
  * @since 1.0.0
  */
 public final class Categories {
@@ -51,6 +52,16 @@ public final class Categories {
 
 	private final boolean violence;
 
+	private final boolean dangerousAndCriminalContent;
+
+	private final boolean health;
+
+	private final boolean financial;
+
+	private final boolean law;
+
+	private final boolean pii;
+
 	private Categories(Builder builder) {
 		this.sexual = builder.sexual;
 		this.hate = builder.hate;
@@ -63,6 +74,11 @@ public final class Categories {
 		this.selfHarmInstructions = builder.selfHarmInstructions;
 		this.harassmentThreatening = builder.harassmentThreatening;
 		this.violence = builder.violence;
+		this.dangerousAndCriminalContent = builder.dangerousAndCriminalContent;
+		this.health = builder.health;
+		this.financial = builder.financial;
+		this.law = builder.law;
+		this.pii = builder.pii;
 	}
 
 	public static Builder builder() {
@@ -113,6 +129,26 @@ public final class Categories {
 		return this.violence;
 	}
 
+	public boolean isDangerousAndCriminalContent() {
+		return this.dangerousAndCriminalContent;
+	}
+
+	public boolean isHealth() {
+		return this.health;
+	}
+
+	public boolean isFinancial() {
+		return this.financial;
+	}
+
+	public boolean isLaw() {
+		return this.law;
+	}
+
+	public boolean isPii() {
+		return this.pii;
+	}
+
 	@Override
 	public boolean equals(Object o) {
 		if (this == o) {
@@ -126,14 +162,17 @@ public final class Categories {
 				&& this.selfHarm == that.selfHarm && this.sexualMinors == that.sexualMinors
 				&& this.hateThreatening == that.hateThreatening && this.violenceGraphic == that.violenceGraphic
 				&& this.selfHarmIntent == that.selfHarmIntent && this.selfHarmInstructions == that.selfHarmInstructions
-				&& this.harassmentThreatening == that.harassmentThreatening && this.violence == that.violence;
+				&& this.harassmentThreatening == that.harassmentThreatening && this.violence == that.violence
+				&& this.dangerousAndCriminalContent == that.dangerousAndCriminalContent && this.health == that.health
+				&& this.financial == that.financial && this.law == that.law && this.pii == that.pii;
 	}
 
 	@Override
 	public int hashCode() {
 		return Objects.hash(this.sexual, this.hate, this.harassment, this.selfHarm, this.sexualMinors,
 				this.hateThreatening, this.violenceGraphic, this.selfHarmIntent, this.selfHarmInstructions,
-				this.harassmentThreatening, this.violence);
+				this.harassmentThreatening, this.violence, this.dangerousAndCriminalContent, this.health,
+				this.financial, this.law, this.pii);
 	}
 
 	@Override
@@ -142,7 +181,9 @@ public final class Categories {
 				+ ", selfHarm=" + this.selfHarm + ", sexualMinors=" + this.sexualMinors + ", hateThreatening="
 				+ this.hateThreatening + ", violenceGraphic=" + this.violenceGraphic + ", selfHarmIntent="
 				+ this.selfHarmIntent + ", selfHarmInstructions=" + this.selfHarmInstructions
-				+ ", harassmentThreatening=" + this.harassmentThreatening + ", violence=" + this.violence + '}';
+				+ ", harassmentThreatening=" + this.harassmentThreatening + ", violence=" + this.violence
+				+ ", dangerousAndCriminalContent=" + this.dangerousAndCriminalContent + ", health=" + this.health
+				+ ", financial=" + this.financial + ", law=" + this.law + ", pii=" + this.pii + '}';
 	}
 
 	public static class Builder {
@@ -169,6 +210,16 @@ public final class Categories {
 
 		private boolean violence;
 
+		private boolean dangerousAndCriminalContent;
+
+		private boolean health;
+
+		private boolean financial;
+
+		private boolean law;
+
+		private boolean pii;
+
 		public Builder sexual(boolean sexual) {
 			this.sexual = sexual;
 			return this;
@@ -224,6 +275,31 @@ public final class Categories {
 			return this;
 		}
 
+		public Builder dangerousAndCriminalContent(boolean dangerousAndCriminalContent) {
+			this.dangerousAndCriminalContent = dangerousAndCriminalContent;
+			return this;
+		}
+
+		public Builder health(boolean health) {
+			this.health = health;
+			return this;
+		}
+
+		public Builder financial(boolean financial) {
+			this.financial = financial;
+			return this;
+		}
+
+		public Builder law(boolean law) {
+			this.law = law;
+			return this;
+		}
+
+		public Builder pii(boolean pii) {
+			this.pii = pii;
+			return this;
+		}
+
 		public Categories build() {
 			return new Categories(this);
 		}

spring-ai-model/src/main/java/org/springframework/ai/moderation/CategoryScores.java

@@ -25,6 +25,7 @@ import java.util.Objects;
  *
  * @author Ahmed Yousri
  * @author Ilayaperumal Gopinathan
+ * @author Ricken Bazolo
  * @since 1.0.0
  */
 public final class CategoryScores {
@@ -51,6 +52,16 @@ public final class CategoryScores {
 
 	private final double violence;
 
+	private final double dangerousAndCriminalContent;
+
+	private final double health;
+
+	private final double financial;
+
+	private final double law;
+
+	private final double pii;
+
 	private CategoryScores(Builder builder) {
 		this.sexual = builder.sexual;
 		this.hate = builder.hate;
@@ -63,6 +74,11 @@ public final class CategoryScores {
 		this.selfHarmInstructions = builder.selfHarmInstructions;
 		this.harassmentThreatening = builder.harassmentThreatening;
 		this.violence = builder.violence;
+		this.dangerousAndCriminalContent = builder.dangerousAndCriminalContent;
+		this.health = builder.health;
+		this.financial = builder.financial;
+		this.law = builder.law;
+		this.pii = builder.pii;
 	}
 
 	public static Builder builder() {
@@ -174,6 +190,16 @@ public final class CategoryScores {
 
 		private double violence;
 
+		private double dangerousAndCriminalContent;
+
+		private double health;
+
+		private double financial;
+
+		private double law;
+
+		private double pii;
+
 		public Builder sexual(double sexual) {
 			this.sexual = sexual;
 			return this;
@@ -229,6 +255,31 @@ public final class CategoryScores {
 			return this;
 		}
 
+		public Builder dangerousAndCriminalContent(double dangerousAndCriminalContent) {
+			this.dangerousAndCriminalContent = dangerousAndCriminalContent;
+			return this;
+		}
+
+		public Builder health(double health) {
+			this.health = health;
+			return this;
+		}
+
+		public Builder financial(double financial) {
+			this.financial = financial;
+			return this;
+		}
+
+		public Builder law(double law) {
+			this.law = law;
+			return this;
+		}
+
+		public Builder pii(double pii) {
+			this.pii = pii;
+			return this;
+		}
+
 		public CategoryScores build() {
 			return new CategoryScores(this);
 		}