aither_openai/
moderation.rs

1use crate::{
2    client::{Config, OpenAI},
3    error::OpenAIError,
4};
5use aither_core::moderation::{Moderation, ModerationCategory, ModerationResult};
6use serde::{Deserialize, Serialize};
7use std::{collections::HashMap, future::Future, sync::Arc};
8use zenwave::{Client, client, header};
9
10impl Moderation for OpenAI {
11    type Error = OpenAIError;
12
13    fn moderate(
14        &self,
15        content: &str,
16    ) -> impl Future<Output = Result<ModerationResult, Self::Error>> + Send {
17        let cfg = self.config();
18        let text = content.to_owned();
19        async move { moderate_once(cfg, text).await }
20    }
21}
22
23async fn moderate_once(cfg: Arc<Config>, content: String) -> Result<ModerationResult, OpenAIError> {
24    let endpoint = cfg.request_url("/moderations");
25    let mut backend = client();
26    let mut builder = backend
27        .post(endpoint)
28        .map_err(OpenAIError::Http)?
29        .header(header::AUTHORIZATION.as_str(), cfg.request_auth())
30        .map_err(OpenAIError::Http)?
31        .header(header::USER_AGENT.as_str(), "aither-openai/0.1")
32        .map_err(OpenAIError::Http)?;
33    if let Some(org) = &cfg.organization {
34        builder = builder
35            .header("OpenAI-Organization", org.clone())
36            .map_err(OpenAIError::Http)?;
37    }
38
39    let request = ModerationRequest {
40        model: &cfg.moderation_model,
41        input: &content,
42    };
43    let response: ModerationResponse = builder
44        .json_body(&request)
45        .map_err(OpenAIError::Http)?
46        .json()
47        .await
48        .map_err(OpenAIError::Http)?;
49
50    response.into_result()
51}
52
53#[derive(Debug, Serialize)]
54struct ModerationRequest<'a> {
55    model: &'a str,
56    input: &'a str,
57}
58
59#[derive(Debug, Deserialize)]
60struct ModerationResponse {
61    results: Vec<ModerationPayload>,
62}
63
64impl ModerationResponse {
65    fn into_result(self) -> Result<ModerationResult, OpenAIError> {
66        let payload = self
67            .results
68            .into_iter()
69            .next()
70            .ok_or_else(|| OpenAIError::Api("moderation response missing results".into()))?;
71        let categories = collect_categories(&payload);
72        Ok(ModerationResult::new(payload.flagged, categories))
73    }
74}
75
76#[derive(Debug, Deserialize)]
77struct ModerationPayload {
78    flagged: bool,
79    #[serde(default)]
80    categories: HashMap<String, bool>,
81    #[serde(default)]
82    category_scores: HashMap<String, f32>,
83}
84
85fn collect_categories(payload: &ModerationPayload) -> Vec<ModerationCategory> {
86    let mut categories = Vec::new();
87    if let Some(score) = select_score(payload, "hate") {
88        categories.push(ModerationCategory::Hate { score });
89    }
90    if let Some(score) = select_score(payload, "hate/threatening") {
91        categories.push(ModerationCategory::HateThreatening { score });
92    }
93    if let Some(score) = select_score(payload, "harassment") {
94        categories.push(ModerationCategory::Harassment { score });
95    }
96    if let Some(score) = select_score(payload, "harassment/threatening") {
97        categories.push(ModerationCategory::HarassmentThreatening { score });
98    }
99    if let Some(score) = select_score(payload, "self-harm") {
100        categories.push(ModerationCategory::SelfHarm { score });
101    }
102    if let Some(score) = select_score(payload, "self-harm/intent") {
103        categories.push(ModerationCategory::SelfHarmIntent { score });
104    }
105    if let Some(score) = select_score(payload, "self-harm/instructions") {
106        categories.push(ModerationCategory::SelfHarmInstructions { score });
107    }
108    if let Some(score) = select_score(payload, "sexual") {
109        categories.push(ModerationCategory::Sexual { score });
110    }
111    if let Some(score) = select_score(payload, "sexual/minors") {
112        categories.push(ModerationCategory::SexualMinors { score });
113    }
114    if let Some(score) = select_score(payload, "violence") {
115        categories.push(ModerationCategory::Violence { score });
116    }
117    if let Some(score) = select_score(payload, "violence/graphic") {
118        categories.push(ModerationCategory::ViolenceGraphic { score });
119    }
120    if let Some(score) = select_score(payload, "illicit") {
121        categories.push(ModerationCategory::Illicit { score });
122    }
123    if let Some(score) = select_score(payload, "illicit/violent") {
124        categories.push(ModerationCategory::IllicitViolent { score });
125    }
126    categories
127}
128
129fn select_score(payload: &ModerationPayload, name: &str) -> Option<f32> {
130    let flagged = payload.categories.get(name).copied().unwrap_or(false);
131    let score = payload.category_scores.get(name).copied();
132    match (flagged, score) {
133        (true, Some(score)) => Some(score),
134        (true, None) => Some(1.0),
135        (false, Some(score)) if score > 0.0 => Some(score),
136        _ => None,
137    }
138}
aither_openai/moderation.rs

aither_openai/
moderation.rs
