use std::sync::{Arc, OnceLock, RwLock};
use crate::storage::schema::Value;
use crate::storage::unified::entity::{EntityData, UnifiedEntity};
use crate::{RedDBError, RedDBResult};
pub const MODERATION_STATUS_FIELD: &str = "__moderation_status";
pub const MODERATION_STATUS_PENDING: &str = "pending";
pub const MODERATION_STATUS_REJECTED: &str = "rejected";
#[inline]
pub fn entity_moderation_hidden(entity: &UnifiedEntity) -> bool {
let EntityData::Row(row) = &entity.data else {
return false;
};
matches!(
row.get_field(MODERATION_STATUS_FIELD),
Some(Value::Text(status))
if status.as_ref() == MODERATION_STATUS_PENDING
|| status.as_ref() == MODERATION_STATUS_REJECTED
)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModerationOutcome {
Allow,
Reject { categories: Vec<String> },
ProviderDown { reason: String },
}
impl ModerationOutcome {
pub fn is_reject(&self) -> bool {
matches!(self, Self::Reject { .. })
}
pub fn is_provider_down(&self) -> bool {
matches!(self, Self::ProviderDown { .. })
}
}
#[derive(Debug, Clone)]
pub struct ModerationRequest {
pub model: String,
pub text: String,
}
pub trait LocalModerationBackend: Send + Sync {
fn moderate(&self, request: &ModerationRequest) -> RedDBResult<ModerationOutcome>;
}
const LOCAL_MODERATION_DISABLED_MESSAGE: &str =
"local moderation requires a backend installed via \
runtime::ai::moderation::install_local_moderation_backend. Alternatively, \
declare a moderation-capable remote provider in the collection's MODERATE \
policy.";
type BackendSlot = Arc<dyn LocalModerationBackend>;
fn backend_slot() -> &'static RwLock<Option<BackendSlot>> {
static SLOT: OnceLock<RwLock<Option<BackendSlot>>> = OnceLock::new();
SLOT.get_or_init(|| RwLock::new(None))
}
pub fn install_local_moderation_backend(backend: Arc<dyn LocalModerationBackend>) {
let mut guard = backend_slot()
.write()
.expect("moderation backend slot poisoned");
*guard = Some(backend);
}
#[doc(hidden)]
pub fn clear_local_moderation_backend_for_tests() {
let mut guard = backend_slot()
.write()
.expect("moderation backend slot poisoned");
*guard = None;
}
fn current_backend() -> Option<BackendSlot> {
backend_slot()
.read()
.expect("moderation backend slot poisoned")
.as_ref()
.map(Arc::clone)
}
pub fn moderate_local(model: &str, text: String) -> RedDBResult<ModerationOutcome> {
let backend = current_backend().ok_or_else(|| {
RedDBError::FeatureNotEnabled(LOCAL_MODERATION_DISABLED_MESSAGE.to_string())
})?;
backend.moderate(&ModerationRequest {
model: model.to_string(),
text,
})
}
#[cfg(test)]
mod tests {
use super::*;
struct AllowAll;
impl LocalModerationBackend for AllowAll {
fn moderate(&self, _request: &ModerationRequest) -> RedDBResult<ModerationOutcome> {
Ok(ModerationOutcome::Allow)
}
}
#[test]
fn missing_backend_is_feature_disabled_error() {
clear_local_moderation_backend_for_tests();
let err = moderate_local("m", "hello".to_string()).expect_err("no backend");
assert!(matches!(err, RedDBError::FeatureNotEnabled(_)));
}
#[test]
fn installed_backend_drives_outcome() {
install_local_moderation_backend(Arc::new(AllowAll));
let outcome = moderate_local("m", "hello".to_string()).expect("allowed");
assert_eq!(outcome, ModerationOutcome::Allow);
clear_local_moderation_backend_for_tests();
}
#[test]
fn outcome_predicates() {
assert!(ModerationOutcome::Reject {
categories: vec!["hate".to_string()],
}
.is_reject());
assert!(ModerationOutcome::ProviderDown {
reason: "timeout".to_string(),
}
.is_provider_down());
assert!(!ModerationOutcome::Allow.is_reject());
}
}