273 lines
8.1 KiB
Rust
273 lines
8.1 KiB
Rust
use std::collections::HashMap;
|
|
|
|
/// Modular system prompt builder inspired by pi's `buildSystemPrompt`.
|
|
///
|
|
/// Supports:
|
|
/// - Base prompt (replaceable or appendable)
|
|
/// - Tool snippets injected into an "Available tools" section
|
|
/// - Project context files (AGENTS.md, etc.)
|
|
/// - Skills injection
|
|
/// - Variable substitution ({{key}})
|
|
/// - Metadata (date)
|
|
///
|
|
/// # Example
|
|
/// ```rust
|
|
/// use ai::agent::prompt_builder::SystemPromptBuilder;
|
|
///
|
|
/// let prompt = SystemPromptBuilder::new()
|
|
/// .base_prompt("You are a helpful assistant.")
|
|
/// .tool_snippet("bash", "Execute shell commands")
|
|
/// .tool_snippet("read", "Read file contents")
|
|
/// .project_context("AGENTS.md", "# Project Rules\n- Follow conventions")
|
|
/// .variable("repo_name", "gitdataai")
|
|
/// .build();
|
|
/// ```
|
|
#[derive(Clone, Debug)]
|
|
pub struct SystemPromptBuilder {
|
|
base_prompt: Option<String>,
|
|
append_prompt: Option<String>,
|
|
tool_snippets: Vec<(String, String)>,
|
|
tool_guidelines: Vec<String>,
|
|
project_contexts: Vec<(String, String)>,
|
|
skills: Vec<String>,
|
|
variables: HashMap<String, String>,
|
|
date: Option<String>,
|
|
custom_sections: Vec<(String, String)>,
|
|
}
|
|
|
|
impl SystemPromptBuilder {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
base_prompt: None,
|
|
append_prompt: None,
|
|
tool_snippets: Vec::new(),
|
|
tool_guidelines: Vec::new(),
|
|
project_contexts: Vec::new(),
|
|
skills: Vec::new(),
|
|
variables: HashMap::new(),
|
|
date: None,
|
|
custom_sections: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Set the base system prompt. Replaces the default prompt.
|
|
pub fn base_prompt(mut self, prompt: impl Into<String>) -> Self {
|
|
self.base_prompt = Some(prompt.into());
|
|
self
|
|
}
|
|
|
|
/// Append additional text to the system prompt after the base.
|
|
pub fn append_prompt(mut self, text: impl Into<String>) -> Self {
|
|
self.append_prompt = Some(text.into());
|
|
self
|
|
}
|
|
|
|
/// Add a one-line tool description snippet.
|
|
pub fn tool_snippet(
|
|
mut self,
|
|
tool_name: impl Into<String>,
|
|
description: impl Into<String>,
|
|
) -> Self {
|
|
self.tool_snippets
|
|
.push((tool_name.into(), description.into()));
|
|
self
|
|
}
|
|
|
|
/// Add a guideline bullet for the tools section.
|
|
pub fn tool_guideline(mut self, guideline: impl Into<String>) -> Self {
|
|
self.tool_guidelines.push(guideline.into());
|
|
self
|
|
}
|
|
|
|
/// Add a project context file (e.g., AGENTS.md content).
|
|
pub fn project_context(
|
|
mut self,
|
|
path: impl Into<String>,
|
|
content: impl Into<String>,
|
|
) -> Self {
|
|
self.project_contexts.push((path.into(), content.into()));
|
|
self
|
|
}
|
|
|
|
/// Add a skill description to inject into the prompt.
|
|
pub fn skill(mut self, skill_description: impl Into<String>) -> Self {
|
|
self.skills.push(skill_description.into());
|
|
self
|
|
}
|
|
|
|
/// Set a variable for {{key}} substitution.
|
|
pub fn variable(
|
|
mut self,
|
|
key: impl Into<String>,
|
|
value: impl Into<String>,
|
|
) -> Self {
|
|
self.variables.insert(key.into(), value.into());
|
|
self
|
|
}
|
|
|
|
/// Set multiple variables from an iterator.
|
|
pub fn variables(
|
|
mut self,
|
|
vars: impl IntoIterator<Item = (String, String)>,
|
|
) -> Self {
|
|
self.variables.extend(vars);
|
|
self
|
|
}
|
|
|
|
/// Set the date metadata (ISO format: YYYY-MM-DD).
|
|
pub fn date(mut self, date: impl Into<String>) -> Self {
|
|
self.date = Some(date.into());
|
|
self
|
|
}
|
|
|
|
/// Add a custom named section to the prompt.
|
|
pub fn custom_section(
|
|
mut self,
|
|
name: impl Into<String>,
|
|
content: impl Into<String>,
|
|
) -> Self {
|
|
self.custom_sections.push((name.into(), content.into()));
|
|
self
|
|
}
|
|
|
|
/// Build the final system prompt string.
|
|
pub fn build(self) -> String {
|
|
let mut parts: Vec<String> = Vec::new();
|
|
|
|
// 1. Base prompt
|
|
if let Some(base) = &self.base_prompt {
|
|
parts.push(base.clone());
|
|
}
|
|
|
|
// 2. Append prompt
|
|
if let Some(append) = &self.append_prompt {
|
|
parts.push(append.clone());
|
|
}
|
|
|
|
// 3. Tool snippets section
|
|
if !self.tool_snippets.is_empty() {
|
|
let mut section = String::from("\n## Available Tools\n");
|
|
for (name, desc) in &self.tool_snippets {
|
|
section.push_str(&format!("- `{name}`: {desc}\n"));
|
|
}
|
|
if !self.tool_guidelines.is_empty() {
|
|
section.push_str("\n### Tool Guidelines\n");
|
|
for guideline in &self.tool_guidelines {
|
|
section.push_str(&format!("- {guideline}\n"));
|
|
}
|
|
}
|
|
parts.push(section);
|
|
}
|
|
|
|
// 4. Project context files
|
|
if !self.project_contexts.is_empty() {
|
|
let mut section = String::from("\n<project_context>\n\n");
|
|
section
|
|
.push_str("Project-specific instructions and guidelines:\n\n");
|
|
for (path, content) in &self.project_contexts {
|
|
section.push_str(&format!("<project_instructions path=\"{path}\">\n{content}\n</project_instructions>\n\n"));
|
|
}
|
|
section.push_str("</project_context>");
|
|
parts.push(section);
|
|
}
|
|
|
|
// 5. Skills section
|
|
if !self.skills.is_empty() {
|
|
let mut section = String::from("\n## Available Skills\n");
|
|
for skill in &self.skills {
|
|
section.push_str(&format!("{skill}\n"));
|
|
}
|
|
parts.push(section);
|
|
}
|
|
|
|
// 6. Custom sections
|
|
for (name, content) in &self.custom_sections {
|
|
parts.push(format!("\n## {name}\n{content}"));
|
|
}
|
|
|
|
// 7. Metadata footer
|
|
if let Some(date) = &self.date {
|
|
parts.push(format!("\nCurrent date: {date}"));
|
|
}
|
|
|
|
let mut result = parts.join("\n");
|
|
|
|
// 8. Variable substitution
|
|
for (key, value) in &self.variables {
|
|
let placeholder = format!("{{{{{}}}}}", key);
|
|
result = result.replace(&placeholder, value);
|
|
}
|
|
|
|
result
|
|
}
|
|
}
|
|
|
|
impl Default for SystemPromptBuilder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_basic_build() {
|
|
let prompt = SystemPromptBuilder::new()
|
|
.base_prompt("You are a helpful assistant.")
|
|
.date("2026-05-29")
|
|
.build();
|
|
|
|
assert!(prompt.contains("You are a helpful assistant."));
|
|
assert!(prompt.contains("Current date: 2026-05-29"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_variable_substitution() {
|
|
let prompt = SystemPromptBuilder::new()
|
|
.base_prompt("Repo: {{repo_name}}, User: {{user}}")
|
|
.variable("repo_name", "gitdataai")
|
|
.variable("user", "zhenyi")
|
|
.build();
|
|
|
|
assert!(prompt.contains("Repo: gitdataai"));
|
|
assert!(prompt.contains("User: zhenyi"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_snippets() {
|
|
let prompt = SystemPromptBuilder::new()
|
|
.base_prompt("Agent prompt.")
|
|
.tool_snippet("bash", "Execute shell commands")
|
|
.tool_snippet("read", "Read file contents")
|
|
.build();
|
|
|
|
assert!(prompt.contains("## Available Tools"));
|
|
assert!(prompt.contains("`bash`: Execute shell commands"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_project_context() {
|
|
let prompt = SystemPromptBuilder::new()
|
|
.base_prompt("Base.")
|
|
.project_context("AGENTS.md", "# Rules\n- Follow conventions")
|
|
.build();
|
|
|
|
assert!(prompt.contains("<project_context>"));
|
|
assert!(prompt.contains("AGENTS.md"));
|
|
assert!(prompt.contains("Follow conventions"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_custom_section() {
|
|
let prompt = SystemPromptBuilder::new()
|
|
.base_prompt("Base.")
|
|
.custom_section("Memory", "Remember: user prefers Rust")
|
|
.build();
|
|
|
|
assert!(prompt.contains("## Memory"));
|
|
assert!(prompt.contains("user prefers Rust"));
|
|
}
|
|
}
|