//! g3-planner: Fast-discovery planner for G3 AI coding agent //! //! This crate provides functionality to generate initial discovery tool calls //! that are injected into the conversation before the first LLM turn. mod code_explore; pub mod prompts; pub use code_explore::explore_codebase; use anyhow::Result; use g3_providers::{CompletionRequest, LLMProvider, Message, MessageRole}; use chrono::Local; use std::fs::{self, OpenOptions}; use std::io::Write; use prompts::{DISCOVERY_REQUIREMENTS_PROMPT, DISCOVERY_SYSTEM_PROMPT}; /// Type alias for a status callback function pub type StatusCallback = Box; /// Generates initial discovery messages for fast codebase exploration. /// /// This function: /// 1. Runs explore_codebase to get a codebase report /// 2. Sends the report to the LLM with DISCOVERY_SYSTEM_PROMPT /// 3. Extracts shell commands from the LLM response /// 4. Returns Assistant messages with tool calls for each command /// /// # Arguments /// /// * `codebase_path` - The path to the codebase to explore /// * `provider` - An LLM provider to query for exploration commands /// * `requirements_text` - Optional requirements text to include in the discovery prompt /// * `status_callback` - Optional callback for status updates /// /// # Returns /// /// A `Result>` containing Assistant messages with JSON tool call strings. pub async fn get_initial_discovery_messages( codebase_path: &str, requirements_text: Option<&str>, provider: &dyn LLMProvider, status_callback: Option<&StatusCallback>, ) -> Result> { // Helper to call status callback if provided let status = |msg: &str| { if let Some(cb) = status_callback { cb(msg); } }; status("🔍 Starting code discovery..."); // Step 1: Run explore_codebase to get the codebase report let codebase_report = explore_codebase(codebase_path); // Write the codebase report to logs directory write_code_report(&codebase_report)?; // Step 2: Build the prompt with the codebase report appended let user_prompt = if let Some(requirements) = requirements_text { format!( "{}\n\n === REQUIREMENTS ===\n\n{}\n\n === CODEBASE REPORT ===\n\n{}", DISCOVERY_REQUIREMENTS_PROMPT, requirements, codebase_report ) } else { format!( "{}\n\n=== CODEBASE REPORT ===\n\n{}", DISCOVERY_REQUIREMENTS_PROMPT, codebase_report ) }; // Step 3: Create messages for the LLM let messages = vec![ Message::new(MessageRole::System, DISCOVERY_SYSTEM_PROMPT.to_string()), Message::new(MessageRole::User, user_prompt), ]; // Step 4: Send to LLM let request = CompletionRequest { messages, max_tokens: Some(provider.max_tokens()), temperature: Some(provider.temperature()), stream: false, tools: None, }; status("🤖 Calling LLM for discovery commands..."); let response = provider.complete(request).await?; // Step 5: Extract shell commands from the response let shell_commands = extract_shell_commands(&response.content); status(&format!("📋 Extracted {} discovery commands", shell_commands.len())); // Write the discovery commands to logs directory write_discovery_commands(&shell_commands)?; // Step 6: Format as tool messages let tool_messages = shell_commands .into_iter() .map(|cmd| create_tool_message("shell", &cmd)) .collect(); Ok(tool_messages) } /// Creates an Assistant message with a tool call in g3's JSON format. pub fn create_tool_message(tool: &str, command: &str) -> Message { let tool_call = serde_json::json!({ "tool": tool, "args": { "command": command } }); Message::new(MessageRole::Assistant, tool_call.to_string()) } /// Extract shell commands from the LLM response. /// Looks for {{CODE EXPLORATION COMMANDS}} section and extracts commands from code blocks. pub fn extract_shell_commands(response: &str) -> Vec { let mut commands = Vec::new(); let section_marker = "{{CODE EXPLORATION COMMANDS}}"; let section_start = match response.find(section_marker) { Some(pos) => pos + section_marker.len(), None => return commands, }; let section_content = &response[section_start..]; let mut in_code_block = false; let mut current_block = String::new(); for line in section_content.lines() { let trimmed = line.trim(); if trimmed.starts_with("```") { if in_code_block { // End of code block - extract commands for cmd_line in current_block.lines() { let cmd = cmd_line.trim(); if !cmd.is_empty() && !cmd.starts_with('#') { commands.push(cmd.to_string()); } } current_block.clear(); } in_code_block = !in_code_block; } else if in_code_block { current_block.push_str(line); current_block.push('\n'); } } commands } /// Extract the summary section from the LLM response pub fn extract_summary(response: &str) -> Option { let section_marker = "{{SUMMARY BASED ON INITIAL INFO}}"; let section_start = match response.find(section_marker) { Some(pos) => pos + section_marker.len(), None => return None, }; let section_content = &response[section_start..]; let section_end = section_content.find("{{").unwrap_or(section_content.len()); let summary = section_content[..section_end].trim().to_string(); if summary.is_empty() { None } else { Some(summary) } } /// Write the codebase report to logs directory fn write_code_report(report: &str) -> Result<()> { // Ensure logs directory exists fs::create_dir_all("logs")?; // Generate timestamp in same format as tool_calls log let timestamp = Local::now().format("%Y%m%d_%H%M%S").to_string(); let filename = format!("logs/code_report_{}.log", timestamp); // Write the report to file let mut file = OpenOptions::new() .create(true) .write(true) .truncate(true) .open(&filename)?; file.write_all(report.as_bytes())?; file.flush()?; Ok(()) } /// Write the discovery commands to logs directory fn write_discovery_commands(commands: &[String]) -> Result<()> { // Ensure logs directory exists fs::create_dir_all("logs")?; // Generate timestamp in same format as tool_calls log let timestamp = Local::now().format("%Y%m%d_%H%M%S").to_string(); let filename = format!("logs/discovery_commands_{}.log", timestamp); // Write the commands to file let mut file = OpenOptions::new() .create(true) .write(true) .truncate(true) .open(&filename)?; // Write header file.write_all(b"# Discovery Commands\n")?; file.write_all(b"# Generated by g3-planner\n\n")?; // Write each command on a separate line for cmd in commands { file.write_all(cmd.as_bytes())?; file.write_all(b"\n")?; } file.flush()?; Ok(()) } #[cfg(test)] mod tests { use super::*; #[test] fn test_create_tool_message_format() { let msg = create_tool_message("shell", "ls -la"); assert!(matches!(msg.role, MessageRole::Assistant)); let parsed: serde_json::Value = serde_json::from_str(&msg.content).unwrap(); assert_eq!(parsed["tool"], "shell"); assert_eq!(parsed["args"]["command"], "ls -la"); } #[test] fn test_extract_shell_commands_basic() { let response = r#" Some text here. {{CODE EXPLORATION COMMANDS}} ```bash ls -la cat README.md rg --files -g '*.rs' ``` More text. "#; let commands = extract_shell_commands(response); assert_eq!(commands.len(), 3); assert_eq!(commands[0], "ls -la"); assert_eq!(commands[1], "cat README.md"); assert_eq!(commands[2], "rg --files -g '*.rs'"); } #[test] fn test_extract_shell_commands_with_comments() { let response = r#" {{CODE EXPLORATION COMMANDS}} ``` # This is a comment ls -la # Another comment cat file.txt ``` "#; let commands = extract_shell_commands(response); assert_eq!(commands.len(), 2); assert_eq!(commands[0], "ls -la"); assert_eq!(commands[1], "cat file.txt"); } #[test] fn test_extract_shell_commands_no_section() { let response = "Some response without the expected section."; let commands = extract_shell_commands(response); assert!(commands.is_empty()); } #[test] fn test_extract_summary() { let response = r#" {{SUMMARY BASED ON INITIAL INFO}} This is a summary of the codebase. It has multiple lines. {{CODE EXPLORATION COMMANDS}} ``` ls -la ``` "#; let summary = extract_summary(response); assert!(summary.is_some()); let summary_text = summary.unwrap(); assert!(summary_text.contains("This is a summary")); assert!(summary_text.contains("multiple lines")); } #[test] fn test_extract_summary_no_section() { let response = "Response without summary section."; let summary = extract_summary(response); assert!(summary.is_none()); } }