mirror of
https://github.com/langgenius/dify.git
synced 2026-05-04 01:18:05 +08:00
234 lines
8.0 KiB
Python
234 lines
8.0 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING
|
|
|
|
from configs import dify_config
|
|
|
|
if TYPE_CHECKING:
|
|
from enums.quota_type import QuotaType
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class QuotaCharge:
|
|
"""
|
|
Result of a quota reservation (Reserve phase).
|
|
|
|
Lifecycle:
|
|
charge = QuotaService.consume(QuotaType.TRIGGER, tenant_id)
|
|
try:
|
|
do_work()
|
|
charge.commit() # Confirm consumption
|
|
except:
|
|
charge.refund() # Release frozen quota
|
|
|
|
If neither commit() nor refund() is called, the billing system's
|
|
cleanup CronJob will auto-release the reservation within ~75 seconds.
|
|
"""
|
|
|
|
success: bool
|
|
charge_id: str | None # reservation_id
|
|
_quota_type: QuotaType
|
|
_tenant_id: str | None = None
|
|
_feature_key: str | None = None
|
|
_amount: int = 0
|
|
_committed: bool = field(default=False, repr=False)
|
|
|
|
def commit(self, actual_amount: int | None = None) -> None:
|
|
"""
|
|
Confirm the consumption with actual amount.
|
|
|
|
Args:
|
|
actual_amount: Actual amount consumed. Defaults to the reserved amount.
|
|
If less than reserved, the difference is refunded automatically.
|
|
"""
|
|
if self._committed or not self.charge_id or not self._tenant_id or not self._feature_key:
|
|
return
|
|
|
|
try:
|
|
from services.billing_service import BillingService
|
|
|
|
amount = actual_amount if actual_amount is not None else self._amount
|
|
BillingService.quota_commit(
|
|
tenant_id=self._tenant_id,
|
|
feature_key=self._feature_key,
|
|
reservation_id=self.charge_id,
|
|
actual_amount=amount,
|
|
)
|
|
self._committed = True
|
|
logger.debug(
|
|
"Committed %s quota for tenant %s, reservation_id: %s, amount: %d",
|
|
self._quota_type,
|
|
self._tenant_id,
|
|
self.charge_id,
|
|
amount,
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to commit quota, reservation_id: %s", self.charge_id)
|
|
|
|
def refund(self) -> None:
|
|
"""
|
|
Release the reserved quota (cancel the charge).
|
|
|
|
Safe to call even if:
|
|
- charge failed or was disabled (charge_id is None)
|
|
- already committed (Release after Commit is a no-op)
|
|
- already refunded (idempotent)
|
|
|
|
This method guarantees no exceptions will be raised.
|
|
"""
|
|
if not self.charge_id or not self._tenant_id or not self._feature_key:
|
|
return
|
|
|
|
QuotaService.release(self._quota_type, self.charge_id, self._tenant_id, self._feature_key)
|
|
|
|
|
|
def unlimited() -> QuotaCharge:
|
|
from enums.quota_type import QuotaType
|
|
|
|
return QuotaCharge(success=True, charge_id=None, _quota_type=QuotaType.UNLIMITED)
|
|
|
|
|
|
class QuotaService:
|
|
"""Orchestrates quota reserve / commit / release lifecycle via BillingService."""
|
|
|
|
@staticmethod
|
|
def consume(quota_type: QuotaType, tenant_id: str, amount: int = 1) -> QuotaCharge:
|
|
"""
|
|
Reserve + immediate Commit (one-shot mode).
|
|
|
|
The returned QuotaCharge supports .refund() which calls Release.
|
|
For two-phase usage (e.g. streaming), use reserve() directly.
|
|
"""
|
|
charge = QuotaService.reserve(quota_type, tenant_id, amount)
|
|
if charge.success and charge.charge_id:
|
|
charge.commit()
|
|
return charge
|
|
|
|
@staticmethod
|
|
def reserve(quota_type: QuotaType, tenant_id: str, amount: int = 1) -> QuotaCharge:
|
|
"""
|
|
Reserve quota before task execution (Reserve phase only).
|
|
|
|
The caller MUST call charge.commit() after the task succeeds,
|
|
or charge.refund() if the task fails.
|
|
|
|
Raises:
|
|
QuotaExceededError: When quota is insufficient
|
|
"""
|
|
from services.billing_service import BillingService
|
|
from services.errors.app import QuotaExceededError
|
|
|
|
if not dify_config.BILLING_ENABLED:
|
|
logger.debug("Billing disabled, allowing request for %s", tenant_id)
|
|
return QuotaCharge(success=True, charge_id=None, _quota_type=quota_type)
|
|
|
|
logger.info("Reserving %d %s quota for tenant %s", amount, quota_type.value, tenant_id)
|
|
|
|
if amount <= 0:
|
|
raise ValueError("Amount to reserve must be greater than 0")
|
|
|
|
request_id = str(uuid.uuid4())
|
|
feature_key = quota_type.billing_key
|
|
|
|
try:
|
|
reserve_resp = BillingService.quota_reserve(
|
|
tenant_id=tenant_id,
|
|
feature_key=feature_key,
|
|
request_id=request_id,
|
|
amount=amount,
|
|
)
|
|
|
|
reservation_id = reserve_resp.get("reservation_id")
|
|
if not reservation_id:
|
|
logger.warning(
|
|
"Reserve returned no reservation_id for %s, feature %s, response: %s",
|
|
tenant_id,
|
|
quota_type.value,
|
|
reserve_resp,
|
|
)
|
|
raise QuotaExceededError(feature=quota_type.value, tenant_id=tenant_id, required=amount)
|
|
|
|
logger.debug(
|
|
"Reserved %d %s quota for tenant %s, reservation_id: %s",
|
|
amount,
|
|
quota_type.value,
|
|
tenant_id,
|
|
reservation_id,
|
|
)
|
|
return QuotaCharge(
|
|
success=True,
|
|
charge_id=reservation_id,
|
|
_quota_type=quota_type,
|
|
_tenant_id=tenant_id,
|
|
_feature_key=feature_key,
|
|
_amount=amount,
|
|
)
|
|
|
|
except QuotaExceededError:
|
|
raise
|
|
except ValueError:
|
|
raise
|
|
except Exception:
|
|
logger.exception("Failed to reserve quota for %s, feature %s", tenant_id, quota_type.value)
|
|
return unlimited()
|
|
|
|
@staticmethod
|
|
def check(quota_type: QuotaType, tenant_id: str, amount: int = 1) -> bool:
|
|
if not dify_config.BILLING_ENABLED:
|
|
return True
|
|
|
|
if amount <= 0:
|
|
raise ValueError("Amount to check must be greater than 0")
|
|
|
|
try:
|
|
remaining = QuotaService.get_remaining(quota_type, tenant_id)
|
|
return remaining >= amount if remaining != -1 else True
|
|
except Exception:
|
|
logger.exception("Failed to check quota for %s, feature %s", tenant_id, quota_type.value)
|
|
return True
|
|
|
|
@staticmethod
|
|
def release(quota_type: QuotaType, reservation_id: str, tenant_id: str, feature_key: str) -> None:
|
|
"""Release a reservation. Guarantees no exceptions."""
|
|
try:
|
|
from services.billing_service import BillingService
|
|
|
|
if not dify_config.BILLING_ENABLED:
|
|
return
|
|
|
|
if not reservation_id:
|
|
return
|
|
|
|
logger.info("Releasing %s quota, reservation_id: %s", quota_type.value, reservation_id)
|
|
BillingService.quota_release(
|
|
tenant_id=tenant_id,
|
|
feature_key=feature_key,
|
|
reservation_id=reservation_id,
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to release quota, reservation_id: %s", reservation_id)
|
|
|
|
@staticmethod
|
|
def get_remaining(quota_type: QuotaType, tenant_id: str) -> int:
|
|
from services.billing_service import BillingService
|
|
|
|
try:
|
|
usage_info = BillingService.get_quota_info(tenant_id)
|
|
if isinstance(usage_info, dict):
|
|
feature_info = usage_info.get(quota_type.billing_key, {})
|
|
if isinstance(feature_info, dict):
|
|
limit = feature_info.get("limit", 0)
|
|
usage = feature_info.get("usage", 0)
|
|
if limit == -1:
|
|
return -1
|
|
return max(0, limit - usage)
|
|
return 0
|
|
except Exception:
|
|
logger.exception("Failed to get remaining quota for %s, feature %s", tenant_id, quota_type.value)
|
|
return -1
|