subx_core/services/ai/
local.rs

1//! Local / OpenAI-compatible LLM provider client.
2//!
3//! `LocalLLMClient` targets any OpenAI-compatible chat-completions endpoint
4//! exposed by a local, LAN, VPN, or remote runtime (Ollama, LM Studio,
5//! llama.cpp `llama-server`, vLLM, text-generation-webui, etc.). It mirrors
6//! the structure of [`crate::services::ai::openrouter::OpenRouterClient`]
7//! but treats the API key as optional and emits actionable, sanitized
8//! error messages when the local endpoint is unreachable, returns a
9//! non-OpenAI response, or signals that the requested model is not loaded.
10
11use crate::Result;
12use crate::error::SubXError;
13use crate::services::ai::AiUsageStats;
14use crate::services::ai::{
15    AIProvider, AnalysisRequest, ConfidenceScore, MatchResult, VerificationRequest,
16};
17use async_trait::async_trait;
18use reqwest::Client;
19use serde_json::{Value, json};
20use std::time::Duration;
21use tokio::time;
22
23use crate::services::ai::prompts::{PromptBuilder, ResponseParser};
24use crate::services::ai::retry::HttpRetryClient;
25
26/// Client for OpenAI-compatible local LLM runtimes.
27///
28/// The struct mirrors the field layout of the hosted-provider clients
29/// (`OpenAIClient`, `OpenRouterClient`) but stores the API key as
30/// `Option<String>` because most local runtimes accept unauthenticated
31/// requests. The `request_timeout_seconds` value is retained so it can be
32/// embedded in timeout error messages.
33pub struct LocalLLMClient {
34    client: Client,
35    api_key: Option<String>,
36    model: String,
37    temperature: f32,
38    max_tokens: u32,
39    retry_attempts: u32,
40    retry_delay_ms: u64,
41    base_url: String,
42    request_timeout_seconds: u64,
43    reporter: std::sync::Arc<dyn crate::core::report::Reporter>,
44}
45
46impl std::fmt::Debug for LocalLLMClient {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("LocalLLMClient")
49            .field("client", &self.client)
50            .field("api_key", &self.api_key.as_ref().map(|_| "[REDACTED]"))
51            .field("model", &self.model)
52            .field("temperature", &self.temperature)
53            .field("max_tokens", &self.max_tokens)
54            .field("retry_attempts", &self.retry_attempts)
55            .field("retry_delay_ms", &self.retry_delay_ms)
56            .field("base_url", &self.base_url)
57            .field("request_timeout_seconds", &self.request_timeout_seconds)
58            .finish()
59    }
60}
61
62impl PromptBuilder for LocalLLMClient {}
63impl ResponseParser for LocalLLMClient {}
64impl HttpRetryClient for LocalLLMClient {
65    fn retry_attempts(&self) -> u32 {
66        self.retry_attempts
67    }
68    fn retry_delay_ms(&self) -> u64 {
69        self.retry_delay_ms
70    }
71}
72
73impl LocalLLMClient {
74    /// Construct a new `LocalLLMClient` with explicit parameters.
75    #[allow(clippy::too_many_arguments)]
76    pub fn new(
77        api_key: Option<String>,
78        model: String,
79        temperature: f32,
80        max_tokens: u32,
81        retry_attempts: u32,
82        retry_delay_ms: u64,
83        base_url: String,
84        request_timeout_seconds: u64,
85    ) -> Self {
86        let client = Client::builder()
87            .timeout(Duration::from_secs(request_timeout_seconds))
88            .build()
89            .expect("Failed to create HTTP client");
90
91        // Normalize an empty/whitespace-only key to `None` so request paths
92        // do not have to repeat the check.
93        let api_key = api_key.and_then(|k| {
94            let trimmed = k.trim().to_string();
95            if trimmed.is_empty() {
96                None
97            } else {
98                Some(trimmed)
99            }
100        });
101
102        Self {
103            client,
104            api_key,
105            model,
106            temperature,
107            max_tokens,
108            retry_attempts,
109            retry_delay_ms,
110            // Trim a single trailing slash so URL joining always yields one
111            // separator. Both `http://h/v1/` and `http://h/v1` collapse to
112            // the same stored form.
113            base_url: base_url.trim_end_matches('/').to_string(),
114            request_timeout_seconds,
115            reporter: crate::core::report::noop(),
116        }
117    }
118
119    /// Attach a reporting sink, consuming and returning the client.
120    ///
121    /// # Arguments
122    ///
123    /// * `reporter` - Sink that receives [`crate::core::report::AiUsage`]
124    ///   after every successful API response.
125    pub fn with_reporter(
126        mut self,
127        reporter: std::sync::Arc<dyn crate::core::report::Reporter>,
128    ) -> Self {
129        self.reporter = reporter;
130        self
131    }
132
133    /// Create a `LocalLLMClient` from the unified `AIConfig`.
134    ///
135    /// Validates that `base_url` is non-empty, emits the shared insecure-HTTP
136    /// warning (which already exempts loopback), and constructs an inner
137    /// `reqwest::Client` honoring `request_timeout_seconds`.
138    pub fn from_config(config: &crate::config::AIConfig) -> Result<Self> {
139        if config.base_url.trim().is_empty() {
140            return Err(SubXError::config(
141                "ai.base_url is required for the local provider",
142            ));
143        }
144
145        // The shared helper warns only when an api_key is configured against
146        // a non-loopback HTTP host. Pass the configured key (or empty string
147        // when absent) so the existing loopback exemption applies.
148        let api_key_for_warning = config.api_key.clone().unwrap_or_default();
149        crate::services::ai::security::warn_on_insecure_http_str(
150            &config.base_url,
151            &api_key_for_warning,
152        );
153
154        Ok(Self::new(
155            config.api_key.clone(),
156            config.model.clone(),
157            config.temperature,
158            config.max_tokens,
159            config.retry_attempts,
160            config.retry_delay_ms,
161            config.base_url.clone(),
162            config.request_timeout_seconds,
163        ))
164    }
165
166    /// Build the `chat/completions` URL by joining the stored `base_url`
167    /// with the path segment using exactly one `/` separator.
168    fn chat_completions_url(&self) -> String {
169        // `base_url` already had any single trailing slash stripped during
170        // construction, so this format never produces `//chat/completions`.
171        format!("{}/chat/completions", self.base_url)
172    }
173
174    /// Issue a chat-completions request, applying retry + actionable error
175    /// mapping. Sends only OpenAI-canonical body fields.
176    pub async fn chat_completion(&self, messages: Vec<Value>) -> Result<String> {
177        let request_body = json!({
178            "model": self.model,
179            "messages": messages,
180            "temperature": self.temperature,
181            "max_tokens": self.max_tokens,
182        });
183
184        let mut builder = self
185            .client
186            .post(self.chat_completions_url())
187            .header("Content-Type", "application/json")
188            .json(&request_body);
189        if let Some(ref key) = self.api_key {
190            builder = builder.header("Authorization", format!("Bearer {}", key));
191        }
192
193        // Note: `local_provider_hint()` is intentionally NOT appended here.
194        // The advisory exists so hosted-provider clients can suggest the
195        // local provider; from this client (which IS the local provider)
196        // it would be tautological.
197
198        let mut response = self.send_with_retry(builder).await?;
199
200        const MAX_AI_RESPONSE_BYTES: u64 = 10 * 1024 * 1024; // 10 MiB
201        if let Some(len) = response.content_length() {
202            if len > MAX_AI_RESPONSE_BYTES {
203                return Err(SubXError::AiService(format!(
204                    "AI response too large: {} bytes (limit: {} bytes)",
205                    len, MAX_AI_RESPONSE_BYTES
206                )));
207            }
208        }
209
210        if !response.status().is_success() {
211            return Err(self.map_http_error(response).await);
212        }
213
214        // Bounded chunked read in case the server omits Content-Length.
215        let mut body = Vec::new();
216        while let Some(chunk) = response
217            .chunk()
218            .await
219            .map_err(|e| self.map_reqwest_error(e))?
220        {
221            body.extend_from_slice(&chunk);
222            if body.len() as u64 > MAX_AI_RESPONSE_BYTES {
223                return Err(SubXError::AiService(format!(
224                    "AI response too large: {} bytes read (limit: {} bytes)",
225                    body.len(),
226                    MAX_AI_RESPONSE_BYTES
227                )));
228            }
229        }
230
231        let response_json: Value = serde_json::from_slice(&body).map_err(|e| {
232            SubXError::AiService(format!(
233                "local LLM response was not OpenAI-compatible JSON: {}",
234                e
235            ))
236        })?;
237
238        let content = response_json["choices"][0]["message"]["content"]
239            .as_str()
240            .ok_or_else(|| {
241                SubXError::AiService(
242                    "local LLM response was not OpenAI-compatible JSON: \
243                     missing choices[0].message.content"
244                        .to_string(),
245                )
246            })?;
247
248        if let Some(usage_obj) = response_json.get("usage") {
249            if let (Some(p), Some(c), Some(t)) = (
250                usage_obj.get("prompt_tokens").and_then(Value::as_u64),
251                usage_obj.get("completion_tokens").and_then(Value::as_u64),
252                usage_obj.get("total_tokens").and_then(Value::as_u64),
253            ) {
254                let stats = AiUsageStats {
255                    model: self.model.clone(),
256                    prompt_tokens: p as u32,
257                    completion_tokens: c as u32,
258                    total_tokens: t as u32,
259                };
260                self.reporter.ai_usage(&stats);
261            }
262        }
263
264        Ok(content.to_string())
265    }
266
267    /// Send `request` with retry, surfacing actionable transport errors.
268    async fn send_with_retry(&self, request: reqwest::RequestBuilder) -> Result<reqwest::Response> {
269        let mut attempts: u32 = 0;
270        loop {
271            let cloned = request.try_clone().ok_or_else(|| {
272                SubXError::AiService("Request body cannot be cloned for retry".to_string())
273            })?;
274            match cloned.send().await {
275                Ok(resp) => {
276                    if resp.status().is_server_error() && attempts < self.retry_attempts {
277                        attempts += 1;
278                        log::warn!(
279                            "Request attempt {} failed with status {}. Retrying in {}ms...",
280                            attempts,
281                            resp.status(),
282                            self.retry_delay_ms
283                        );
284                        time::sleep(Duration::from_millis(self.retry_delay_ms)).await;
285                        continue;
286                    }
287                    return Ok(resp);
288                }
289                Err(e) if attempts < self.retry_attempts => {
290                    attempts += 1;
291                    log::warn!(
292                        "Request attempt {} failed: {}. Retrying in {}ms...",
293                        attempts,
294                        e,
295                        self.retry_delay_ms
296                    );
297                    time::sleep(Duration::from_millis(self.retry_delay_ms)).await;
298                    continue;
299                }
300                Err(e) => return Err(self.map_reqwest_error(e)),
301            }
302        }
303    }
304
305    /// Map a low-level `reqwest::Error` into an actionable `SubXError`.
306    fn map_reqwest_error(&self, err: reqwest::Error) -> SubXError {
307        let url = sanitize_base_url(&self.base_url);
308        if err.is_timeout() {
309            return SubXError::AiService(format!(
310                "local LLM endpoint timed out after {}s: {}",
311                self.request_timeout_seconds, url
312            ));
313        }
314        if err.is_connect() {
315            return SubXError::AiService(format!("local LLM endpoint unreachable: {}", url));
316        }
317        // Fall back to the generic conversion (which already strips query
318        // strings from any embedded URLs via `error_sanitizer`).
319        err.into()
320    }
321
322    /// Map a non-2xx HTTP response into an actionable `SubXError`. Reads
323    /// the body, sanitizes it, and detects the "model not found" pattern.
324    async fn map_http_error(&self, response: reqwest::Response) -> SubXError {
325        let status = response.status();
326        let body_text = response.text().await.unwrap_or_default();
327        let safe_body = crate::services::ai::error_sanitizer::sanitize_url_in_error(
328            &crate::services::ai::error_sanitizer::truncate_error_body(
329                &body_text,
330                crate::services::ai::error_sanitizer::DEFAULT_ERROR_BODY_MAX_LEN,
331            ),
332        );
333
334        if status.as_u16() == 404 || body_indicates_model_missing(&body_text) {
335            return SubXError::AiService(format!("local LLM model not found: {}", self.model));
336        }
337
338        SubXError::AiService(format!(
339            "local LLM endpoint returned HTTP {}: {}",
340            status, safe_body
341        ))
342    }
343}
344
345/// Detect common "model not loaded / no such model" body patterns emitted
346/// by Ollama, LM Studio, llama.cpp, and vLLM.
347fn body_indicates_model_missing(body: &str) -> bool {
348    let lower = body.to_ascii_lowercase();
349    let mentions_model = lower.contains("model");
350    if !mentions_model {
351        return false;
352    }
353    lower.contains("not found")
354        || lower.contains("not loaded")
355        || lower.contains("no such model")
356        || lower.contains("unknown model")
357}
358
359/// Strip userinfo, query strings, and fragments from a base URL so it can
360/// be safely embedded in user-facing error messages.
361///
362/// Returns only `scheme://host[:port]/path`. Trailing slashes on the path
363/// are preserved as authored. If the input cannot be parsed as a URL the
364/// helper falls back to `"<unparseable URL>"` so credentials hidden in a
365/// malformed string can never leak.
366pub(crate) fn sanitize_base_url(input: &str) -> String {
367    match url::Url::parse(input) {
368        Ok(mut url) => {
369            // Wipe credentials. The setters only fail when the URL cannot
370            // carry userinfo; ignore any error and continue.
371            let _ = url.set_username("");
372            let _ = url.set_password(None);
373            url.set_query(None);
374            url.set_fragment(None);
375
376            let scheme = url.scheme();
377            // Render IPv6 hosts with brackets so the result remains a
378            // syntactically valid URL (e.g. `http://[::1]:11434/v1`).
379            let host_display = match url.host() {
380                Some(url::Host::Ipv6(addr)) => format!("[{}]", addr),
381                Some(_) => url.host_str().unwrap_or_default().to_string(),
382                None => return "<unparseable URL>".to_string(),
383            };
384            let path = url.path();
385            match url.port() {
386                Some(port) => format!("{}://{}:{}{}", scheme, host_display, port, path),
387                None => format!("{}://{}{}", scheme, host_display, path),
388            }
389        }
390        Err(_) => "<unparseable URL>".to_string(),
391    }
392}
393
394#[async_trait]
395impl AIProvider for LocalLLMClient {
396    async fn analyze_content(&self, request: AnalysisRequest) -> Result<MatchResult> {
397        let prompt = self.build_analysis_prompt(&request);
398        let messages = vec![
399            json!({"role": "system", "content": "You are a professional subtitle matching assistant that can analyze the correspondence between video and subtitle files."}),
400            json!({"role": "user", "content": prompt}),
401        ];
402        let response = self.chat_completion(messages).await?;
403        self.parse_match_result(&response)
404    }
405
406    async fn verify_match(&self, verification: VerificationRequest) -> Result<ConfidenceScore> {
407        let prompt = self.build_verification_prompt(&verification);
408        let messages = vec![
409            json!({"role": "system", "content": "Please evaluate the confidence level of subtitle matching and provide a score between 0-1."}),
410            json!({"role": "user", "content": prompt}),
411        ];
412        let response = self.chat_completion(messages).await?;
413        self.parse_confidence_score(&response)
414    }
415
416    async fn chat_completion(&self, messages: Vec<Value>) -> Result<String> {
417        LocalLLMClient::chat_completion(self, messages).await
418    }
419}
420
421#[cfg(test)]
422mod tests {
423    use super::*;
424
425    fn make_client(base_url: &str, api_key: Option<&str>) -> LocalLLMClient {
426        LocalLLMClient::new(
427            api_key.map(|s| s.to_string()),
428            "llama3.1:8b-instruct".to_string(),
429            0.3,
430            1024,
431            1,
432            10,
433            base_url.to_string(),
434            120,
435        )
436    }
437
438    #[test]
439    fn debug_redacts_api_key() {
440        let client = make_client("http://localhost:11434/v1", Some("super-secret-token"));
441        let rendered = format!("{:?}", client);
442        assert!(
443            rendered.contains("[REDACTED]"),
444            "Debug output should redact api_key, got: {rendered}"
445        );
446        assert!(!rendered.contains("super-secret-token"));
447    }
448
449    #[test]
450    fn debug_marks_missing_api_key_as_none() {
451        let client = make_client("http://localhost:11434/v1", None);
452        let rendered = format!("{:?}", client);
453        assert!(rendered.contains("api_key: None"), "got: {rendered}");
454    }
455
456    #[test]
457    fn url_join_with_trailing_slash() {
458        let client = make_client("http://localhost:11434/v1/", None);
459        assert_eq!(
460            client.chat_completions_url(),
461            "http://localhost:11434/v1/chat/completions"
462        );
463        assert!(!client.chat_completions_url().contains("//chat"));
464    }
465
466    #[test]
467    fn url_join_without_trailing_slash() {
468        let client = make_client("http://localhost:11434/v1", None);
469        assert_eq!(
470            client.chat_completions_url(),
471            "http://localhost:11434/v1/chat/completions"
472        );
473    }
474
475    #[test]
476    fn url_join_root_base_url() {
477        let client = make_client("http://localhost:11434", None);
478        assert_eq!(
479            client.chat_completions_url(),
480            "http://localhost:11434/chat/completions"
481        );
482    }
483
484    #[test]
485    fn sanitize_base_url_strips_userinfo_query_and_fragment() {
486        assert_eq!(
487            sanitize_base_url("http://user:secret@127.0.0.1:11434/v1?token=abc#frag"),
488            "http://127.0.0.1:11434/v1"
489        );
490    }
491
492    #[test]
493    fn sanitize_base_url_preserves_plain_localhost() {
494        assert_eq!(
495            sanitize_base_url("http://localhost:11434/v1"),
496            "http://localhost:11434/v1"
497        );
498    }
499
500    #[test]
501    fn sanitize_base_url_preserves_trailing_slash() {
502        // Trailing slashes on the path are preserved as authored.
503        assert_eq!(
504            sanitize_base_url("https://host:8080/api/v1/"),
505            "https://host:8080/api/v1/"
506        );
507    }
508
509    #[test]
510    fn sanitize_base_url_handles_unparseable_input() {
511        assert_eq!(sanitize_base_url("not a url"), "<unparseable URL>");
512        assert_eq!(sanitize_base_url(""), "<unparseable URL>");
513    }
514
515    #[test]
516    fn sanitize_base_url_strips_password_only() {
517        assert_eq!(
518            sanitize_base_url("https://:pwd@host:8080/v1"),
519            "https://host:8080/v1"
520        );
521    }
522
523    #[test]
524    fn sanitize_base_url_preserves_ipv6_brackets() {
525        assert_eq!(
526            sanitize_base_url("http://[::1]:11434/v1"),
527            "http://[::1]:11434/v1"
528        );
529        assert_eq!(
530            sanitize_base_url("https://[fd00::1]:8443/v1/"),
531            "https://[fd00::1]:8443/v1/"
532        );
533        // Userinfo on IPv6 host must still be stripped while brackets remain.
534        assert_eq!(
535            sanitize_base_url("http://user:pwd@[::1]:11434/v1?token=secret"),
536            "http://[::1]:11434/v1"
537        );
538    }
539
540    #[test]
541    fn body_indicates_model_missing_detects_common_patterns() {
542        assert!(body_indicates_model_missing(
543            "{\"error\":\"model 'foo' not found, try pulling it first\"}"
544        ));
545        assert!(body_indicates_model_missing(
546            "{\"error\":\"Model not loaded\"}"
547        ));
548        assert!(body_indicates_model_missing(
549            "{\"detail\":\"no such model: bar\"}"
550        ));
551        assert!(body_indicates_model_missing(
552            "{\"error\":\"unknown model llama99\"}"
553        ));
554        assert!(!body_indicates_model_missing(
555            "{\"error\":\"server overloaded\"}"
556        ));
557        assert!(!body_indicates_model_missing(""));
558    }
559
560    fn make_config(base_url: &str, api_key: Option<&str>) -> crate::config::AIConfig {
561        crate::config::AIConfig {
562            provider: "local".to_string(),
563            api_key: api_key.map(|s| s.to_string()),
564            model: "llama3.1:8b-instruct".to_string(),
565            base_url: base_url.to_string(),
566            max_sample_length: 500,
567            temperature: 0.3,
568            max_tokens: 1024,
569            retry_attempts: 2,
570            retry_delay_ms: 100,
571            request_timeout_seconds: 120,
572            api_version: None,
573        }
574    }
575
576    #[test]
577    fn from_config_rejects_empty_base_url() {
578        let config = make_config("", None);
579        let err = LocalLLMClient::from_config(&config).unwrap_err();
580        assert!(
581            err.to_string().contains("ai.base_url is required"),
582            "unexpected error: {err}"
583        );
584    }
585
586    #[test]
587    fn from_config_rejects_whitespace_base_url() {
588        let config = make_config("   ", None);
589        assert!(LocalLLMClient::from_config(&config).is_err());
590    }
591
592    #[test]
593    fn from_config_accepts_loopback_http() {
594        let config = make_config("http://localhost:11434/v1", None);
595        let client = LocalLLMClient::from_config(&config).expect("should accept loopback HTTP");
596        assert!(client.api_key.is_none());
597        assert_eq!(client.base_url, "http://localhost:11434/v1");
598    }
599
600    #[test]
601    fn from_config_accepts_lan_http() {
602        let config = make_config("http://192.168.1.50:11434/v1", None);
603        let client = LocalLLMClient::from_config(&config).expect("LAN HTTP must be accepted");
604        assert_eq!(client.base_url, "http://192.168.1.50:11434/v1");
605    }
606
607    #[test]
608    fn from_config_accepts_https() {
609        let config = make_config("https://ollama.tailnet.ts.net/v1", Some("vllm-token"));
610        let client = LocalLLMClient::from_config(&config).expect("HTTPS must be accepted");
611        assert_eq!(client.base_url, "https://ollama.tailnet.ts.net/v1");
612        assert_eq!(client.api_key.as_deref(), Some("vllm-token"));
613    }
614
615    #[test]
616    fn from_config_normalizes_empty_api_key_to_none() {
617        let config = make_config("http://localhost:11434/v1", Some(""));
618        let client = LocalLLMClient::from_config(&config).unwrap();
619        assert!(
620            client.api_key.is_none(),
621            "empty api_key should normalize to None"
622        );
623    }
624
625    #[test]
626    fn from_config_trims_trailing_slash_in_base_url() {
627        let config = make_config("http://localhost:11434/v1/", None);
628        let client = LocalLLMClient::from_config(&config).unwrap();
629        assert_eq!(client.base_url, "http://localhost:11434/v1");
630        assert_eq!(
631            client.chat_completions_url(),
632            "http://localhost:11434/v1/chat/completions"
633        );
634    }
635}
subx_core/services/ai/local.rs

subx_core/services/ai/
local.rs
