subx_core/services/ai/
azure_openai.rs

1use crate::error::SubXError;
2use crate::services::ai::hosted_hint::{append_local_hint, maybe_attach_local_hint};
3use crate::services::ai::prompts::{PromptBuilder, ResponseParser};
4use crate::services::ai::retry::HttpRetryClient;
5use crate::services::ai::{
6    AIProvider, AnalysisRequest, ConfidenceScore, MatchResult, VerificationRequest,
7};
8use async_trait::async_trait;
9use reqwest::Client;
10use serde_json::{Value, json};
11use std::time::Duration;
12use tokio::time;
13use url::{ParseError, Url};
14
15/// Azure OpenAI client implementation
16pub struct AzureOpenAIClient {
17    client: Client,
18    api_key: String,
19    model: String,
20    base_url: String,
21    api_version: String,
22    temperature: f32,
23    max_tokens: u32,
24    retry_attempts: u32,
25    retry_delay_ms: u64,
26    request_timeout_seconds: u64,
27    reporter: std::sync::Arc<dyn crate::core::report::Reporter>,
28}
29
30impl std::fmt::Debug for AzureOpenAIClient {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("AzureOpenAIClient")
33            .field("client", &self.client)
34            .field("api_key", &"[REDACTED]")
35            .field("model", &self.model)
36            .field("base_url", &self.base_url)
37            .field("api_version", &self.api_version)
38            .field("temperature", &self.temperature)
39            .field("max_tokens", &self.max_tokens)
40            .field("retry_attempts", &self.retry_attempts)
41            .field("retry_delay_ms", &self.retry_delay_ms)
42            .field("request_timeout_seconds", &self.request_timeout_seconds)
43            .finish()
44    }
45}
46
47const DEFAULT_AZURE_API_VERSION: &str = "2025-04-01-preview";
48
49impl AzureOpenAIClient {
50    /// Create a new AzureOpenAIClient with full configuration
51    #[allow(clippy::too_many_arguments)]
52    pub fn new_with_all(
53        api_key: String,
54        model: String,
55        base_url: String,
56        api_version: String,
57        temperature: f32,
58        max_tokens: u32,
59        retry_attempts: u32,
60        retry_delay_ms: u64,
61        request_timeout_seconds: u64,
62    ) -> Self {
63        let client = Client::builder()
64            .timeout(Duration::from_secs(request_timeout_seconds))
65            .build()
66            .expect("Failed to create HTTP client");
67        AzureOpenAIClient {
68            client,
69            api_key,
70            model,
71            base_url: base_url.trim_end_matches('/').to_string(),
72            api_version,
73            temperature,
74            max_tokens,
75            retry_attempts,
76            retry_delay_ms,
77            request_timeout_seconds,
78            reporter: crate::core::report::noop(),
79        }
80    }
81
82    /// Attach a reporting sink, consuming and returning the client.
83    ///
84    /// # Arguments
85    ///
86    /// * `reporter` - Sink that receives [`crate::core::report::AiUsage`]
87    ///   after every successful API response.
88    pub fn with_reporter(
89        mut self,
90        reporter: std::sync::Arc<dyn crate::core::report::Reporter>,
91    ) -> Self {
92        self.reporter = reporter;
93        self
94    }
95
96    /// Create client from AIConfig
97    pub fn from_config(config: &crate::config::AIConfig) -> crate::Result<Self> {
98        let api_key = config
99            .api_key
100            .as_ref()
101            .filter(|key| !key.trim().is_empty())
102            .ok_or_else(|| SubXError::config("Missing Azure OpenAI API Key".to_string()))?
103            .clone();
104        // Use the model value as the deployment identifier; ensure it's provided
105        let deployment_name = config.model.clone();
106        if deployment_name.trim().is_empty() {
107            return Err(SubXError::config(
108                "Missing Azure OpenAI deployment name in model field".to_string(),
109            ));
110        }
111        let api_version = config
112            .api_version
113            .clone()
114            .unwrap_or_else(|| DEFAULT_AZURE_API_VERSION.to_string());
115
116        // Validate base URL format, handle missing host specially
117        let parsed = match Url::parse(&config.base_url) {
118            Ok(u) => u,
119            Err(ParseError::EmptyHost) => {
120                return Err(SubXError::config(
121                    "Azure OpenAI endpoint missing host".to_string(),
122                ));
123            }
124            Err(e) => {
125                return Err(SubXError::config(format!(
126                    "Invalid Azure OpenAI endpoint: {}",
127                    e
128                )));
129            }
130        };
131        if !matches!(parsed.scheme(), "http" | "https") {
132            return Err(SubXError::config(
133                "Azure OpenAI endpoint must use http or https".to_string(),
134            ));
135        }
136        crate::services::ai::security::warn_on_insecure_http(&parsed, &api_key);
137
138        Ok(Self::new_with_all(
139            api_key,
140            config.model.clone(),
141            config.base_url.clone(),
142            api_version,
143            config.temperature,
144            config.max_tokens,
145            config.retry_attempts,
146            config.retry_delay_ms,
147            config.request_timeout_seconds,
148        ))
149    }
150
151    async fn make_request_with_retry(
152        &self,
153        request: reqwest::RequestBuilder,
154    ) -> crate::Result<reqwest::Response> {
155        let mut attempts = 0;
156        loop {
157            let cloned = request.try_clone().ok_or_else(|| {
158                crate::error::SubXError::AiService(
159                    "Request body cannot be cloned for retry".to_string(),
160                )
161            })?;
162            match cloned.send().await {
163                Ok(resp) => {
164                    if attempts > 0 {
165                        log::info!("Request succeeded after {} retry attempts", attempts);
166                    }
167                    return Ok(resp);
168                }
169                Err(e) if (attempts as u32) < self.retry_attempts => {
170                    attempts += 1;
171                    log::warn!(
172                        "Request attempt {} failed: {}. Retrying in {}ms...",
173                        attempts,
174                        e,
175                        self.retry_delay_ms
176                    );
177                    if e.is_timeout() {
178                        log::warn!(
179                            "This appears to be a timeout error. Consider increasing 'ai.request_timeout_seconds' in config."
180                        );
181                    }
182                    time::sleep(Duration::from_millis(self.retry_delay_ms)).await;
183                }
184                Err(e) => {
185                    log::error!(
186                        "Request failed after {} attempts. Final error: {}",
187                        attempts + 1,
188                        e
189                    );
190                    if e.is_timeout() {
191                        log::error!(
192                            "AI service error: Request timed out after multiple attempts. Try increasing 'ai.request_timeout_seconds' configuration."
193                        );
194                    } else if e.is_connect() {
195                        log::error!(
196                            "AI service error: Connection failed. Check network connection and Azure OpenAI endpoint settings."
197                        );
198                    }
199                    return Err(e.into());
200                }
201            }
202        }
203    }
204
205    /// Send a raw chat completion request to the Azure OpenAI Chat Completions API.
206    pub async fn chat_completion(&self, messages: Vec<Value>) -> crate::Result<String> {
207        let url = format!(
208            "{}/openai/deployments/{}/chat/completions?api-version={}",
209            self.base_url, self.model, self.api_version
210        );
211        let mut req = self
212            .client
213            .post(url)
214            .header("Content-Type", "application/json");
215        if self.api_key.to_lowercase().starts_with("bearer ") {
216            req = req.header("Authorization", self.api_key.clone());
217        } else {
218            req = req.header("api-key", self.api_key.clone());
219        }
220        let body = json!({
221            "messages": messages,
222            "temperature": self.temperature,
223            "max_tokens": self.max_tokens,
224            "stream": false
225        });
226        let request = req.json(&body);
227        let mut response = match self.make_request_with_retry(request).await {
228            Ok(r) => r,
229            Err(e) => return Err(maybe_attach_local_hint(e, &self.base_url)),
230        };
231
232        const MAX_AI_RESPONSE_BYTES: u64 = 10 * 1024 * 1024; // 10 MiB
233        if let Some(len) = response.content_length() {
234            if len > MAX_AI_RESPONSE_BYTES {
235                return Err(SubXError::AiService(format!(
236                    "AI response too large: {} bytes (limit: {} bytes)",
237                    len, MAX_AI_RESPONSE_BYTES
238                )));
239            }
240        }
241
242        if !response.status().is_success() {
243            let status = response.status();
244            let text = response.text().await?;
245            let safe_body = crate::services::ai::error_sanitizer::sanitize_url_in_error(
246                &crate::services::ai::error_sanitizer::truncate_error_body(
247                    &text,
248                    crate::services::ai::error_sanitizer::DEFAULT_ERROR_BODY_MAX_LEN,
249                ),
250            );
251            return Err(SubXError::AiService(format!(
252                "Azure OpenAI API error {}: {}",
253                status, safe_body
254            )));
255        }
256        // Bounded chunked read to guard against oversized responses when
257        // content_length() is not reported by the server.
258        let mut body = Vec::new();
259        while let Some(chunk) = response.chunk().await? {
260            body.extend_from_slice(&chunk);
261            if body.len() as u64 > MAX_AI_RESPONSE_BYTES {
262                return Err(SubXError::AiService(format!(
263                    "AI response too large: {} bytes read (limit: {} bytes)",
264                    body.len(),
265                    MAX_AI_RESPONSE_BYTES
266                )));
267            }
268        }
269        let resp_json: Value = serde_json::from_slice(&body)
270            .map_err(|e| SubXError::AiService(format!("Failed to parse AI response: {}", e)))?;
271        if let Some(usage) = resp_json.get("usage") {
272            if let (Some(p), Some(c), Some(t)) = (
273                usage.get("prompt_tokens").and_then(Value::as_u64),
274                usage.get("completion_tokens").and_then(Value::as_u64),
275                usage.get("total_tokens").and_then(Value::as_u64),
276            ) {
277                // Get model from response JSON, fallback to self.model if missing
278                let model = resp_json
279                    .get("model")
280                    .and_then(Value::as_str)
281                    .unwrap_or(self.model.as_str())
282                    .to_string();
283                let stats = crate::services::ai::AiUsageStats {
284                    model,
285                    prompt_tokens: p as u32,
286                    completion_tokens: c as u32,
287                    total_tokens: t as u32,
288                };
289                self.reporter.ai_usage(&stats);
290            }
291        }
292        let content = resp_json["choices"][0]["message"]["content"]
293            .as_str()
294            .ok_or_else(|| {
295                SubXError::AiService(append_local_hint("Invalid API response format"))
296            })?;
297        Ok(content.to_string())
298    }
299}
300
301impl PromptBuilder for AzureOpenAIClient {}
302impl ResponseParser for AzureOpenAIClient {}
303impl HttpRetryClient for AzureOpenAIClient {
304    fn retry_attempts(&self) -> u32 {
305        self.retry_attempts
306    }
307
308    fn retry_delay_ms(&self) -> u64 {
309        self.retry_delay_ms
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::config::Config;
317
318    #[test]
319    fn test_azure_openai_from_config_and_url_construction() {
320        let mut config = Config::default();
321        config.ai.provider = "azure-openai".to_string();
322        config.ai.api_key = Some("test-api-key".to_string());
323        config.ai.model = "deployment-name".to_string();
324        config.ai.base_url = "https://example.openai.azure.com".to_string();
325        config.ai.api_version = Some("2025-04-01-preview".to_string());
326
327        let client = AzureOpenAIClient::from_config(&config.ai).unwrap();
328        let url = format!(
329            "{}/openai/deployments/{}/chat/completions?api-version={}",
330            client.base_url, client.model, client.api_version
331        );
332        assert!(url.contains("deployment-name"));
333    }
334
335    #[test]
336    fn test_missing_model_error() {
337        let mut config = Config::default();
338        config.ai.provider = "azure-openai".to_string();
339        config.ai.api_key = Some("test-api-key".to_string());
340        config.ai.model = "".to_string();
341        config.ai.base_url = "https://example.openai.azure.com".to_string();
342
343        let err = AzureOpenAIClient::from_config(&config.ai)
344            .unwrap_err()
345            .to_string();
346        assert!(err.contains("Missing Azure OpenAI deployment name in model field"));
347    }
348
349    #[test]
350    fn test_azure_openai_client_creation_with_defaults() {
351        let mut config = Config::default();
352        config.ai.provider = "azure-openai".to_string();
353        config.ai.api_key = Some("test-api-key".to_string());
354        config.ai.model = "deployment-name".to_string();
355        config.ai.base_url = "https://example.openai.azure.com".to_string();
356        // api_version defaults to DEFAULT_AZURE_API_VERSION
357
358        let client = AzureOpenAIClient::from_config(&config.ai).unwrap();
359        assert_eq!(
360            client.api_version,
361            super::DEFAULT_AZURE_API_VERSION.to_string()
362        );
363    }
364
365    #[test]
366    fn test_azure_openai_client_missing_api_key() {
367        let mut config = Config::default();
368        config.ai.provider = "azure-openai".to_string();
369        config.ai.api_key = None;
370        config.ai.model = "deployment-name".to_string();
371        config.ai.base_url = "https://example.openai.azure.com".to_string();
372
373        let result = AzureOpenAIClient::from_config(&config.ai);
374        let err = result.unwrap_err().to_string();
375        assert!(err.contains("Missing Azure OpenAI API Key"));
376    }
377
378    #[test]
379    fn test_azure_openai_client_invalid_base_url() {
380        let mut config = Config::default();
381        config.ai.provider = "azure-openai".to_string();
382        config.ai.api_key = Some("test-api-key".to_string());
383        config.ai.model = "deployment-name".to_string();
384        config.ai.base_url = "invalid-url".to_string();
385
386        let result = AzureOpenAIClient::from_config(&config.ai);
387        let err = result.unwrap_err().to_string();
388        assert!(err.contains("Invalid Azure OpenAI endpoint"));
389    }
390
391    #[test]
392    fn test_azure_openai_client_invalid_url_scheme() {
393        let mut config = Config::default();
394        config.ai.provider = "azure-openai".to_string();
395        config.ai.api_key = Some("test-api-key".to_string());
396        config.ai.model = "deployment-name".to_string();
397        config.ai.base_url = "ftp://example.openai.azure.com".to_string();
398
399        let result = AzureOpenAIClient::from_config(&config.ai);
400        let err = result.unwrap_err().to_string();
401        assert!(err.contains("must use http or https"));
402    }
403
404    #[test]
405    fn test_azure_openai_client_url_without_host() {
406        let mut config = Config::default();
407        config.ai.provider = "azure-openai".to_string();
408        config.ai.api_key = Some("test-api-key".to_string());
409        config.ai.model = "deployment-name".to_string();
410        config.ai.base_url = "https://".to_string();
411
412        let result = AzureOpenAIClient::from_config(&config.ai);
413        let err = result.unwrap_err().to_string();
414        assert!(err.contains("missing host"));
415    }
416
417    #[test]
418    fn test_azure_openai_with_custom_model_and_version() {
419        let mock_model = "custom-model-123";
420        let mock_version = "2023-12-01-preview";
421
422        let mut config = Config::default();
423        config.ai.provider = "azure-openai".to_string();
424        config.ai.api_key = Some("test-api-key".to_string());
425        config.ai.model = mock_model.to_string();
426        config.ai.base_url = "https://custom.openai.azure.com".to_string();
427        config.ai.api_version = Some(mock_version.to_string());
428
429        let client = AzureOpenAIClient::from_config(&config.ai).unwrap();
430        assert_eq!(client.model, mock_model);
431        assert_eq!(client.api_version, mock_version);
432    }
433
434    #[test]
435    fn test_azure_openai_with_trailing_slash_in_url() {
436        let mut config = Config::default();
437        config.ai.provider = "azure-openai".to_string();
438        config.ai.api_key = Some("test-api-key".to_string());
439        config.ai.model = "deployment-name".to_string();
440        config.ai.base_url = "https://example.openai.azure.com/".to_string(); // Trailing slash
441
442        let client = AzureOpenAIClient::from_config(&config.ai).unwrap();
443        assert_eq!(
444            client.base_url,
445            "https://example.openai.azure.com".to_string()
446        );
447    }
448
449    #[test]
450    fn test_azure_openai_with_custom_temperature_and_tokens() {
451        let mut config = Config::default();
452        config.ai.provider = "azure-openai".to_string();
453        config.ai.api_key = Some("test-api-key".to_string());
454        config.ai.model = "deployment-name".to_string();
455        config.ai.base_url = "https://example.openai.azure.com".to_string();
456        config.ai.temperature = 0.8;
457        config.ai.max_tokens = 2000;
458
459        let client = AzureOpenAIClient::from_config(&config.ai).unwrap();
460        assert!((client.temperature - 0.8).abs() < f32::EPSILON);
461        assert_eq!(client.max_tokens, 2000);
462    }
463
464    #[test]
465    fn test_azure_openai_with_custom_retry_and_timeout() {
466        let mut config = Config::default();
467        config.ai.provider = "azure-openai".to_string();
468        config.ai.api_key = Some("test-api-key".to_string());
469        config.ai.model = "deployment-name".to_string();
470        config.ai.base_url = "https://example.openai.azure.com".to_string();
471        config.ai.retry_attempts = 5;
472        config.ai.retry_delay_ms = 2000;
473        config.ai.request_timeout_seconds = 180;
474
475        let client = AzureOpenAIClient::from_config(&config.ai).unwrap();
476        assert_eq!(client.retry_attempts, 5);
477        assert_eq!(client.retry_delay_ms, 2000);
478        assert_eq!(client.request_timeout_seconds, 180);
479    }
480
481    #[test]
482    fn test_azure_openai_new_with_all_parameters() {
483        let client = AzureOpenAIClient::new_with_all(
484            "test-api-key".to_string(),
485            "gpt-test".to_string(),
486            "https://example.openai.azure.com".to_string(),
487            "2025-04-01-preview".to_string(),
488            0.7,
489            4000,
490            3,
491            1000,
492            120,
493        );
494        assert!(format!("{:?}", client).contains("AzureOpenAIClient"));
495    }
496
497    #[test]
498    fn test_azure_openai_error_handling_empty_api_key() {
499        let mut config = Config::default();
500        config.ai.provider = "azure-openai".to_string();
501        config.ai.api_key = Some("".to_string()); // Empty string
502        config.ai.model = "deployment-name".to_string();
503        config.ai.base_url = "https://example.openai.azure.com".to_string();
504
505        let err = AzureOpenAIClient::from_config(&config.ai)
506            .unwrap_err()
507            .to_string();
508        assert!(err.contains("Missing Azure OpenAI API Key"));
509    }
510
511    /// §3.6 — connection refused against `127.0.0.1` MUST surface the hint.
512    #[tokio::test]
513    async fn test_hosted_hint_connection_refused_loopback() {
514        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
515        let port = listener.local_addr().unwrap().port();
516        drop(listener);
517        let client = AzureOpenAIClient::new_with_all(
518            "k".into(),
519            "dep".into(),
520            format!("http://127.0.0.1:{}", port),
521            "2025-04-01-preview".into(),
522            0.0,
523            16,
524            0,
525            0,
526            1,
527        );
528        let err = client
529            .chat_completion(vec![json!({"role":"user","content":"x"})])
530            .await
531            .unwrap_err();
532        let msg = err.to_string();
533        assert!(
534            msg.contains("ollama") && msg.contains("ai.provider"),
535            "expected local-provider hint: {msg}"
536        );
537    }
538
539    /// §3.6 — HTTP 200 with non-OpenAI body must surface the hint via the
540    /// parse-shape branch.
541    #[tokio::test]
542    async fn test_hosted_hint_http_200_non_openai_body() {
543        use wiremock::matchers::{method, path};
544        use wiremock::{Mock, MockServer, ResponseTemplate};
545        let server = MockServer::start().await;
546        Mock::given(method("POST"))
547            // Azure URL shape:
548            // {base}/openai/deployments/{model}/chat/completions
549            .and(path("/openai/deployments/dep/chat/completions"))
550            .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "hello": "world" })))
551            .mount(&server)
552            .await;
553        let client = AzureOpenAIClient::new_with_all(
554            "k".into(),
555            "dep".into(),
556            server.uri(),
557            "2025-04-01-preview".into(),
558            0.0,
559            16,
560            0,
561            0,
562            5,
563        );
564        let err = client
565            .chat_completion(vec![json!({"role":"user","content":"x"})])
566            .await
567            .unwrap_err();
568        let msg = err.to_string();
569        assert!(
570            msg.contains("Invalid API response format")
571                && msg.contains("ollama")
572                && msg.contains("ai.provider"),
573            "expected hint-bearing parse-shape error: {msg}"
574        );
575    }
576
577    /// §3.6 negative — a public host MUST NOT surface the hint. We use
578    /// TEST-NET-1 (RFC 5737) so the test is hermetic.
579    #[tokio::test]
580    async fn test_hosted_hint_not_emitted_for_public_host() {
581        let client = AzureOpenAIClient::new_with_all(
582            "k".into(),
583            "dep".into(),
584            "https://192.0.2.1".into(),
585            "2025-04-01-preview".into(),
586            0.0,
587            16,
588            0,
589            0,
590            1,
591        );
592        let err = client
593            .chat_completion(vec![json!({"role":"user","content":"x"})])
594            .await
595            .unwrap_err();
596        let msg = err.to_string();
597        assert!(
598            !msg.contains("ollama"),
599            "public-host failure must NOT carry the hint: {msg}"
600        );
601    }
602}
603
604#[async_trait]
605impl AIProvider for AzureOpenAIClient {
606    async fn analyze_content(&self, request: AnalysisRequest) -> crate::Result<MatchResult> {
607        let prompt = self.build_analysis_prompt(&request);
608        let messages = vec![
609            json!({"role": "system", "content": "You are a professional subtitle matching assistant that can analyze the correspondence between video and subtitle files."}),
610            json!({"role": "user", "content": prompt}),
611        ];
612        let resp = self.chat_completion(messages).await?;
613        self.parse_match_result(&resp)
614    }
615
616    async fn verify_match(
617        &self,
618        verification: VerificationRequest,
619    ) -> crate::Result<ConfidenceScore> {
620        let prompt = self.build_verification_prompt(&verification);
621        let messages = vec![
622            json!({"role": "system", "content": "Please evaluate the confidence level of subtitle matching and provide a score between 0-1."}),
623            json!({"role": "user", "content": prompt}),
624        ];
625        let resp = self.chat_completion(messages).await?;
626        self.parse_confidence_score(&resp)
627    }
628
629    async fn chat_completion(&self, messages: Vec<Value>) -> crate::Result<String> {
630        AzureOpenAIClient::chat_completion(self, messages).await
631    }
632}
subx_core/services/ai/azure_openai.rs

subx_core/services/ai/
azure_openai.rs
