lc_rag/loaders/
sitemap.rs

1//! Sitemap loader
2//!
3//! Parses a sitemap.xml and crawls the page contents in bulk.
4
5use std::collections::HashMap;
6use std::time::Duration;
7
8use async_trait::async_trait;
9use regex::Regex;
10use std::sync::LazyLock;
11
12use super::{DocumentLoader, LoaderError};
13use lc_vector_stores::Document;
14
15/// M59: compile regex once using LazyLock instead of on every call
16static LOC_REGEX: LazyLock<Regex> =
17    LazyLock::new(|| Regex::new(r"<loc>\s*(.*?)\s*</loc>").unwrap());
18
19/// H8: default per-HTTP-request timeout — a hung target site will not block forever.
20const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30);
21
22/// Sitemap loader
23///
24/// Loads pages from a sitemap.xml URL or content, extracting the body text.
25pub struct SitemapLoader {
26    /// Sitemap URL or XML content
27    source: SitemapSource,
28    /// Maximum number of pages to crawl
29    max_pages: usize,
30    /// H8: per-HTTP-request timeout, preventing the crawler from blocking forever on a hung target site
31    timeout: Duration,
32}
33
34/// Sitemap source
35enum SitemapSource {
36    /// Fetched from a URL
37    Url(String),
38    /// Direct XML content
39    Xml(String),
40}
41
42impl SitemapLoader {
43    /// Creates from a sitemap URL
44    pub fn from_url(url: impl Into<String>) -> Self {
45        Self {
46            source: SitemapSource::Url(url.into()),
47            max_pages: 100,
48            timeout: DEFAULT_HTTP_TIMEOUT,
49        }
50    }
51
52    /// Creates from sitemap XML content
53    pub fn from_xml(xml: impl Into<String>) -> Self {
54        Self {
55            source: SitemapSource::Xml(xml.into()),
56            max_pages: 100,
57            timeout: DEFAULT_HTTP_TIMEOUT,
58        }
59    }
60
61    /// Sets the maximum number of pages to crawl
62    pub fn with_max_pages(mut self, max: usize) -> Self {
63        self.max_pages = max;
64        self
65    }
66
67    /// Sets the per-HTTP-request timeout (H8, default 30s)
68    pub fn with_timeout(mut self, timeout: Duration) -> Self {
69        self.timeout = timeout;
70        self
71    }
72
73    /// Parses sitemap XML, extracting the URL list
74    fn parse_urls(xml: &str) -> Vec<String> {
75        LOC_REGEX
76            .captures_iter(xml)
77            .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
78            .collect()
79    }
80
81    /// Crawls a single page through the shared SSRF-hardened helper
82    async fn fetch_page(url: &str, timeout: Duration) -> Result<String, LoaderError> {
83        let (_final_url, body) = super::guarded_fetch(url, timeout).await?;
84        Ok(body)
85    }
86}
87
88#[async_trait]
89impl DocumentLoader for SitemapLoader {
90    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
91        // Fetch the sitemap XML
92        let xml = match &self.source {
93            SitemapSource::Url(url) => Self::fetch_page(url, self.timeout).await?,
94            SitemapSource::Xml(content) => content.clone(),
95        };
96
97        // Parse the URL list
98        let urls = Self::parse_urls(&xml);
99        let mut documents = Vec::new();
100
101        for url in urls.iter().take(self.max_pages) {
102            match Self::fetch_page(url, self.timeout).await {
103                Ok(html) => {
104                    let text = super::HTMLLoader::extract_text(&html);
105                    let mut metadata = HashMap::new();
106                    metadata.insert("format".to_string(), "html".to_string().into());
107                    metadata.insert("source".to_string(), url.clone().into());
108
109                    documents.push(Document {
110                        content: text,
111                        metadata,
112                        id: None,
113                    });
114                }
115                Err(e) => {
116                    log::warn!("Failed to crawl {} (skipped from results): {}", url, e);
117                    continue;
118                }
119            }
120        }
121
122        Ok(documents)
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn test_parse_urls_simple() {
132        let xml = r#"<?xml version="1.0" encoding="UTF-8"?>
133        <urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
134            <url><loc>https://example.com/</loc></url>
135            <url><loc>https://example.com/about</loc></url>
136            <url><loc>https://example.com/contact</loc></url>
137        </urlset>"#;
138        let urls = SitemapLoader::parse_urls(xml);
139        assert_eq!(urls.len(), 3);
140        assert_eq!(urls[0], "https://example.com/");
141        assert_eq!(urls[1], "https://example.com/about");
142    }
143
144    #[test]
145    fn test_parse_urls_with_whitespace() {
146        let xml = r#"<urlset>
147            <url><loc>  https://example.com/page1  </loc></url>
148            <url><loc>https://example.com/page2</loc></url>
149        </urlset>"#;
150        let urls = SitemapLoader::parse_urls(xml);
151        assert_eq!(urls.len(), 2);
152        assert_eq!(urls[0], "https://example.com/page1");
153    }
154
155    #[test]
156    fn test_parse_urls_empty() {
157        let xml = r#"<?xml version="1.0"?><urlset></urlset>"#;
158        let urls = SitemapLoader::parse_urls(xml);
159        assert!(urls.is_empty());
160    }
161
162    #[test]
163    fn test_from_url() {
164        let loader = SitemapLoader::from_url("https://example.com/sitemap.xml");
165        assert_eq!(loader.max_pages, 100);
166    }
167
168    #[test]
169    fn test_with_max_pages() {
170        let loader = SitemapLoader::from_url("https://example.com/sitemap.xml").with_max_pages(5);
171        assert_eq!(loader.max_pages, 5);
172    }
173
174    #[tokio::test]
175    async fn test_load_from_xml() {
176        let xml = r#"<?xml version="1.0"?>
177        <urlset>
178            <url><loc>https://example.com/</loc></url>
179        </urlset>"#;
180        // This test would actually crawl, but max_pages=0 skips the crawl
181        let loader = SitemapLoader::from_xml(xml).with_max_pages(0);
182        let docs = loader.load().await.unwrap();
183        assert!(docs.is_empty()); // max_pages=0, no pages are crawled
184    }
185}
lc_rag/loaders/sitemap.rs

lc_rag/loaders/
sitemap.rs
