Feat/billing enhancement (#2239)

Co-authored-by: takatost <takatost@gmail.com>

api/config.py

@@ -40,17 +40,11 @@ DEFAULTS = {
     'HOSTED_OPENAI_QUOTA_LIMIT': 200,
     'HOSTED_OPENAI_TRIAL_ENABLED': 'False',
     'HOSTED_OPENAI_PAID_ENABLED': 'False',
-    'HOSTED_OPENAI_PAID_INCREASE_QUOTA': 1,
-    'HOSTED_OPENAI_PAID_MIN_QUANTITY': 1,
-    'HOSTED_OPENAI_PAID_MAX_QUANTITY': 1,
     'HOSTED_AZURE_OPENAI_ENABLED': 'False',
     'HOSTED_AZURE_OPENAI_QUOTA_LIMIT': 200,
     'HOSTED_ANTHROPIC_QUOTA_LIMIT': 600000,
     'HOSTED_ANTHROPIC_TRIAL_ENABLED': 'False',
     'HOSTED_ANTHROPIC_PAID_ENABLED': 'False',
-    'HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA': 1,
-    'HOSTED_ANTHROPIC_PAID_MIN_QUANTITY': 1,
-    'HOSTED_ANTHROPIC_PAID_MAX_QUANTITY': 1,
     'HOSTED_MODERATION_ENABLED': 'False',
     'HOSTED_MODERATION_PROVIDERS': '',
     'CLEAN_DAY_SETTING': 30,
@@ -262,10 +256,6 @@ class Config:
         self.HOSTED_OPENAI_TRIAL_ENABLED = get_bool_env('HOSTED_OPENAI_TRIAL_ENABLED')
         self.HOSTED_OPENAI_QUOTA_LIMIT = int(get_env('HOSTED_OPENAI_QUOTA_LIMIT'))
         self.HOSTED_OPENAI_PAID_ENABLED = get_bool_env('HOSTED_OPENAI_PAID_ENABLED')
-        self.HOSTED_OPENAI_PAID_STRIPE_PRICE_ID = get_env('HOSTED_OPENAI_PAID_STRIPE_PRICE_ID')
-        self.HOSTED_OPENAI_PAID_INCREASE_QUOTA = int(get_env('HOSTED_OPENAI_PAID_INCREASE_QUOTA'))
-        self.HOSTED_OPENAI_PAID_MIN_QUANTITY = int(get_env('HOSTED_OPENAI_PAID_MIN_QUANTITY'))
-        self.HOSTED_OPENAI_PAID_MAX_QUANTITY = int(get_env('HOSTED_OPENAI_PAID_MAX_QUANTITY'))
 
         self.HOSTED_AZURE_OPENAI_ENABLED = get_bool_env('HOSTED_AZURE_OPENAI_ENABLED')
         self.HOSTED_AZURE_OPENAI_API_KEY = get_env('HOSTED_AZURE_OPENAI_API_KEY')
@@ -277,10 +267,6 @@ class Config:
         self.HOSTED_ANTHROPIC_TRIAL_ENABLED = get_bool_env('HOSTED_ANTHROPIC_TRIAL_ENABLED')
         self.HOSTED_ANTHROPIC_QUOTA_LIMIT = int(get_env('HOSTED_ANTHROPIC_QUOTA_LIMIT'))
         self.HOSTED_ANTHROPIC_PAID_ENABLED = get_bool_env('HOSTED_ANTHROPIC_PAID_ENABLED')
-        self.HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID = get_env('HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID')
-        self.HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA = int(get_env('HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA'))
-        self.HOSTED_ANTHROPIC_PAID_MIN_QUANTITY = int(get_env('HOSTED_ANTHROPIC_PAID_MIN_QUANTITY'))
-        self.HOSTED_ANTHROPIC_PAID_MAX_QUANTITY = int(get_env('HOSTED_ANTHROPIC_PAID_MAX_QUANTITY'))
 
         self.HOSTED_MINIMAX_ENABLED = get_bool_env('HOSTED_MINIMAX_ENABLED')
         self.HOSTED_SPARK_ENABLED = get_bool_env('HOSTED_SPARK_ENABLED')

api/controllers/console/billing/billing.py

@@ -20,7 +20,7 @@ class Subscription(Resource):
         parser.add_argument('interval', type=str, required=True, location='args', choices=['month', 'year'])
         args = parser.parse_args()
 
-        BillingService.is_tenant_owner(current_user)
+        BillingService.is_tenant_owner_or_admin(current_user)
 
         return BillingService.get_subscription(args['plan'],
                                                args['interval'],
@@ -35,8 +35,8 @@ class Invoices(Resource):
     @account_initialization_required
     @only_edition_cloud
     def get(self):
-        BillingService.is_tenant_owner(current_user)
-        return BillingService.get_invoices(current_user.email)
+        BillingService.is_tenant_owner_or_admin(current_user)
+        return BillingService.get_invoices(current_user.email, current_user.current_tenant_id)
 
 
 api.add_resource(Subscription, '/billing/subscription')

api/controllers/console/workspace/model_providers.py

@@ -186,10 +186,11 @@ class ModelProviderPaymentCheckoutUrlApi(Resource):
     def get(self, provider: str):
         if provider != 'anthropic':
             raise ValueError(f'provider name {provider} is invalid')
-
+        BillingService.is_tenant_owner_or_admin(current_user)
         data = BillingService.get_model_provider_payment_link(provider_name=provider,
                                                               tenant_id=current_user.current_tenant_id,
-                                                              account_id=current_user.id)
+                                                              account_id=current_user.id,
+                                                              prefilled_email=current_user.email)
         return data
 
 

api/core/entities/provider_entities.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel
 class QuotaUnit(Enum):
     TIMES = 'times'
     TOKENS = 'tokens'
+    CREDITS = 'credits'
 
 
 class SystemConfigurationStatus(Enum):

api/core/hosting_configuration.py

@@ -20,10 +20,6 @@ class TrialHostingQuota(HostingQuota):
 
 class PaidHostingQuota(HostingQuota):
     quota_type: ProviderQuotaType = ProviderQuotaType.PAID
-    stripe_price_id: str = None
-    increase_quota: int = 1
-    min_quantity: int = 20
-    max_quantity: int = 100
 
 
 class FreeHostingQuota(HostingQuota):
@@ -102,7 +98,7 @@ class HostingConfiguration:
         )
 
     def init_openai(self, app_config: Config) -> HostingProvider:
-        quota_unit = QuotaUnit.TIMES
+        quota_unit = QuotaUnit.CREDITS
         quotas = []
 
         if app_config.get("HOSTED_OPENAI_TRIAL_ENABLED"):
@@ -114,6 +110,8 @@ class HostingConfiguration:
                     RestrictModel(model="gpt-3.5-turbo-1106", model_type=ModelType.LLM),
                     RestrictModel(model="gpt-3.5-turbo-instruct", model_type=ModelType.LLM),
                     RestrictModel(model="gpt-3.5-turbo-16k", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-16k-0613", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-0613", model_type=ModelType.LLM),
                     RestrictModel(model="text-davinci-003", model_type=ModelType.LLM),
                     RestrictModel(model="whisper-1", model_type=ModelType.SPEECH2TEXT),
                 ]
@@ -122,10 +120,20 @@ class HostingConfiguration:
 
         if app_config.get("HOSTED_OPENAI_PAID_ENABLED"):
             paid_quota = PaidHostingQuota(
-                stripe_price_id=app_config.get("HOSTED_OPENAI_PAID_STRIPE_PRICE_ID"),
-                increase_quota=int(app_config.get("HOSTED_OPENAI_PAID_INCREASE_QUOTA", "1")),
-                min_quantity=int(app_config.get("HOSTED_OPENAI_PAID_MIN_QUANTITY", "1")),
-                max_quantity=int(app_config.get("HOSTED_OPENAI_PAID_MAX_QUANTITY", "1"))
+                restrict_models=[
+                    RestrictModel(model="gpt-4", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4-turbo-preview", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4-32k", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4-1106-preview", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-16k", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-16k-0613", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-1106", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-4-0125-preview", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-0613", model_type=ModelType.LLM),
+                    RestrictModel(model="gpt-3.5-turbo-instruct", model_type=ModelType.LLM),
+                    RestrictModel(model="text-davinci-003", model_type=ModelType.LLM),
+                ]
             )
             quotas.append(paid_quota)
 
@@ -164,12 +172,7 @@ class HostingConfiguration:
             quotas.append(trial_quota)
 
         if app_config.get("HOSTED_ANTHROPIC_PAID_ENABLED"):
-            paid_quota = PaidHostingQuota(
-                stripe_price_id=app_config.get("HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID"),
-                increase_quota=int(app_config.get("HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA", "1000000")),
-                min_quantity=int(app_config.get("HOSTED_ANTHROPIC_PAID_MIN_QUANTITY", "20")),
-                max_quantity=int(app_config.get("HOSTED_ANTHROPIC_PAID_MAX_QUANTITY", "100"))
-            )
+            paid_quota = PaidHostingQuota()
             quotas.append(paid_quota)
 
         if len(quotas) > 0:

api/core/model_runtime/model_providers/openai/llm/text-davinci-003.yaml

@@ -26,3 +26,4 @@ pricing:
   output: '0.002'
   unit: '0.001'
   currency: USD
+deprecated: true

api/events/event_handlers/deduct_quota_when_messaeg_created.py

@@ -33,6 +33,11 @@ def handle(sender, **kwargs):
     if quota_unit:
         if quota_unit == QuotaUnit.TOKENS:
             used_quota = message.message_tokens + message.answer_tokens
+        elif quota_unit == QuotaUnit.CREDITS:
+            used_quota = 1
+
+            if 'gpt-4' in model_config.model:
+                used_quota = 20
         else:
             used_quota = 1
 

api/services/billing_service.py

@@ -34,17 +34,22 @@ class BillingService:
     def get_model_provider_payment_link(cls,
                                         provider_name: str,
                                         tenant_id: str,
-                                        account_id: str):
+                                        account_id: str,
+                                        prefilled_email: str):
         params = {
             'provider_name': provider_name,
             'tenant_id': tenant_id,
-            'account_id': account_id
+            'account_id': account_id,
+            'prefilled_email': prefilled_email
         }
         return cls._send_request('GET', '/model-provider/payment-link', params=params)
 
     @classmethod
-    def get_invoices(cls, prefilled_email: str = ''):
-        params = {'prefilled_email': prefilled_email}
+    def get_invoices(cls, prefilled_email: str = '', tenant_id: str = ''):
+        params = {
+            'prefilled_email': prefilled_email,
+            'tenant_id': tenant_id
+        }
         return cls._send_request('GET', '/invoices', params=params)
 
     @classmethod
@@ -60,7 +65,7 @@ class BillingService:
         return response.json()
 
     @staticmethod
-    def is_tenant_owner(current_user):
+    def is_tenant_owner_or_admin(current_user):
         tenant_id = current_user.current_tenant_id
 
         join = db.session.query(TenantAccountJoin).filter(
@@ -68,5 +73,5 @@ class BillingService:
             TenantAccountJoin.account_id == current_user.id
         ).first()
 
-        if join.role != 'owner':
-            raise ValueError('Only tenant owner can perform this action')
+        if join.role not in ['owner', 'admin']:
+            raise ValueError('Only team owner or team admin can perform this action')