feat: Integrate WaterCrawl.dev as a new knowledge base provider (#16396)

Co-authored-by: crazywoola <427733928@qq.com>

api/services/auth/api_key_auth_factory.py

@@ -17,6 +17,10 @@ class ApiKeyAuthFactory:
                 from services.auth.firecrawl.firecrawl import FirecrawlAuth
 
                 return FirecrawlAuth
+            case AuthType.WATERCRAWL:
+                from services.auth.watercrawl.watercrawl import WatercrawlAuth
+
+                return WatercrawlAuth
             case AuthType.JINA:
                 from services.auth.jina.jina import JinaAuth
 

api/services/auth/auth_type.py

@@ -3,4 +3,5 @@ from enum import StrEnum
 
 class AuthType(StrEnum):
     FIRECRAWL = "firecrawl"
+    WATERCRAWL = "watercrawl"
     JINA = "jinareader"

api/services/auth/watercrawl/watercrawl.py

@@ -0,0 +1,44 @@
+import json
+from urllib.parse import urljoin
+
+import requests
+
+from services.auth.api_key_auth_base import ApiKeyAuthBase
+
+
+class WatercrawlAuth(ApiKeyAuthBase):
+    def __init__(self, credentials: dict):
+        super().__init__(credentials)
+        auth_type = credentials.get("auth_type")
+        if auth_type != "x-api-key":
+            raise ValueError("Invalid auth type, WaterCrawl auth type must be x-api-key")
+        self.api_key = credentials.get("config", {}).get("api_key", None)
+        self.base_url = credentials.get("config", {}).get("base_url", "https://app.watercrawl.dev")
+
+        if not self.api_key:
+            raise ValueError("No API key provided")
+
+    def validate_credentials(self):
+        headers = self._prepare_headers()
+        url = urljoin(self.base_url, "/api/v1/core/crawl-requests/")
+        response = self._get_request(url, headers)
+        if response.status_code == 200:
+            return True
+        else:
+            self._handle_error(response)
+
+    def _prepare_headers(self):
+        return {"Content-Type": "application/json", "X-API-KEY": self.api_key}
+
+    def _get_request(self, url, headers):
+        return requests.get(url, headers=headers)
+
+    def _handle_error(self, response):
+        if response.status_code in {402, 409, 500}:
+            error_message = response.json().get("error", "Unknown error occurred")
+            raise Exception(f"Failed to authorize. Status code: {response.status_code}. Error: {error_message}")
+        else:
+            if response.text:
+                error_message = json.loads(response.text).get("error", "Unknown error occurred")
+                raise Exception(f"Failed to authorize. Status code: {response.status_code}. Error: {error_message}")
+            raise Exception(f"Unexpected error occurred while trying to authorize. Status code: {response.status_code}")

api/services/website_service.py

@@ -7,6 +7,7 @@ from flask_login import current_user  # type: ignore
 
 from core.helper import encrypter
 from core.rag.extractor.firecrawl.firecrawl_app import FirecrawlApp
+from core.rag.extractor.watercrawl.provider import WaterCrawlProvider
 from extensions.ext_redis import redis_client
 from extensions.ext_storage import storage
 from services.auth.api_key_auth_service import ApiKeyAuthService
@@ -59,6 +60,13 @@ class WebsiteService:
             time = str(datetime.datetime.now().timestamp())
             redis_client.setex(website_crawl_time_cache_key, 3600, time)
             return {"status": "active", "job_id": job_id}
+        elif provider == "watercrawl":
+            # decrypt api_key
+            api_key = encrypter.decrypt_token(
+                tenant_id=current_user.current_tenant_id, token=credentials.get("config").get("api_key")
+            )
+            return WaterCrawlProvider(api_key, credentials.get("config").get("base_url", None)).crawl_url(url, options)
+
         elif provider == "jinareader":
             api_key = encrypter.decrypt_token(
                 tenant_id=current_user.current_tenant_id, token=credentials.get("config").get("api_key")
@@ -116,6 +124,14 @@ class WebsiteService:
                     time_consuming = abs(end_time - float(start_time))
                     crawl_status_data["time_consuming"] = f"{time_consuming:.2f}"
                     redis_client.delete(website_crawl_time_cache_key)
+        elif provider == "watercrawl":
+            # decrypt api_key
+            api_key = encrypter.decrypt_token(
+                tenant_id=current_user.current_tenant_id, token=credentials.get("config").get("api_key")
+            )
+            crawl_status_data = WaterCrawlProvider(
+                api_key, credentials.get("config").get("base_url", None)
+            ).get_crawl_status(job_id)
         elif provider == "jinareader":
             api_key = encrypter.decrypt_token(
                 tenant_id=current_user.current_tenant_id, token=credentials.get("config").get("api_key")
@@ -180,6 +196,11 @@ class WebsiteService:
                     if item.get("source_url") == url:
                         return dict(item)
             return None
+        elif provider == "watercrawl":
+            api_key = encrypter.decrypt_token(tenant_id=tenant_id, token=credentials.get("config").get("api_key"))
+            return WaterCrawlProvider(api_key, credentials.get("config").get("base_url", None)).get_crawl_url_data(
+                job_id, url
+            )
         elif provider == "jinareader":
             if not job_id:
                 response = requests.get(
@@ -223,5 +244,8 @@ class WebsiteService:
             params = {"onlyMainContent": only_main_content}
             result = firecrawl_app.scrape_url(url, params)
             return result
+        elif provider == "watercrawl":
+            api_key = encrypter.decrypt_token(tenant_id=tenant_id, token=credentials.get("config").get("api_key"))
+            return WaterCrawlProvider(api_key, credentials.get("config").get("base_url", None)).scrape_url(url)
         else:
             raise ValueError("Invalid provider")