psyche_subtitle_toolkit/translation/
openrouter.rs

1use serde::{Deserialize, Serialize};
2
3use crate::error::{Result, SubtitleToolkitError};
4
5use super::{TranslationRequest, Translator};
6
7/// Translator backend that calls the [OpenRouter](https://openrouter.ai) `/api/v1/chat/completions` endpoint.
8///
9/// OpenRouter provides an OpenAI-compatible API to 400+ models, including
10/// free models (no credit card required). Model slugs use the format
11/// `provider/model-name` (e.g. `meta-llama/llama-3.3-70b-instruct:free`).
12///
13/// # Example
14///
15/// ```no_run
16/// # async fn example() -> psyche_subtitle_toolkit::Result<()> {
17/// use psyche_subtitle_toolkit::OpenRouterTranslator;
18///
19/// let translator = OpenRouterTranslator::new("your-api-key", "meta-llama/llama-3.3-70b-instruct:free")?;
20/// // let result = translator.translate(request).await?;
21/// # Ok(())
22/// # }
23/// ```
24#[derive(Debug, Clone)]
25pub struct OpenRouterTranslator {
26    client: reqwest::Client,
27    base_url: String,
28    api_key: String,
29    model: String,
30}
31
32impl OpenRouterTranslator {
33    /// Create a new translator targeting the default OpenRouter API (`https://openrouter.ai/api`).
34    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Result<Self> {
35        Self::with_base_url("https://openrouter.ai/api", api_key, model)
36    }
37
38    /// Create a new translator with a custom base URL, API key, and model.
39    pub fn with_base_url(
40        base_url: impl Into<String>,
41        api_key: impl Into<String>,
42        model: impl Into<String>,
43    ) -> Result<Self> {
44        let client = reqwest::Client::builder()
45            .timeout(std::time::Duration::from_secs(120))
46            .build()
47            .map_err(SubtitleToolkitError::Http)?;
48        Ok(Self {
49            client,
50            base_url: base_url.into().trim_end_matches('/').to_string(),
51            api_key: api_key.into(),
52            model: model.into(),
53        })
54    }
55}
56
57#[async_trait::async_trait]
58impl Translator for OpenRouterTranslator {
59    async fn translate(&self, request: TranslationRequest<'_>) -> Result<String> {
60        let messages = build_messages(&request);
61
62        let response = self
63            .client
64            .post(format!("{}/v1/chat/completions", self.base_url))
65            .header("Authorization", format!("Bearer {}", self.api_key))
66            .header("HTTP-Referer", "https://github.com/Gitlawb/psyche-subtitle-toolkit")
67            .header("X-Title", "psyche-subtitle-toolkit")
68            .json(&ChatCompletionRequest {
69                model: &self.model,
70                messages,
71                stream: false,
72            })
73            .send()
74            .await?;
75
76        if !response.status().is_success() {
77            return Err(SubtitleToolkitError::Translation {
78                provider: "openrouter",
79                message: response
80                    .text()
81                    .await
82                    .unwrap_or_else(|_| "request failed".into()),
83            });
84        }
85
86        let body = response.json::<ChatCompletionResponse>().await?;
87        let content = body
88            .choices
89            .first()
90            .ok_or_else(|| SubtitleToolkitError::Translation {
91                provider: "openrouter",
92                message: "response contained no choices".into(),
93            })?
94            .message
95            .content
96            .trim()
97            .to_string();
98
99        Ok(content)
100    }
101}
102
103fn build_messages(request: &TranslationRequest<'_>) -> Vec<ChatMessage> {
104    let source_hint = match request.source_language {
105        Some(lang) => format!("The source language is {lang}.\n"),
106        None => String::new(),
107    };
108    vec![
109        ChatMessage {
110            role: "system",
111            content: "You are a subtitle translator. You translate subtitle dialogue while \
112                      preserving numbered tags exactly. Return only the translated lines. \
113                      Do not add explanations, markdown, notes, or code fences. \
114                      Do not add curly-brace commands or backslash formatting."
115                .to_string(),
116        },
117        ChatMessage {
118            role: "user",
119            content: format!(
120                "Translate the following subtitle dialogue to {target_language}.\n\
121                 {source_hint}\n\
122                 Preserve every numeric tag exactly, like <1>, <2>, <3>.\n\
123                 Keep line breaks inside each subtitle when needed.\n\n\
124                 Subtitle dialogue:\n\
125                 {source_text}",
126                target_language = request.target_language,
127                source_text = request.source_text,
128            ),
129        },
130    ]
131}
132
133#[derive(Debug, Serialize)]
134struct ChatMessage {
135    role: &'static str,
136    content: String,
137}
138
139#[derive(Debug, Serialize)]
140struct ChatCompletionRequest<'a> {
141    model: &'a str,
142    messages: Vec<ChatMessage>,
143    stream: bool,
144}
145
146#[derive(Debug, Deserialize)]
147struct ChatCompletionResponse {
148    choices: Vec<ChatChoice>,
149}
150
151#[derive(Debug, Deserialize)]
152struct ChatChoice {
153    message: ChatChoiceMessage,
154}
155
156#[derive(Debug, Deserialize)]
157struct ChatChoiceMessage {
158    content: String,
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use wiremock::matchers::{header, method, path};
165    use wiremock::{Mock, MockServer, ResponseTemplate};
166
167    #[tokio::test]
168    async fn translates_numbered_text() {
169        let server = MockServer::start().await;
170
171        Mock::given(method("POST"))
172            .and(path("/v1/chat/completions"))
173            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
174                "choices": [{
175                    "message": { "content": "<1> Olá\n<2> mundo" }
176                }]
177            })))
178            .mount(&server)
179            .await;
180
181        let translator = OpenRouterTranslator::with_base_url(
182            server.uri(),
183            "sk-test",
184            "meta-llama/llama-3.3-70b-instruct:free",
185        )
186        .unwrap();
187        let result = translator
188            .translate(TranslationRequest {
189                source_text: "<1> hello\n<2> world",
190                target_language: "pt-BR",
191            source_language: None,
192            })
193            .await
194            .unwrap();
195
196        assert_eq!(result, "<1> Olá\n<2> mundo");
197    }
198
199    #[tokio::test]
200    async fn sends_bearer_auth_and_attribution_headers() {
201        let server = MockServer::start().await;
202
203        Mock::given(method("POST"))
204            .and(header("Authorization", "Bearer sk-my-key"))
205            .and(header("HTTP-Referer", "https://github.com/Gitlawb/psyche-subtitle-toolkit"))
206            .and(header("X-Title", "psyche-subtitle-toolkit"))
207            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
208                "choices": [{ "message": { "content": "<1> ok" } }]
209            })))
210            .mount(&server)
211            .await;
212
213        let translator =
214            OpenRouterTranslator::with_base_url(server.uri(), "sk-my-key", "some/model").unwrap();
215        translator
216            .translate(TranslationRequest {
217                source_text: "<1> test",
218                target_language: "en",
219            source_language: None,
220            })
221            .await
222            .unwrap();
223    }
224
225    #[tokio::test]
226    async fn error_on_non_200() {
227        let server = MockServer::start().await;
228
229        Mock::given(method("POST"))
230            .respond_with(
231                ResponseTemplate::new(401).set_body_string(r#"{"error": "invalid api key"}"#),
232            )
233            .mount(&server)
234            .await;
235
236        let translator =
237            OpenRouterTranslator::with_base_url(server.uri(), "sk-bad", "some/model").unwrap();
238        let err = translator
239            .translate(TranslationRequest {
240                source_text: "<1> hello",
241                target_language: "en",
242            source_language: None,
243            })
244            .await
245            .unwrap_err();
246
247        assert!(err.to_string().contains("openrouter"));
248    }
249
250    #[tokio::test]
251    async fn error_on_empty_choices() {
252        let server = MockServer::start().await;
253
254        Mock::given(method("POST"))
255            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
256                "choices": []
257            })))
258            .mount(&server)
259            .await;
260
261        let translator =
262            OpenRouterTranslator::with_base_url(server.uri(), "sk-test", "some/model").unwrap();
263        let err = translator
264            .translate(TranslationRequest {
265                source_text: "<1> hello",
266                target_language: "en",
267            source_language: None,
268            })
269            .await
270            .unwrap_err();
271
272        assert!(err.to_string().contains("no choices"));
273    }
274}
psyche_subtitle_toolkit/translation/openrouter.rs

psyche_subtitle_toolkit/translation/
openrouter.rs
