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/// Refine a synthesized observation summary with an LLM (#802). Opt-in via
124/// `llm.enabled`; the deterministic input is always a valid result and is returned
125/// unchanged when LLM is disabled, errors, times out, or yields something unusable.
126/// The model is constrained to the supplied notes (no invention). Because the
127/// result is stored once and recalled byte-stably, enabling this does not break
128/// hot-path read determinism (#498) — only the one-time stored value differs.
129pub fn enhance_observation(entity: &str, deterministic: &str) -> String {
130    let cfg = crate::core::config::Config::load().llm;
131    if !cfg.enabled {
132        return deterministic.to_string();
133    }
134
135    let prompt = format!(
136        "Summarize what is known about `{entity}` in ONE concise, factual sentence. \
137         Use ONLY these notes; do not invent. Return ONLY the sentence.\n\
138         Notes: {deterministic}"
139    );
140
141    match call_llm(&cfg, &prompt) {
142        Ok(text) => {
143            let cleaned = text.trim();
144            // Reject empty or runaway output; the deterministic digest is the floor.
145            if cleaned.is_empty() || cleaned.len() > deterministic.len() * 4 {
146                deterministic.to_string()
147            } else {
148                format!("{entity} — {cleaned}")
149            }
150        }
151        Err(_) => deterministic.to_string(),
152    }
153}
154
155/// Low-level LLM call. Supports Ollama, OpenRouter, and Anthropic.
156fn call_llm(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
157    let truncated = if prompt.len() > MAX_PROMPT_CHARS {
158        &prompt[..prompt.floor_char_boundary(MAX_PROMPT_CHARS)]
159    } else {
160        prompt
161    };
162
163    match cfg.backend {
164        LlmBackend::Ollama => call_ollama(cfg, truncated),
165        LlmBackend::OpenRouter => call_openai_compatible(cfg, truncated),
166        LlmBackend::Anthropic => call_anthropic(cfg, truncated),
167    }
168}
169
170fn make_agent(cfg: &LlmConfig) -> ureq::Agent {
171    ureq::Agent::new_with_config(
172        ureq::config::Config::builder()
173            .timeout_global(Some(cfg.timeout()))
174            .build(),
175    )
176}
177
178fn call_ollama(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
179    let url = format!("{}/api/generate", cfg.effective_base_url());
180    let body = serde_json::json!({
181        "model": cfg.model,
182        "prompt": prompt,
183        "stream": false,
184        "options": { "num_predict": 100 }
185    });
186
187    let agent = make_agent(cfg);
188    let payload = serde_json::to_vec(&body).map_err(|e| format!("json: {e}"))?;
189    let resp = agent
190        .post(&url)
191        .header("Content-Type", "application/json")
192        .send(payload.as_slice())
193        .map_err(|e| format!("ollama: {e}"))?;
194
195    let text = resp
196        .into_body()
197        .read_to_string()
198        .map_err(|e| format!("read: {e}"))?;
199    let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
200    json.get("response")
201        .and_then(|v| v.as_str())
202        .map(str::to_string)
203        .ok_or_else(|| "no response field".to_string())
204}
205
206fn call_openai_compatible(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
207    let key = cfg.api_key().ok_or("OPENROUTER_API_KEY not set")?;
208    let url = format!("{}/v1/chat/completions", cfg.effective_base_url());
209    let body = serde_json::json!({
210        "model": cfg.model,
211        "messages": [{"role": "user", "content": prompt}],
212        "max_tokens": 100
213    });
214
215    let agent = make_agent(cfg);
216    let payload = serde_json::to_vec(&body).map_err(|e| format!("json: {e}"))?;
217    let resp = agent
218        .post(&url)
219        .header("Authorization", &format!("Bearer {key}"))
220        .header("Content-Type", "application/json")
221        .send(payload.as_slice())
222        .map_err(|e| format!("openrouter: {e}"))?;
223
224    let text = resp
225        .into_body()
226        .read_to_string()
227        .map_err(|e| format!("read: {e}"))?;
228    let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
229    json.pointer("/choices/0/message/content")
230        .and_then(|v| v.as_str())
231        .map(str::to_string)
232        .ok_or_else(|| "no content in response".to_string())
233}
234
235fn call_anthropic(cfg: &LlmConfig, prompt: &str) -> Result<String, String> {
236    let key = cfg.api_key().ok_or("ANTHROPIC_API_KEY not set")?;
237    let url = format!("{}/v1/messages", cfg.effective_base_url());
238    let body = serde_json::json!({
239        "model": cfg.model,
240        "max_tokens": 100,
241        "messages": [{"role": "user", "content": prompt}]
242    });
243
244    let agent = make_agent(cfg);
245    let payload = serde_json::to_vec(&body).map_err(|e| format!("json: {e}"))?;
246    let resp = agent
247        .post(&url)
248        .header("x-api-key", &key)
249        .header("anthropic-version", "2023-06-01")
250        .header("Content-Type", "application/json")
251        .send(payload.as_slice())
252        .map_err(|e| format!("anthropic: {e}"))?;
253
254    let text = resp
255        .into_body()
256        .read_to_string()
257        .map_err(|e| format!("read: {e}"))?;
258    let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| format!("parse: {e}"))?;
259    json.pointer("/content/0/text")
260        .and_then(|v| v.as_str())
261        .map(str::to_string)
262        .ok_or_else(|| "no text in response".to_string())
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn default_config_disabled() {
271        let cfg = LlmConfig::default();
272        assert!(!cfg.enabled);
273        assert!(matches!(cfg.backend, LlmBackend::Ollama));
274    }
275
276    #[test]
277    fn expand_query_passthrough_when_disabled() {
278        let result = expand_query("test query");
279        assert_eq!(result, "test query");
280    }
281
282    #[test]
283    fn deterministic_contradiction_format() {
284        let result = deterministic_contradiction("A is true", "A is false");
285        assert!(result.contains("Conflict"));
286        assert!(result.contains("A is true"));
287    }
288
289    #[test]
290    fn effective_base_url_defaults() {
291        let cfg = LlmConfig::default();
292        assert!(cfg.effective_base_url().contains("11434"));
293
294        let cfg = LlmConfig {
295            backend: LlmBackend::OpenRouter,
296            ..Default::default()
297        };
298        assert!(cfg.effective_base_url().contains("openrouter"));
299    }
300}