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

lc_tools/
url_fetch.rs
