494 lines · 15.8 KB · Rust
RawDownload
use std::collections::HashMap; use serde::{Deserialize, Serialize}; use crate::{    crawl::{CrawlErrorsResponse, CrawlStatus, CrawlStatusTypes},    scrape::ScrapeOptions,    FirecrawlApp, FirecrawlError, API_VERSION,}; #[serde_with::skip_serializing_none]#[derive(Deserialize, Serialize, Debug, Default, Clone)]#[serde(rename_all = "camelCase")]pub struct BatchScrapeParams {    /// List of URLs to scrape    pub urls: Vec<String>,    /// Scrape options to apply to all URLs    #[serde(flatten)]    pub options: Option<ScrapeOptions>,    /// Whether to ignore invalid URLs    #[serde(rename = "ignoreInvalidURLs")]    pub ignore_invalid_urls: bool,    /// ID of an existing job to append these URLs to    pub append_to_id: Option<String>,    /// Webhook configuration    pub webhook: Option<WebhookOptions>,     /// Idempotency key to send to the crawl endpoint.    #[serde(skip)]    pub idempotency_key: Option<String>,} /// Options for webhook notifications#[serde_with::skip_serializing_none]#[derive(Deserialize, Serialize, Debug, Default, Clone)]#[serde(rename_all = "camelCase")]pub struct WebhookOptions {    /// URL to send webhook notifications to    pub url: String,    /// Custom headers to include in webhook requests    pub headers: Option<HashMap<String, String>>,    /// Authentication token for the webhook    pub auth_token: Option<String>,} impl From<&str> for WebhookOptions {    fn from(url: &str) -> Self {        Self {            url: url.to_string(),            headers: None,            auth_token: None,        }    }} /// Response from initiating a batch scrape job#[derive(Deserialize, Serialize, Debug, Clone)]#[serde(rename_all = "camelCase")]pub struct BatchScrapeResponse {    /// Whether the request was successful    pub success: bool,    /// The ID of the batch scrape job    pub id: String,    /// URL to get the status of the batch scrape job    pub url: String,    /// List of URLs that were invalid and could not be processed    pub invalid_urls: Option<Vec<String>>,} impl From<String> for WebhookOptions {    fn from(url: String) -> Self {        Self {            url,            headers: None,            auth_token: None,        }    }} impl FirecrawlApp {    /// Initiates an asynchronous batch scrape job    pub async fn async_batch_scrape_urls(        &self,        params: BatchScrapeParams,    ) -> Result<BatchScrapeResponse, FirecrawlError> {        let headers = self.prepare_headers(params.idempotency_key.as_ref());         let response = self            .client            .post(format!("{}{}/batch/scrape", self.api_url, API_VERSION))            .headers(headers)            .json(&params)            .send()            .await            .map_err(|e| FirecrawlError::HttpError("Initiating batch scrape job".to_string(), e))?;         self.handle_response(response, "initiate batch scrape job")            .await    }     /// Initiates a batch scrape job and waits for completion    pub async fn batch_scrape_urls(        &self,        params: BatchScrapeParams,        poll_interval: Option<u64>,    ) -> Result<CrawlStatus, FirecrawlError> {        let poll_interval_ms = poll_interval.unwrap_or(2000);         let response = self.async_batch_scrape_urls(params).await?;         self.monitor_batch_job_status(&response.id, poll_interval_ms)            .await    }     /// Checks the status of a batch scrape job    pub async fn check_batch_scrape_status(        &self,        id: impl AsRef<str>,    ) -> Result<CrawlStatus, FirecrawlError> {        let response = self            .client            .get(format!(                "{}{}/batch/scrape/{}",                self.api_url,                API_VERSION,                id.as_ref()            ))            .headers(self.prepare_headers(None))            .send()            .await            .map_err(|e| {                FirecrawlError::HttpError(                    format!("Checking status of batch scrape {}", id.as_ref()),                    e,                )            })?;         let mut status: CrawlStatus = self            .handle_response(                response,                format!("Checking status of batch scrape {}", id.as_ref()),            )            .await?;         if status.status == CrawlStatusTypes::Completed {            while let Some(next) = status.next.clone() {                let new_status = self.check_batch_scrape_status_next(next).await?;                status.data.extend_from_slice(&new_status.data);                status.next = new_status.next;            }        }         Ok(status)    }     /// Helper function to paginate through batch scrape status results    async fn check_batch_scrape_status_next(        &self,        next: impl AsRef<str>,    ) -> Result<CrawlStatus, FirecrawlError> {        let response = self            .client            .get(next.as_ref())            .headers(self.prepare_headers(None))            .send()            .await            .map_err(|e| {                FirecrawlError::HttpError(                    format!("Paginating batch scrape using URL {:?}", next.as_ref()),                    e,                )            })?;         self.handle_response(            response,            format!("Paginating batch scrape using URL {:?}", next.as_ref()),        )        .await    }     /// Check for errors in a batch scrape job    pub async fn check_batch_scrape_errors(        &self,        id: impl AsRef<str>,    ) -> Result<CrawlErrorsResponse, FirecrawlError> {        let response = self            .client            .get(format!(                "{}{}/batch/scrape/{}/errors",                self.api_url,                API_VERSION,                id.as_ref()            ))            .headers(self.prepare_headers(None))            .send()            .await            .map_err(|e| {                FirecrawlError::HttpError(                    format!("Checking errors for batch scrape {}", id.as_ref()),                    e,                )            })?;         self.handle_response(            response,            format!("Checking errors for batch scrape {}", id.as_ref()),        )        .await    }     /// Helper function to poll for batch job status until completion    async fn monitor_batch_job_status(        &self,        id: &str,        poll_interval: u64,    ) -> Result<CrawlStatus, FirecrawlError> {        loop {            let status_data = self.check_batch_scrape_status(id).await?;            match status_data.status {                CrawlStatusTypes::Completed => {                    break Ok(status_data);                }                CrawlStatusTypes::Scraping => {                    tokio::time::sleep(tokio::time::Duration::from_millis(poll_interval)).await;                }                CrawlStatusTypes::Failed => {                    break Err(FirecrawlError::CrawlJobFailed(                        "Batch scrape job failed".into(),                        status_data,                    ));                }                CrawlStatusTypes::Cancelled => {                    break Err(FirecrawlError::CrawlJobFailed(                        "Batch scrape job was cancelled".into(),                        status_data,                    ));                }            }        }    }} #[cfg(test)]mod tests {    use super::*;    use serde_json::json;     #[tokio::test]    #[ignore = "Makes real network request"]    async fn test_real_batch_scrape() {        let api_url = std::env::var("FIRECRAWL_API_URL")            .expect("Please set the FIRECRAWL_API_URL environment variable");        let app = FirecrawlApp::new_selfhosted(api_url, None::<&str>).unwrap();         // Start a batch scrape job        let params = BatchScrapeParams {            urls: vec![                "https://example.com".to_string(),                "https://example.org".to_string(),            ],            ignore_invalid_urls: true,            ..Default::default()        };         let response = app.async_batch_scrape_urls(params).await.unwrap();         assert!(response.success);        assert!(!response.id.is_empty());        assert!(!response.url.is_empty());    }     #[tokio::test]    async fn test_async_batch_scrape_with_mock() {        let mut server = mockito::Server::new_async().await;         // Set up the mock        let mock = server            .mock("POST", "/v1/batch/scrape")            // Remove the match_body expectation which might be causing issues            .with_status(200)            .with_header("content-type", "application/json")            .with_body(                json!({                    "success": true,                    "id": "batch-123",                    "url": "https://api.example.com/v1/batch/batch-123",                    "invalidUrls": []                })                .to_string(),            )            .create();         let app = FirecrawlApp::new_selfhosted(server.url(), Some("test_key")).unwrap();         let params = BatchScrapeParams {            urls: vec![                "https://example.com".to_string(),                "https://example.org".to_string(),            ],            ignore_invalid_urls: true,            ..Default::default()        };         let response = app.async_batch_scrape_urls(params).await.unwrap();         assert!(response.success);        assert_eq!(response.id, "batch-123");        assert_eq!(response.url, "https://api.example.com/v1/batch/batch-123");        assert!(response.invalid_urls.unwrap_or_default().is_empty());        mock.assert();    }     #[tokio::test]    async fn test_batch_scrape_with_webhook() {        let mut server = mockito::Server::new_async().await;         let mock = server            .mock("POST", "/v1/batch/scrape")            // Remove the match_body expectation to simplify            .with_status(200)            .with_header("content-type", "application/json")            .with_body(                json!({                    "success": true,                    "id": "batch-123",                    "url": "https://api.example.com/v1/batch/batch-123"                })                .to_string(),            )            .create();         let app = FirecrawlApp::new_selfhosted(server.url(), Some("test_key")).unwrap();         let params = BatchScrapeParams {            urls: vec!["https://example.com".to_string()],            webhook: Some("https://webhook.example.com/notify".into()),            ..Default::default()        };         let response = app.async_batch_scrape_urls(params).await.unwrap();         assert!(response.success);        assert_eq!(response.id, "batch-123");        mock.assert();    }     #[tokio::test]    async fn test_check_batch_scrape_status_with_mock() {        let mut server = mockito::Server::new_async().await;         let mock = server            .mock("GET", "/v1/batch/scrape/batch-123")            .with_status(200)            .with_header("content-type", "application/json")            .with_body(                json!({                    "success": true,                    "status": "completed",                    "total": 2,                    "completed": 2,                    "creditsUsed": 2,                    "expiresAt": "2023-12-31T23:59:59Z",                    "data": [                        {                            "metadata": {                                "sourceURL": "https://example.com",                                "statusCode": 200                            },                            "markdown": "Example Domain content"                        },                        {                            "metadata": {                                "sourceURL": "https://example.org",                                "statusCode": 200                            },                            "markdown": "Another example content"                        }                    ]                })                .to_string(),            )            .create();         let app = FirecrawlApp::new_selfhosted(server.url(), Some("test_key")).unwrap();        let status = app.check_batch_scrape_status("batch-123").await.unwrap();         assert_eq!(status.total, 2);        assert_eq!(status.completed, 2);        assert_eq!(status.data.len(), 2);        assert_eq!(status.data[0].metadata.source_url, "https://example.com");        assert_eq!(status.data[1].metadata.source_url, "https://example.org");        mock.assert();    }     #[tokio::test]    async fn test_check_batch_scrape_errors_with_mock() {        let mut server = mockito::Server::new_async().await;         let mock = server            .mock("GET", "/v1/batch/scrape/batch-123/errors")            .with_status(200)            .with_header("content-type", "application/json")            .with_body(                json!({                    "success": true,                    "errors": [                        {                            "id": "error1",                            "timestamp": "2023-01-01T00:00:00Z",                            "url": "https://invalid.example.com",                            "error": "Failed to load page"                        }                    ],                    "robotsBlocked": [                        "https://example.com/admin"                    ]                })                .to_string(),            )            .create_async()            .await;         let app = FirecrawlApp::new_selfhosted(server.url(), Some("test_key")).unwrap();        let errors = app.check_batch_scrape_errors("batch-123").await.unwrap();         assert_eq!(errors.errors.len(), 1);        assert_eq!(errors.errors[0].url, "https://invalid.example.com");        assert_eq!(errors.robots_blocked.len(), 1);        assert_eq!(errors.robots_blocked[0], "https://example.com/admin");        mock.assert();    }     #[tokio::test]    async fn test_batch_scrape_with_invalid_urls() {        let mut server = mockito::Server::new_async().await;         let mock = server            .mock("POST", "/v1/batch/scrape")            // Remove the match_body expectation            .with_status(200)            .with_header("content-type", "application/json")            .with_body(                json!({                    "success": true,                    "id": "batch-123",                    "url": "https://api.example.com/v1/batch/batch-123",                    "invalidUrls": ["invalid-url"]                })                .to_string(),            )            .create();         let app = FirecrawlApp::new_selfhosted(server.url(), Some("test_key")).unwrap();         let params = BatchScrapeParams {            urls: vec!["https://example.com".to_string(), "invalid-url".to_string()],            ignore_invalid_urls: true,            ..Default::default()        };         let response = app.async_batch_scrape_urls(params).await.unwrap();         assert!(response.success);        assert_eq!(response.id, "batch-123");        assert_eq!(response.invalid_urls, Some(vec!["invalid-url".to_string()]));        mock.assert();    }     #[tokio::test]    async fn test_batch_scrape_error_response() {        let mut server = mockito::Server::new_async().await;         let mock = server            .mock("POST", "/v1/batch/scrape")            .with_status(400)            .with_header("content-type", "application/json")            .with_body(                json!({                    "success": false,                    "error": "No valid URLs provided"                })                .to_string(),            )            .create();         let app = FirecrawlApp::new_selfhosted(server.url(), Some("test_key")).unwrap();         let params = BatchScrapeParams::default();        let result = app.async_batch_scrape_urls(params).await;         assert!(result.is_err());        mock.assert();    }}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t