Underground-Digital/Workflow-Engine

1import json2 3import requests4 5from services.auth.api_key_auth_base import ApiKeyAuthBase6 7 8class FirecrawlAuth(ApiKeyAuthBase):9    def __init__(self, credentials: dict):10        super().__init__(credentials)11        auth_type = credentials.get("auth_type")12        if auth_type != "bearer":13            raise ValueError("Invalid auth type, Firecrawl auth type must be Bearer")14        self.api_key = credentials.get("config").get("api_key", None)15        self.base_url = credentials.get("config").get("base_url", "https://api.firecrawl.dev")16 17        if not self.api_key:18            raise ValueError("No API key provided")19 20    def validate_credentials(self):21        headers = self._prepare_headers()22        options = {23            "url": "https://example.com",24            "crawlerOptions": {"excludes": [], "includes": [], "limit": 1},25            "pageOptions": {"onlyMainContent": True},26        }27        response = self._post_request(f"{self.base_url}/v0/crawl", options, headers)28        if response.status_code == 200:29            return True30        else:31            self._handle_error(response)32 33    def _prepare_headers(self):34        return {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}35 36    def _post_request(self, url, data, headers):37        return requests.post(url, headers=headers, json=data)38 39    def _handle_error(self, response):40        if response.status_code in {402, 409, 500}:41            error_message = response.json().get("error", "Unknown error occurred")42            raise Exception(f"Failed to authorize. Status code: {response.status_code}. Error: {error_message}")43        else:44            if response.text:45                error_message = json.loads(response.text).get("error", "Unknown error occurred")46                raise Exception(f"Failed to authorize. Status code: {response.status_code}. Error: {error_message}")47            raise Exception(f"Unexpected error occurred while trying to authorize. Status code: {response.status_code}")48