aurum_core/cleanup/
openrouter.rs

1//! LLM-assisted cleanup via OpenRouter chat completions.
2
3use super::{CleanupProviderKind, CleanupResult, CleanupStyle, TextCleanup};
4use crate::error::{ProviderError, Result, UserError};
5use crate::postprocess::truncate_chars;
6use async_trait::async_trait;
7use serde::Deserialize;
8use serde_json::json;
9use std::time::Duration;
10
11const DEFAULT_BASE_URL: &str = "https://openrouter.ai/api/v1";
12const DEFAULT_MODEL: &str = "google/gemini-2.5-flash";
13
14/// OpenRouter-backed text cleanup (explicit opt-in; not local-first).
15pub struct OpenRouterCleanup {
16    api_key: String,
17    base_url: String,
18    model: String,
19    http: reqwest::Client,
20}
21
22impl OpenRouterCleanup {
23    pub fn new(
24        api_key: Option<String>,
25        base_url: Option<String>,
26        model: Option<String>,
27    ) -> Result<Self> {
28        let api_key = api_key
29            .map(|s| s.trim().to_string())
30            .filter(|s| !s.is_empty())
31            .ok_or(UserError::MissingApiKey)?;
32        let base_url = base_url
33            .map(|s| s.trim_end_matches('/').to_string())
34            .filter(|s| !s.is_empty())
35            .unwrap_or_else(|| DEFAULT_BASE_URL.to_string());
36        let model = model
37            .filter(|s| !s.trim().is_empty())
38            .unwrap_or_else(|| DEFAULT_MODEL.to_string());
39        let http = reqwest::Client::builder()
40            .user_agent(concat!("aurum-core/", env!("CARGO_PKG_VERSION")))
41            .timeout(Duration::from_secs(120))
42            .build()
43            .map_err(|e| ProviderError::Network {
44                provider: "openrouter-cleanup".into(),
45                reason: e.to_string(),
46            })?;
47        Ok(Self {
48            api_key,
49            base_url,
50            model,
51            http,
52        })
53    }
54
55    fn prompt(style: CleanupStyle, text: &str) -> String {
56        let instruction = match style {
57            CleanupStyle::Raw => "Return the text unchanged.",
58            CleanupStyle::Clean => {
59                "Clean up this speech transcript: remove filler words (um, uh, you know), \
60                 fix spacing/punctuation, keep meaning verbatim. Reply with ONLY the cleaned text."
61            }
62            CleanupStyle::Bullets => {
63                "Turn this transcript into a concise bullet list. One bullet per idea. \
64                 Reply with ONLY the bullets, each starting with • ."
65            }
66            CleanupStyle::Professional => {
67                "Rewrite this transcript in clear professional prose. Expand casualisms, \
68                 keep facts. Reply with ONLY the rewritten text."
69            }
70            CleanupStyle::Summary => {
71                "Summarize this transcript in 1-3 short sentences. Reply with ONLY the summary."
72            }
73        };
74        format!("{instruction}\n\n---\n{text}\n---")
75    }
76}
77
78#[async_trait]
79impl TextCleanup for OpenRouterCleanup {
80    fn name(&self) -> &'static str {
81        "openrouter"
82    }
83
84    fn kind(&self) -> CleanupProviderKind {
85        CleanupProviderKind::OpenRouter
86    }
87
88    async fn cleanup(&self, text: &str, style: CleanupStyle) -> Result<CleanupResult> {
89        let original = text.to_string();
90        if text.trim().is_empty() || matches!(style, CleanupStyle::Raw) {
91            return Ok(CleanupResult {
92                text: text.trim().to_string(),
93                style,
94                provider: CleanupProviderKind::OpenRouter,
95                original_text: original,
96            });
97        }
98
99        let body = json!({
100            "model": self.model,
101            "messages": [
102                {"role": "user", "content": Self::prompt(style, text)}
103            ],
104            "temperature": 0.2,
105        });
106
107        let url = format!("{}/chat/completions", self.base_url);
108        let response = self
109            .http
110            .post(&url)
111            .header("Authorization", format!("Bearer {}", self.api_key))
112            .header("Content-Type", "application/json")
113            .header("HTTP-Referer", "https://github.com/joe-broadhead/aurum")
114            .header("X-Title", "Aurum Cleanup")
115            .json(&body)
116            .send()
117            .await
118            .map_err(|e| ProviderError::Network {
119                provider: "openrouter-cleanup".into(),
120                reason: e.to_string(),
121            })?;
122
123        let status = response.status();
124        let body_text = response.text().await.map_err(|e| ProviderError::Network {
125            provider: "openrouter-cleanup".into(),
126            reason: e.to_string(),
127        })?;
128
129        if status.as_u16() == 401 || status.as_u16() == 403 {
130            return Err(ProviderError::Auth {
131                provider: "openrouter-cleanup".into(),
132                reason: truncate_chars(&body_text, 300),
133            }
134            .into());
135        }
136        if status.as_u16() == 429 {
137            return Err(ProviderError::RateLimited {
138                provider: "openrouter-cleanup".into(),
139            }
140            .into());
141        }
142        if !status.is_success() {
143            return Err(ProviderError::Remote {
144                provider: "openrouter-cleanup".into(),
145                reason: format!("HTTP {status}: {}", truncate_chars(&body_text, 400)),
146            }
147            .into());
148        }
149
150        let parsed: ChatResponse =
151            serde_json::from_str(&body_text).map_err(|e| ProviderError::Remote {
152                provider: "openrouter-cleanup".into(),
153                reason: format!("invalid JSON: {e}"),
154            })?;
155        let content = parsed
156            .choices
157            .first()
158            .and_then(|c| c.message.content.as_deref())
159            .unwrap_or("")
160            .trim()
161            .to_string();
162
163        if content.is_empty() {
164            return Err(ProviderError::TranscriptionFailed {
165                reason: "cleanup model returned empty text".into(),
166            }
167            .into());
168        }
169
170        Ok(CleanupResult {
171            text: content,
172            style,
173            provider: CleanupProviderKind::OpenRouter,
174            original_text: original,
175        })
176    }
177}
178
179#[derive(Debug, Deserialize)]
180struct ChatResponse {
181    choices: Vec<Choice>,
182}
183#[derive(Debug, Deserialize)]
184struct Choice {
185    message: Msg,
186}
187#[derive(Debug, Deserialize)]
188struct Msg {
189    content: Option<String>,
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use serde_json::json;
196    use wiremock::matchers::{method, path};
197    use wiremock::{Mock, MockServer, ResponseTemplate};
198
199    #[tokio::test]
200    async fn missing_key() {
201        assert!(OpenRouterCleanup::new(None, None, None).is_err());
202    }
203
204    #[tokio::test]
205    async fn cleans_via_mock() {
206        let server = MockServer::start().await;
207        Mock::given(method("POST"))
208            .and(path("/chat/completions"))
209            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
210                "choices": [{"message": {"content": "Hello world."}}]
211            })))
212            .mount(&server)
213            .await;
214
215        let c = OpenRouterCleanup::new(
216            Some("k".into()),
217            Some(server.uri()),
218            Some("test/model".into()),
219        )
220        .unwrap();
221        let out = c
222            .cleanup("um hello world", CleanupStyle::Clean)
223            .await
224            .unwrap();
225        assert_eq!(out.text, "Hello world.");
226        assert_eq!(out.provider, CleanupProviderKind::OpenRouter);
227    }
228}
aurum_core/cleanup/openrouter.rs

aurum_core/cleanup/
openrouter.rs
