From f0e921bde409e4bdb8fe8187a17bf47830c73ef6 Mon Sep 17 00:00:00 2001
From: "lau.td" <lau.td.developer@gmail.com>
Date: Tue, 1 Oct 2024 15:53:19 +0700
Subject: [PATCH] feat: implement custom add firecrawl datasource

---
 api/controllers/codelight/__init__.py         |  2 +-
 .../internal_api/data_source_bearer_auth.py   | 37 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 api/controllers/codelight/internal_api/data_source_bearer_auth.py

diff --git a/api/controllers/codelight/__init__.py b/api/controllers/codelight/__init__.py
index 1012a01b9ba310..7e7fc8d66cd3ad 100644
--- a/api/controllers/codelight/__init__.py
+++ b/api/controllers/codelight/__init__.py
@@ -6,5 +6,5 @@
 
 
 from .console_api import dataset, message, conversation
-from .internal_api import account, app, auth, model_config, model_providers, tenant
+from .internal_api import account, app, auth, model_config, model_providers, tenant, data_source_bearer_auth
 from .web import passport, conversation
diff --git a/api/controllers/codelight/internal_api/data_source_bearer_auth.py b/api/controllers/codelight/internal_api/data_source_bearer_auth.py
new file mode 100644
index 00000000000000..1d74b72ae3449c
--- /dev/null
+++ b/api/controllers/codelight/internal_api/data_source_bearer_auth.py
@@ -0,0 +1,37 @@
+import uuid
+from flask_restful import Resource, reqparse
+
+from controllers.codelight import api
+from controllers.console.auth.error import ApiKeyAuthFailedError
+from services.auth.api_key_auth_service import ApiKeyAuthService
+from controllers.console.setup import setup_required
+from controllers.inner_api.wraps import inner_api_only
+
+
+class CodelightApiKeyAuthDataSourceBinding(Resource):
+    @setup_required
+    @inner_api_only
+    def post(self, tenant_id: uuid.UUID):
+        parser = reqparse.RequestParser()
+        parser.add_argument(
+            "category", type=str, required=True, nullable=False, location="json"
+        )
+        parser.add_argument(
+            "provider", type=str, required=True, nullable=False, location="json"
+        )
+        parser.add_argument(
+            "credentials", type=dict, required=True, nullable=False, location="json"
+        )
+        args = parser.parse_args()
+        ApiKeyAuthService.validate_api_key_auth_args(args)
+
+        try:
+            ApiKeyAuthService.create_provider_auth(tenant_id, args)
+        except Exception as e:
+            raise ApiKeyAuthFailedError(str(e))
+        return {"result": "success"}, 200
+
+
+api.add_resource(
+    CodelightApiKeyAuthDataSourceBinding, "/tenants/<uuid:tenant_id>/data-sources"
+)