lc_tools/
wikipedia.rs

1// lc-tools/src/wikipedia.rs
2//! Wikipedia search tool
3//!
4//! Searches and fetches encyclopedia entry content via the 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 tool input
13#[derive(Debug, Deserialize, JsonSchema)]
14pub struct WikipediaInput {
15    /// The search query
16    pub query: String,
17    /// Number of results to return (default: 3)
18    pub top_k: Option<usize>,
19    /// Language (default: zh, supports en/zh/ja, etc.)
20    pub lang: Option<String>,
21    /// Whether to fetch full content (default false, summary only)
22    pub full_content: Option<bool>,
23}
24
25/// Wikipedia tool output
26#[derive(Debug, Serialize)]
27pub struct WikipediaOutput {
28    /// The query
29    pub query: String,
30    /// The result list
31    pub results: Vec<WikipediaResult>,
32    /// Number of results
33    pub total: usize,
34}
35
36#[derive(Debug, Serialize)]
37pub struct WikipediaResult {
38    /// Title
39    pub title: String,
40    /// Snippet or content
41    pub snippet: String,
42    /// Full page URL
43    pub url: String,
44}
45
46/// Wikipedia search tool
47pub struct WikipediaTool {
48    client: reqwest::Client,
49}
50
51impl WikipediaTool {
52    /// Creates a Wikipedia search tool.
53    pub fn new() -> Self {
54        Self {
55            client: reqwest::Client::builder()
56                .timeout(std::time::Duration::from_secs(15))
57                .user_agent("LangChainRust/0.1 (Wikipedia Tool)")
58                .build()
59                .unwrap_or_else(|_| reqwest::Client::new()),
60        }
61    }
62}
63
64impl Default for WikipediaTool {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70impl WikipediaTool {
71    /// Searches Wikipedia entries
72    async fn search(
73        &self,
74        query: &str,
75        top_k: usize,
76        lang: &str,
77    ) -> Result<WikipediaOutput, ToolError> {
78        let search_url = format!(
79            "https://{}.wikipedia.org/w/api.php?action=query&list=search&srsearch={}&format=json&srlimit={}",
80            lang, urlencoding(query), top_k
81        );
82
83        let response =
84            self.client.get(&search_url).send().await.map_err(|e| {
85                ToolError::ExecutionFailed(format!("Wikipedia search failed: {}", e))
86            })?;
87
88        let body: serde_json::Value = response.json().await.map_err(|e| {
89            ToolError::ExecutionFailed(format!("failed to parse search results: {}", e))
90        })?;
91
92        let search_results = body["query"]["search"]
93            .as_array()
94            .map(|arr| arr.to_vec())
95            .unwrap_or_default();
96
97        let mut results = Vec::new();
98
99        for item in search_results.iter().take(top_k) {
100            let title = item["title"].as_str().unwrap_or("").to_string();
101            let snippet_html = item["snippet"].as_str().unwrap_or("").to_string();
102            let snippet = strip_html(&snippet_html);
103            let page_url = format!(
104                "https://{}.wikipedia.org/wiki/{}",
105                lang,
106                urlencoding(&title)
107            );
108
109            results.push(WikipediaResult {
110                title,
111                snippet,
112                url: page_url,
113            });
114        }
115
116        Ok(WikipediaOutput {
117            query: query.to_string(),
118            total: results.len(),
119            results,
120        })
121    }
122
123    /// Fetches the full content of an entry
124    async fn get_full_content(&self, title: &str, lang: &str) -> Result<String, ToolError> {
125        let url = format!(
126            "https://{}.wikipedia.org/w/api.php?action=query&prop=extracts&exintro&explaintext&titles={}&format=json",
127            lang, urlencoding(title)
128        );
129
130        let response = self.client.get(&url).send().await.map_err(|e| {
131            ToolError::ExecutionFailed(format!("failed to fetch page content: {}", e))
132        })?;
133
134        let body: serde_json::Value = response.json().await.map_err(|e| {
135            ToolError::ExecutionFailed(format!("failed to parse page content: {}", e))
136        })?;
137
138        let pages = body["query"]["pages"]
139            .as_object()
140            .cloned()
141            .unwrap_or_default();
142        for (_, page) in pages {
143            if let Some(extract) = page["extract"].as_str() {
144                if extract.len() > 5000 {
145                    return Ok(
146                        extract.chars().take(5000).collect::<String>() + "\n... [内容已截断]"
147                    );
148                }
149                return Ok(extract.to_string());
150            }
151        }
152
153        Err(ToolError::ExecutionFailed(
154            "page content not found".to_string(),
155        ))
156    }
157}
158
159/// Strips HTML tags
160fn strip_html(html: &str) -> String {
161    static TAG_RE: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
162        regex::Regex::new(r"<[^>]+>").expect("static regex literal must compile")
163    });
164    static WHITESPACE_RE: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
165        regex::Regex::new(r"\s+").expect("static regex literal must compile")
166    });
167
168    let result = TAG_RE.replace_all(html, "");
169    WHITESPACE_RE.replace_all(&result, " ").trim().to_string()
170}
171
172/// URL encoding (using the urlencoding crate for complete encoding)
173fn urlencoding(s: &str) -> String {
174    urlencoding::encode(s).to_string()
175}
176
177#[async_trait]
178impl Tool for WikipediaTool {
179    type Input = WikipediaInput;
180    type Output = WikipediaOutput;
181
182    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
183        let top_k = input.top_k.unwrap_or(3);
184        let lang = input.lang.as_deref().unwrap_or("zh");
185        let full = input.full_content.unwrap_or(false);
186
187        if input.query.trim().is_empty() {
188            return Err(ToolError::InvalidInput(
189                "query must not be empty".to_string(),
190            ));
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 parse failed: {}", 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}
lc_tools/wikipedia.rs

lc_tools/
wikipedia.rs
