1use serde_json::{json, Value};
2use std::time::Duration;
3
4use crate::embedding::EmbeddingProvider;
5use crate::errors::{InnateError, Result};
6use crate::refine::{DistillProvenance, DistilledChunk, Distiller, Reranker};
7use crate::settings::{EmbeddingConfig, LlmConfig};
8
9const DISTILL_PROMPT_VERSION: &str = "4";
14
15fn safe_prompt_field(value: Option<&str>) -> String {
16 let value = value.unwrap_or("");
17 let (cleaned, action) = crate::utils::sanitize(value);
18 match action {
19 crate::utils::SanitizeAction::Discard => "[removed unsafe content]".to_string(),
20 _ => cleaned,
21 }
22}
23
24fn build_distill_prompt(log: &Value) -> String {
25 let query = safe_prompt_field(log.get("query").and_then(Value::as_str));
26 let output = safe_prompt_field(log.get("output").and_then(Value::as_str));
27 let output_summary = safe_prompt_field(log.get("output_summary").and_then(Value::as_str));
28 let nomination = safe_prompt_field(log.get("nomination").and_then(Value::as_str));
29 let outcome = safe_prompt_field(log.get("outcome").and_then(Value::as_str));
30
31 let mut context_parts = vec![];
32 if !query.is_empty() {
33 context_parts.push(format!("Query: {query}"));
34 }
35 if !nomination.is_empty() {
36 context_parts.push(format!("Nominated insight: {nomination}"));
37 }
38 if !output_summary.is_empty() {
39 context_parts.push(format!("Summary: {output_summary}"));
40 }
41 if !output.is_empty() {
42 let truncated: String = output.chars().take(1500).collect();
43 context_parts.push(format!("Output (truncated): {truncated}"));
44 }
45 if !outcome.is_empty() {
46 context_parts.push(format!("Outcome: {outcome}"));
47 }
48
49 let context = context_parts.join("\n");
50
51 format!(
52 r#"You are a knowledge distillation assistant. Given an agent interaction log, \
53extract zero or more independent reusable procedural principles. Favor GENERAL, \
54transferable skills, methods, and techniques over project-specific facts.
55
56Agent interaction:
57{context}
58
59Output a JSON array. Each item has:
60{{
61 "skill_name": "<1-3 word skill/topic label for this principle>",
62 "content": "<principle; when it applies; what to avoid>",
63 "trigger_desc": "<2-6 word canonical phrase>",
64 "anti_trigger_desc": "<when NOT to apply this, or null>"
65}}
66Return [] if nothing is worth keeping.
67
68Rules:
69- skill_name is a short human label (1-3 words) naming the skill/topic, e.g.
70 "error handling", "git rebase", "async retries"; not a sentence
71- content must be self-contained and actionable for a future agent reading cold
72- Prefer transferable methods and techniques; a principle that helps across many
73 projects is worth far more than one tied to this codebase
74- Abstract away project-specific detail: strip repo/file/function/path/variable names
75 and one-off identifiers, and rephrase the lesson as a general principle whoever the
76 next project is. Keep concrete project-specific detail ONLY when the lesson genuinely
77 cannot be generalized without losing its meaning
78- trigger_desc must match the vocabulary a future agent would use in a search query;
79 prefer general, technology- or domain-level phrasing over project-name phrasing
80- Never store conversation text verbatim; always distil to reusable principle form
81- If outcome is "fail", focus on what to avoid
82- Keep principles independent; do not combine unrelated lessons"#
83 )
84}
85
86fn build_distill_prompt_with_related(log: &Value, logs: &[Value]) -> String {
87 let mut prompt = build_distill_prompt(log);
88 let log_id = log.get("id").and_then(Value::as_str).unwrap_or("");
89 let context_key = log.get("context_key").and_then(Value::as_str);
90 let related: Vec<String> = logs
91 .iter()
92 .filter(|other| other.get("id").and_then(Value::as_str).unwrap_or("") != log_id)
93 .filter(|other| {
94 context_key.is_some() && other.get("context_key").and_then(Value::as_str) == context_key
95 })
96 .take(4)
97 .map(|other| {
98 let query = safe_prompt_field(other.get("query").and_then(Value::as_str));
99 let summary = safe_prompt_field(other.get("output_summary").and_then(Value::as_str));
100 let outcome = safe_prompt_field(other.get("outcome").and_then(Value::as_str));
101 format!("- Query: {query}; outcome: {outcome}; summary: {summary}")
102 })
103 .collect();
104 if !related.is_empty() {
105 prompt.push_str(
106 "\n\nRelated recent interactions (use only to identify repeated patterns or conflicts):\n",
107 );
108 prompt.push_str(&related.join("\n"));
109 }
110 prompt
111}
112
113const HTTP_MAX_ATTEMPTS: u32 = 3;
119const HTTP_TIMEOUT: Duration = Duration::from_secs(30);
121
122fn post_json_retry(
127 url: &str,
128 headers: &[(&str, &str)],
129 body: &Value,
130 label: &str,
131) -> Result<Value> {
132 let start = std::time::Instant::now();
136 let mut attempt = 0;
137 let outcome: Result<Value> = loop {
138 attempt += 1;
139 let mut req = ureq::post(url)
144 .config()
145 .timeout_global(Some(HTTP_TIMEOUT))
146 .http_status_as_error(false)
147 .build()
148 .header("Content-Type", "application/json");
149 for (k, v) in headers {
150 req = req.header(*k, *v);
151 }
152 match req.send_json(body) {
153 Ok(mut response) => {
154 let code = response.status().as_u16();
155 if (200..300).contains(&code) {
156 break response.body_mut().read_json::<Value>().map_err(|e| {
157 InnateError::Other(format!("{label} response parse error: {e}"))
158 });
159 }
160 let retry_after = response
161 .headers()
162 .get("retry-after")
163 .and_then(|h| h.to_str().ok())
164 .and_then(|s| s.trim().parse::<u64>().ok());
165 if status_is_retryable(code) && attempt < HTTP_MAX_ATTEMPTS {
166 std::thread::sleep(backoff_delay(attempt, retry_after));
167 continue;
168 }
169 let detail = response.body_mut().read_to_string().unwrap_or_default();
172 break Err(InnateError::Other(format!(
173 "{label} HTTP error: status: {code} {detail}"
174 )));
175 }
176 Err(err) => {
177 if attempt < HTTP_MAX_ATTEMPTS {
178 std::thread::sleep(backoff_delay(attempt, None));
179 continue;
180 }
181 break Err(InnateError::Other(format!(
183 "{label} HTTP error: transport: {err}"
184 )));
185 }
186 }
187 };
188 crate::llm_trace::record(label, url, body, &outcome, attempt, start.elapsed());
189 outcome
190}
191
192fn status_is_retryable(code: u16) -> bool {
194 code == 429 || (500..=599).contains(&code)
195}
196
197fn backoff_delay(attempt: u32, retry_after_secs: Option<u64>) -> Duration {
200 if let Some(secs) = retry_after_secs {
201 return Duration::from_secs(secs.min(30));
202 }
203 let shift = (attempt - 1).min(6);
204 Duration::from_millis(250u64.saturating_mul(1 << shift))
205}
206
207pub struct HttpDistiller {
215 config: LlmConfig,
216}
217
218impl HttpDistiller {
219 pub fn new(config: LlmConfig) -> Self {
220 Self { config }
221 }
222
223 pub fn call(&self, prompt: &str) -> Result<String> {
226 if self.config.provider == "anthropic" {
227 self.call_anthropic(prompt)
228 } else {
229 self.call_openai(prompt)
230 }
231 }
232
233 fn call_openai(&self, prompt: &str) -> Result<String> {
234 let api_key = self
235 .config
236 .resolved_api_key()
237 .ok_or_else(|| InnateError::Other("LLM API key not configured".into()))?;
238
239 let base = self.config.resolved_base_url();
240 let url = format!("{base}/chat/completions");
241
242 let body = json!({
243 "model": self.config.model_id,
244 "messages": [{"role": "user", "content": prompt}],
245 "max_tokens": 800,
246 "temperature": 0.2,
247 });
248
249 let auth = format!("Bearer {api_key}");
250 let resp_json = post_json_retry(&url, &[("Authorization", &auth)], &body, "LLM")?;
251
252 resp_json
253 .pointer("/choices/0/message/content")
254 .and_then(Value::as_str)
255 .map(str::to_string)
256 .ok_or_else(|| InnateError::Other("unexpected LLM response shape".into()))
257 }
258
259 fn call_anthropic(&self, prompt: &str) -> Result<String> {
260 let api_key = self
261 .config
262 .resolved_api_key()
263 .ok_or_else(|| InnateError::Other("Anthropic API key not configured".into()))?;
264
265 let base = self.config.resolved_base_url();
266 let url = format!("{base}/v1/messages");
267
268 let body = json!({
269 "model": self.config.model_id,
270 "max_tokens": 800,
271 "messages": [{"role": "user", "content": prompt}],
272 });
273
274 let resp_json = post_json_retry(
275 &url,
276 &[("x-api-key", &api_key), ("anthropic-version", "2023-06-01")],
277 &body,
278 "Anthropic",
279 )?;
280
281 resp_json
282 .pointer("/content/0/text")
283 .and_then(Value::as_str)
284 .map(str::to_string)
285 .ok_or_else(|| InnateError::Other("unexpected Anthropic response shape".into()))
286 }
287}
288
289impl Distiller for HttpDistiller {
290 fn distill(&self, log_entries: &[Value]) -> crate::errors::Result<Vec<DistilledChunk>> {
291 distill_with(log_entries, |prompt| self.call(prompt))
292 }
293
294 fn distill_with_context(
295 &self,
296 primary: &Value,
297 related_logs: &[Value],
298 ) -> crate::errors::Result<Vec<DistilledChunk>> {
299 distill_entry_with(primary, related_logs, |prompt| self.call(prompt))
300 }
301
302 fn provenance(&self) -> DistillProvenance {
303 DistillProvenance {
304 provider: Some(self.config.provider.clone()),
305 model: Some(self.config.model_id.clone()),
306 prompt_version: Some(DISTILL_PROMPT_VERSION.to_string()),
307 }
308 }
309}
310
311fn distill_with(
316 log_entries: &[Value],
317 call: impl Fn(&str) -> Result<String> + Copy,
318) -> Result<Vec<DistilledChunk>> {
319 let mut out = Vec::new();
320 for entry in log_entries {
321 out.extend(distill_entry_with(entry, log_entries, call)?);
322 }
323 Ok(out)
324}
325
326fn distill_entry_with(
327 entry: &Value,
328 related_logs: &[Value],
329 call: impl Fn(&str) -> Result<String>,
330) -> Result<Vec<DistilledChunk>> {
331 let log_id = entry["id"].as_str().unwrap_or("").to_string();
332 let prompt = build_distill_prompt_with_related(entry, related_logs);
333 let mut raw = call(&prompt)?;
334 let mut parsed = parse_distill_response(&raw);
335 if parsed.is_err() {
336 raw = call(&format!(
337 "{prompt}\n\nYour previous response was invalid. Return only a valid JSON array."
338 ))?;
339 parsed = parse_distill_response(&raw);
340 }
341 let items = parsed.map_err(|error| {
342 InnateError::Other(format!("LLM distillation response invalid: {error}"))
343 })?;
344 let mut out = Vec::new();
345 for parsed in items {
346 let content = parsed
347 .get("content")
348 .and_then(Value::as_str)
349 .map(str::trim)
350 .filter(|s| !s.is_empty());
351 let Some(content) = content else { continue };
352 let skill_name = parsed
353 .get("skill_name")
354 .and_then(Value::as_str)
355 .map(|s| {
356 s.split_whitespace()
357 .take(3)
358 .collect::<Vec<_>>()
359 .join(" ")
360 })
361 .filter(|s| !s.is_empty() && s.to_lowercase() != "null");
362 let trigger_desc = parsed
363 .get("trigger_desc")
364 .and_then(Value::as_str)
365 .map(str::to_string)
366 .filter(|s| !s.is_empty());
367 let anti_trigger_desc = parsed
368 .get("anti_trigger_desc")
369 .and_then(Value::as_str)
370 .map(str::to_string)
371 .filter(|s| !s.is_empty() && s.to_lowercase() != "null");
372 out.push(DistilledChunk {
373 content: content.to_string(),
374 skill_name,
375 trigger_desc,
376 anti_trigger_desc,
377 source_log_id: log_id.clone(),
378 nomination: entry
379 .get("nomination")
380 .and_then(Value::as_str)
381 .map(str::to_string),
382 provider_override: None,
383 });
384 }
385 Ok(out)
386}
387
388fn parse_distill_response(raw: &str) -> std::result::Result<Vec<Value>, String> {
389 let json_str = extract_json(raw);
390 let parsed: Value = serde_json::from_str(json_str.trim()).map_err(|e| e.to_string())?;
391 if parsed.get("skip").and_then(Value::as_bool) == Some(true) {
392 return Ok(vec![]);
393 }
394 match parsed {
395 Value::Array(items) => Ok(items),
396 Value::Object(_) => Ok(vec![parsed]),
397 _ => Err("expected a JSON object or array".to_string()),
398 }
399}
400
401fn extract_json(text: &str) -> &str {
402 let stripped = text.trim();
404 if let Some(inner) = stripped
405 .strip_prefix("```json")
406 .or_else(|| stripped.strip_prefix("```"))
407 {
408 if let Some(end) = inner.rfind("```") {
409 return inner[..end].trim();
410 }
411 }
412 if let (Some(start), Some(end)) = (stripped.find('['), stripped.rfind(']')) {
413 return &stripped[start..=end];
414 }
415 if let (Some(start), Some(end)) = (stripped.find('{'), stripped.rfind('}')) {
417 return &stripped[start..=end];
418 }
419 stripped
420}
421
422pub fn build_distiller(config: &LlmConfig) -> std::sync::Arc<dyn Distiller + Send + Sync> {
427 std::sync::Arc::new(HttpDistiller::new(config.clone()))
428}
429
430pub struct LlmReranker {
437 inner: HttpDistiller,
438}
439
440impl LlmReranker {
441 pub fn new(config: LlmConfig) -> Self {
442 Self {
443 inner: HttpDistiller::new(config),
444 }
445 }
446}
447
448impl Reranker for LlmReranker {
449 fn rerank(&self, query: &str, candidates: &[Value]) -> Result<Vec<String>> {
450 if candidates.is_empty() {
451 return Ok(Vec::new());
452 }
453 let mut list = String::new();
454 for c in candidates {
455 let id = c.get("id").and_then(Value::as_str).unwrap_or("");
456 let trig = c.get("trigger_desc").and_then(Value::as_str).unwrap_or("");
457 let content: String = c
458 .get("content")
459 .and_then(Value::as_str)
460 .unwrap_or("")
461 .chars()
462 .take(280)
463 .collect();
464 list.push_str(&format!("- id={id} | when={trig} | {content}\n"));
465 }
466 let prompt = format!(
467 "You are reranking knowledge snippets by how directly each one helps with the QUERY. \
468 Consider the snippet's `when` (trigger) and content. Return ONLY a JSON array of the \
469 ids, most relevant first, no prose, no ids that are not listed.\n\n\
470 QUERY: {query}\n\nCANDIDATES:\n{list}"
471 );
472 let resp = self.inner.call(&prompt)?;
473 parse_id_array(&resp)
474 .ok_or_else(|| InnateError::Other("reranker: no id array in LLM response".into()))
475 }
476}
477
478fn parse_id_array(resp: &str) -> Option<Vec<String>> {
480 let start = resp.find('[')?;
481 let end = resp.rfind(']')?;
482 if end <= start {
483 return None;
484 }
485 let arr: Value = serde_json::from_str(&resp[start..=end]).ok()?;
486 let ids: Vec<String> = arr
487 .as_array()?
488 .iter()
489 .filter_map(|v| v.as_str().map(str::to_string))
490 .collect();
491 if ids.is_empty() {
492 None
493 } else {
494 Some(ids)
495 }
496}
497
498pub struct LlmEmbeddingProvider {
503 config: EmbeddingConfig,
504}
505
506#[cfg(test)]
507#[allow(clippy::items_after_test_module)]
508mod tests {
509 use std::cell::Cell;
510
511 use serde_json::json;
512
513 use std::time::Duration;
514
515 use super::{
516 backoff_delay, build_distill_prompt, distill_entry_with, distill_with,
517 parse_distill_response, parse_embedding_response, status_is_retryable,
518 };
519
520 #[test]
521 fn embedding_response_is_parsed_fail_closed() {
522 let resp = json!({"data": [{"embedding": [0.1, 0.2, 0.3]}]});
524 assert_eq!(
525 parse_embedding_response(&resp, 3).unwrap(),
526 vec![0.1f32, 0.2, 0.3]
527 );
528
529 assert!(parse_embedding_response(&resp, 4).is_err());
531
532 let bad = json!({"data": [{"embedding": [0.1, "oops", 0.3]}]});
534 assert!(parse_embedding_response(&bad, 3).is_err());
535
536 let shape = json!({"data": []});
538 assert!(parse_embedding_response(&shape, 3).is_err());
539 }
540
541 #[test]
542 fn only_rate_limit_and_5xx_are_retryable() {
543 assert!(status_is_retryable(429));
544 assert!(status_is_retryable(500));
545 assert!(status_is_retryable(503));
546 assert!(status_is_retryable(599));
547 assert!(!status_is_retryable(400));
548 assert!(!status_is_retryable(401));
549 assert!(!status_is_retryable(404));
550 assert!(!status_is_retryable(200));
551 }
552
553 #[test]
554 fn backoff_is_exponential_and_honors_retry_after() {
555 assert_eq!(backoff_delay(1, None), Duration::from_millis(250));
557 assert_eq!(backoff_delay(2, None), Duration::from_millis(500));
558 assert_eq!(backoff_delay(3, None), Duration::from_millis(1000));
559 assert_eq!(backoff_delay(1, Some(5)), Duration::from_secs(5));
561 assert_eq!(backoff_delay(1, Some(120)), Duration::from_secs(30));
562 }
563
564 #[test]
565 fn prompt_redacts_secrets_before_external_llm_call() {
566 let prompt = build_distill_prompt(&json!({
567 "query": "debug sk-12345678901234567890",
568 "output_summary": "Authorization: Bearer secret-token-value"
569 }));
570 assert!(!prompt.contains("sk-12345678901234567890"));
571 assert!(!prompt.contains("secret-token-value"));
572 assert!(prompt.contains("[REDACTED]"));
573 }
574
575 #[test]
576 fn malformed_response_is_retried_instead_of_silently_skipped() {
577 let calls = Cell::new(0);
578 let chunks = distill_with(&[json!({"id": "log-1", "query": "q"})], |_| {
579 calls.set(calls.get() + 1);
580 if calls.get() == 1 {
581 Ok("not json".to_string())
582 } else {
583 Ok(r#"[{"content":"retry worked","trigger_desc":"retry"}]"#.to_string())
584 }
585 })
586 .unwrap();
587 assert_eq!(calls.get(), 2);
588 assert_eq!(chunks.len(), 1);
589 assert_eq!(chunks[0].content, "retry worked");
590 }
591
592 #[test]
593 fn parser_accepts_multiple_distilled_chunks() {
594 let parsed = parse_distill_response(
595 r#"[{"content":"one"},{"content":"two","anti_trigger_desc":"never"}]"#,
596 )
597 .unwrap();
598 assert_eq!(parsed.len(), 2);
599 }
600
601 #[test]
602 fn nomination_is_distilled_instead_of_bypassing_the_model() {
603 let prompt_seen = Cell::new(false);
604 let entry = json!({
605 "id": "log-1",
606 "query": "original query",
607 "nomination": "raw agent nomination",
608 "output_summary": "summary",
609 "outcome": "ok"
610 });
611 let chunks = distill_entry_with(&entry, std::slice::from_ref(&entry), |prompt| {
612 prompt_seen.set(prompt.contains("raw agent nomination"));
613 Ok(
614 r#"[{"content":"generalized principle","trigger_desc":"generalize","anti_trigger_desc":null}]"#
615 .to_string(),
616 )
617 })
618 .unwrap();
619
620 assert!(prompt_seen.get());
621 assert_eq!(chunks[0].content, "generalized principle");
622 assert_eq!(
623 chunks[0].nomination.as_deref(),
624 Some("raw agent nomination")
625 );
626 }
627}
628
629impl LlmEmbeddingProvider {
630 pub fn new(config: EmbeddingConfig) -> Self {
631 Self { config }
632 }
633
634 fn embed(&self, text: &str) -> Result<Vec<f32>> {
635 let api_key = self
636 .config
637 .resolved_api_key()
638 .ok_or_else(|| InnateError::Other("Embedding API key not configured".into()))?;
639
640 let base = self.config.resolved_base_url();
641 let url = format!("{base}/embeddings");
642
643 let body = json!({
644 "input": text,
645 "model": self.config.model_id,
646 });
647
648 let auth = format!("Bearer {api_key}");
649 let resp_json = post_json_retry(&url, &[("Authorization", &auth)], &body, "Embedding")?;
650
651 parse_embedding_response(&resp_json, self.config.dim)
652 }
653}
654
655fn parse_embedding_response(resp_json: &Value, expected_dim: usize) -> Result<Vec<f32>> {
661 let embedding = resp_json
662 .pointer("/data/0/embedding")
663 .and_then(Value::as_array)
664 .ok_or_else(|| InnateError::Other("unexpected embedding response shape".into()))?;
665 let vec: Vec<f32> = embedding
666 .iter()
667 .map(|v| {
668 v.as_f64().map(|x| x as f32).ok_or_else(|| {
669 InnateError::Other("embedding response contains a non-numeric element".into())
670 })
671 })
672 .collect::<Result<Vec<f32>>>()?;
673 if vec.len() != expected_dim {
674 return Err(InnateError::Other(format!(
675 "embedding dimension mismatch: provider returned {}, expected {expected_dim} (check embedding.dim)",
676 vec.len(),
677 )));
678 }
679 Ok(vec)
680}
681
682impl EmbeddingProvider for LlmEmbeddingProvider {
683 fn model_name(&self) -> &'static str {
684 "llm-embedding"
685 }
686
687 fn content_dim(&self) -> usize {
688 self.config.dim
689 }
690
691 fn trigger_dim(&self) -> usize {
692 self.config.dim
693 }
694
695 fn embed_content(&self, text: &str) -> Result<Vec<f32>> {
696 self.embed(text)
697 }
698
699 fn embed_trigger(&self, text: &str) -> Result<Vec<f32>> {
700 self.embed(text)
701 }
702
703 fn embed_both(&self, text: &str) -> Result<(Vec<f32>, Vec<f32>)> {
706 let v = self.embed(text)?;
707 Ok((v.clone(), v))
708 }
709}
710
711pub fn test_llm(config: &LlmConfig) -> Result<String> {
717 let distiller = build_distiller(config);
718 let dummy_log = json!({
719 "id": "test",
720 "query": "connection test",
721 "output_summary": "test",
722 "outcome": "ok"
723 });
724 distiller.distill(&[dummy_log])?;
726 Ok(format!("OK — model: {}", config.model_id))
727}
728
729pub fn test_embedding(config: &EmbeddingConfig) -> Result<usize> {
731 let provider = LlmEmbeddingProvider::new(config.clone());
732 let vec = provider.embed("connection test")?;
733 Ok(vec.len())
734}