From 983834cd52181f18a99d6e5c76e54adc29952058 Mon Sep 17 00:00:00 2001
From: takatost <takatost@users.noreply.github.com>
Date: Mon, 11 Sep 2023 17:31:03 +0800
Subject: [PATCH] feat: spark check (#1134)

---
 .../console/workspace/model_providers.py      | 16 ++++++++
 api/services/provider_service.py              | 39 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/api/controllers/console/workspace/model_providers.py b/api/controllers/console/workspace/model_providers.py
index 08fdb4f117..7498072e3e 100644
--- a/api/controllers/console/workspace/model_providers.py
+++ b/api/controllers/console/workspace/model_providers.py
@@ -285,6 +285,20 @@ class ModelProviderFreeQuotaSubmitApi(Resource):
         return result
 
 
+class ModelProviderFreeQuotaQualificationVerifyApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider_name: str):
+        provider_service = ProviderService()
+        result = provider_service.free_quota_qualification_verify(
+            tenant_id=current_user.current_tenant_id,
+            provider_name=provider_name
+        )
+
+        return result
+
+
 api.add_resource(ModelProviderListApi, '/workspaces/current/model-providers')
 api.add_resource(ModelProviderValidateApi, '/workspaces/current/model-providers/<string:provider_name>/validate')
 api.add_resource(ModelProviderUpdateApi, '/workspaces/current/model-providers/<string:provider_name>')
@@ -300,3 +314,5 @@ api.add_resource(ModelProviderPaymentCheckoutUrlApi,
                  '/workspaces/current/model-providers/<string:provider_name>/checkout-url')
 api.add_resource(ModelProviderFreeQuotaSubmitApi,
                  '/workspaces/current/model-providers/<string:provider_name>/free-quota-submit')
+api.add_resource(ModelProviderFreeQuotaQualificationVerifyApi,
+                 '/workspaces/current/model-providers/<string:provider_name>/free-quota-qualification-verify')
diff --git a/api/services/provider_service.py b/api/services/provider_service.py
index 703ec51b6b..a0eed4d273 100644
--- a/api/services/provider_service.py
+++ b/api/services/provider_service.py
@@ -518,7 +518,8 @@ class ProviderService:
 
     def free_quota_submit(self, tenant_id: str, provider_name: str):
         api_key = os.environ.get("FREE_QUOTA_APPLY_API_KEY")
-        api_url = os.environ.get("FREE_QUOTA_APPLY_URL")
+        api_base_url = os.environ.get("FREE_QUOTA_APPLY_BASE_URL")
+        api_url = api_base_url + '/api/v1/providers/apply'
 
         headers = {
             'Content-Type': 'application/json',
@@ -546,3 +547,39 @@ class ProviderService:
                 'type': rst['type'],
                 'result': 'success'
             }
+
+    def free_quota_qualification_verify(self, tenant_id: str, provider_name: str):
+        api_key = os.environ.get("FREE_QUOTA_APPLY_API_KEY")
+        api_base_url = os.environ.get("FREE_QUOTA_APPLY_BASE_URL")
+        api_url = api_base_url + '/api/v1/providers/qualification-verify'
+
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': f"Bearer {api_key}"
+        }
+        response = requests.post(api_url, headers=headers,
+                                 json={'workspace_id': tenant_id, 'provider_name': provider_name})
+        if not response.ok:
+            logging.error(f"Request FREE QUOTA APPLY SERVER Error: {response.status_code} ")
+            raise ValueError(f"Error: {response.status_code} ")
+
+        rst = response.json()
+        if rst["code"] != 'success':
+            raise ValueError(
+                f"error: {rst['message']}"
+            )
+
+        data = rst['data']
+        if data['qualified'] is True:
+            return {
+                'result': 'success',
+                'provider_name': provider_name,
+                'flag': True
+            }
+        else:
+            return {
+                'result': 'success',
+                'provider_name': provider_name,
+                'flag': False,
+                'reason': data['reason']
+            }