refactor: Add an enumeration type and use the factory pattern to obtain the corresponding class (#9356)

api/core/rag/datasource/keyword/keyword_factory.py

@@ -1,8 +1,8 @@
 from typing import Any
 
 from configs import dify_config
-from core.rag.datasource.keyword.jieba.jieba import Jieba
 from core.rag.datasource.keyword.keyword_base import BaseKeyword
+from core.rag.datasource.keyword.keyword_type import KeyWordType
 from core.rag.models.document import Document
 from models.dataset import Dataset
 
@@ -13,16 +13,19 @@ class Keyword:
         self._keyword_processor = self._init_keyword()
 
     def _init_keyword(self) -> BaseKeyword:
-        config = dify_config
-        keyword_type = config.KEYWORD_STORE
+        keyword_type = dify_config.KEYWORD_STORE
+        keyword_factory = self.get_keyword_factory(keyword_type)
+        return keyword_factory(self._dataset)
 
-        if not keyword_type:
-            raise ValueError("Keyword store must be specified.")
+    @staticmethod
+    def get_keyword_factory(keyword_type: str) -> type[BaseKeyword]:
+        match keyword_type:
+            case KeyWordType.JIEBA:
+                from core.rag.datasource.keyword.jieba.jieba import Jieba
 
-        if keyword_type == "jieba":
-            return Jieba(dataset=self._dataset)
-        else:
-            raise ValueError(f"Keyword store {keyword_type} is not supported.")
+                return Jieba
+            case _:
+                raise ValueError(f"Keyword store {keyword_type} is not supported.")
 
     def create(self, texts: list[Document], **kwargs):
         self._keyword_processor.create(texts, **kwargs)

api/core/rag/datasource/keyword/keyword_type.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class KeyWordType(str, Enum):
+    JIEBA = "jieba"

api/services/auth/api_key_auth_factory.py

@@ -1,15 +1,25 @@
-from services.auth.firecrawl import FirecrawlAuth
-from services.auth.jina import JinaAuth
+from services.auth.api_key_auth_base import ApiKeyAuthBase
+from services.auth.auth_type import AuthType
 
 
 class ApiKeyAuthFactory:
     def __init__(self, provider: str, credentials: dict):
-        if provider == "firecrawl":
-            self.auth = FirecrawlAuth(credentials)
-        elif provider == "jinareader":
-            self.auth = JinaAuth(credentials)
-        else:
-            raise ValueError("Invalid provider")
+        auth_factory = self.get_apikey_auth_factory(provider)
+        self.auth = auth_factory(credentials)
 
     def validate_credentials(self):
         return self.auth.validate_credentials()
+
+    @staticmethod
+    def get_apikey_auth_factory(provider: str) -> type[ApiKeyAuthBase]:
+        match provider:
+            case AuthType.FIRECRAWL:
+                from services.auth.firecrawl.firecrawl import FirecrawlAuth
+
+                return FirecrawlAuth
+            case AuthType.JINA:
+                from services.auth.jina.jina import JinaAuth
+
+                return JinaAuth
+            case _:
+                raise ValueError("Invalid provider")

api/services/auth/auth_type.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class AuthType(str, Enum):
+    FIRECRAWL = "firecrawl"
+    JINA = "jinareader"

api/services/auth/firecrawl/firecrawl.py

@@ -0,0 +1,47 @@
+import json
+
+import requests
+
+from services.auth.api_key_auth_base import ApiKeyAuthBase
+
+
+class FirecrawlAuth(ApiKeyAuthBase):
+    def __init__(self, credentials: dict):
+        super().__init__(credentials)
+        auth_type = credentials.get("auth_type")
+        if auth_type != "bearer":
+            raise ValueError("Invalid auth type, Firecrawl auth type must be Bearer")
+        self.api_key = credentials.get("config").get("api_key", None)
+        self.base_url = credentials.get("config").get("base_url", "https://api.firecrawl.dev")
+
+        if not self.api_key:
+            raise ValueError("No API key provided")
+
+    def validate_credentials(self):
+        headers = self._prepare_headers()
+        options = {
+            "url": "https://example.com",
+            "crawlerOptions": {"excludes": [], "includes": [], "limit": 1},
+            "pageOptions": {"onlyMainContent": True},
+        }
+        response = self._post_request(f"{self.base_url}/v0/crawl", options, headers)
+        if response.status_code == 200:
+            return True
+        else:
+            self._handle_error(response)
+
+    def _prepare_headers(self):
+        return {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
+
+    def _post_request(self, url, data, headers):
+        return requests.post(url, headers=headers, json=data)
+
+    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/auth/jina/jina.py

@@ -0,0 +1,44 @@
+import json
+
+import requests
+
+from services.auth.api_key_auth_base import ApiKeyAuthBase
+
+
+class JinaAuth(ApiKeyAuthBase):
+    def __init__(self, credentials: dict):
+        super().__init__(credentials)
+        auth_type = credentials.get("auth_type")
+        if auth_type != "bearer":
+            raise ValueError("Invalid auth type, Jina Reader auth type must be Bearer")
+        self.api_key = credentials.get("config").get("api_key", None)
+
+        if not self.api_key:
+            raise ValueError("No API key provided")
+
+    def validate_credentials(self):
+        headers = self._prepare_headers()
+        options = {
+            "url": "https://example.com",
+        }
+        response = self._post_request("https://r.jina.ai", options, headers)
+        if response.status_code == 200:
+            return True
+        else:
+            self._handle_error(response)
+
+    def _prepare_headers(self):
+        return {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
+
+    def _post_request(self, url, data, headers):
+        return requests.post(url, headers=headers, json=data)
+
+    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}")