lc_tools/
wikipedia.rs

1// lc-tools/src/wikipedia.rs
2//! Wikipedia 搜索工具
3//!
4//! 通过 Wikipedia API 搜索和获取百科条目内容。
5
6use async_trait::async_trait;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10use lc_core::tools::{BaseTool, Tool, ToolError};
11
12/// Wikipedia 工具输入
13#[derive(Debug, Deserialize, JsonSchema)]
14pub struct WikipediaInput {
15    /// 搜索查询
16    pub query: String,
17    /// 返回结果数量（默认 3）
18    pub top_k: Option<usize>,
19    /// 语言（默认 zh，支持 en/zh/ja 等）
20    pub lang: Option<String>,
21    /// 是否获取完整内容（默认 false，只获取摘要）
22    pub full_content: Option<bool>,
23}
24
25/// Wikipedia 工具输出
26#[derive(Debug, Serialize)]
27pub struct WikipediaOutput {
28    /// 查询
29    pub query: String,
30    /// 结果列表
31    pub results: Vec<WikipediaResult>,
32    /// 结果数量
33    pub total: usize,
34}
35
36#[derive(Debug, Serialize)]
37pub struct WikipediaResult {
38    /// 标题
39    pub title: String,
40    /// 摘要或内容
41    pub snippet: String,
42    /// 完整页面 URL
43    pub url: String,
44}
45
46/// Wikipedia 搜索工具
47pub struct WikipediaTool {
48    client: reqwest::Client,
49}
50
51impl WikipediaTool {
52    pub fn new() -> Self {
53        Self {
54            client: reqwest::Client::builder()
55                .timeout(std::time::Duration::from_secs(15))
56                .user_agent("LangChainRust/0.1 (Wikipedia Tool)")
57                .build()
58                .unwrap_or_else(|_| reqwest::Client::new()),
59        }
60    }
61}
62
63impl Default for WikipediaTool {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69impl WikipediaTool {
70    /// 搜索 Wikipedia 条目
71    async fn search(
72        &self,
73        query: &str,
74        top_k: usize,
75        lang: &str,
76    ) -> Result<WikipediaOutput, ToolError> {
77        let search_url = format!(
78            "https://{}.wikipedia.org/w/api.php?action=query&list=search&srsearch={}&format=json&srlimit={}",
79            lang, urlencoding(query), top_k
80        );
81
82        let response = self
83            .client
84            .get(&search_url)
85            .send()
86            .await
87            .map_err(|e| ToolError::ExecutionFailed(format!("Wikipedia 搜索失败: {}", e)))?;
88
89        let body: serde_json::Value = response
90            .json()
91            .await
92            .map_err(|e| ToolError::ExecutionFailed(format!("解析搜索结果失败: {}", e)))?;
93
94        let search_results = body["query"]["search"]
95            .as_array()
96            .map(|arr| arr.to_vec())
97            .unwrap_or_default();
98
99        let mut results = Vec::new();
100
101        for item in search_results.iter().take(top_k) {
102            let title = item["title"].as_str().unwrap_or("").to_string();
103            let snippet_html = item["snippet"].as_str().unwrap_or("").to_string();
104            let snippet = strip_html(&snippet_html);
105            let page_url = format!(
106                "https://{}.wikipedia.org/wiki/{}",
107                lang,
108                urlencoding(&title)
109            );
110
111            results.push(WikipediaResult {
112                title,
113                snippet,
114                url: page_url,
115            });
116        }
117
118        Ok(WikipediaOutput {
119            query: query.to_string(),
120            total: results.len(),
121            results,
122        })
123    }
124
125    /// 获取条目完整内容
126    async fn get_full_content(&self, title: &str, lang: &str) -> Result<String, ToolError> {
127        let url = format!(
128            "https://{}.wikipedia.org/w/api.php?action=query&prop=extracts&exintro&explaintext&titles={}&format=json",
129            lang, urlencoding(title)
130        );
131
132        let response = self
133            .client
134            .get(&url)
135            .send()
136            .await
137            .map_err(|e| ToolError::ExecutionFailed(format!("获取页面内容失败: {}", e)))?;
138
139        let body: serde_json::Value = response
140            .json()
141            .await
142            .map_err(|e| ToolError::ExecutionFailed(format!("解析页面内容失败: {}", e)))?;
143
144        let pages = body["query"]["pages"]
145            .as_object()
146            .cloned()
147            .unwrap_or_default();
148        for (_, page) in pages {
149            if let Some(extract) = page["extract"].as_str() {
150                if extract.len() > 5000 {
151                    return Ok(
152                        extract.chars().take(5000).collect::<String>() + "\n... [内容已截断]"
153                    );
154                }
155                return Ok(extract.to_string());
156            }
157        }
158
159        Err(ToolError::ExecutionFailed("未找到页面内容".to_string()))
160    }
161}
162
163/// 去除 HTML 标签
164fn strip_html(html: &str) -> String {
165    static TAG_RE: std::sync::LazyLock<regex::Regex> =
166        std::sync::LazyLock::new(|| regex::Regex::new(r"<[^>]+>").unwrap());
167    static WHITESPACE_RE: std::sync::LazyLock<regex::Regex> =
168        std::sync::LazyLock::new(|| regex::Regex::new(r"\s+").unwrap());
169
170    let result = TAG_RE.replace_all(html, "");
171    WHITESPACE_RE.replace_all(&result, " ").trim().to_string()
172}
173
174/// URL 编码 (using urlencoding crate for complete encoding)
175fn urlencoding(s: &str) -> String {
176    urlencoding::encode(s).to_string()
177}
178
179#[async_trait]
180impl Tool for WikipediaTool {
181    type Input = WikipediaInput;
182    type Output = WikipediaOutput;
183
184    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
185        let top_k = input.top_k.unwrap_or(3);
186        let lang = input.lang.as_deref().unwrap_or("zh");
187        let full = input.full_content.unwrap_or(false);
188
189        if input.query.trim().is_empty() {
190            return Err(ToolError::InvalidInput("查询不能为空".to_string()));
191        }
192
193        let mut output = self.search(&input.query, top_k, lang).await?;
194
195        if full {
196            for result in &mut output.results {
197                if let Ok(content) = self.get_full_content(&result.title, lang).await {
198                    result.snippet = content;
199                }
200            }
201        }
202
203        Ok(output)
204    }
205}
206
207#[async_trait]
208impl BaseTool for WikipediaTool {
209    fn name(&self) -> &str {
210        "wikipedia"
211    }
212
213    fn description(&self) -> &str {
214        "Wikipedia 百科搜索工具。搜索 Wikipedia 百科条目并返回摘要或完整内容。
215
216参数:
217- query: 搜索关键词
218- top_k: 返回结果数量（默认 3）
219- lang: 语言代码，如 zh/en/ja（默认 zh）
220- full_content: 是否获取完整内容（默认 false）
221
222示例:
223- 搜索百科: {\"query\": \"Rust\", \"lang\": \"zh\"}
224- 获取详细内容: {\"query\": \"Rust\", \"lang\": \"en\", \"full_content\": true}"
225    }
226
227    async fn run(&self, input: String) -> Result<String, ToolError> {
228        let parsed: WikipediaInput = serde_json::from_str(&input)
229            .map_err(|e| ToolError::InvalidInput(format!("JSON 解析失败: {}", e)))?;
230
231        let output = self.invoke(parsed).await?;
232
233        let mut text = format!("Wikipedia 搜索结果 (查询: {})\n\n", output.query);
234        for (i, result) in output.results.iter().enumerate() {
235            text.push_str(&format!("{}. {}\n", i + 1, result.title));
236            text.push_str(&format!("   {}\n", result.snippet));
237            text.push_str(&format!("   URL: {}\n\n", result.url));
238        }
239        text.push_str(&format!("共 {} 条结果", output.total));
240
241        Ok(text)
242    }
243
244    fn args_schema(&self) -> Option<serde_json::Value> {
245        use schemars::schema_for;
246        serde_json::to_value(schema_for!(WikipediaInput)).ok()
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253
254    #[test]
255    fn test_wikipedia_tool_properties() {
256        let tool = WikipediaTool::new();
257        assert_eq!(tool.name(), "wikipedia");
258        assert!(tool.description().contains("Wikipedia"));
259        assert!(BaseTool::args_schema(&tool).is_some());
260    }
261
262    #[tokio::test]
263    async fn test_wikipedia_empty_query() {
264        let tool = WikipediaTool::new();
265        let result = tool.run(r#"{"query": ""}"#.to_string()).await;
266        assert!(result.is_err());
267    }
268
269    #[test]
270    fn test_strip_html() {
271        let html = "<p>Hello <b>World</b></p>";
272        assert_eq!(strip_html(html), "Hello World");
273    }
274
275    #[test]
276    fn test_urlencoding() {
277        let encoded = urlencoding("Rust programming");
278        assert_eq!(encoded, "Rust%20programming");
279        assert!(urlencoding("a&b=c#d?e").contains("%26"));
280        assert!(urlencoding("a&b=c#d?e").contains("%3D"));
281    }
282
283    #[tokio::test]
284    #[ignore = "需要网络连接"]
285    async fn test_wikipedia_search_real() {
286        let tool = WikipediaTool::new();
287        let result = tool
288            .invoke(WikipediaInput {
289                query: "Rust".to_string(),
290                top_k: Some(2),
291                lang: Some("en".into()),
292                full_content: Some(false),
293            })
294            .await
295            .unwrap();
296
297        assert!(!result.results.is_empty());
298        assert!(result.results[0].title.to_lowercase().contains("rust"));
299    }
300}
lc_tools/wikipedia.rs

lc_tools/
wikipedia.rs
