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> =
162        std::sync::LazyLock::new(|| regex::Regex::new(r"<[^>]+>").unwrap());
163    static WHITESPACE_RE: std::sync::LazyLock<regex::Regex> =
164        std::sync::LazyLock::new(|| regex::Regex::new(r"\s+").unwrap());
165
166    let result = TAG_RE.replace_all(html, "");
167    WHITESPACE_RE.replace_all(&result, " ").trim().to_string()
168}
169
170/// URL encoding (using the urlencoding crate for complete encoding)
171fn urlencoding(s: &str) -> String {
172    urlencoding::encode(s).to_string()
173}
174
175#[async_trait]
176impl Tool for WikipediaTool {
177    type Input = WikipediaInput;
178    type Output = WikipediaOutput;
179
180    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
181        let top_k = input.top_k.unwrap_or(3);
182        let lang = input.lang.as_deref().unwrap_or("zh");
183        let full = input.full_content.unwrap_or(false);
184
185        if input.query.trim().is_empty() {
186            return Err(ToolError::InvalidInput(
187                "query must not be empty".to_string(),
188            ));
189        }
190
191        let mut output = self.search(&input.query, top_k, lang).await?;
192
193        if full {
194            for result in &mut output.results {
195                if let Ok(content) = self.get_full_content(&result.title, lang).await {
196                    result.snippet = content;
197                }
198            }
199        }
200
201        Ok(output)
202    }
203}
204
205#[async_trait]
206impl BaseTool for WikipediaTool {
207    fn name(&self) -> &str {
208        "wikipedia"
209    }
210
211    fn description(&self) -> &str {
212        "Wikipedia 百科搜索工具。搜索 Wikipedia 百科条目并返回摘要或完整内容。
213
214参数:
215- query: 搜索关键词
216- top_k: 返回结果数量（默认 3）
217- lang: 语言代码，如 zh/en/ja（默认 zh）
218- full_content: 是否获取完整内容（默认 false）
219
220示例:
221- 搜索百科: {\"query\": \"Rust\", \"lang\": \"zh\"}
222- 获取详细内容: {\"query\": \"Rust\", \"lang\": \"en\", \"full_content\": true}"
223    }
224
225    async fn run(&self, input: String) -> Result<String, ToolError> {
226        let parsed: WikipediaInput = serde_json::from_str(&input)
227            .map_err(|e| ToolError::InvalidInput(format!("JSON parse failed: {}", e)))?;
228
229        let output = self.invoke(parsed).await?;
230
231        let mut text = format!("Wikipedia 搜索结果 (查询: {})\n\n", output.query);
232        for (i, result) in output.results.iter().enumerate() {
233            text.push_str(&format!("{}. {}\n", i + 1, result.title));
234            text.push_str(&format!("   {}\n", result.snippet));
235            text.push_str(&format!("   URL: {}\n\n", result.url));
236        }
237        text.push_str(&format!("共 {} 条结果", output.total));
238
239        Ok(text)
240    }
241
242    fn args_schema(&self) -> Option<serde_json::Value> {
243        use schemars::schema_for;
244        serde_json::to_value(schema_for!(WikipediaInput)).ok()
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn test_wikipedia_tool_properties() {
254        let tool = WikipediaTool::new();
255        assert_eq!(tool.name(), "wikipedia");
256        assert!(tool.description().contains("Wikipedia"));
257        assert!(BaseTool::args_schema(&tool).is_some());
258    }
259
260    #[tokio::test]
261    async fn test_wikipedia_empty_query() {
262        let tool = WikipediaTool::new();
263        let result = tool.run(r#"{"query": ""}"#.to_string()).await;
264        assert!(result.is_err());
265    }
266
267    #[test]
268    fn test_strip_html() {
269        let html = "<p>Hello <b>World</b></p>";
270        assert_eq!(strip_html(html), "Hello World");
271    }
272
273    #[test]
274    fn test_urlencoding() {
275        let encoded = urlencoding("Rust programming");
276        assert_eq!(encoded, "Rust%20programming");
277        assert!(urlencoding("a&b=c#d?e").contains("%26"));
278        assert!(urlencoding("a&b=c#d?e").contains("%3D"));
279    }
280}
lc_tools/wikipedia.rs

lc_tools/
wikipedia.rs
