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    vec![
105        ChatMessage {
106            role: "system",
107            content: "You are a subtitle translator. You translate subtitle dialogue while \
108                      preserving numbered tags exactly. Return only the translated lines. \
109                      Do not add explanations, markdown, notes, or code fences. \
110                      Do not add curly-brace commands or backslash formatting."
111                .to_string(),
112        },
113        ChatMessage {
114            role: "user",
115            content: format!(
116                "Translate the following subtitle dialogue to {target_language}.\n\n\
117                 Preserve every numeric tag exactly, like <1>, <2>, <3>.\n\
118                 Keep line breaks inside each subtitle when needed.\n\n\
119                 Subtitle dialogue:\n\
120                 {source_text}",
121                target_language = request.target_language,
122                source_text = request.source_text,
123            ),
124        },
125    ]
126}
127
128#[derive(Debug, Serialize)]
129struct ChatMessage {
130    role: &'static str,
131    content: String,
132}
133
134#[derive(Debug, Serialize)]
135struct ChatCompletionRequest<'a> {
136    model: &'a str,
137    messages: Vec<ChatMessage>,
138    stream: bool,
139}
140
141#[derive(Debug, Deserialize)]
142struct ChatCompletionResponse {
143    choices: Vec<ChatChoice>,
144}
145
146#[derive(Debug, Deserialize)]
147struct ChatChoice {
148    message: ChatChoiceMessage,
149}
150
151#[derive(Debug, Deserialize)]
152struct ChatChoiceMessage {
153    content: String,
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use wiremock::matchers::{header, method, path};
160    use wiremock::{Mock, MockServer, ResponseTemplate};
161
162    #[tokio::test]
163    async fn translates_numbered_text() {
164        let server = MockServer::start().await;
165
166        Mock::given(method("POST"))
167            .and(path("/v1/chat/completions"))
168            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
169                "choices": [{
170                    "message": { "content": "<1> Olá\n<2> mundo" }
171                }]
172            })))
173            .mount(&server)
174            .await;
175
176        let translator = OpenRouterTranslator::with_base_url(
177            server.uri(),
178            "sk-test",
179            "meta-llama/llama-3.3-70b-instruct:free",
180        )
181        .unwrap();
182        let result = translator
183            .translate(TranslationRequest {
184                source_text: "<1> hello\n<2> world",
185                target_language: "pt-BR",
186            })
187            .await
188            .unwrap();
189
190        assert_eq!(result, "<1> Olá\n<2> mundo");
191    }
192
193    #[tokio::test]
194    async fn sends_bearer_auth_and_attribution_headers() {
195        let server = MockServer::start().await;
196
197        Mock::given(method("POST"))
198            .and(header("Authorization", "Bearer sk-my-key"))
199            .and(header("HTTP-Referer", "https://github.com/Gitlawb/psyche-subtitle-toolkit"))
200            .and(header("X-Title", "psyche-subtitle-toolkit"))
201            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
202                "choices": [{ "message": { "content": "<1> ok" } }]
203            })))
204            .mount(&server)
205            .await;
206
207        let translator =
208            OpenRouterTranslator::with_base_url(server.uri(), "sk-my-key", "some/model").unwrap();
209        translator
210            .translate(TranslationRequest {
211                source_text: "<1> test",
212                target_language: "en",
213            })
214            .await
215            .unwrap();
216    }
217
218    #[tokio::test]
219    async fn error_on_non_200() {
220        let server = MockServer::start().await;
221
222        Mock::given(method("POST"))
223            .respond_with(
224                ResponseTemplate::new(401).set_body_string(r#"{"error": "invalid api key"}"#),
225            )
226            .mount(&server)
227            .await;
228
229        let translator =
230            OpenRouterTranslator::with_base_url(server.uri(), "sk-bad", "some/model").unwrap();
231        let err = translator
232            .translate(TranslationRequest {
233                source_text: "<1> hello",
234                target_language: "en",
235            })
236            .await
237            .unwrap_err();
238
239        assert!(err.to_string().contains("openrouter"));
240    }
241
242    #[tokio::test]
243    async fn error_on_empty_choices() {
244        let server = MockServer::start().await;
245
246        Mock::given(method("POST"))
247            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
248                "choices": []
249            })))
250            .mount(&server)
251            .await;
252
253        let translator =
254            OpenRouterTranslator::with_base_url(server.uri(), "sk-test", "some/model").unwrap();
255        let err = translator
256            .translate(TranslationRequest {
257                source_text: "<1> hello",
258                target_language: "en",
259            })
260            .await
261            .unwrap_err();
262
263        assert!(err.to_string().contains("no choices"));
264    }
265}
psyche_subtitle_toolkit/translation/openrouter.rs

psyche_subtitle_toolkit/translation/
openrouter.rs
