mirror of
https://github.com/langgenius/dify.git
synced 2026-05-13 08:57:28 +08:00
94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
import logging
|
|
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from configs import dify_config
|
|
from core.errors.error import QuotaExceededError
|
|
from extensions.ext_database import db
|
|
from models import TenantCreditPool
|
|
from models.enums import ProviderQuotaType
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CreditPoolService:
|
|
@classmethod
|
|
def create_default_pool(cls, tenant_id: str) -> TenantCreditPool:
|
|
"""create default credit pool for new tenant"""
|
|
credit_pool = TenantCreditPool(
|
|
tenant_id=tenant_id,
|
|
quota_limit=dify_config.HOSTED_POOL_CREDITS,
|
|
quota_used=0,
|
|
pool_type=ProviderQuotaType.TRIAL,
|
|
)
|
|
db.session.add(credit_pool)
|
|
db.session.commit()
|
|
return credit_pool
|
|
|
|
@classmethod
|
|
def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | None:
|
|
"""get tenant credit pool"""
|
|
with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
|
|
return session.scalar(
|
|
select(TenantCreditPool)
|
|
.where(
|
|
TenantCreditPool.tenant_id == tenant_id,
|
|
TenantCreditPool.pool_type == pool_type,
|
|
)
|
|
.limit(1)
|
|
)
|
|
|
|
@classmethod
|
|
def check_credits_available(
|
|
cls,
|
|
tenant_id: str,
|
|
credits_required: int,
|
|
pool_type: str = "trial",
|
|
) -> bool:
|
|
"""check if credits are available without deducting"""
|
|
pool = cls.get_pool(tenant_id, pool_type)
|
|
if not pool:
|
|
return False
|
|
return pool.remaining_credits >= credits_required
|
|
|
|
@classmethod
|
|
def check_and_deduct_credits(
|
|
cls,
|
|
tenant_id: str,
|
|
credits_required: int,
|
|
pool_type: str = "trial",
|
|
) -> int:
|
|
"""Deduct credits, depleting the pool before raising if the balance is insufficient."""
|
|
|
|
pool = cls.get_pool(tenant_id, pool_type)
|
|
if not pool:
|
|
raise QuotaExceededError("Credit pool not found")
|
|
|
|
if pool.remaining_credits <= 0:
|
|
raise QuotaExceededError("No credits remaining")
|
|
|
|
remaining_credits = pool.remaining_credits
|
|
actual_credits = min(credits_required, remaining_credits)
|
|
quota_exceeded = actual_credits < credits_required
|
|
|
|
try:
|
|
with sessionmaker(db.engine).begin() as session:
|
|
stmt = (
|
|
update(TenantCreditPool)
|
|
.where(
|
|
TenantCreditPool.tenant_id == tenant_id,
|
|
TenantCreditPool.pool_type == pool_type,
|
|
)
|
|
.values(quota_used=TenantCreditPool.quota_used + actual_credits)
|
|
)
|
|
session.execute(stmt)
|
|
except Exception:
|
|
logger.exception("Failed to deduct credits for tenant %s", tenant_id)
|
|
raise QuotaExceededError("Failed to deduct credits")
|
|
|
|
if quota_exceeded:
|
|
raise QuotaExceededError("Insufficient credits remaining")
|
|
|
|
return actual_credits
|