battlecommand_forge/
memory.rs

1//! Self-improving memory system.
2//!
3//! After successful missions: distill learnings + save few-shot examples.
4//! Before missions: load relevant context from learnings + examples.
5
6use anyhow::Result;
7use std::fs;
8use std::path::{Path, PathBuf};
9
10use crate::llm::LlmClient;
11
12const LEARNINGS_FILE: &str = ".battlecommand/learnings.md";
13const EXAMPLES_DIR: &str = ".battlecommand/examples";
14const CONTEXT_FILE: &str = ".battlecommand/context.md";
15const FAILURES_FILE: &str = ".battlecommand/failure_patterns.md";
16const MAX_EXAMPLES: usize = 100;
17const MAX_FAILURE_PATTERNS: usize = 30;
18
19/// Load relevant context for a mission prompt.
20/// Combines: context.md + matching learnings + relevant few-shot examples.
21pub fn load_context(prompt: &str) -> String {
22    let mut context = String::new();
23
24    // Always load context.md
25    if let Ok(ctx) = fs::read_to_string(CONTEXT_FILE) {
26        context.push_str(&ctx);
27        context.push_str("\n\n");
28    }
29
30    // Search learnings for relevant entries
31    if let Ok(learnings) = fs::read_to_string(LEARNINGS_FILE) {
32        let relevant = find_relevant_sections(&learnings, prompt, 3);
33        if !relevant.is_empty() {
34            context.push_str("## Relevant Learnings from Past Missions\n\n");
35            context.push_str(&relevant);
36            context.push_str("\n\n");
37        }
38    }
39
40    // Find relevant few-shot examples
41    let examples = find_relevant_examples(prompt, 2);
42    if !examples.is_empty() {
43        context.push_str("## Reference Examples (follow this style)\n\n");
44        for (name, content) in &examples {
45            context.push_str(&format!("### Example: {}\n```\n{}\n```\n\n", name, content));
46        }
47    }
48
49    context
50}
51
52/// Save a learning after a successful mission.
53pub async fn distill_and_save(
54    llm: &LlmClient,
55    prompt: &str,
56    code_summary: &str,
57    score: f32,
58) -> Result<()> {
59    // Distill the mission into a learning
60    let system = "You are a knowledge distiller. Extract the key patterns, decisions, and \
61                  pitfalls from this successful mission into 2-3 bullet points. \
62                  Focus on reusable patterns, not specifics. Output ONLY bullet points.";
63    let user_prompt = format!(
64        "Mission: {}\nScore: {:.1}/10\nCode summary:\n{}",
65        prompt, score, code_summary
66    );
67
68    let learning = llm
69        .generate("DISTILL", system, &user_prompt)
70        .await
71        .unwrap_or_else(|_| format!("- Completed: {}", prompt));
72
73    // Append to learnings file
74    let timestamp = chrono::Utc::now().format("%Y-%m-%d %H:%M");
75    let entry = format!(
76        "\n## [{}] {}\nScore: {:.1}/10\n{}\n",
77        timestamp,
78        prompt.chars().take(80).collect::<String>(),
79        score,
80        learning.trim()
81    );
82
83    let mut existing = fs::read_to_string(LEARNINGS_FILE).unwrap_or_default();
84    existing.push_str(&entry);
85    fs::write(LEARNINGS_FILE, existing)?;
86
87    Ok(())
88}
89
90/// Save successful output files as a few-shot example.
91pub fn save_example(prompt: &str, output_dir: &Path, language: &str) -> Result<()> {
92    let examples_dir = Path::new(EXAMPLES_DIR);
93    fs::create_dir_all(examples_dir)?;
94
95    // Create example directory name
96    let slug: String = prompt
97        .to_lowercase()
98        .chars()
99        .filter(|c| c.is_alphanumeric() || *c == ' ')
100        .collect::<String>()
101        .split_whitespace()
102        .take(4)
103        .collect::<Vec<_>>()
104        .join("_");
105
106    let example_dir = examples_dir.join(format!("{}_{}", language, slug));
107    fs::create_dir_all(&example_dir)?;
108
109    // Copy key source files (not __init__.py, not boilerplate)
110    for entry in walkdir_source_files(output_dir, language)? {
111        let relative = entry.strip_prefix(output_dir).unwrap_or(&entry);
112        let dest = example_dir.join(relative);
113        if let Some(parent) = dest.parent() {
114            fs::create_dir_all(parent)?;
115        }
116        let _ = fs::copy(&entry, &dest);
117    }
118
119    // Enforce max examples limit
120    enforce_example_limit(examples_dir)?;
121
122    Ok(())
123}
124
125/// Find sections in learnings.md that match the prompt keywords.
126fn find_relevant_sections(learnings: &str, prompt: &str, max: usize) -> String {
127    let lowered = prompt.to_lowercase();
128    let keywords: Vec<&str> = lowered.split_whitespace().filter(|w| w.len() > 3).collect();
129
130    let mut sections: Vec<(usize, &str)> = Vec::new();
131    let mut current_section = String::new();
132    let mut section_start = 0;
133
134    for (i, line) in learnings.lines().enumerate() {
135        if line.starts_with("## ") {
136            if !current_section.is_empty() {
137                let score = keywords
138                    .iter()
139                    .filter(|k| current_section.to_lowercase().contains(*k))
140                    .count();
141                if score > 0 {
142                    let start = section_start;
143                    let end = i;
144                    let section_text = learnings
145                        .lines()
146                        .skip(start)
147                        .take(end - start)
148                        .collect::<Vec<_>>()
149                        .join("\n");
150                    sections.push((score, Box::leak(section_text.into_boxed_str())));
151                }
152            }
153            current_section = line.to_string();
154            section_start = i;
155        } else {
156            current_section.push('\n');
157            current_section.push_str(line);
158        }
159    }
160
161    sections.sort_by(|a, b| b.0.cmp(&a.0));
162    sections
163        .into_iter()
164        .take(max)
165        .map(|(_, s)| s)
166        .collect::<Vec<_>>()
167        .join("\n\n")
168}
169
170/// Find relevant few-shot examples by matching directory names to prompt keywords.
171fn find_relevant_examples(prompt: &str, max: usize) -> Vec<(String, String)> {
172    let examples_dir = Path::new(EXAMPLES_DIR);
173    if !examples_dir.exists() {
174        return vec![];
175    }
176
177    let keywords: Vec<String> = prompt
178        .to_lowercase()
179        .split_whitespace()
180        .filter(|w| w.len() > 3)
181        .map(String::from)
182        .collect();
183
184    let mut matches: Vec<(usize, String, String)> = Vec::new();
185
186    if let Ok(entries) = fs::read_dir(examples_dir) {
187        for entry in entries.flatten() {
188            if !entry.path().is_dir() {
189                continue;
190            }
191            let name = entry.file_name().to_string_lossy().to_string();
192            let score = keywords
193                .iter()
194                .filter(|k| name.contains(k.as_str()))
195                .count();
196
197            if score > 0 {
198                // Read the main source file from the example
199                let content = read_example_summary(&entry.path());
200                if !content.is_empty() {
201                    matches.push((score, name, content));
202                }
203            }
204        }
205    }
206
207    matches.sort_by(|a, b| b.0.cmp(&a.0));
208    matches
209        .into_iter()
210        .take(max)
211        .map(|(_, name, content)| (name, content))
212        .collect()
213}
214
215/// Read a summary of an example (first source file, truncated).
216fn read_example_summary(dir: &Path) -> String {
217    let extensions = ["py", "ts", "js", "rs", "go"];
218    if let Ok(entries) = fs::read_dir(dir) {
219        for entry in entries.flatten() {
220            let path = entry.path();
221            if path.is_file() {
222                if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
223                    if extensions.contains(&ext) {
224                        if let Ok(content) = fs::read_to_string(&path) {
225                            // Truncate to first 50 lines
226                            return content.lines().take(50).collect::<Vec<_>>().join("\n");
227                        }
228                    }
229                }
230            }
231        }
232    }
233    String::new()
234}
235
236/// Walk directory for source files (skip __init__.py, boilerplate).
237fn walkdir_source_files(dir: &Path, language: &str) -> Result<Vec<PathBuf>> {
238    let source_exts: Vec<&str> = match language {
239        "python" => vec!["py"],
240        "typescript" => vec!["ts", "tsx"],
241        "javascript" => vec!["js", "jsx"],
242        "rust" => vec!["rs"],
243        "go" => vec!["go"],
244        _ => vec!["py"],
245    };
246
247    let mut files = Vec::new();
248    if dir.is_dir() {
249        for entry in fs::read_dir(dir)? {
250            let entry = entry?;
251            let path = entry.path();
252            if path.is_dir() {
253                files.extend(walkdir_source_files(&path, language)?);
254            } else if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
255                if source_exts.contains(&ext) {
256                    let name = path.file_name().unwrap_or_default().to_string_lossy();
257                    if name != "__init__.py" {
258                        files.push(path);
259                    }
260                }
261            }
262        }
263    }
264    Ok(files)
265}
266
267/// Remove oldest examples if we exceed the limit.
268fn enforce_example_limit(examples_dir: &Path) -> Result<()> {
269    let mut entries: Vec<_> = fs::read_dir(examples_dir)?
270        .flatten()
271        .filter(|e| e.path().is_dir())
272        .collect();
273
274    if entries.len() <= MAX_EXAMPLES {
275        return Ok(());
276    }
277
278    // Sort by modification time (oldest first)
279    entries.sort_by_key(|e| {
280        e.metadata()
281            .and_then(|m| m.modified())
282            .unwrap_or(std::time::SystemTime::UNIX_EPOCH)
283    });
284
285    // Remove oldest
286    let to_remove = entries.len() - MAX_EXAMPLES;
287    for entry in entries.into_iter().take(to_remove) {
288        let _ = fs::remove_dir_all(entry.path());
289    }
290
291    Ok(())
292}
293
294/// Save failure patterns from a failed mission for future runs.
295/// Extracts error categories from verifier output and stores them.
296pub fn save_failure_patterns(language: &str, errors: &[String], score: f32) {
297    if errors.is_empty() {
298        return;
299    }
300
301    let timestamp = chrono::Utc::now().format("%Y-%m-%d %H:%M");
302    let mut entry = format!("\n## [{}] {} (score: {:.1})\n", timestamp, language, score);
303    for err in errors.iter().take(10) {
304        // Normalize errors into reusable patterns
305        let pattern = normalize_error_pattern(err);
306        if !pattern.is_empty() {
307            entry.push_str(&format!("- {}\n", pattern));
308        }
309    }
310
311    let mut existing = fs::read_to_string(FAILURES_FILE).unwrap_or_default();
312    existing.push_str(&entry);
313
314    // Trim to max patterns (keep most recent)
315    let sections: Vec<&str> = existing.split("\n## ").collect();
316    if sections.len() > MAX_FAILURE_PATTERNS {
317        let kept: String = sections
318            .iter()
319            .skip(sections.len() - MAX_FAILURE_PATTERNS)
320            .map(|s| format!("\n## {}", s))
321            .collect();
322        let _ = fs::write(FAILURES_FILE, kept.trim_start());
323    } else {
324        let _ = fs::write(FAILURES_FILE, &existing);
325    }
326}
327
328/// Load failure patterns relevant to a language.
329pub fn load_failure_patterns(language: &str) -> String {
330    let content = match fs::read_to_string(FAILURES_FILE) {
331        Ok(c) => c,
332        Err(_) => return String::new(),
333    };
334
335    let mut patterns: Vec<String> = Vec::new();
336    let mut in_matching_section = false;
337
338    for line in content.lines() {
339        if line.starts_with("## ") {
340            in_matching_section = line.to_lowercase().contains(language);
341        } else if in_matching_section && line.starts_with("- ") {
342            let pattern = line[2..].trim().to_string();
343            if !patterns.contains(&pattern) {
344                patterns.push(pattern);
345            }
346        }
347    }
348
349    if patterns.is_empty() {
350        return String::new();
351    }
352
353    // Deduplicate and limit
354    patterns.truncate(15);
355    let mut result =
356        String::from("## Patterns from previous failed runs (DO NOT repeat these mistakes):\n");
357    for p in &patterns {
358        result.push_str(&format!("- {}\n", p));
359    }
360    result
361}
362
363/// Normalize an error message into a reusable pattern.
364fn normalize_error_pattern(error: &str) -> String {
365    let lower = error.to_lowercase();
366
367    // Python-specific normalizations
368    if lower.contains("modulenotfounderror") || lower.contains("importerror") {
369        if lower.contains("pydantic") && lower.contains("basesettings") {
370            return "Use pydantic_settings.BaseSettings, not pydantic.BaseSettings (Pydantic v2)"
371                .to_string();
372        }
373        if lower.contains("no module named") {
374            return format!(
375                "Missing import: {}",
376                error
377                    .split("named")
378                    .last()
379                    .unwrap_or("")
380                    .trim()
381                    .trim_matches('\'')
382            );
383        }
384    }
385    if lower.contains("nameerror") {
386        return format!(
387            "Undefined name: {}",
388            error
389                .split("name")
390                .last()
391                .unwrap_or("")
392                .trim()
393                .trim_matches('\'')
394        );
395    }
396    if lower.contains("attributeerror") {
397        return format!("Wrong attribute/method: {}", error.trim());
398    }
399    if lower.contains("syntax error") || lower.contains("syntaxerror") {
400        return "Python syntax error in generated code".to_string();
401    }
402    if lower.contains("hardcoded secret") || lower.contains("hardcoded") {
403        return "Hardcoded secrets — use environment variables".to_string();
404    }
405
406    // Generic: keep short errors as-is
407    if error.len() < 100 {
408        error.trim().to_string()
409    } else {
410        error.trim().chars().take(100).collect::<String>()
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    #[test]
419    fn test_load_context_missing_files() {
420        // Should not panic even if files don't exist
421        let ctx = load_context("build something");
422        // May or may not have content depending on if .battlecommand exists
423        let _ = ctx; // just verify it doesn't panic
424    }
425
426    #[test]
427    fn test_find_relevant_examples_empty() {
428        let examples = find_relevant_examples("nonexistent prompt xyz", 3);
429        // May be empty if no examples dir exists
430        assert!(examples.len() <= 3);
431    }
432}
battlecommand_forge/memory.rs

battlecommand_forge/
memory.rs
