refactor: extract template and tools modules from agent
Split template selection and external tool management into dedicated modules for better separation of concerns.
This commit is contained in:
Fam Zheng
244
src/template.rs
Normal file
244
src/template.rs
Normal file
@@ -0,0 +1,244 @@
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::llm::{ChatMessage, LlmClient};
|
||||
use crate::tools::ExternalToolManager;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct TemplateInfo {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub match_hint: String,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub struct LoadedTemplate {
|
||||
pub id: String,
|
||||
pub info: TemplateInfo,
|
||||
pub instructions: String,
|
||||
pub external_tools: ExternalToolManager,
|
||||
pub kb_files: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
pub fn templates_dir() -> &'static str {
|
||||
if Path::new("/app/templates").is_dir() {
|
||||
"/app/templates"
|
||||
} else {
|
||||
"app-templates"
|
||||
}
|
||||
}
|
||||
|
||||
/// Scan available templates and ask LLM to pick one (or none).
|
||||
pub async fn select_template(llm: &LlmClient, requirement: &str) -> Option<String> {
|
||||
let base = Path::new(templates_dir());
|
||||
let mut entries = match tokio::fs::read_dir(base).await {
|
||||
Ok(e) => e,
|
||||
Err(_) => return None,
|
||||
};
|
||||
|
||||
let mut templates: Vec<(String, TemplateInfo)> = Vec::new();
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
// Use metadata() instead of file_type() to follow symlinks
|
||||
let is_dir = tokio::fs::metadata(entry.path())
|
||||
.await
|
||||
.map(|m| m.is_dir())
|
||||
.unwrap_or(false);
|
||||
if !is_dir {
|
||||
continue;
|
||||
}
|
||||
let id = entry.file_name().to_string_lossy().to_string();
|
||||
let meta_path = entry.path().join("template.json");
|
||||
if let Ok(data) = tokio::fs::read_to_string(&meta_path).await {
|
||||
if let Ok(info) = serde_json::from_str::<TemplateInfo>(&data) {
|
||||
templates.push((id, info));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if templates.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let listing: String = templates
|
||||
.iter()
|
||||
.map(|(id, info)| {
|
||||
format!(
|
||||
"- id: {}\n 名称: {}\n 描述: {}\n 适用场景: {}",
|
||||
id, info.name, info.description, info.match_hint
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let prompt = format!(
|
||||
"以下是可用的项目模板：\n{}\n\n用户需求：{}\n\n选择最匹配的模板 ID，如果都不合适则回复 none。只回复模板 ID 或 none，不要其他内容。",
|
||||
listing, requirement
|
||||
);
|
||||
|
||||
let response = llm
|
||||
.chat(vec![
|
||||
ChatMessage::system("你是一个模板选择助手。根据用户需求选择最合适的项目模板。只回复模板 ID 或 none。"),
|
||||
ChatMessage::user(&prompt),
|
||||
])
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
let answer = response.trim().to_lowercase();
|
||||
tracing::info!("Template selection LLM response: '{}' (available: {:?})",
|
||||
answer, templates.iter().map(|(id, _)| id.as_str()).collect::<Vec<_>>());
|
||||
if answer == "none" {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Verify the answer matches an actual template ID
|
||||
let result = templates
|
||||
.iter()
|
||||
.find(|(id, _)| id == &answer)
|
||||
.map(|(id, _)| id.clone());
|
||||
if result.is_none() {
|
||||
tracing::warn!("Template selection: LLM returned '{}' which doesn't match any template ID", answer);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Copy template contents to workdir (excluding template.json, tools/, kb/).
|
||||
pub async fn apply_template(template_id: &str, workdir: &str) -> anyhow::Result<()> {
|
||||
let src = Path::new(templates_dir()).join(template_id);
|
||||
if !src.is_dir() {
|
||||
anyhow::bail!("Template directory not found: {}", template_id);
|
||||
}
|
||||
copy_dir_recursive(&src, Path::new(workdir)).await
|
||||
}
|
||||
|
||||
/// Recursively copy directory contents, skipping template.json/tools/kb at the top level.
|
||||
async fn copy_dir_recursive(src: &Path, dst: &Path) -> anyhow::Result<()> {
|
||||
let mut stack: Vec<(std::path::PathBuf, std::path::PathBuf, bool)> =
|
||||
vec![(src.to_path_buf(), dst.to_path_buf(), true)];
|
||||
|
||||
while let Some((src_dir, dst_dir, top_level)) = stack.pop() {
|
||||
tokio::fs::create_dir_all(&dst_dir).await?;
|
||||
let mut entries = tokio::fs::read_dir(&src_dir).await?;
|
||||
while let Some(entry) = entries.next_entry().await? {
|
||||
let name = entry.file_name();
|
||||
let name_str = name.to_string_lossy();
|
||||
|
||||
// Skip metadata/convention dirs at top level
|
||||
if top_level && (name_str == "template.json" || name_str == "tools" || name_str == "kb")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let src_path = entry.path();
|
||||
let dst_path = dst_dir.join(&name);
|
||||
|
||||
if entry.file_type().await?.is_dir() {
|
||||
stack.push((src_path, dst_path, false));
|
||||
} else {
|
||||
tokio::fs::copy(&src_path, &dst_path).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl LoadedTemplate {
|
||||
/// Load a template: discover tools, read KB files, read instructions.
|
||||
pub async fn load(template_id: &str) -> anyhow::Result<Self> {
|
||||
let base = Path::new(templates_dir()).join(template_id);
|
||||
if !base.is_dir() {
|
||||
anyhow::bail!("Template directory not found: {}", template_id);
|
||||
}
|
||||
|
||||
// Read template.json
|
||||
let meta_path = base.join("template.json");
|
||||
let meta_data = tokio::fs::read_to_string(&meta_path).await?;
|
||||
let info: TemplateInfo = serde_json::from_str(&meta_data)?;
|
||||
|
||||
// Read INSTRUCTIONS.md
|
||||
let instructions_path = base.join("INSTRUCTIONS.md");
|
||||
let instructions = tokio::fs::read_to_string(&instructions_path)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
// Discover external tools
|
||||
let tools_dir = base.join("tools");
|
||||
let external_tools = ExternalToolManager::discover(&tools_dir).await;
|
||||
tracing::info!(
|
||||
"Template '{}': discovered {} external tools",
|
||||
template_id,
|
||||
external_tools.len()
|
||||
);
|
||||
|
||||
// Scan KB files
|
||||
let kb_dir = base.join("kb");
|
||||
let kb_files = scan_kb_files(&kb_dir).await;
|
||||
tracing::info!(
|
||||
"Template '{}': found {} KB files",
|
||||
template_id,
|
||||
kb_files.len()
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
id: template_id.to_string(),
|
||||
info,
|
||||
instructions,
|
||||
external_tools,
|
||||
kb_files,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Scan kb/ directory for .md files. Returns (title, content) pairs.
|
||||
/// Title is extracted from the first `# heading` line, or falls back to the filename.
|
||||
async fn scan_kb_files(kb_dir: &Path) -> Vec<(String, String)> {
|
||||
let mut results = Vec::new();
|
||||
|
||||
let mut entries = match tokio::fs::read_dir(kb_dir).await {
|
||||
Ok(e) => e,
|
||||
Err(_) => return results,
|
||||
};
|
||||
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
let path = entry.path();
|
||||
|
||||
// Resolve symlinks
|
||||
let real_path = match tokio::fs::canonicalize(&path).await {
|
||||
Ok(p) => p,
|
||||
Err(_) => path.clone(),
|
||||
};
|
||||
|
||||
// Only process .md files
|
||||
let ext = real_path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or("");
|
||||
if ext != "md" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let content = match tokio::fs::read_to_string(&real_path).await {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to read KB file {}: {}", real_path.display(), e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Extract title: first `# heading` line, or filename without extension
|
||||
let title = content
|
||||
.lines()
|
||||
.find(|l| l.starts_with("# "))
|
||||
.map(|l| l.trim_start_matches("# ").trim().to_string())
|
||||
.unwrap_or_else(|| {
|
||||
path.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or("untitled")
|
||||
.to_string()
|
||||
});
|
||||
|
||||
results.push((title, content));
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
Reference in New Issue
Block a user