crawlkit_extensions/firecrawl/
mod.rs

1//! Firecrawl API 客户端
2//!
3//! 通过 [Firecrawl SDK](https://firecrawl.dev) 进行云端网页抓取，支持 JS 渲染。
4//! 适用于 SPA 站点、需要 JavaScript 执行的页面。
5//!
6//! # 示例
7//!
8//! ```rust,no_run
9//! use crawlkit_core::HttpClient;
10//! use crawlkit_extensions::firecrawl::FirecrawlClient;
11//!
12//! # #[tokio::main]
13//! # async fn main() -> anyhow::Result<()> {
14//! let client = FirecrawlClient::builder()
15//!     .with_tokens(&["fc-api-token-1", "fc-api-token-2"])
16//!     .build()?;
17//!
18//! let response = client.get("https://example.com", &Default::default()).await?;
19//! println!("{}", response.body);
20//! # Ok(())
21//! # }
22//! ```
23
24use std::collections::HashMap;
25
26use async_trait::async_trait;
27use backon::{ExponentialBuilder, Retryable};
28use crawlkit_core::{CrawlError, HttpClient, Response};
29use rand::Rng;
30
31/// Firecrawl API 客户端
32///
33/// 支持多 token 轮换，适合免费套餐多账号场景。
34pub struct FirecrawlClient {
35    name: String,
36    api_tokens: Vec<String>,
37    max_retries: usize,
38}
39
40impl FirecrawlClient {
41    /// 创建 Builder
42    pub fn builder() -> FirecrawlClientBuilder {
43        FirecrawlClientBuilder::default()
44    }
45}
46
47/// FirecrawlClient 构建器
48pub struct FirecrawlClientBuilder {
49    name: String,
50    api_tokens: Vec<String>,
51    max_retries: usize,
52}
53
54impl Default for FirecrawlClientBuilder {
55    fn default() -> Self {
56        Self {
57            name: "firecrawl".to_string(),
58            api_tokens: Vec::new(),
59            max_retries: 3,
60        }
61    }
62}
63
64impl FirecrawlClientBuilder {
65    /// 设置客户端名称（用于日志和 CompositeFetcher 识别，默认 "firecrawl"）
66    pub fn with_name(mut self, name: impl Into<String>) -> Self {
67        self.name = name.into();
68        self
69    }
70
71    /// 设置 API token 列表（每次请求随机选择一个）
72    pub fn with_tokens(mut self, tokens: &[impl Into<String> + Clone]) -> Self {
73        self.api_tokens = tokens.iter().map(|t| t.clone().into()).collect();
74        self
75    }
76
77    /// 从环境变量读取 token（逗号分隔多个）
78    pub fn with_tokens_from_env(mut self, env_var: &str) -> Self {
79        if let Ok(val) = std::env::var(env_var) {
80            self.api_tokens = val.split(',').map(|s| s.trim().to_string()).filter(|s| !s.is_empty()).collect();
81        }
82        self
83    }
84
85    /// 设置最大重试次数（默认 3）
86    pub fn with_max_retries(mut self, max_retries: usize) -> Self {
87        self.max_retries = max_retries;
88        self
89    }
90
91    /// 构建 FirecrawlClient
92    ///
93    /// # 错误
94    ///
95    /// 如果未设置任何 token，返回错误。
96    pub fn build(self) -> anyhow::Result<FirecrawlClient> {
97        if self.api_tokens.is_empty() {
98            anyhow::bail!("Firecrawl 至少需要一个 API Token");
99        }
100        Ok(FirecrawlClient {
101            name: self.name,
102            api_tokens: self.api_tokens,
103            max_retries: self.max_retries,
104        })
105    }
106}
107
108#[async_trait]
109impl HttpClient for FirecrawlClient {
110    async fn get(&self, url: &str, _headers: &HashMap<String, String>) -> crawlkit_core::Result<Response> {
111        let tokens = self.api_tokens.clone();
112        let max_retries = self.max_retries;
113
114        let fetch = || {
115            let tokens = tokens.clone();
116            async move {
117                let token = {
118                    let mut rng = rand::rng();
119                    let index = rng.random_range(0..tokens.len());
120                    tokens[index].clone()
121                };
122
123                let client = firecrawl2::Client::new(&token)
124                    .map_err(|e: firecrawl2::error::FirecrawlError| anyhow::anyhow!("Firecrawl 客户端创建失败: {e}"))?;
125
126                let options = firecrawl2::ScrapeOptions {
127                    formats: Some(vec![firecrawl2::Format::Html]),
128                    only_main_content: Some(true),
129                    ..Default::default()
130                };
131
132                let doc = client.scrape(url, options).await
133                    .map_err(|e: firecrawl2::error::FirecrawlError| anyhow::anyhow!("Firecrawl 请求失败: {e}"))?;
134
135                let html = doc.html
136                    .ok_or_else(|| anyhow::anyhow!("Firecrawl 返回 HTML 内容为空"))?;
137
138                Ok::<(String, String), anyhow::Error>((html, url.to_string()))
139            }
140        };
141
142        let (body, final_url) = fetch
143            .retry(&ExponentialBuilder::default().with_max_times(max_retries))
144            .await
145            .map_err(|e| CrawlError::Http(format!("Firecrawl 请求失败（重试 {max_retries} 次）: {e}")))?;
146
147        Ok(Response {
148            url: final_url,
149            status: 200,
150            headers: Default::default(),
151            body,
152        })
153    }
154
155    async fn post(
156        &self,
157        url: &str,
158        headers: &HashMap<String, String>,
159        _body: Vec<u8>,
160    ) -> crawlkit_core::Result<Response> {
161        self.get(url, headers).await
162    }
163
164    fn name(&self) -> &str {
165        &self.name
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn test_firecrawl_builder_defaults() {
175        let builder = FirecrawlClientBuilder::default();
176        assert_eq!(builder.name, "firecrawl");
177        assert!(builder.api_tokens.is_empty());
178        assert_eq!(builder.max_retries, 3);
179    }
180
181    #[test]
182    fn test_firecrawl_builder_requires_token() {
183        let result = FirecrawlClientBuilder::default().build();
184        assert!(result.is_err());
185    }
186
187    #[test]
188    fn test_firecrawl_builder_with_tokens() {
189        let client = FirecrawlClient::builder()
190            .with_tokens(&["token1", "token2"])
191            .build()
192            .unwrap();
193        assert_eq!(client.api_tokens.len(), 2);
194    }
195}
crawlkit_extensions/firecrawl/mod.rs

crawlkit_extensions/firecrawl/
mod.rs
