Skip to main content

hyprcorrect_core/
llm.rs

1//! LLM-backed correction provider (M4).
2//!
3//! Currently wires the Anthropic Messages API only — the `backend`
4//! field on [`crate::LlmConfig`] is read but only `"anthropic"` is
5//! implemented. Synchronous on purpose: the daemon's main loop calls
6//! this from the trigger handler and we expect ~1s round-trip; an
7//! async runtime would be overkill.
8//!
9//! Construction reads the API key out of the OS keychain via
10//! [`crate::secrets`]. Missing key → `Err(LlmError::NoApiKey)` — the
11//! daemon falls back to the offline provider so the trigger never
12//! silently no-ops.
13
14use std::time::Duration;
15
16use crate::secrets;
17
18/// The keyring entry name the prefs UI writes to and the daemon
19/// reads from. Kept in lock-step with the prefs constant.
20const ANTHROPIC_KEY_NAME: &str = "llm.anthropic";
21
22const ANTHROPIC_URL: &str = "https://api.anthropic.com/v1/messages";
23const ANTHROPIC_VERSION: &str = "2023-06-01";
24const DEFAULT_MAX_TOKENS: u32 = 1024;
25
26const SYSTEM_PROMPT: &str = "You are a spelling, typo, and minor-grammar corrector. Return ONLY the \
27     corrected version of the user's text — no preamble, no commentary, no \
28     quotation marks. Preserve the user's voice, register, and punctuation \
29     style. If the text is already fine, return it unchanged.";
30
31const WORD_SYSTEM_PROMPT: &str = "You correct ONE word at a time using sentence context. The \
32     user gives you a SENTENCE and one WORD from it to correct. Return ONLY the corrected \
33     version of that word — nothing else: no quotes, no punctuation, no commentary, no rest \
34     of the sentence. Use the rest of the sentence to disambiguate homophones \
35     (their/there/they're, its/it's, your/you're, etc.) and to pick the right fix for typos. \
36     Preserve the original casing of the word's first letter. If the word is already correct \
37     in context, return it unchanged.";
38
39/// Errors from an LLM correction request.
40#[derive(Debug, thiserror::Error)]
41pub enum LlmError {
42    /// No API key is stored in the OS keychain under the expected entry.
43    #[error("no API key for the LLM provider — set one in Preferences → Providers")]
44    NoApiKey,
45    /// The keychain itself returned an error.
46    #[error("keychain: {0}")]
47    Keychain(String),
48    /// The configured backend ID isn't one we support yet.
49    #[error("unsupported LLM backend: {0}")]
50    UnsupportedBackend(String),
51    /// The network request itself failed (DNS / TLS / non-2xx, …).
52    #[error("LLM request failed: {0}")]
53    Request(String),
54    /// We reached the API but couldn't read what came back.
55    #[error("LLM response was unparseable: {0}")]
56    Response(String),
57}
58
59/// The LLM correction provider.
60pub struct LlmProvider {
61    api_key: String,
62    model: String,
63}
64
65impl std::fmt::Debug for LlmProvider {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        // Never print the API key — Debug is used in tests and logs.
68        f.debug_struct("LlmProvider")
69            .field("model", &self.model)
70            .field("api_key", &"[redacted]")
71            .finish()
72    }
73}
74
75impl LlmProvider {
76    /// Build the provider from the user's [`crate::LlmConfig`] —
77    /// reads the API key out of the OS keychain.
78    ///
79    /// # Errors
80    ///
81    /// See [`LlmError`].
82    pub fn from_config(llm: &crate::LlmConfig) -> Result<Self, LlmError> {
83        if llm.backend != "anthropic" {
84            return Err(LlmError::UnsupportedBackend(llm.backend.clone()));
85        }
86        let api_key = secrets::get(ANTHROPIC_KEY_NAME)
87            .map_err(|e| LlmError::Keychain(e.to_string()))?
88            .ok_or(LlmError::NoApiKey)?;
89        Ok(Self {
90            api_key,
91            model: llm.model.clone(),
92        })
93    }
94
95    /// Rewrite `text` into its corrected form. Returns the corrected
96    /// string verbatim; callers compare against the input to decide
97    /// whether an edit is needed.
98    ///
99    /// # Errors
100    ///
101    /// See [`LlmError`].
102    pub fn rewrite(&self, text: &str) -> Result<String, LlmError> {
103        if text.trim().is_empty() {
104            return Ok(text.to_string());
105        }
106        self.request(SYSTEM_PROMPT, text.to_string())
107    }
108
109    /// Correct a single word using the surrounding sentence as
110    /// context. The LLM is told to return ONLY the corrected word,
111    /// not the rest of the sentence — callers splice it back in at
112    /// the caret. Good for homophones and context-dependent typos
113    /// where the offline spellbook either can't see the error
114    /// (their/there) or picks the wrong nearest neighbor.
115    ///
116    /// # Errors
117    ///
118    /// See [`LlmError`].
119    pub fn fix_word_in_context(&self, sentence: &str, word: &str) -> Result<String, LlmError> {
120        if word.trim().is_empty() {
121            return Ok(word.to_string());
122        }
123        let content = format!("SENTENCE: {sentence}\nWORD: {word}");
124        let corrected = self.request(WORD_SYSTEM_PROMPT, content)?;
125        // Defensive: strip any wrapping whitespace or quotation
126        // marks the LLM may include despite the system prompt
127        // telling it not to.
128        Ok(corrected
129            .trim()
130            .trim_matches(|c: char| c == '"' || c == '\'')
131            .to_string())
132    }
133
134    fn request(&self, system: &str, content: String) -> Result<String, LlmError> {
135        let agent: ureq::Agent = ureq::AgentBuilder::new()
136            .timeout(Duration::from_secs(20))
137            .build();
138        let body = serde_json::json!({
139            "model": self.model,
140            "max_tokens": DEFAULT_MAX_TOKENS,
141            "system": system,
142            "messages": [{
143                "role": "user",
144                "content": content,
145            }],
146        });
147        let response = agent
148            .post(ANTHROPIC_URL)
149            .set("x-api-key", &self.api_key)
150            .set("anthropic-version", ANTHROPIC_VERSION)
151            .set("content-type", "application/json")
152            .send_json(body)
153            .map_err(|e| LlmError::Request(e.to_string()))?;
154        let json: serde_json::Value = response
155            .into_json()
156            .map_err(|e| LlmError::Response(e.to_string()))?;
157        // Anthropic's response: { "content": [ { "type": "text", "text": "..." }, ... ], ... }
158        let corrected = json["content"]
159            .as_array()
160            .and_then(|parts| {
161                parts
162                    .iter()
163                    .filter_map(|p| p.get("text").and_then(|t| t.as_str()))
164                    .next()
165            })
166            .ok_or_else(|| LlmError::Response("no `content[*].text` in response".into()))?;
167        Ok(corrected.trim_end_matches('\n').to_string())
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use crate::LlmConfig;
175
176    #[test]
177    fn unsupported_backend_is_rejected_cleanly() {
178        let cfg = LlmConfig {
179            backend: "openai".into(),
180            model: "gpt-5".into(),
181        };
182        match LlmProvider::from_config(&cfg) {
183            Err(LlmError::UnsupportedBackend(name)) => assert_eq!(name, "openai"),
184            other => panic!("expected UnsupportedBackend, got {other:?}"),
185        }
186    }
187}