lc_tools/
url_fetch.rs

1// lc-tools/src/url_fetch.rs
2//! 网页抓取工具
3//!
4//! 提供网页内容抓取和解析功能。
5
6use async_trait::async_trait;
7use regex::Regex;
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10
11use crate::ssrf::guarded_get;
12use lc_core::tools::{BaseTool, Tool, ToolError};
13
14static SCRIPT_REGEX: std::sync::LazyLock<Regex> =
15    std::sync::LazyLock::new(|| Regex::new(r"<script[^>]*>.*?</script>").unwrap());
16
17static STYLE_REGEX: std::sync::LazyLock<Regex> =
18    std::sync::LazyLock::new(|| Regex::new(r"<style[^>]*>.*?</style>").unwrap());
19
20static TAG_REGEX: std::sync::LazyLock<Regex> =
21    std::sync::LazyLock::new(|| Regex::new(r"<[^>]+>").unwrap());
22
23static WHITESPACE_REGEX: std::sync::LazyLock<Regex> =
24    std::sync::LazyLock::new(|| Regex::new(r"\s+").unwrap());
25
26static LINK_REGEX: std::sync::LazyLock<Regex> =
27    std::sync::LazyLock::new(|| Regex::new(r#"<a[^>]+href\s*=\s*['"]([^'"]+)['"][^>]*>"#).unwrap());
28
29static IMG_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
30    Regex::new(r#"<img[^>]+src\s*=\s*['"]([^'"]+)['"][^>]*>"#).unwrap()
31});
32
33static TITLE_REGEX: std::sync::LazyLock<Regex> =
34    std::sync::LazyLock::new(|| Regex::new(r"<title[^>]*>(.*?)</title>").unwrap());
35
36static DESC_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
37    Regex::new(r#"<meta[^>]+name\s*=\s*['"]description['"][^>]+content\s*=\s*['"]([^'"]+)['"]"#)
38        .unwrap()
39});
40
41static KW_REGEX: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
42    Regex::new(r#"<meta[^>]+name\s*=\s*['"]keywords['"][^>]+content\s*=\s*['"]([^'"]+)['"]"#)
43        .unwrap()
44});
45
46/// 从 HTML 中提取链接并去重(保留首次出现顺序)。
47///
48/// 返回 `(去重后的链接, 原始条数)`。原始条数用于 details 展示"去重前后对比"。
49/// Q4: 旧实现只给 Vec 换了个名 `unique_links`,并没有真正去重。
50fn extract_unique_links(html: &str) -> (Vec<String>, usize) {
51    let raw: Vec<String> = LINK_REGEX
52        .captures_iter(html)
53        .map(|cap| cap[1].to_string())
54        .collect();
55    let raw_count = raw.len();
56    let mut seen = std::collections::HashSet::new();
57    let unique: Vec<String> = raw
58        .into_iter()
59        .filter(|link| seen.insert(link.clone()))
60        .collect();
61    (unique, raw_count)
62}
63
64/// URLFetch 工具输入
65#[derive(Debug, Deserialize, JsonSchema)]
66pub struct URLFetchInput {
67    /// 操作类型: "fetch", "extract_text", "extract_links", "extract_images", "metadata"
68    pub operation: String,
69
70    /// URL 地址
71    pub url: String,
72
73    /// 是否包含头部信息（用于 fetch 操作）
74    pub include_headers: Option<bool>,
75
76    /// 最大内容长度（字节）
77    pub max_length: Option<usize>,
78}
79
80/// URLFetch 工具输出
81#[derive(Debug, Serialize)]
82pub struct URLFetchOutput {
83    /// 操作结果
84    pub result: String,
85
86    /// 操作类型
87    pub operation: String,
88
89    /// URL
90    pub url: String,
91
92    /// 内容长度
93    pub content_length: usize,
94
95    /// 额外信息
96    pub details: Option<String>,
97}
98
99/// 网页抓取工具
100pub struct URLFetchTool {
101    /// HTTP 客户端
102    client: reqwest::Client,
103    /// 是否允许访问内网 IP（默认 false）
104    allow_private_ips: bool,
105}
106
107impl URLFetchTool {
108    pub fn new() -> Self {
109        Self {
110            client: reqwest::Client::builder()
111                .timeout(std::time::Duration::from_secs(30))
112                .user_agent("LangChainRust/0.1 (URL Fetch Tool)")
113                // SSRF: 禁用自动重定向,由 guarded_get 逐跳重查
114                .redirect(reqwest::redirect::Policy::none())
115                .build()
116                .unwrap_or_else(|_| reqwest::Client::new()),
117            allow_private_ips: false,
118        }
119    }
120
121    /// Allow requests to private/internal IP addresses (SSRF opt-in).
122    pub fn with_allow_private_ips(mut self, allow: bool) -> Self {
123        self.allow_private_ips = allow;
124        self
125    }
126
127    /// 抓取网页内容
128    async fn fetch_url(
129        &self,
130        url: &str,
131        max_length: Option<usize>,
132        include_headers: Option<bool>,
133    ) -> Result<URLFetchOutput, ToolError> {
134        if !url.starts_with("http://") && !url.starts_with("https://") {
135            return Err(ToolError::InvalidInput(
136                "URL 必须以 http:// 或 https:// 开头".to_string(),
137            ));
138        }
139
140        // SSRF: guarded_get 逐跳检查并手动跟随重定向(首跳与每一跳都会重查内网地址)
141        let response = guarded_get(&self.client, url, !self.allow_private_ips).await?;
142
143        let status = response.status();
144        if !status.is_success() {
145            return Err(ToolError::ExecutionFailed(format!(
146                "HTTP 错误: {} - {}",
147                status.as_u16(),
148                status.canonical_reason().unwrap_or("未知")
149            )));
150        }
151
152        // Q3: include_headers = Some(true) 时,把响应头并入输出(details),不再静默忽略。
153        // 头必须在消费 response 前读出(reqwest 的 headers() 借用、text() 消费)。
154        let header_block: String = if include_headers.unwrap_or(false) {
155            let mut lines: Vec<String> = response
156                .headers()
157                .iter()
158                .map(|(name, value)| format!("{}: {}", name, value.to_str().unwrap_or("<非UTF-8>")))
159                .collect();
160            lines.sort();
161            let mut block = String::from("响应头:\n");
162            for line in lines {
163                block.push_str(&line);
164                block.push('\n');
165            }
166            block
167        } else {
168            String::new()
169        };
170
171        let content = response
172            .text()
173            .await
174            .map_err(|e| ToolError::ExecutionFailed(format!("读取响应失败: {}", e)))?;
175
176        let max_len = max_length.unwrap_or(50000);
177        let content_len = content.len();
178        let truncated = content_len > max_len;
179        let result = if truncated {
180            content.chars().take(max_len).collect::<String>() + "\n... [内容已截断]"
181        } else {
182            content
183        };
184
185        let mut details = format!(
186            "状态码: {}, 内容长度: {} 字节{}",
187            status.as_u16(),
188            content_len,
189            if truncated { " (已截断)" } else { "" }
190        );
191        if !header_block.is_empty() {
192            details.push('\n');
193            details.push_str(&header_block);
194        }
195
196        Ok(URLFetchOutput {
197            result,
198            operation: "fetch".to_string(),
199            url: url.to_string(),
200            content_length: content_len,
201            details: Some(details),
202        })
203    }
204
205    /// 提取纯文本内容
206    async fn extract_text(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
207        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
208        let html = &fetch_result.result;
209
210        let html = SCRIPT_REGEX.replace_all(html, "");
211        let html = STYLE_REGEX.replace_all(&html, "");
212
213        let text = TAG_REGEX.replace_all(&html, "");
214
215        let clean_text = WHITESPACE_REGEX.replace_all(&text, " ").trim().to_string();
216
217        let max_len = 5000;
218        let clean_len = clean_text.len();
219        let result = if clean_len > max_len {
220            clean_text.chars().take(max_len).collect::<String>() + "..."
221        } else {
222            clean_text
223        };
224
225        Ok(URLFetchOutput {
226            result,
227            operation: "extract_text".to_string(),
228            url: url.to_string(),
229            content_length: clean_len,
230            details: Some(format!("提取了 {} 字符的纯文本", clean_len)),
231        })
232    }
233
234    /// 提取链接(去重,保留首次出现顺序)
235    async fn extract_links(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
236        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
237        let html = &fetch_result.result;
238
239        let (unique_links, raw_count) = extract_unique_links(html);
240        let result = unique_links.join("\n");
241
242        Ok(URLFetchOutput {
243            result,
244            operation: "extract_links".to_string(),
245            url: url.to_string(),
246            content_length: html.len(), // Q4: 真实正文长度,而非链接条数
247            details: Some(format!(
248                "找到 {} 个唯一链接(原始 {} 个)",
249                unique_links.len(),
250                raw_count
251            )),
252        })
253    }
254
255    /// 提取图片链接
256    async fn extract_images(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
257        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
258        let html = &fetch_result.result;
259
260        let images: Vec<String> = IMG_REGEX
261            .captures_iter(html)
262            .map(|cap| cap[1].to_string())
263            .collect();
264
265        let result = images.join("\n");
266
267        Ok(URLFetchOutput {
268            result,
269            operation: "extract_images".to_string(),
270            url: url.to_string(),
271            content_length: images.len(),
272            details: Some(format!("找到 {} 张图片", images.len())),
273        })
274    }
275
276    /// 提取元数据
277    async fn extract_metadata(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
278        let fetch_result = self.fetch_url(url, Some(50000), None).await?;
279        let html = &fetch_result.result;
280
281        let title = TITLE_REGEX
282            .captures(html)
283            .map(|cap| cap[1].trim().to_string())
284            .unwrap_or_default();
285
286        let description = DESC_REGEX
287            .captures(html)
288            .map(|cap| cap[1].to_string())
289            .unwrap_or_default();
290
291        let keywords = KW_REGEX
292            .captures(html)
293            .map(|cap| cap[1].to_string())
294            .unwrap_or_default();
295
296        let result = format!(
297            "标题: {}\n描述: {}\n关键词: {}",
298            title, description, keywords
299        );
300
301        Ok(URLFetchOutput {
302            result,
303            operation: "metadata".to_string(),
304            url: url.to_string(),
305            content_length: title.len() + description.len() + keywords.len(),
306            details: Some("提取了网页元数据".to_string()),
307        })
308    }
309}
310
311impl Default for URLFetchTool {
312    fn default() -> Self {
313        Self::new()
314    }
315}
316
317#[async_trait]
318impl Tool for URLFetchTool {
319    type Input = URLFetchInput;
320    type Output = URLFetchOutput;
321
322    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
323        match input.operation.as_str() {
324            "fetch" => {
325                self.fetch_url(&input.url, input.max_length, input.include_headers)
326                    .await
327            }
328            "extract_text" => self.extract_text(&input.url).await,
329            "extract_links" => self.extract_links(&input.url).await,
330            "extract_images" => self.extract_images(&input.url).await,
331            "metadata" => self.extract_metadata(&input.url).await,
332            _ => Err(ToolError::InvalidInput(
333                format!("不支持的操作: {}，请使用: fetch, extract_text, extract_links, extract_images, metadata", input.operation)
334            )),
335        }
336    }
337}
338
339#[async_trait]
340impl BaseTool for URLFetchTool {
341    fn name(&self) -> &str {
342        "url_fetch"
343    }
344
345    fn description(&self) -> &str {
346        "网页抓取工具。支持多种操作：
347
348操作类型:
349- fetch: 抓取完整网页内容
350- extract_text: 提取纯文本内容（去除HTML标签）
351- extract_links: 提取所有链接
352- extract_images: 提取所有图片链接
353- metadata: 提取网页元数据（标题、描述、关键词）
354
355参数:
356- url: 网页地址（必须以 http:// 或 https:// 开头）
357- max_length: 最大内容长度（可选，默认50KB）
358- include_headers: 是否包含头部信息（可选）
359
360示例:
361- 抓取网页: {\"operation\": \"fetch\", \"url\": \"https://example.com\"}
362- 提取文本: {\"operation\": \"extract_text\", \"url\": \"https://example.com\"}
363- 提取链接: {\"operation\": \"extract_links\", \"url\": \"https://example.com\"}"
364    }
365
366    async fn run(&self, input: String) -> Result<String, ToolError> {
367        let parsed: URLFetchInput = serde_json::from_str(&input)
368            .map_err(|e| ToolError::InvalidInput(format!("JSON 解析失败: {}", e)))?;
369
370        let output = self.invoke(parsed).await?;
371
372        Ok(format!(
373            "URL: {}\n操作: {}\n内容长度: {} 字节\n\n{}\n详细信息: {}",
374            output.url,
375            output.operation,
376            output.content_length,
377            output.result,
378            output.details.unwrap_or_default()
379        ))
380    }
381
382    fn args_schema(&self) -> Option<serde_json::Value> {
383        use schemars::schema_for;
384        serde_json::to_value(schema_for!(URLFetchInput)).ok()
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391
392    #[test]
393    fn test_url_validation() {
394        let valid_url = "https://example.com";
395        assert!(valid_url.starts_with("http://") || valid_url.starts_with("https://"));
396
397        let valid_url2 = "http://example.org";
398        assert!(valid_url2.starts_with("http://") || valid_url2.starts_with("https://"));
399    }
400
401    #[tokio::test]
402    async fn test_url_fetch_invalid_url() {
403        let tool = URLFetchTool::new();
404
405        let input = URLFetchInput {
406            operation: "fetch".to_string(),
407            url: "invalid-url".to_string(),
408            include_headers: None,
409            max_length: None,
410        };
411
412        let result = tool.invoke(input).await;
413        assert!(result.is_err());
414        assert!(result.unwrap_err().to_string().contains("http://"));
415    }
416
417    /// Q4: extract_links 真正去重且保留首次出现顺序;content_length 是正文长度。
418    #[test]
419    fn test_extract_unique_links_dedups() {
420        let html = r#"
421            <a href="https://a.com/1">first</a>
422            <a href="https://a.com/1">dup</a>
423            <a href="https://a.com/2">second</a>
424            <a href="https://a.com/1">dup2</a>
425        "#;
426        let (unique, raw) = extract_unique_links(html);
427        assert_eq!(raw, 4, "原始链接数应为 4");
428        assert_eq!(
429            unique,
430            vec!["https://a.com/1".to_string(), "https://a.com/2".to_string()],
431            "应去重且保持首次出现顺序"
432        );
433    }
434
435    /// Q1: SSRF 抽取公共模块后,URLFetch 默认仍拦截内网地址。
436    #[tokio::test]
437    async fn test_url_fetch_blocks_localhost_by_default() {
438        let tool = URLFetchTool::new();
439        let input = URLFetchInput {
440            operation: "fetch".to_string(),
441            url: "http://127.0.0.1:6379/".to_string(),
442            include_headers: None,
443            max_length: None,
444        };
445        let result = tool.invoke(input).await;
446        assert!(result.is_err());
447        let err = result.unwrap_err().to_string();
448        assert!(err.contains("SSRF"), "expected SSRF error, got: {}", err);
449    }
450
451    /// Q3: include_headers = Some(true) 时,响应头并入输出(details)。
452    #[tokio::test]
453    async fn test_fetch_include_headers() {
454        use tokio::io::{AsyncReadExt, AsyncWriteExt};
455        use tokio::net::TcpListener;
456
457        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
458        let addr = listener.local_addr().unwrap();
459
460        let server = tokio::spawn(async move {
461            let (mut socket, _) = listener.accept().await.unwrap();
462            let mut buf = [0u8; 4096];
463            let _ = socket.read(&mut buf).await;
464            let body = "hello world";
465            let resp = format!(
466                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: text/plain\r\nX-Test-Header: hello\r\n\r\n{}",
467                body.len(),
468                body
469            );
470            socket.write_all(resp.as_bytes()).await.unwrap();
471            drop(socket);
472        });
473
474        let tool = URLFetchTool::new().with_allow_private_ips(true);
475        let input = URLFetchInput {
476            operation: "fetch".to_string(),
477            url: format!("http://{}/", addr),
478            include_headers: Some(true),
479            max_length: None,
480        };
481
482        let output = tool.invoke(input).await.unwrap();
483        let details = output.details.unwrap();
484        assert!(details.contains("响应头:"), "details: {}", details);
485        // reqwest 会把响应头名规范化为小写（HTTP 头名不区分大小写）
486        assert!(
487            details.contains("x-test-header: hello"),
488            "details: {}",
489            details
490        );
491        assert!(output.result.contains("hello world"));
492        server.await.unwrap();
493    }
494
495    /// Q3: include_headers = false/None 时不返回响应头。
496    #[tokio::test]
497    async fn test_fetch_without_include_headers() {
498        use tokio::io::{AsyncReadExt, AsyncWriteExt};
499        use tokio::net::TcpListener;
500
501        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
502        let addr = listener.local_addr().unwrap();
503
504        let server = tokio::spawn(async move {
505            let (mut socket, _) = listener.accept().await.unwrap();
506            let mut buf = [0u8; 4096];
507            let _ = socket.read(&mut buf).await;
508            let body = "hello world";
509            let resp = format!(
510                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nX-Test-Header: hello\r\n\r\n{}",
511                body.len(),
512                body
513            );
514            socket.write_all(resp.as_bytes()).await.unwrap();
515            drop(socket);
516        });
517
518        let tool = URLFetchTool::new().with_allow_private_ips(true);
519        let input = URLFetchInput {
520            operation: "fetch".to_string(),
521            url: format!("http://{}/", addr),
522            include_headers: Some(false),
523            max_length: None,
524        };
525
526        let output = tool.invoke(input).await.unwrap();
527        let details = output.details.unwrap();
528        assert!(
529            !details.contains("X-Test-Header"),
530            "不应包含响应头, got: {}",
531            details
532        );
533        server.await.unwrap();
534    }
535
536    #[test]
537    fn test_tool_properties() {
538        let tool = URLFetchTool::new();
539
540        assert_eq!(tool.name(), "url_fetch");
541        assert!(tool.description().contains("fetch"));
542        assert!(BaseTool::args_schema(&tool).is_some());
543    }
544}
lc_tools/url_fetch.rs

lc_tools/
url_fetch.rs
