Skip to main content

batuta/agent/
code_prompts.rs

1//! System prompts and exit codes for `apr code` / `batuta code`.
2//!
3//! Extracted from code.rs to keep module under 500-line threshold.
4
5/// Compact system prompt for -p mode (PMAT-197).
6/// Minimal context to avoid overwhelming small models (Qwen3 1.7B).
7/// The full CODE_SYSTEM_PROMPT causes Qwen3 1.7B to loop on `</think>` tags
8/// when combined with 9 tool JSON schemas consuming most of the context window.
9pub(super) const COMPACT_SYSTEM_PROMPT: &str = "\
10Answer the question. Be direct.\
11";
12
13/// System prompt โ€” PMAT-168: optimized for 1.5B-7B with explicit tool table.
14pub(super) const CODE_SYSTEM_PROMPT: &str = "\
15You are apr code, a sovereign AI coding assistant. All inference runs locally โ€” \
16no data ever leaves the machine.
17
18## Tools
19
20You have 9 tools. To use one, emit a <tool_call> block:
21
22<tool_call>
23{\"name\": \"tool_name\", \"input\": {\"param\": \"value\"}}
24</tool_call>
25
26| Tool | Use for | Example input |
27|------|---------|---------------|
28| file_read | Read a file | {\"path\": \"src/main.rs\"} |
29| file_write | Create/overwrite file | {\"path\": \"new.rs\", \"content\": \"fn main() {}\"} |
30| file_edit | Replace text in file | {\"path\": \"src/lib.rs\", \"old\": \"foo\", \"new\": \"bar\"} |
31| glob | Find files by pattern | {\"pattern\": \"src/**/*.rs\"} |
32| grep | Search file contents | {\"pattern\": \"TODO\", \"path\": \"src/\"} |
33| shell | Run a command | {\"command\": \"cargo test --lib\"} |
34| memory | Remember/recall facts | {\"action\": \"remember\", \"key\": \"bug\", \"value\": \"off-by-one\"} |
35| pmat_query | Search code by intent | {\"query\": \"error handling\", \"limit\": 5} |
36| rag | Search project docs | {\"query\": \"authentication flow\"} |
37
38## Guidelines
39
40- Read files before editing โ€” understand first
41- Use file_edit for changes, file_write only for new files
42- Run tests after changes: shell with cargo test
43- Use pmat_query for code search (returns quality-graded functions), glob for files, grep for text
44- Be concise
45";
46
47/// Exit codes for non-interactive mode (spec ยง9.1).
48pub mod exit_code {
49    pub const SUCCESS: i32 = 0;
50    pub const AGENT_ERROR: i32 = 1;
51    pub const BUDGET_EXHAUSTED: i32 = 2;
52    pub const MAX_TURNS: i32 = 3;
53    pub const SANDBOX_VIOLATION: i32 = 4;
54    pub const NO_MODEL: i32 = 5;
55}
56
57/// Mid-size system prompt for 2-7B models (PMAT-198).
58/// Includes tool names and format but shorter examples than CODE_SYSTEM_PROMPT.
59const MID_SYSTEM_PROMPT: &str = "\
60You are apr code, a sovereign AI coding assistant. All inference runs locally.
61
62To use a tool: <tool_call>{\"name\": \"tool\", \"input\": {...}}</tool_call>
63
64Tools: file_read, file_write, file_edit, glob, grep, shell, memory, pmat_query, rag
65
66- Read before editing. Use file_edit for changes, file_write for new files.
67- Run tests: shell with cargo test
68- Be concise
69";
70
71/// Estimate model parameter count (billions) from filename.
72///
73/// Parses patterns like `Qwen3-1.7B`, `llama-8b`, `phi-3-mini-3.8b`.
74/// Returns 0.0 if no size found (caller should assume large model).
75pub(super) fn estimate_model_params_from_name(path: &std::path::Path) -> f64 {
76    let name = path.file_stem().and_then(|s| s.to_str()).unwrap_or("").to_lowercase();
77
78    // Match patterns: "1.7b", "8b", "0.6b", "70b", "3.8b"
79    // Look for a number followed by 'b' (case-insensitive, already lowered)
80    let mut best: f64 = 0.0;
81    let chars: Vec<char> = name.chars().collect();
82    let mut i = 0;
83    while i < chars.len() {
84        // Find start of a number
85        if chars[i].is_ascii_digit() {
86            let start = i;
87            // Consume digits and optional decimal
88            while i < chars.len() && (chars[i].is_ascii_digit() || chars[i] == '.') {
89                i += 1;
90            }
91            // Check if followed by 'b'
92            if i < chars.len() && chars[i] == 'b' {
93                if let Ok(n) = name[start..i].parse::<f64>() {
94                    if n > best {
95                        best = n;
96                    }
97                }
98            }
99        }
100        i += 1;
101    }
102    best
103}
104
105/// Select system prompt based on model parameter count (PMAT-198).
106///
107/// | Size | Prompt | Rationale |
108/// |------|--------|-----------|
109/// | <2B  | COMPACT | Avoids thinking loops, keeps tool format |
110/// | 2-7B | MID | Tool names + format, no example JSON |
111/// | 7B+  | FULL | Full table with examples + guidelines |
112pub(super) fn scale_prompt_for_model(params_b: f64) -> String {
113    if params_b < 2.0 {
114        COMPACT_SYSTEM_PROMPT.to_string()
115    } else if params_b < 7.0 {
116        MID_SYSTEM_PROMPT.to_string()
117    } else {
118        CODE_SYSTEM_PROMPT.to_string()
119    }
120}
121
122pub(super) fn map_error_to_exit_code(e: &crate::agent::result::AgentError) -> i32 {
123    use crate::agent::result::AgentError;
124    match e {
125        AgentError::CircuitBreak(_) => exit_code::BUDGET_EXHAUSTED,
126        AgentError::MaxIterationsReached => exit_code::MAX_TURNS,
127        AgentError::CapabilityDenied { .. } => exit_code::SANDBOX_VIOLATION,
128        _ => exit_code::AGENT_ERROR,
129    }
130}