use crate::generator::compose::types::AgentType;
use crate::generator::{compose::memory::MemoryScope, context::GeneratorContext};
use anyhow::Result;
use std::collections::HashMap;
use std::fs;
pub mod summary_generator;
pub mod summary_outlet;
pub mod fixer;
pub use summary_outlet::SummaryOutlet;
pub use fixer::MermaidFixer;
pub trait Outlet {
async fn save(&self, context: &GeneratorContext) -> Result<()>;
}
pub struct DocTree {
structure: HashMap<String, String>,
}
impl DocTree {
pub fn insert(&mut self, scoped_key: &str, relative_path: &str) {
self.structure
.insert(scoped_key.to_string(), relative_path.to_string());
}
}
impl Default for DocTree {
fn default() -> Self {
let structure = HashMap::from([
(
AgentType::Overview.to_string(),
"1、项目概述.md".to_string(),
),
(
AgentType::Architecture.to_string(),
"2、架构概览.md".to_string(),
),
(
AgentType::Workflow.to_string(),
"3、工作流程.md".to_string(),
),
(
AgentType::Boundary.to_string(),
"5、边界调用.md".to_string(),
),
]);
Self { structure }
}
}
pub struct DiskOutlet {
doc_tree: DocTree,
}
impl DiskOutlet {
pub fn new(doc_tree: DocTree) -> Self {
Self { doc_tree }
}
}
impl Outlet for DiskOutlet {
async fn save(&self, context: &GeneratorContext) -> Result<()> {
println!("\n🖊️ 文档存储中...");
let output_dir = &context.config.output_path;
if output_dir.exists() {
fs::remove_dir_all(output_dir)?;
}
fs::create_dir_all(output_dir)?;
for (scoped_key, relative_path) in &self.doc_tree.structure {
if let Some(doc_markdown) = context
.get_from_memory::<String>(MemoryScope::DOCUMENTATION, scoped_key)
.await
{
let output_file_path = output_dir.join(relative_path);
if let Some(parent_dir) = output_file_path.parent() {
if !parent_dir.exists() {
fs::create_dir_all(parent_dir)?;
}
}
fs::write(&output_file_path, doc_markdown)?;
println!("💾 已保存文档: {}", output_file_path.display());
} else {
eprintln!("⚠️ 警告: 未找到文档内容,键: {}", scoped_key);
}
}
println!("💾 文档保存完成,输出目录: {}", output_dir.display());
if let Err(e) = MermaidFixer::auto_fix_after_output(context).await {
eprintln!("⚠️ mermaid图表修复过程中出现错误: {}", e);
eprintln!("💡 这不会影响文档生成的主要流程");
}
Ok(())
}
}