feat: server multi models support (#799)

api/services/provider_checkout_service.py

@@ -0,0 +1,158 @@
+import datetime
+import logging
+
+import stripe
+from flask import current_app
+
+from core.model_providers.model_provider_factory import ModelProviderFactory
+from extensions.ext_database import db
+from models.account import Account
+from models.provider import ProviderOrder, ProviderOrderPaymentStatus, ProviderType, Provider, ProviderQuotaType
+
+
+class ProviderCheckout:
+    def __init__(self, stripe_checkout_session):
+        self.stripe_checkout_session = stripe_checkout_session
+
+    def get_checkout_url(self):
+        return self.stripe_checkout_session.url
+
+
+class ProviderCheckoutService:
+    def create_checkout(self, tenant_id: str, provider_name: str, account: Account) -> ProviderCheckout:
+        # check provider name is valid
+        model_provider_rules = ModelProviderFactory.get_provider_rules()
+        if provider_name not in model_provider_rules:
+            raise ValueError(f'provider name {provider_name} is invalid')
+
+        model_provider_rule = model_provider_rules[provider_name]
+
+        # check provider name can be paid
+        self._check_provider_payable(provider_name, model_provider_rule)
+
+        # get stripe checkout product id
+        paid_provider = self._get_paid_provider(tenant_id, provider_name)
+        model_provider_class = ModelProviderFactory.get_model_provider_class(provider_name)
+        model_provider = model_provider_class(provider=paid_provider)
+        payment_info = model_provider.get_payment_info()
+        if not payment_info:
+            raise ValueError(f'provider name {provider_name} not support payment')
+
+        payment_product_id = payment_info['product_id']
+
+        # create provider order
+        provider_order = ProviderOrder(
+            tenant_id=tenant_id,
+            provider_name=provider_name,
+            account_id=account.id,
+            payment_product_id=payment_product_id,
+            quantity=1,
+            payment_status=ProviderOrderPaymentStatus.WAIT_PAY.value
+        )
+
+        db.session.add(provider_order)
+        db.session.flush()
+
+        try:
+            # create stripe checkout session
+            checkout_session = stripe.checkout.Session.create(
+                line_items=[
+                    {
+                        'price': f'{payment_product_id}',
+                        'quantity': 1,
+                    },
+                ],
+                mode='payment',
+                success_url=current_app.config.get("CONSOLE_WEB_URL") + '?provider_payment=succeeded',
+                cancel_url=current_app.config.get("CONSOLE_WEB_URL") + '?provider_payment=cancelled',
+                automatic_tax={'enabled': True},
+            )
+        except Exception as e:
+            logging.exception(e)
+            raise ValueError(f'provider name {provider_name} create checkout session failed, please try again later')
+
+        provider_order.payment_id = checkout_session.id
+        db.session.commit()
+
+        return ProviderCheckout(checkout_session)
+
+    def fulfill_provider_order(self, event):
+        provider_order = db.session.query(ProviderOrder) \
+            .filter(ProviderOrder.payment_id == event['data']['object']['id']) \
+            .first()
+
+        if not provider_order:
+            raise ValueError(f'provider order not found, payment id: {event["data"]["object"]["id"]}')
+
+        if provider_order.payment_status != ProviderOrderPaymentStatus.WAIT_PAY.value:
+            raise ValueError(f'provider order payment status is not wait pay, payment id: {event["data"]["object"]["id"]}')
+
+        provider_order.transaction_id = event['data']['object']['payment_intent']
+        provider_order.currency = event['data']['object']['currency']
+        provider_order.total_amount = event['data']['object']['amount_subtotal']
+        provider_order.payment_status = ProviderOrderPaymentStatus.PAID.value
+        provider_order.paid_at = datetime.datetime.utcnow()
+        provider_order.updated_at = provider_order.paid_at
+
+        # update provider quota
+        provider = db.session.query(Provider).filter(
+            Provider.tenant_id == provider_order.tenant_id,
+            Provider.provider_name == provider_order.provider_name,
+            Provider.provider_type == ProviderType.SYSTEM.value,
+            Provider.quota_type == ProviderQuotaType.PAID.value
+        ).first()
+
+        if not provider:
+            raise ValueError(f'provider not found, tenant id: {provider_order.tenant_id}, '
+                             f'provider name: {provider_order.provider_name}')
+
+        model_provider_class = ModelProviderFactory.get_model_provider_class(provider_order.provider_name)
+        model_provider = model_provider_class(provider=provider)
+        payment_info = model_provider.get_payment_info()
+
+        if not payment_info:
+            increase_quota = 0
+        else:
+            increase_quota = int(payment_info['increase_quota'])
+
+        if increase_quota > 0:
+            provider.quota_limit += increase_quota
+            provider.is_valid = True
+
+        db.session.commit()
+
+    def _check_provider_payable(self, provider_name: str, model_provider_rule: dict):
+        if ProviderType.SYSTEM.value not in model_provider_rule['support_provider_types']:
+            raise ValueError(f'provider name {provider_name} not support payment')
+
+        if 'system_config' not in model_provider_rule:
+            raise ValueError(f'provider name {provider_name} not support payment')
+
+        if 'supported_quota_types' not in model_provider_rule['system_config']:
+            raise ValueError(f'provider name {provider_name} not support payment')
+
+        if 'paid' not in model_provider_rule['system_config']['supported_quota_types']:
+            raise ValueError(f'provider name {provider_name} not support payment')
+
+    def _get_paid_provider(self, tenant_id: str, provider_name: str):
+        paid_provider = db.session.query(Provider) \
+            .filter(
+            Provider.tenant_id == tenant_id,
+            Provider.provider_name == provider_name,
+            Provider.provider_type == ProviderType.SYSTEM.value,
+            Provider.quota_type == ProviderQuotaType.PAID.value,
+        ).first()
+
+        if not paid_provider:
+            paid_provider = Provider(
+                tenant_id=tenant_id,
+                provider_name=provider_name,
+                provider_type=ProviderType.SYSTEM.value,
+                quota_type=ProviderQuotaType.PAID.value,
+                quota_limit=0,
+                quota_used=0,
+            )
+            db.session.add(paid_provider)
+            db.session.commit()
+
+        return paid_provider