aither_gemini/
moderation.rs

1use aither_core::moderation::{Moderation, ModerationCategory, ModerationResult};
2
3use crate::{
4    client::call_generate,
5    config::Gemini,
6    error::GeminiError,
7    types::{
8        Candidate, GeminiContent, GenerateContentRequest, GenerationConfig, SafetyRating,
9        SafetySetting,
10    },
11};
12
13impl Moderation for Gemini {
14    type Error = GeminiError;
15
16    fn moderate(
17        &self,
18        content: &str,
19    ) -> impl core::future::Future<Output = Result<ModerationResult, Self::Error>> + Send {
20        let cfg = self.config();
21        let text = content.to_owned();
22        async move {
23            let request = GenerateContentRequest {
24                system_instruction: None,
25                contents: vec![GeminiContent::text("user", text)],
26                generation_config: Some(GenerationConfig {
27                    response_mime_type: Some("text/plain".into()),
28                    ..GenerationConfig::default()
29                }),
30                tools: Vec::new(),
31                tool_config: None,
32                safety_settings: default_safety_settings(),
33                cached_content: None,
34            };
35            let model_id = cfg.text_model.clone();
36            let response = call_generate(cfg, &model_id, request).await?;
37            let mut flagged = false;
38            let mut categories = Vec::new();
39            if let Some(feedback) = response.prompt_feedback {
40                for rating in feedback.safety_ratings {
41                    if let Some(cat) = to_moderation_category(&rating) {
42                        flagged |= rating.blocked.unwrap_or(false);
43                        categories.push(cat);
44                    }
45                }
46                flagged |= feedback.block_reason.is_some();
47            }
48            for candidate in response.candidates {
49                categories.extend(ratings_to_categories(&candidate));
50                flagged |= candidate
51                    .safety_ratings
52                    .iter()
53                    .any(|rating| rating.blocked.unwrap_or(false));
54            }
55            Ok(ModerationResult::new(flagged, categories))
56        }
57    }
58}
59
60fn default_safety_settings() -> Vec<SafetySetting> {
61    vec![
62        SafetySetting::new("HARM_CATEGORY_HATE_SPEECH"),
63        SafetySetting::new("HARM_CATEGORY_SEXUALLY_EXPLICIT"),
64        SafetySetting::new("HARM_CATEGORY_DANGEROUS_CONTENT"),
65        SafetySetting::new("HARM_CATEGORY_HARASSMENT"),
66        SafetySetting::new("HARM_CATEGORY_CIVIC_INTEGRITY"),
67    ]
68}
69
70fn ratings_to_categories(candidate: &Candidate) -> Vec<ModerationCategory> {
71    candidate
72        .safety_ratings
73        .iter()
74        .filter_map(to_moderation_category)
75        .collect()
76}
77
78fn to_moderation_category(rating: &SafetyRating) -> Option<ModerationCategory> {
79    let score = probability_to_score(rating.probability.as_deref());
80    match rating.category.as_str() {
81        "HARM_CATEGORY_HATE_SPEECH" => Some(ModerationCategory::Hate { score }),
82        "HARM_CATEGORY_HARASSMENT" => Some(ModerationCategory::Harassment { score }),
83        "HARM_CATEGORY_SEXUALLY_EXPLICIT" => Some(ModerationCategory::Sexual { score }),
84        "HARM_CATEGORY_DANGEROUS_CONTENT" => Some(ModerationCategory::Violence { score }),
85        _ => None,
86    }
87}
88
89fn probability_to_score(label: Option<&str>) -> f32 {
90    match label {
91        Some("NEGLIGIBLE") => 0.0,
92        Some("VERY_LOW") => 0.1,
93        Some("LOW") => 0.3,
94        Some("MEDIUM") => 0.6,
95        Some("HIGH") => 0.9,
96        _ => 0.5,
97    }
98}
aither_gemini/moderation.rs

aither_gemini/
moderation.rs
