Skip to main content

lean_ctx/core/
llm_enhance.rs

1//! Optional LLM enhancement layer.
2//!
3//! Deterministic by default — LLM calls are opt-in and always fall back to
4//! the deterministic pipeline on failure, timeout, or when disabled.
5//!
6//! Supported backends: Ollama (local), OpenRouter, Claude (Anthropic).
7
8use std::time::Duration;
9
10const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
11const MAX_PROMPT_CHARS: usize = 2000;
12
13#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
14#[serde(default)]
15pub struct LlmConfig {
16    pub enabled: bool,
17    pub backend: LlmBackend,
18    pub model: String,
19    pub timeout_secs: u64,
20    pub base_url: Option<String>,
21}
22
23impl Default for LlmConfig {
24    fn default() -> Self {
25        Self {
26            enabled: false,
27            backend: LlmBackend::Ollama,
28            model: "qwen2.5-coder:1.5b".to_string(),
29            timeout_secs: 10,
30            base_url: None,
31        }
32    }
33}
34
35#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
36#[serde(rename_all = "lowercase")]
37#[derive(Default)]
38pub enum LlmBackend {
39    #[default]
40    Ollama,
41    OpenRouter,
42    Anthropic,
43}
44
45impl LlmConfig {
46    fn effective_base_url(&self) -> String {
47        if let Some(ref url) = self.base_url {
48            return url.clone();
49        }
50        match self.backend {
51            LlmBackend::Ollama => "http://localhost:11434".to_string(),
52            LlmBackend::OpenRouter => "https://openrouter.ai/api".to_string(),
53            LlmBackend::Anthropic => "https://api.anthropic.com".to_string(),
54        }
55    }
56
57    fn api_key(&self) -> Option<String> {
58        match self.backend {
59            LlmBackend::Ollama => None,
60            LlmBackend::OpenRouter => std::env::var("OPENROUTER_API_KEY").ok(),
61            LlmBackend::Anthropic => std::env::var("ANTHROPIC_API_KEY").ok(),
62        }
63    }
64
65    fn timeout(&self) -> Duration {
66        if self.timeout_secs > 0 {
67            Duration::from_secs(self.timeout_secs)
68        } else {
69            DEFAULT_TIMEOUT
70        }
71    }
72}
73
74/// Expand a search query using LLM. Falls back to the original query on failure.
75pub fn expand_query(query: &str) -> String {
76    let cfg = crate::core::config::Config::load().llm;
77    if !cfg.enabled {
78        return query.to_string();
79    }
80
81    let prompt = format!(
82        "Expand this code search query with 2-3 related terms. \
83         Return ONLY the expanded query, no explanation.\n\
84         Query: {query}"
85    );
86
87    match call_llm(&cfg, &prompt) {
88        Ok(expanded) => {
89            let cleaned = expanded.trim().to_string();
90            if cleaned.is_empty() || cleaned.len() > query.len() * 5 {
91                query.to_string()
92            } else {
93                cleaned
94            }
95        }
96        Err(_) => query.to_string(),
97    }
98}
99
100/// Generate a human-readable explanation for a knowledge contradiction.
101/// Falls back to a simple diff-style description.
102pub fn explain_contradiction(fact_a: &str, fact_b: &str) -> String {
103    let cfg = crate::core::config::Config::load().llm;
104    if !cfg.enabled {
105        return deterministic_contradiction(fact_a, fact_b);
106    }
107
108    let prompt = format!(
109        "These two facts contradict. Explain the conflict in one sentence:\n\
110         A: {fact_a}\nB: {fact_b}"
111    );
112
113    match call_llm(&cfg, &prompt) {
114        Ok(explanation) => explanation.trim().to_string(),
115        Err(_) => deterministic_contradiction(fact_a, fact_b),
116    }
117}
118
119fn deterministic_contradiction(a: &str, b: &str) -> String {
120    format!("Conflict: \"{a}\" vs \"{b}\"")
121}
122
123/// Low-level LLM call. Supports Ollama, OpenRouter, and Anthropic.
124fn call_llm(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
125    let truncated = if prompt.len() > MAX_PROMPT_CHARS {
126        &prompt[..prompt.floor_char_boundary(MAX_PROMPT_CHARS)]
127    } else {
128        prompt
129    };
130
131    match cfg.backend {
132        LlmBackend::Ollama => call_ollama(cfg, truncated),
133        LlmBackend::OpenRouter => call_openai_compatible(cfg, truncated),
134        LlmBackend::Anthropic => call_anthropic(cfg, truncated),
135    }
136}
137
138fn make_agent(cfg: &LlmConfig) -> ureq::Agent {
139    ureq::Agent::new_with_config(
140        ureq::config::Config::builder()
141            .timeout_global(Some(cfg.timeout()))
142            .build(),
143    )
144}
145
146fn call_ollama(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
147    let url = format!("{}/api/generate", cfg.effective_base_url());
148    let body = serde_json::json!({
149        "model": cfg.model,
150        "prompt": prompt,
151        "stream": false,
152        "options": { "num_predict": 100 }
153    });
154
155    let agent = make_agent(cfg);
156    let payload = serde_json::to_vec(&body).map_err(|e| format!("json: {e}"))?;
157    let resp = agent
158        .post(&url)
159        .header("Content-Type", "application/json")
160        .send(payload.as_slice())
161        .map_err(|e| format!("ollama: {e}"))?;
162
163    let text = resp
164        .into_body()
165        .read_to_string()
166        .map_err(|e| format!("read: {e}"))?;
167    let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
168    json.get("response")
169        .and_then(|v| v.as_str())
170        .map(str::to_string)
171        .ok_or_else(|| "no response field".to_string())
172}
173
174fn call_openai_compatible(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
175    let key = cfg.api_key().ok_or("OPENROUTER_API_KEY not set")?;
176    let url = format!("{}/v1/chat/completions", cfg.effective_base_url());
177    let body = serde_json::json!({
178        "model": cfg.model,
179        "messages": [{"role": "user", "content": prompt}],
180        "max_tokens": 100
181    });
182
183    let agent = make_agent(cfg);
184    let payload = serde_json::to_vec(&body).map_err(|e| format!("json: {e}"))?;
185    let resp = agent
186        .post(&url)
187        .header("Authorization", &format!("Bearer {key}"))
188        .header("Content-Type", "application/json")
189        .send(payload.as_slice())
190        .map_err(|e| format!("openrouter: {e}"))?;
191
192    let text = resp
193        .into_body()
194        .read_to_string()
195        .map_err(|e| format!("read: {e}"))?;
196    let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
197    json.pointer("/choices/0/message/content")
198        .and_then(|v| v.as_str())
199        .map(str::to_string)
200        .ok_or_else(|| "no content in response".to_string())
201}
202
203fn call_anthropic(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
204    let key = cfg.api_key().ok_or("ANTHROPIC_API_KEY not set")?;
205    let url = format!("{}/v1/messages", cfg.effective_base_url());
206    let body = serde_json::json!({
207        "model": cfg.model,
208        "max_tokens": 100,
209        "messages": [{"role": "user", "content": prompt}]
210    });
211
212    let agent = make_agent(cfg);
213    let payload = serde_json::to_vec(&body).map_err(|e| format!("json: {e}"))?;
214    let resp = agent
215        .post(&url)
216        .header("x-api-key", &key)
217        .header("anthropic-version", "2023-06-01")
218        .header("Content-Type", "application/json")
219        .send(payload.as_slice())
220        .map_err(|e| format!("anthropic: {e}"))?;
221
222    let text = resp
223        .into_body()
224        .read_to_string()
225        .map_err(|e| format!("read: {e}"))?;
226    let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
227    json.pointer("/content/0/text")
228        .and_then(|v| v.as_str())
229        .map(str::to_string)
230        .ok_or_else(|| "no text in response".to_string())
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    #[test]
238    fn default_config_disabled() {
239        let cfg = LlmConfig::default();
240        assert!(!cfg.enabled);
241        assert!(matches!(cfg.backend, LlmBackend::Ollama));
242    }
243
244    #[test]
245    fn expand_query_passthrough_when_disabled() {
246        let result = expand_query("test query");
247        assert_eq!(result, "test query");
248    }
249
250    #[test]
251    fn deterministic_contradiction_format() {
252        let result = deterministic_contradiction("A is true", "A is false");
253        assert!(result.contains("Conflict"));
254        assert!(result.contains("A is true"));
255    }
256
257    #[test]
258    fn effective_base_url_defaults() {
259        let cfg = LlmConfig::default();
260        assert!(cfg.effective_base_url().contains("11434"));
261
262        let cfg = LlmConfig {
263            backend: LlmBackend::OpenRouter,
264            ..Default::default()
265        };
266        assert!(cfg.effective_base_url().contains("openrouter"));
267    }
268}