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
123pub 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 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
155fn 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}