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
211
src/agent.rs
211
src/agent.rs
@@ -8,6 +8,8 @@ use tokio::sync::{mpsc, RwLock, broadcast};
|
||||
|
||||
use crate::llm::{LlmClient, ChatMessage, Tool, ToolFunction};
|
||||
use crate::exec::LocalExecutor;
|
||||
use crate::template::{self, LoadedTemplate};
|
||||
use crate::tools::ExternalToolManager;
|
||||
use crate::LlmConfig;
|
||||
|
||||
use crate::state::{AgentState, AgentPhase, Step, StepStatus};
|
||||
@@ -131,115 +133,7 @@ impl AgentManager {
|
||||
}
|
||||
}
|
||||
|
||||
// --- Template system ---
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TemplateInfo {
|
||||
name: String,
|
||||
description: String,
|
||||
match_hint: String,
|
||||
}
|
||||
|
||||
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).
|
||||
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 {
|
||||
if !entry.file_type().await.map(|t| t.is_dir()).unwrap_or(false) {
|
||||
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();
|
||||
if answer == "none" {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Verify the answer matches an actual template ID
|
||||
templates.iter().find(|(id, _)| id == &answer).map(|(id, _)| id.clone())
|
||||
}
|
||||
|
||||
/// Copy template contents to workdir (excluding template.json).
|
||||
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 at the top level.
|
||||
async fn copy_dir_recursive(src: &Path, dst: &Path) -> anyhow::Result<()> {
|
||||
// Use a stack to avoid async recursion
|
||||
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();
|
||||
|
||||
if top_level && name_str == "template.json" {
|
||||
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(())
|
||||
}
|
||||
// Template system is in crate::template
|
||||
|
||||
/// Read INSTRUCTIONS.md from workdir if it exists.
|
||||
async fn read_instructions(workdir: &str) -> String {
|
||||
@@ -312,25 +206,90 @@ async fn agent_loop(
|
||||
.await;
|
||||
|
||||
// Template selection + workspace setup
|
||||
let template_id = select_template(&llm, &requirement).await;
|
||||
if let Some(ref tid) = template_id {
|
||||
let template_id = template::select_template(&llm, &requirement).await;
|
||||
let loaded_template = if let Some(ref tid) = template_id {
|
||||
tracing::info!("Template selected for workflow {}: {}", workflow_id, tid);
|
||||
let _ = tokio::fs::create_dir_all(&workdir).await;
|
||||
if let Err(e) = apply_template(tid, &workdir).await {
|
||||
if let Err(e) = template::apply_template(tid, &workdir).await {
|
||||
tracing::error!("Failed to apply template {}: {}", tid, e);
|
||||
}
|
||||
match LoadedTemplate::load(tid).await {
|
||||
Ok(t) => Some(t),
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to load template {}: {}", tid, e);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Import KB files from template
|
||||
if let Some(ref t) = loaded_template {
|
||||
if let Some(ref kb) = mgr.kb {
|
||||
let mut batch_items: Vec<(String, String)> = Vec::new();
|
||||
for (title, content) in &t.kb_files {
|
||||
// Check if article already exists by title
|
||||
let existing: Option<String> = sqlx::query_scalar(
|
||||
"SELECT id FROM kb_articles WHERE title = ?"
|
||||
)
|
||||
.bind(title)
|
||||
.fetch_optional(&pool)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
|
||||
let article_id = if let Some(id) = existing {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE kb_articles SET content = ?, updated_at = datetime('now') WHERE id = ?"
|
||||
)
|
||||
.bind(content)
|
||||
.bind(&id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
id
|
||||
} else {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO kb_articles (id, title, content) VALUES (?, ?, ?)"
|
||||
)
|
||||
.bind(&id)
|
||||
.bind(title)
|
||||
.bind(content)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
id
|
||||
};
|
||||
|
||||
batch_items.push((article_id, content.clone()));
|
||||
}
|
||||
// Batch index: single embed.py call for all articles
|
||||
if !batch_items.is_empty() {
|
||||
if let Err(e) = kb.index_batch(&batch_items).await {
|
||||
tracing::warn!("Failed to batch index KB articles: {}", e);
|
||||
}
|
||||
}
|
||||
tracing::info!("Imported {} KB articles from template", t.kb_files.len());
|
||||
}
|
||||
}
|
||||
|
||||
ensure_workspace(&exec, &workdir).await;
|
||||
let _ = tokio::fs::write(format!("{}/requirement.md", workdir), &requirement).await;
|
||||
|
||||
let instructions = read_instructions(&workdir).await;
|
||||
let instructions = if let Some(ref t) = loaded_template {
|
||||
t.instructions.clone()
|
||||
} else {
|
||||
read_instructions(&workdir).await
|
||||
};
|
||||
|
||||
let ext_tools = loaded_template.as_ref().map(|t| &t.external_tools);
|
||||
|
||||
tracing::info!("Starting agent loop for workflow {}", workflow_id);
|
||||
// Run tool-calling agent loop
|
||||
let result = run_agent_loop(
|
||||
&llm, &exec, &pool, &broadcast_tx,
|
||||
&project_id, &workflow_id, &requirement, &workdir, &mgr,
|
||||
&instructions, None,
|
||||
&instructions, None, ext_tools,
|
||||
).await;
|
||||
|
||||
let final_status = if result.is_ok() { "done" } else { "failed" };
|
||||
@@ -433,10 +392,12 @@ async fn agent_loop(
|
||||
|
||||
let instructions = read_instructions(&workdir).await;
|
||||
|
||||
// Try to detect which template was used (check for tools/ in workdir parent template)
|
||||
// For comments, we don't re-load the template — external tools are not available in feedback resume
|
||||
let result = run_agent_loop(
|
||||
&llm, &exec, &pool, &broadcast_tx,
|
||||
&project_id, &workflow_id, &wf.requirement, &workdir, &mgr,
|
||||
&instructions, Some(state),
|
||||
&instructions, Some(state), None,
|
||||
).await;
|
||||
|
||||
let final_status = if result.is_ok() { "done" } else { "failed" };
|
||||
@@ -990,9 +951,13 @@ async fn run_agent_loop(
|
||||
mgr: &Arc<AgentManager>,
|
||||
instructions: &str,
|
||||
initial_state: Option<AgentState>,
|
||||
external_tools: Option<&ExternalToolManager>,
|
||||
) -> anyhow::Result<()> {
|
||||
let planning_tools = build_planning_tools();
|
||||
let execution_tools = build_execution_tools();
|
||||
let mut execution_tools = build_execution_tools();
|
||||
if let Some(ext) = external_tools {
|
||||
execution_tools.extend(ext.tool_definitions());
|
||||
}
|
||||
|
||||
let mut state = initial_state.unwrap_or_else(AgentState::new);
|
||||
|
||||
@@ -1247,6 +1212,20 @@ async fn run_agent_loop(
|
||||
state.current_step_chat_history.push(ChatMessage::tool_result(&tc.id, &result));
|
||||
}
|
||||
|
||||
// External tools (convention-based)
|
||||
name if external_tools.as_ref().is_some_and(|e| e.has_tool(name)) => {
|
||||
let result = match external_tools.unwrap().invoke(name, &tc.function.arguments, workdir).await {
|
||||
Ok(output) => {
|
||||
let truncated = truncate_str(&output, 8192);
|
||||
truncated.to_string()
|
||||
}
|
||||
Err(e) => format!("Tool error: {}", e),
|
||||
};
|
||||
let status = if result.starts_with("Tool error:") { "failed" } else { "done" };
|
||||
log_execution(pool, broadcast_tx, workflow_id, cur, &tc.function.name, &tc.function.arguments, &result, status).await;
|
||||
state.current_step_chat_history.push(ChatMessage::tool_result(&tc.id, &result));
|
||||
}
|
||||
|
||||
// IO tools: execute, read_file, write_file, list_files
|
||||
_ => {
|
||||
let result = execute_tool(&tc.function.name, &tc.function.arguments, workdir, exec).await;
|
||||
|
||||
Reference in New Issue
Block a user