lean_ctx/core/
llm_enhance.rs1use 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
74pub 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
100pub 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
123fn 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}