probe_code/extract/
prompts.rs

1//! Prompt templates for LLM models.
2//!
3//! This module provides functionality for loading and formatting prompt templates
4//! for use with LLM models. It supports built-in templates and loading from files.
5
6use anyhow::{Context, Result};
7use std::fs;
8use std::path::Path;
9
10/// Built-in engineer prompt template
11pub const ENGINEER_PROMPT: &str = r#"As a senior software engineer, your task is providing explicit, actionable code adjustments. For each required change:
12
131. Clearly indicate:
14   - File path and name.
15   - Function or class to modify.
16   - Type of modification: add, modify, or remove.
17
182. Provide complete code blocks only for:
19   - Newly created functions or methods.
20   - Entire modified functions.
21   - Updated class definitions.
22   - Adjusted configuration segments.
23
243. Structure responses precisely as follows:
25
26   File: path/filename.ext  
27   Change: Concise description of change made  
28   ```language  
29   [Complete code block of the specified change] "#;
30
31/// Built-in architect prompt template
32pub const ARCHITECT_PROMPT: &str = r#"You are a senior software architect with expertise in designing code structures and planning implementations. Your responsibilities include:
33
341. Analyzing requested modifications, clearly identifying actionable steps.
352. Preparing a thorough implementation plan, detailing:
36   - Files requiring modification.
37   - Specific code segments to update.
38   - New functions, classes, or methods to introduce.
39   - Dependencies and import statements to revise.
40   - Adjustments to data structures.
41   - Changes in interfaces.
42   - Updates needed in configuration files.
43
44For every change, explicitly state:
45- Exact location within the codebase.
46- Reasoning and logic behind each decision.
47- Example code signatures with parameters and return types.
48- Possible side effects or impacts on existing code.
49- Essential architectural choices that must be addressed.
50
51Brief code snippets may be included for clarity, but do not produce a full implementation.
52
53Your analysis should strictly cover the technical implementation plan, excluding deployment, testing, or validation unless explicitly tied to architectural impact."#;
54
55/// Enum representing different prompt template sources
56#[derive(Debug, Clone)]
57pub enum PromptTemplate {
58    /// Built-in engineer template
59    Engineer,
60    /// Built-in architect template
61    Architect,
62    /// Custom template loaded from a file
63    Custom(String),
64}
65
66impl PromptTemplate {
67    /// Parse a prompt template string into a PromptTemplate enum
68    #[allow(clippy::should_implement_trait)]
69    pub fn from_str(template_str: &str) -> Result<Self> {
70        match template_str.to_lowercase().as_str() {
71            "engineer" => Ok(PromptTemplate::Engineer),
72            "architect" => Ok(PromptTemplate::Architect),
73            path => {
74                // Check if the string is a valid file path
75                let path_obj = Path::new(path);
76                if path_obj.exists() && path_obj.is_file() {
77                    Ok(PromptTemplate::Custom(path.to_string()))
78                } else {
79                    Err(anyhow::anyhow!(
80                        "Invalid prompt template: '{}'. Use 'engineer', 'architect', or a valid file path.",
81                        template_str
82                    ))
83                }
84            }
85        }
86    }
87
88    /// Get the content of the prompt template
89    pub fn get_content(&self) -> Result<String> {
90        match self {
91            PromptTemplate::Engineer => Ok(ENGINEER_PROMPT.to_string()),
92            PromptTemplate::Architect => Ok(ARCHITECT_PROMPT.to_string()),
93            PromptTemplate::Custom(path) => fs::read_to_string(path)
94                .with_context(|| format!("Failed to read prompt file: {path}")),
95        }
96    }
97}