Underground-Digital/Workflow-Engine

1import json2import time3 4import requests5 6from extensions.ext_storage import storage7 8 9class FirecrawlApp:10    def __init__(self, api_key=None, base_url=None):11        self.api_key = api_key12        self.base_url = base_url or "https://api.firecrawl.dev"13        if self.api_key is None and self.base_url == "https://api.firecrawl.dev":14            raise ValueError("No API key provided")15 16    def scrape_url(self, url, params=None) -> dict:17        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}18        json_data = {"url": url}19        if params:20            json_data.update(params)21        response = requests.post(f"{self.base_url}/v0/scrape", headers=headers, json=json_data)22        if response.status_code == 200:23            response = response.json()24            if response["success"] == True:25                data = response["data"]26                return {27                    "title": data.get("metadata").get("title"),28                    "description": data.get("metadata").get("description"),29                    "source_url": data.get("metadata").get("sourceURL"),30                    "markdown": data.get("markdown"),31                }32            else:33                raise Exception(f'Failed to scrape URL. Error: {response["error"]}')34 35        elif response.status_code in {402, 409, 500}:36            error_message = response.json().get("error", "Unknown error occurred")37            raise Exception(f"Failed to scrape URL. Status code: {response.status_code}. Error: {error_message}")38        else:39            raise Exception(f"Failed to scrape URL. Status code: {response.status_code}")40 41    def crawl_url(self, url, params=None) -> str:42        headers = self._prepare_headers()43        json_data = {"url": url}44        if params:45            json_data.update(params)46        response = self._post_request(f"{self.base_url}/v0/crawl", json_data, headers)47        if response.status_code == 200:48            job_id = response.json().get("jobId")49            return job_id50        else:51            self._handle_error(response, "start crawl job")52 53    def check_crawl_status(self, job_id) -> dict:54        headers = self._prepare_headers()55        response = self._get_request(f"{self.base_url}/v0/crawl/status/{job_id}", headers)56        if response.status_code == 200:57            crawl_status_response = response.json()58            if crawl_status_response.get("status") == "completed":59                total = crawl_status_response.get("total", 0)60                if total == 0:61                    raise Exception("Failed to check crawl status. Error: No page found")62                data = crawl_status_response.get("data", [])63                url_data_list = []64                for item in data:65                    if isinstance(item, dict) and "metadata" in item and "markdown" in item:66                        url_data = {67                            "title": item.get("metadata").get("title"),68                            "description": item.get("metadata").get("description"),69                            "source_url": item.get("metadata").get("sourceURL"),70                            "markdown": item.get("markdown"),71                        }72                        url_data_list.append(url_data)73                if url_data_list:74                    file_key = "website_files/" + job_id + ".txt"75                    if storage.exists(file_key):76                        storage.delete(file_key)77                    storage.save(file_key, json.dumps(url_data_list).encode("utf-8"))78                return {79                    "status": "completed",80                    "total": crawl_status_response.get("total"),81                    "current": crawl_status_response.get("current"),82                    "data": url_data_list,83                }84 85            else:86                return {87                    "status": crawl_status_response.get("status"),88                    "total": crawl_status_response.get("total"),89                    "current": crawl_status_response.get("current"),90                    "data": [],91                }92 93        else:94            self._handle_error(response, "check crawl status")95 96    def _prepare_headers(self):97        return {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}98 99    def _post_request(self, url, data, headers, retries=3, backoff_factor=0.5):100        for attempt in range(retries):101            response = requests.post(url, headers=headers, json=data)102            if response.status_code == 502:103                time.sleep(backoff_factor * (2**attempt))104            else:105                return response106        return response107 108    def _get_request(self, url, headers, retries=3, backoff_factor=0.5):109        for attempt in range(retries):110            response = requests.get(url, headers=headers)111            if response.status_code == 502:112                time.sleep(backoff_factor * (2**attempt))113            else:114                return response115        return response116 117    def _handle_error(self, response, action):118        error_message = response.json().get("error", "Unknown error occurred")119        raise Exception(f"Failed to {action}. Status code: {response.status_code}. Error: {error_message}")120