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

matrixcode_core/tools/workflow/
content.rs
