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::url_points_to_private_ip;
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                .build()
114                .unwrap_or_else(|_| reqwest::Client::new()),
115            allow_private_ips: false,
116        }
117    }
118
119    /// Allow requests to private/internal IP addresses (SSRF opt-in).
120    pub fn with_allow_private_ips(mut self, allow: bool) -> Self {
121        self.allow_private_ips = allow;
122        self
123    }
124
125    /// Check SSRF protection before making a request.
126    async fn check_ssrf(&self, url: &str) -> Result<(), ToolError> {
127        if self.allow_private_ips {
128            return Ok(());
129        }
130        if url_points_to_private_ip(url).await? {
131            return Err(ToolError::ExecutionFailed(
132                "Request to private/internal IP address is blocked by SSRF protection. \
133                 Call .with_allow_private_ips(true) to allow."
134                    .to_string(),
135            ));
136        }
137        Ok(())
138    }
139
140    /// 抓取网页内容
141    async fn fetch_url(
142        &self,
143        url: &str,
144        max_length: Option<usize>,
145        include_headers: Option<bool>,
146    ) -> Result<URLFetchOutput, ToolError> {
147        if !url.starts_with("http://") && !url.starts_with("https://") {
148            return Err(ToolError::InvalidInput(
149                "URL 必须以 http:// 或 https:// 开头".to_string(),
150            ));
151        }
152
153        self.check_ssrf(url).await?;
154
155        let response = self
156            .client
157            .get(url)
158            .send()
159            .await
160            .map_err(|e| ToolError::ExecutionFailed(format!("HTTP 请求失败: {}", e)))?;
161
162        let status = response.status();
163        if !status.is_success() {
164            return Err(ToolError::ExecutionFailed(format!(
165                "HTTP 错误: {} - {}",
166                status.as_u16(),
167                status.canonical_reason().unwrap_or("未知")
168            )));
169        }
170
171        // Q3: include_headers = Some(true) 时,把响应头并入输出(details),不再静默忽略。
172        // 头必须在消费 response 前读出(reqwest 的 headers() 借用、text() 消费)。
173        let header_block: String = if include_headers.unwrap_or(false) {
174            let mut lines: Vec<String> = response
175                .headers()
176                .iter()
177                .map(|(name, value)| format!("{}: {}", name, value.to_str().unwrap_or("<非UTF-8>")))
178                .collect();
179            lines.sort();
180            let mut block = String::from("响应头:\n");
181            for line in lines {
182                block.push_str(&line);
183                block.push('\n');
184            }
185            block
186        } else {
187            String::new()
188        };
189
190        let content = response
191            .text()
192            .await
193            .map_err(|e| ToolError::ExecutionFailed(format!("读取响应失败: {}", e)))?;
194
195        let max_len = max_length.unwrap_or(50000);
196        let content_len = content.len();
197        let truncated = content_len > max_len;
198        let result = if truncated {
199            content.chars().take(max_len).collect::<String>() + "\n... [内容已截断]"
200        } else {
201            content
202        };
203
204        let mut details = format!(
205            "状态码: {}, 内容长度: {} 字节{}",
206            status.as_u16(),
207            content_len,
208            if truncated { " (已截断)" } else { "" }
209        );
210        if !header_block.is_empty() {
211            details.push('\n');
212            details.push_str(&header_block);
213        }
214
215        Ok(URLFetchOutput {
216            result,
217            operation: "fetch".to_string(),
218            url: url.to_string(),
219            content_length: content_len,
220            details: Some(details),
221        })
222    }
223
224    /// 提取纯文本内容
225    async fn extract_text(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
226        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
227        let html = &fetch_result.result;
228
229        let html = SCRIPT_REGEX.replace_all(html, "");
230        let html = STYLE_REGEX.replace_all(&html, "");
231
232        let text = TAG_REGEX.replace_all(&html, "");
233
234        let clean_text = WHITESPACE_REGEX.replace_all(&text, " ").trim().to_string();
235
236        let max_len = 5000;
237        let clean_len = clean_text.len();
238        let result = if clean_len > max_len {
239            clean_text.chars().take(max_len).collect::<String>() + "..."
240        } else {
241            clean_text
242        };
243
244        Ok(URLFetchOutput {
245            result,
246            operation: "extract_text".to_string(),
247            url: url.to_string(),
248            content_length: clean_len,
249            details: Some(format!("提取了 {} 字符的纯文本", clean_len)),
250        })
251    }
252
253    /// 提取链接(去重,保留首次出现顺序)
254    async fn extract_links(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
255        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
256        let html = &fetch_result.result;
257
258        let (unique_links, raw_count) = extract_unique_links(html);
259        let result = unique_links.join("\n");
260
261        Ok(URLFetchOutput {
262            result,
263            operation: "extract_links".to_string(),
264            url: url.to_string(),
265            content_length: html.len(), // Q4: 真实正文长度,而非链接条数
266            details: Some(format!(
267                "找到 {} 个唯一链接(原始 {} 个)",
268                unique_links.len(),
269                raw_count
270            )),
271        })
272    }
273
274    /// 提取图片链接
275    async fn extract_images(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
276        let fetch_result = self.fetch_url(url, Some(100000), None).await?;
277        let html = &fetch_result.result;
278
279        let images: Vec<String> = IMG_REGEX
280            .captures_iter(html)
281            .map(|cap| cap[1].to_string())
282            .collect();
283
284        let result = images.join("\n");
285
286        Ok(URLFetchOutput {
287            result,
288            operation: "extract_images".to_string(),
289            url: url.to_string(),
290            content_length: images.len(),
291            details: Some(format!("找到 {} 张图片", images.len())),
292        })
293    }
294
295    /// 提取元数据
296    async fn extract_metadata(&self, url: &str) -> Result<URLFetchOutput, ToolError> {
297        let fetch_result = self.fetch_url(url, Some(50000), None).await?;
298        let html = &fetch_result.result;
299
300        let title = TITLE_REGEX
301            .captures(html)
302            .map(|cap| cap[1].trim().to_string())
303            .unwrap_or_default();
304
305        let description = DESC_REGEX
306            .captures(html)
307            .map(|cap| cap[1].to_string())
308            .unwrap_or_default();
309
310        let keywords = KW_REGEX
311            .captures(html)
312            .map(|cap| cap[1].to_string())
313            .unwrap_or_default();
314
315        let result = format!(
316            "标题: {}\n描述: {}\n关键词: {}",
317            title, description, keywords
318        );
319
320        Ok(URLFetchOutput {
321            result,
322            operation: "metadata".to_string(),
323            url: url.to_string(),
324            content_length: title.len() + description.len() + keywords.len(),
325            details: Some("提取了网页元数据".to_string()),
326        })
327    }
328}
329
330impl Default for URLFetchTool {
331    fn default() -> Self {
332        Self::new()
333    }
334}
335
336#[async_trait]
337impl Tool for URLFetchTool {
338    type Input = URLFetchInput;
339    type Output = URLFetchOutput;
340
341    async fn invoke(&self, input: Self::Input) -> Result<Self::Output, ToolError> {
342        match input.operation.as_str() {
343            "fetch" => {
344                self.fetch_url(&input.url, input.max_length, input.include_headers)
345                    .await
346            }
347            "extract_text" => self.extract_text(&input.url).await,
348            "extract_links" => self.extract_links(&input.url).await,
349            "extract_images" => self.extract_images(&input.url).await,
350            "metadata" => self.extract_metadata(&input.url).await,
351            _ => Err(ToolError::InvalidInput(
352                format!("不支持的操作: {}，请使用: fetch, extract_text, extract_links, extract_images, metadata", input.operation)
353            )),
354        }
355    }
356}
357
358#[async_trait]
359impl BaseTool for URLFetchTool {
360    fn name(&self) -> &str {
361        "url_fetch"
362    }
363
364    fn description(&self) -> &str {
365        "网页抓取工具。支持多种操作：
366
367操作类型:
368- fetch: 抓取完整网页内容
369- extract_text: 提取纯文本内容（去除HTML标签）
370- extract_links: 提取所有链接
371- extract_images: 提取所有图片链接
372- metadata: 提取网页元数据（标题、描述、关键词）
373
374参数:
375- url: 网页地址（必须以 http:// 或 https:// 开头）
376- max_length: 最大内容长度（可选，默认50KB）
377- include_headers: 是否包含头部信息（可选）
378
379示例:
380- 抓取网页: {\"operation\": \"fetch\", \"url\": \"https://example.com\"}
381- 提取文本: {\"operation\": \"extract_text\", \"url\": \"https://example.com\"}
382- 提取链接: {\"operation\": \"extract_links\", \"url\": \"https://example.com\"}"
383    }
384
385    async fn run(&self, input: String) -> Result<String, ToolError> {
386        let parsed: URLFetchInput = serde_json::from_str(&input)
387            .map_err(|e| ToolError::InvalidInput(format!("JSON 解析失败: {}", e)))?;
388
389        let output = self.invoke(parsed).await?;
390
391        Ok(format!(
392            "URL: {}\n操作: {}\n内容长度: {} 字节\n\n{}\n详细信息: {}",
393            output.url,
394            output.operation,
395            output.content_length,
396            output.result,
397            output.details.unwrap_or_default()
398        ))
399    }
400
401    fn args_schema(&self) -> Option<serde_json::Value> {
402        use schemars::schema_for;
403        serde_json::to_value(schema_for!(URLFetchInput)).ok()
404    }
405}
406
407#[cfg(test)]
408mod tests {
409    use super::*;
410
411    #[test]
412    fn test_url_validation() {
413        let valid_url = "https://example.com";
414        assert!(valid_url.starts_with("http://") || valid_url.starts_with("https://"));
415
416        let valid_url2 = "http://example.org";
417        assert!(valid_url2.starts_with("http://") || valid_url2.starts_with("https://"));
418    }
419
420    #[tokio::test]
421    async fn test_url_fetch_invalid_url() {
422        let tool = URLFetchTool::new();
423
424        let input = URLFetchInput {
425            operation: "fetch".to_string(),
426            url: "invalid-url".to_string(),
427            include_headers: None,
428            max_length: None,
429        };
430
431        let result = tool.invoke(input).await;
432        assert!(result.is_err());
433        assert!(result.unwrap_err().to_string().contains("http://"));
434    }
435
436    /// Q4: extract_links 真正去重且保留首次出现顺序;content_length 是正文长度。
437    #[test]
438    fn test_extract_unique_links_dedups() {
439        let html = r#"
440            <a href="https://a.com/1">first</a>
441            <a href="https://a.com/1">dup</a>
442            <a href="https://a.com/2">second</a>
443            <a href="https://a.com/1">dup2</a>
444        "#;
445        let (unique, raw) = extract_unique_links(html);
446        assert_eq!(raw, 4, "原始链接数应为 4");
447        assert_eq!(
448            unique,
449            vec!["https://a.com/1".to_string(), "https://a.com/2".to_string()],
450            "应去重且保持首次出现顺序"
451        );
452    }
453
454    /// Q1: SSRF 抽取公共模块后,URLFetch 默认仍拦截内网地址。
455    #[tokio::test]
456    async fn test_url_fetch_blocks_localhost_by_default() {
457        let tool = URLFetchTool::new();
458        let input = URLFetchInput {
459            operation: "fetch".to_string(),
460            url: "http://127.0.0.1:6379/".to_string(),
461            include_headers: None,
462            max_length: None,
463        };
464        let result = tool.invoke(input).await;
465        assert!(result.is_err());
466        let err = result.unwrap_err().to_string();
467        assert!(err.contains("SSRF"), "expected SSRF error, got: {}", err);
468    }
469
470    /// Q3: include_headers = Some(true) 时,响应头并入输出(details)。
471    #[tokio::test]
472    async fn test_fetch_include_headers() {
473        use tokio::io::{AsyncReadExt, AsyncWriteExt};
474        use tokio::net::TcpListener;
475
476        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
477        let addr = listener.local_addr().unwrap();
478
479        let server = tokio::spawn(async move {
480            let (mut socket, _) = listener.accept().await.unwrap();
481            let mut buf = [0u8; 4096];
482            let _ = socket.read(&mut buf).await;
483            let body = "hello world";
484            let resp = format!(
485                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: text/plain\r\nX-Test-Header: hello\r\n\r\n{}",
486                body.len(),
487                body
488            );
489            socket.write_all(resp.as_bytes()).await.unwrap();
490            drop(socket);
491        });
492
493        let tool = URLFetchTool::new().with_allow_private_ips(true);
494        let input = URLFetchInput {
495            operation: "fetch".to_string(),
496            url: format!("http://{}/", addr),
497            include_headers: Some(true),
498            max_length: None,
499        };
500
501        let output = tool.invoke(input).await.unwrap();
502        let details = output.details.unwrap();
503        assert!(details.contains("响应头:"), "details: {}", details);
504        // reqwest 会把响应头名规范化为小写（HTTP 头名不区分大小写）
505        assert!(
506            details.contains("x-test-header: hello"),
507            "details: {}",
508            details
509        );
510        assert!(output.result.contains("hello world"));
511        server.await.unwrap();
512    }
513
514    /// Q3: include_headers = false/None 时不返回响应头。
515    #[tokio::test]
516    async fn test_fetch_without_include_headers() {
517        use tokio::io::{AsyncReadExt, AsyncWriteExt};
518        use tokio::net::TcpListener;
519
520        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
521        let addr = listener.local_addr().unwrap();
522
523        let server = tokio::spawn(async move {
524            let (mut socket, _) = listener.accept().await.unwrap();
525            let mut buf = [0u8; 4096];
526            let _ = socket.read(&mut buf).await;
527            let body = "hello world";
528            let resp = format!(
529                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nX-Test-Header: hello\r\n\r\n{}",
530                body.len(),
531                body
532            );
533            socket.write_all(resp.as_bytes()).await.unwrap();
534            drop(socket);
535        });
536
537        let tool = URLFetchTool::new().with_allow_private_ips(true);
538        let input = URLFetchInput {
539            operation: "fetch".to_string(),
540            url: format!("http://{}/", addr),
541            include_headers: Some(false),
542            max_length: None,
543        };
544
545        let output = tool.invoke(input).await.unwrap();
546        let details = output.details.unwrap();
547        assert!(
548            !details.contains("X-Test-Header"),
549            "不应包含响应头, got: {}",
550            details
551        );
552        server.await.unwrap();
553    }
554
555    #[tokio::test]
556    #[ignore = "需要网络连接"]
557    async fn test_url_fetch_real() {
558        let tool = URLFetchTool::new();
559
560        let input = URLFetchInput {
561            operation: "fetch".to_string(),
562            url: "https://example.com".to_string(),
563            include_headers: None,
564            max_length: Some(5000),
565        };
566
567        let result = tool.invoke(input).await.unwrap();
568        assert!(result.result.contains("example"));
569        assert!(result.content_length > 0);
570    }
571
572    #[tokio::test]
573    #[ignore = "需要网络连接"]
574    async fn test_url_extract_text_real() {
575        let tool = URLFetchTool::new();
576
577        let input = URLFetchInput {
578            operation: "extract_text".to_string(),
579            url: "https://example.com".to_string(),
580            include_headers: None,
581            max_length: None,
582        };
583
584        let result = tool.invoke(input).await.unwrap();
585        assert!(!result.result.contains("<"));
586    }
587
588    #[tokio::test]
589    #[ignore = "需要网络连接"]
590    async fn test_url_extract_links_real() {
591        let tool = URLFetchTool::new();
592
593        let input = URLFetchInput {
594            operation: "extract_links".to_string(),
595            url: "https://example.com".to_string(),
596            include_headers: None,
597            max_length: None,
598        };
599
600        let result = tool.invoke(input).await.unwrap();
601        assert!(result.details.unwrap().contains("链接"));
602    }
603
604    #[tokio::test]
605    #[ignore = "需要网络连接"]
606    async fn test_url_extract_metadata_real() {
607        let tool = URLFetchTool::new();
608
609        let input = URLFetchInput {
610            operation: "metadata".to_string(),
611            url: "https://example.com".to_string(),
612            include_headers: None,
613            max_length: None,
614        };
615
616        let result = tool.invoke(input).await.unwrap();
617        assert!(result.result.contains("标题"));
618    }
619
620    #[test]
621    fn test_tool_properties() {
622        let tool = URLFetchTool::new();
623
624        assert_eq!(tool.name(), "url_fetch");
625        assert!(tool.description().contains("fetch"));
626        assert!(BaseTool::args_schema(&tool).is_some());
627    }
628}
lc_tools/url_fetch.rs

lc_tools/
url_fetch.rs
