matrixcode_core/tools/workflow/
content.rs

1//! Content Generation Tool
2//!
3//! AI 内容生成工具，用于工作流的最终输出节点
4
5use anyhow::Result;
6use async_trait::async_trait;
7use serde_json::{Value, json};
8
9use crate::tools::{Tool, ToolDefinition};
10use crate::providers::Provider;
11use std::sync::Arc;
12
13/// AI 内容生成工具
14pub struct ContentGenerationTool {
15    provider: Arc<dyn Provider>,
16}
17
18impl ContentGenerationTool {
19    pub fn new(provider: Arc<dyn Provider>) -> Self {
20        Self { provider }
21    }
22
23    /// 从参数中提取主题
24    fn extract_topic<'a>(&self, params: &'a Value) -> Result<&'a str> {
25        params.get("topic")
26            .and_then(|v| v.as_str())
27            .ok_or_else(|| anyhow::anyhow!("缺少 topic 参数"))
28    }
29
30    /// 构建 AI prompt
31    fn build_prompt(&self, params: &Value) -> String {
32        let topic = params.get("topic")
33            .and_then(|v| v.as_str())
34            .unwrap_or("");
35        let style = params.get("style")
36            .and_then(|v| v.as_str())
37            .unwrap_or("informative");
38
39        let mut prompt = format!("主题: {}\n\n风格: {}\n\n", topic, style);
40
41        if let Some(research) = params.get("research_data").and_then(|v| v.as_str()) {
42            prompt.push_str(&format!("参考资料:\n{}\n\n", research));
43        }
44
45        self.append_image_instructions(&mut prompt, params);
46        prompt.push_str("\n请生成一篇完整的图文文章，**必须包含图片**。");
47
48        prompt
49    }
50
51    /// 添加图片插入指令
52    fn append_image_instructions(&self, prompt: &mut String, params: &Value) {
53        if let Some(images) = params.get("image_urls")
54            && let Some(arr) = images.as_array()
55                && !arr.is_empty() {
56                    prompt.push_str("\n**重要：请在文章中插入以下图片**（使用 Markdown 图片格式 `![描述](URL)`）：\n");
57                    for (idx, img) in arr.iter().enumerate() {
58                        let (url, desc) = self.extract_image_info(img, idx);
59                        prompt.push_str(&format!("{}. ![{}]({})\n", idx + 1, desc, url));
60                    }
61                    prompt.push_str("\n请将图片插入到文章的合适位置，使文章更加生动。\n");
62                }
63    }
64
65    /// 从图片参数中提取 URL 和描述
66    fn extract_image_info(&self, img: &Value, idx: usize) -> (String, String) {
67        if let Some(url_str) = img.as_str() {
68            (url_str.to_string(), format!("图片{}", idx + 1))
69        } else if let Some(obj) = img.as_object() {
70            let url = obj.get("url")
71                .and_then(|u| u.as_str())
72                .unwrap_or("")
73                .to_string();
74            let desc = obj.get("description")
75                .and_then(|d| d.as_str())
76                .unwrap_or("配图")
77                .to_string();
78            (url, desc)
79        } else {
80            ("".to_string(), format!("图片{}", idx + 1))
81        }
82    }
83
84    /// 构建图片画廊
85    fn build_image_gallery(&self, params: &Value) -> String {
86        if let Some(images) = params.get("image_urls")
87            && let Some(arr) = images.as_array()
88                && !arr.is_empty() {
89                    let mut gallery = String::from("\n## 📷 配图\n\n");
90                    for (idx, img) in arr.iter().enumerate().take(5) {
91                        let (url, desc) = self.extract_image_info(img, idx);
92                        if !url.is_empty() {
93                            gallery.push_str(&format!("![{}]({})\n\n", desc, url));
94                        }
95                    }
96                    return gallery;
97                }
98        String::new()
99    }
100
101    /// 调用 AI 生成内容
102    async fn call_ai(&self, prompt: String) -> Result<String> {
103        log::info!("ContentGenerationTool: starting for model {}", self.provider.model_name());
104
105        let request = crate::providers::ChatRequest {
106            messages: vec![crate::providers::Message {
107                role: crate::providers::Role::User,
108                content: crate::providers::MessageContent::Text(prompt),
109            }],
110            system: Some("你是一个专业的内容创作者。".to_string()),
111            tools: vec![],
112            think: false,
113            max_tokens: 4096,
114            server_tools: vec![],
115            enable_caching: false,
116        };
117
118        match self.provider.chat(request).await {
119            Ok(response) => {
120                log::info!("ContentGenerationTool: received successful response");
121                Ok(self.extract_text_content(response))
122            }
123            Err(e) => {
124                log::error!("ContentGenerationTool: AI call failed: {:?}", e);
125                Err(e)
126            }
127        }
128    }
129
130    /// 从响应中提取文本内容
131    fn extract_text_content(&self, response: crate::providers::ChatResponse) -> String {
132        response.content.iter()
133            .filter_map(|block| {
134                match block {
135                    crate::providers::ContentBlock::Text { text } => Some(text.clone()),
136                    _ => None,
137                }
138            })
139            .collect::<Vec<_>>()
140            .join("\n")
141    }
142
143    /// 格式化最终响应
144    fn format_response(&self, content: String, topic: &str, style: &str, params: &Value) -> String {
145        let gallery = self.build_image_gallery(params);
146        let final_content = if gallery.is_empty() {
147            content
148        } else {
149            format!("{}\n\n{}", gallery, content)
150        };
151
152        json!({
153            "content": final_content,
154            "topic": topic,
155            "style": style,
156            "word_count": final_content.chars().count()
157        }).to_string()
158    }
159
160    /// 格式化错误响应
161    fn format_error_response(&self, topic: &str, style: &str, error: anyhow::Error) -> String {
162        json!({
163            "content": format!("主题《{}》的内容生成失败: {}", topic, error),
164            "topic": topic,
165            "style": style,
166            "error": true
167        }).to_string()
168    }
169}
170
171#[async_trait]
172impl Tool for ContentGenerationTool {
173    fn definition(&self) -> ToolDefinition {
174        ToolDefinition {
175            name: "content_generation".to_string(),
176            description: "使用 AI 生成内容（文章等）。需要 Provider 支持。".to_string(),
177            parameters: json!({
178                "type": "object",
179                "properties": {
180                    "topic": {
181                        "type": "string",
182                        "description": "主题"
183                    },
184                    "research_data": {
185                        "type": "string",
186                        "description": "研究资料（可选）"
187                    },
188                    "image_urls": {
189                        "type": "array",
190                        "description": "图片 URL 列表（可选）",
191                        "items": {"type": "string"}
192                    },
193                    "style": {
194                        "type": "string",
195                        "description": "写作风格（informative/casual/professional）",
196                        "default": "informative"
197                    }
198                },
199                "required": ["topic"]
200            }),
201            ..Default::default()
202        }
203    }
204
205    async fn execute(&self, params: Value) -> Result<String> {
206        let topic = self.extract_topic(&params)?;
207        let style = params.get("style")
208            .and_then(|v| v.as_str())
209            .unwrap_or("informative");
210
211        let prompt = self.build_prompt(&params);
212
213        match self.call_ai(prompt).await {
214            Ok(content) => Ok(self.format_response(content, topic, style, &params)),
215            Err(e) => Ok(self.format_error_response(topic, style, e)),
216        }
217    }
218}
matrixcode_core/tools/workflow/content.rs

matrixcode_core/tools/workflow/
content.rs
