From 2b51fc23d9231f7aa96cc97fe2d7934d84da613f Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 10:43:34 +0800 Subject: [PATCH 01/13] add credit pool sys --- .../feature/hosted_service/__init__.py | 122 ++++++++++++++++-- .../console/workspace/workspace.py | 2 + api/core/hosting_configuration.py | 90 ++++++++++++- api/core/provider_manager.py | 40 ++++-- .../update_provider_when_message_created.py | 36 ++++-- ...1520-58a70d22fdbd_add_table_credit_pool.py | 96 ++++++++++++++ api/models/__init__.py | 2 + api/models/model.py | 28 +++- api/services/account_service.py | 5 + api/services/credit_pool_service.py | 107 +++++++++++++++ api/services/workspace_service.py | 6 + 11 files changed, 493 insertions(+), 41 deletions(-) create mode 100644 api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py create mode 100644 api/services/credit_pool_service.py diff --git a/api/configs/feature/hosted_service/__init__.py b/api/configs/feature/hosted_service/__init__.py index 4ad30014c7..538c55d931 100644 --- a/api/configs/feature/hosted_service/__init__.py +++ b/api/configs/feature/hosted_service/__init__.py @@ -8,6 +8,11 @@ class HostedCreditConfig(BaseSettings): default="", ) + HOSTED_POOL_CREDITS: int = Field( + description="Pool credits for hosted service", + default=200, + ) + def get_model_credits(self, model_name: str) -> int: """ Get credit value for a specific model name. @@ -70,11 +75,6 @@ class HostedOpenAiConfig(BaseSettings): "text-davinci-003", ) - HOSTED_OPENAI_QUOTA_LIMIT: NonNegativeInt = Field( - description="Quota limit for hosted OpenAI service usage", - default=200, - ) - HOSTED_OPENAI_PAID_ENABLED: bool = Field( description="Enable paid access to hosted OpenAI service", default=False, @@ -98,6 +98,99 @@ class HostedOpenAiConfig(BaseSettings): ) +class HostedGeminiConfig(BaseSettings): + """ + Configuration for fetching Gemini service + """ + + HOSTED_GEMINI_API_KEY: str | None = Field( + description="API key for hosted Gemini service", + default=None, + ) + + HOSTED_GEMINI_API_BASE: str | None = Field( + description="Base URL for hosted Gemini API", + default=None, + ) + + HOSTED_GEMINI_API_ORGANIZATION: str | None = Field( + description="Organization ID for hosted Gemini service", + default=None, + ) + + HOSTED_GEMINI_TRIAL_ENABLED: bool = Field( + description="Enable trial access to hosted Gemini service", + default=False, + ) + + HOSTED_GEMINI_TRIAL_MODELS: str = Field( + description="Comma-separated list of available models for trial access", + default="gemini-2.5-flash,gemini-2.0-flash,gemini-2.0-flash-lite,", + ) + + +class HostedXAIConfig(BaseSettings): + """ + Configuration for fetching XAI service + """ + + HOSTED_XAI_API_KEY: str | None = Field( + description="API key for hosted XAI service", + default=None, + ) + + HOSTED_XAI_API_BASE: str | None = Field( + description="Base URL for hosted XAI API", + default=None, + ) + + HOSTED_XAI_API_ORGANIZATION: str | None = Field( + description="Organization ID for hosted XAI service", + default=None, + ) + + HOSTED_XAI_TRIAL_ENABLED: bool = Field( + description="Enable trial access to hosted XAI service", + default=False, + ) + + HOSTED_XAI_TRIAL_MODELS: str = Field( + description="Comma-separated list of available models for trial access", + default="grok-3,grok-3-mini,grok-3-mini-fast", + ) + + +class HostedDeepseekConfig(BaseSettings): + """ + Configuration for fetching Deepseek service + """ + + HOSTED_DEEPSEEK_API_KEY: str | None = Field( + description="API key for hosted Deepseek service", + default=None, + ) + + HOSTED_DEEPSEEK_API_BASE: str | None = Field( + description="Base URL for hosted Deepseek API", + default=None, + ) + + HOSTED_DEEPSEEK_API_ORGANIZATION: str | None = Field( + description="Organization ID for hosted Deepseek service", + default=None, + ) + + HOSTED_DEEPSEEK_TRIAL_ENABLED: bool = Field( + description="Enable trial access to hosted Deepseek service", + default=False, + ) + + HOSTED_DEEPSEEK_TRIAL_MODELS: str = Field( + description="Comma-separated list of available models for trial access", + default="deepseek-chat,deepseek-reasoner", + ) + + class HostedAzureOpenAiConfig(BaseSettings): """ Configuration for hosted Azure OpenAI service @@ -144,16 +237,22 @@ class HostedAnthropicConfig(BaseSettings): default=False, ) - HOSTED_ANTHROPIC_QUOTA_LIMIT: NonNegativeInt = Field( - description="Quota limit for hosted Anthropic service usage", - default=600000, - ) - HOSTED_ANTHROPIC_PAID_ENABLED: bool = Field( description="Enable paid access to hosted Anthropic service", default=False, ) + HOSTED_ANTHROPIC_TRIAL_MODELS: str = Field( + description="Comma-separated list of available models for paid access", + default="claude-opus-4-20250514," + "claude-opus-4-20250514," + "claude-sonnet-4-20250514," + "claude-3-5-haiku-20241022," + "claude-3-opus-20240229," + "claude-3-7-sonnet-20250219," + "claude-3-haiku-20240307", + ) + class HostedMinmaxConfig(BaseSettings): """ @@ -250,5 +349,8 @@ class HostedServiceConfig( HostedModerationConfig, # credit config HostedCreditConfig, + HostedGeminiConfig, + HostedXAIConfig, + HostedDeepseekConfig, ): pass diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index 6bec70b5da..5242e6e04c 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -51,6 +51,8 @@ tenant_fields = { "in_trial": fields.Boolean, "trial_end_reason": fields.String, "custom_config": fields.Raw(attribute="custom_config"), + "trial_credits": fields.Integer, + "trial_credits_used": fields.Integer, } tenants_fields = { diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py index af860a1070..7aafa4bc80 100644 --- a/api/core/hosting_configuration.py +++ b/api/core/hosting_configuration.py @@ -56,6 +56,9 @@ class HostingConfiguration: self.provider_map[f"{DEFAULT_PLUGIN_ID}/minimax/minimax"] = self.init_minimax() self.provider_map[f"{DEFAULT_PLUGIN_ID}/spark/spark"] = self.init_spark() self.provider_map[f"{DEFAULT_PLUGIN_ID}/zhipuai/zhipuai"] = self.init_zhipuai() + self.provider_map[f"{DEFAULT_PLUGIN_ID}/gemini/google"] = self.init_gemini() + self.provider_map[f"{DEFAULT_PLUGIN_ID}/x/x"] = self.init_xai() + self.provider_map[f"{DEFAULT_PLUGIN_ID}/deepseek/deepseek"] = self.init_deepseek() self.moderation_config = self.init_moderation_config() @@ -128,7 +131,7 @@ class HostingConfiguration: quotas: list[HostingQuota] = [] if dify_config.HOSTED_OPENAI_TRIAL_ENABLED: - hosted_quota_limit = dify_config.HOSTED_OPENAI_QUOTA_LIMIT + hosted_quota_limit = 0 trial_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_TRIAL_MODELS") trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models) quotas.append(trial_quota) @@ -156,14 +159,39 @@ class HostingConfiguration: quota_unit=quota_unit, ) - @staticmethod - def init_anthropic() -> HostingProvider: - quota_unit = QuotaUnit.TOKENS + def init_gemini(self) -> HostingProvider: + quota_unit = QuotaUnit.CREDITS + quotas: list[HostingQuota] = [] + + if dify_config.HOSTED_GEMINI_TRIAL_ENABLED: + hosted_quota_limit = 0 + trial_models = self.parse_restrict_models_from_env("HOSTED_GEMINI_TRIAL_MODELS") + trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models) + quotas.append(trial_quota) + + if len(quotas) > 0: + credentials = { + "google_api_key": dify_config.HOSTED_GEMINI_API_KEY, + } + + if dify_config.HOSTED_GEMINI_API_BASE: + credentials["google_base_url"] = dify_config.HOSTED_GEMINI_API_BASE + + return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + def init_anthropic(self) -> HostingProvider: + quota_unit = QuotaUnit.CREDITS quotas: list[HostingQuota] = [] if dify_config.HOSTED_ANTHROPIC_TRIAL_ENABLED: - hosted_quota_limit = dify_config.HOSTED_ANTHROPIC_QUOTA_LIMIT - trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit) + hosted_quota_limit = 0 + trail_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_TRIAL_MODELS") + trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models) quotas.append(trial_quota) if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED: @@ -185,6 +213,56 @@ class HostingConfiguration: quota_unit=quota_unit, ) + def init_xai(self) -> HostingProvider: + quota_unit = QuotaUnit.CREDITS + quotas: list[HostingQuota] = [] + + if dify_config.HOSTED_XAI_TRIAL_ENABLED: + hosted_quota_limit = 0 + trail_models = self.parse_restrict_models_from_env("HOSTED_XAI_TRIAL_MODELS") + trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models) + quotas.append(trial_quota) + + if len(quotas) > 0: + credentials = { + "api_key": dify_config.HOSTED_XAI_API_KEY, + } + + if dify_config.HOSTED_XAI_API_BASE: + credentials["endpoint_url"] = dify_config.HOSTED_XAI_API_BASE + + return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + + def init_deepseek(self) -> HostingProvider: + quota_unit = QuotaUnit.CREDITS + quotas: list[HostingQuota] = [] + + if dify_config.HOSTED_DEEPSEEK_TRIAL_ENABLED: + hosted_quota_limit = 0 + trail_models = self.parse_restrict_models_from_env("HOSTED_DEEPSEEK_TRIAL_MODELS") + trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models) + quotas.append(trial_quota) + + if len(quotas) > 0: + credentials = { + "api_key": dify_config.HOSTED_DEEPSEEK_API_KEY, + } + + if dify_config.HOSTED_DEEPSEEK_API_BASE: + credentials["endpoint_url"] = dify_config.HOSTED_DEEPSEEK_API_BASE + + return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas) + + return HostingProvider( + enabled=False, + quota_unit=quota_unit, + ) + @staticmethod def init_minimax() -> HostingProvider: quota_unit = QuotaUnit.TOKENS diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 499d39bd5d..1ac02d9b6a 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -629,7 +629,7 @@ class ProviderManager: provider_name=ModelProviderID(provider_name).provider_name, provider_type=ProviderType.SYSTEM.value, quota_type=ProviderQuotaType.TRIAL.value, - quota_limit=quota.quota_limit, # type: ignore + quota_limit=0, # type: ignore quota_used=0, is_valid=True, ) @@ -912,6 +912,16 @@ class ProviderManager: provider_record ) quota_configurations = [] + + if dify_config.EDITION == "CLOUD": + from services.credit_pool_service import CreditPoolService + + pool = CreditPoolService.get_or_create_pool( + tenant_id=tenant_id, + ) + else: + pool = None + for provider_quota in provider_hosting_configuration.quotas: if provider_quota.quota_type not in quota_type_to_provider_records_dict: if provider_quota.quota_type == ProviderQuotaType.FREE: @@ -932,16 +942,26 @@ class ProviderManager: raise ValueError("quota_used is None") if provider_record.quota_limit is None: raise ValueError("quota_limit is None") + if provider_quota.quota_type == ProviderQuotaType.TRIAL and pool is not None: + quota_configuration = QuotaConfiguration( + quota_type=provider_quota.quota_type, + quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, + quota_used=pool.quota_used, + quota_limit=pool.quota_limit, + is_valid=pool.quota_limit > pool.quota_used or pool.quota_limit == -1, + restrict_models=provider_quota.restrict_models, + ) - quota_configuration = QuotaConfiguration( - quota_type=provider_quota.quota_type, - quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, - quota_used=provider_record.quota_used, - quota_limit=provider_record.quota_limit, - is_valid=provider_record.quota_limit > provider_record.quota_used - or provider_record.quota_limit == -1, - restrict_models=provider_quota.restrict_models, - ) + else: + quota_configuration = QuotaConfiguration( + quota_type=provider_quota.quota_type, + quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, + quota_used=provider_record.quota_used, + quota_limit=provider_record.quota_limit, + is_valid=provider_record.quota_limit > provider_record.quota_used + or provider_record.quota_limit == -1, + restrict_models=provider_quota.restrict_models, + ) quota_configurations.append(quota_configuration) diff --git a/api/events/event_handlers/update_provider_when_message_created.py b/api/events/event_handlers/update_provider_when_message_created.py index 27efa539dc..d787b23fac 100644 --- a/api/events/event_handlers/update_provider_when_message_created.py +++ b/api/events/event_handlers/update_provider_when_message_created.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Session from configs import dify_config from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity -from core.entities.provider_entities import QuotaUnit, SystemConfiguration +from core.entities.provider_entities import ProviderQuotaType, QuotaUnit, SystemConfiguration from events.message_event import message_was_created from extensions.ext_database import db from extensions.ext_redis import redis_client, redis_fallback @@ -135,20 +135,28 @@ def handle(sender: Message, **kwargs): ) if used_quota is not None: - quota_update = _ProviderUpdateOperation( - filters=_ProviderUpdateFilters( + if provider_configuration.system_configuration.current_quota_type == ProviderQuotaType.TRIAL: + from services.credit_pool_service import CreditPoolService + + CreditPoolService.check_and_deduct_credits( tenant_id=tenant_id, - provider_name=ModelProviderID(model_config.provider).provider_name, - provider_type=ProviderType.SYSTEM.value, - quota_type=provider_configuration.system_configuration.current_quota_type.value, - ), - values=_ProviderUpdateValues(quota_used=Provider.quota_used + used_quota, last_used=current_time), - additional_filters=_ProviderUpdateAdditionalFilters( - quota_limit_check=True # Provider.quota_limit > Provider.quota_used - ), - description="quota_deduction_update", - ) - updates_to_perform.append(quota_update) + credits_required=used_quota, + ) + else: + quota_update = _ProviderUpdateOperation( + filters=_ProviderUpdateFilters( + tenant_id=tenant_id, + provider_name=ModelProviderID(model_config.provider).provider_name, + provider_type=ProviderType.SYSTEM.value, + quota_type=provider_configuration.system_configuration.current_quota_type.value, + ), + values=_ProviderUpdateValues(quota_used=Provider.quota_used + used_quota, last_used=current_time), + additional_filters=_ProviderUpdateAdditionalFilters( + quota_limit_check=True # Provider.quota_limit > Provider.quota_used + ), + description="quota_deduction_update", + ) + updates_to_perform.append(quota_update) # Execute all updates start_time = time_module.perf_counter() diff --git a/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py b/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py new file mode 100644 index 0000000000..d298d885f4 --- /dev/null +++ b/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py @@ -0,0 +1,96 @@ +"""add table credit pool + +Revision ID: 58a70d22fdbd +Revises: 68519ad5cd18 +Create Date: 2025-09-25 15:20:40.367078 + +""" +from alembic import op +import models as models +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '58a70d22fdbd' +down_revision = '68519ad5cd18' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('tenant_credit_pools', + sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('pool_type', sa.String(length=40), nullable=False), + sa.Column('quota_limit', sa.BigInteger(), nullable=False), + sa.Column('quota_used', sa.BigInteger(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.PrimaryKeyConstraint('id', name='tenant_credit_pool_pkey') + ) + with op.batch_alter_table('tenant_credit_pools', schema=None) as batch_op: + batch_op.create_index('tenant_credit_pool_pool_type_idx', ['pool_type'], unique=False) + batch_op.create_index('tenant_credit_pool_tenant_id_idx', ['tenant_id'], unique=False) + # Data migration: Move trial quota data from providers to tenant_credit_pools + migrate_trial_quota_data() + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('tenant_credit_pools', schema=None) as batch_op: + batch_op.drop_index('tenant_credit_pool_tenant_id_idx') + batch_op.drop_index('tenant_credit_pool_pool_type_idx') + + op.drop_table('tenant_credit_pools') + # ### end Alembic commands ### + + +def migrate_trial_quota_data(): + """ + Migrate quota data from providers table to tenant_credit_pools table + for providers with quota_type='trial', provider_name='openai', provider_type='system' + """ + # Create connection + bind = op.get_bind() + + # Query providers that match the criteria + select_sql = sa.text(""" + SELECT tenant_id, quota_limit, quota_used + FROM providers + WHERE quota_type = 'trial' + AND provider_name = 'openai' + AND provider_type = 'system' + AND quota_limit IS NOT NULL + """) + + result = bind.execute(select_sql) + providers_data = result.fetchall() + + # Insert data into tenant_credit_pools + for provider_data in providers_data: + tenant_id, quota_limit, quota_used = provider_data + + # Check if credit pool already exists for this tenant + check_sql = sa.text(""" + SELECT COUNT(*) + FROM tenant_credit_pools + WHERE tenant_id = :tenant_id AND pool_type = 'trial' + """) + + existing_count = bind.execute(check_sql, {"tenant_id": tenant_id}).scalar() + + if existing_count == 0: + # Insert new credit pool record + insert_sql = sa.text(""" + INSERT INTO tenant_credit_pools (tenant_id, pool_type, quota_limit, quota_used, created_at, updated_at) + VALUES (:tenant_id, 'trial', :quota_limit, :quota_used, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + """) + + bind.execute(insert_sql, { + "tenant_id": tenant_id, + "quota_limit": quota_limit or 0, + "quota_used": quota_used or 0 + }) diff --git a/api/models/__init__.py b/api/models/__init__.py index 779484283f..6cdb7529e3 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -53,6 +53,7 @@ from .model import ( Site, Tag, TagBinding, + TenantCreditPool, TraceAppConfig, UploadFile, ) @@ -159,6 +160,7 @@ __all__ = [ "Tenant", "TenantAccountJoin", "TenantAccountRole", + "TenantCreditPool", "TenantDefaultModel", "TenantPreferredModelProvider", "TenantStatus", diff --git a/api/models/model.py b/api/models/model.py index a8218c3a4e..30ec03de97 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, cast import sqlalchemy as sa from flask import request from flask_login import UserMixin # type: ignore[import-untyped] -from sqlalchemy import Float, Index, PrimaryKeyConstraint, String, exists, func, select, text +from sqlalchemy import BigInteger, Float, Index, PrimaryKeyConstraint, String, exists, func, select, text from sqlalchemy.orm import Mapped, Session, mapped_column from configs import dify_config @@ -1944,3 +1944,29 @@ class TraceAppConfig(Base): "created_at": str(self.created_at) if self.created_at else None, "updated_at": str(self.updated_at) if self.updated_at else None, } + + +class TenantCreditPool(Base): + __tablename__ = "tenant_credit_pools" + __table_args__ = ( + sa.PrimaryKeyConstraint("id", name="tenant_credit_pool_pkey"), + sa.Index("tenant_credit_pool_tenant_id_idx", "tenant_id"), + sa.Index("tenant_credit_pool_pool_type_idx", "pool_type"), + ) + + id = mapped_column(StringUUID, primary_key=True, server_default=text("uuid_generate_v4()")) + tenant_id = mapped_column(StringUUID, nullable=False) + pool_type = mapped_column(String(40), nullable=False, default="trial", server_default="trial") + quota_limit = mapped_column(BigInteger, nullable=False, default=0) + quota_used = mapped_column(BigInteger, nullable=False, default=0) + created_at = mapped_column(sa.DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP")) + updated_at = mapped_column( + sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp() + ) + + @property + def remaining_credits(self) -> int: + return max(0, self.quota_limit - self.quota_used) + + def has_sufficient_credits(self, required_credits: int) -> bool: + return self.remaining_credits >= required_credits diff --git a/api/services/account_service.py b/api/services/account_service.py index 0e699d16da..21637a69e5 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -993,6 +993,11 @@ class TenantService: tenant.encrypt_public_key = generate_key_pair(tenant.id) db.session.commit() + + from services.credit_pool_service import CreditPoolService + + CreditPoolService.create_default_pool(tenant.id) + return tenant @staticmethod diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py new file mode 100644 index 0000000000..2bf9f4118f --- /dev/null +++ b/api/services/credit_pool_service.py @@ -0,0 +1,107 @@ +from typing import Optional + +from sqlalchemy import update + +from configs import dify_config +from core.errors.error import QuotaExceededError +from extensions.ext_database import db +from models import TenantCreditPool + + +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="trial" + ) + db.session.add(credit_pool) + db.session.commit() + return credit_pool + + @classmethod + def get_pool(cls, tenant_id: str) -> Optional[TenantCreditPool]: + """get tenant credit pool""" + return ( + db.session.query(TenantCreditPool) + .filter_by( + tenant_id=tenant_id, + ) + .first() + ) + + @classmethod + def get_or_create_pool(cls, tenant_id: str) -> TenantCreditPool: + """get or create credit pool""" + # First try to get existing pool + pool = cls.get_pool(tenant_id) + if pool: + return pool + + # Create new pool if not exists, handle race condition + try: + # Double-check in case another thread created it + pool = ( + db.session.query(TenantCreditPool) + .filter_by( + tenant_id=tenant_id, + ) + .first() + ) + if pool: + return pool + + # Create new pool + pool = TenantCreditPool( + tenant_id=tenant_id, quota_limit=dify_config.HOSTED_POOL_CREDITS, quota_used=0, pool_type="trial" + ) + db.session.add(pool) + db.session.commit() + + except Exception: + # If creation fails (e.g., due to race condition), rollback and try to get existing one + db.session.rollback() + pool = cls.get_pool(tenant_id) + if not pool: + raise + + return pool + + @classmethod + def check_and_deduct_credits( + cls, + tenant_id: str, + credits_required: int, + ) -> bool: + """check and deduct credits""" + pool = cls.get_pool(tenant_id) + if not pool: + raise QuotaExceededError("Credit pool not found") + + if pool.remaining_credits < credits_required: + raise QuotaExceededError( + f"Insufficient credits. Required: {credits_required}, Available: {pool.remaining_credits}" + ) + + with db.session.begin(): + update_values = {"quota_used": pool.quota_used + credits_required} + + where_conditions = [ + TenantCreditPool.tenant_id == tenant_id, + TenantCreditPool.quota_used + credits_required <= TenantCreditPool.quota_limit, + ] + stmt = update(TenantCreditPool).where(*where_conditions).values(**update_values) + db.session.execute(stmt) + + return True + + @classmethod + def check_deduct_credits(cls, tenant_id: str, credits_required: int) -> bool: + """check and deduct credits""" + pool = cls.get_pool(tenant_id) + if not pool: + return False + + if pool.remaining_credits < credits_required: + return False + return True diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index 292ac6e008..a21aac1984 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -46,5 +46,11 @@ class WorkspaceService: "remove_webapp_brand": remove_webapp_brand, "replace_webapp_logo": replace_webapp_logo, } + if dify_config.EDITION == "CLOUD": + from services.credit_pool_service import CreditPoolService + + pool = CreditPoolService.get_or_create_pool(tenant_id=tenant.id) + tenant_info["trial_credits"] = pool.quota_limit + tenant_info["trial_credits_used"] = pool.quota_used return tenant_info From db0780cfa8390498197d7bc71ed38edcd513fbfd Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 11:18:28 +0800 Subject: [PATCH 02/13] add:log --- .../update_provider_when_message_created.py | 4 +++- api/services/credit_pool_service.py | 8 +++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/api/events/event_handlers/update_provider_when_message_created.py b/api/events/event_handlers/update_provider_when_message_created.py index d787b23fac..9efe9a79af 100644 --- a/api/events/event_handlers/update_provider_when_message_created.py +++ b/api/events/event_handlers/update_provider_when_message_created.py @@ -133,9 +133,10 @@ def handle(sender: Message, **kwargs): system_configuration=system_configuration, model_name=model_config.model, ) - + logger.info("used_quota: %s", used_quota) if used_quota is not None: if provider_configuration.system_configuration.current_quota_type == ProviderQuotaType.TRIAL: + logger.info("deduct credits") from services.credit_pool_service import CreditPoolService CreditPoolService.check_and_deduct_credits( @@ -143,6 +144,7 @@ def handle(sender: Message, **kwargs): credits_required=used_quota, ) else: + logger.info("update provider quota") quota_update = _ProviderUpdateOperation( filters=_ProviderUpdateFilters( tenant_id=tenant_id, diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index 2bf9f4118f..f72686b9ff 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -1,3 +1,4 @@ +import logging from typing import Optional from sqlalchemy import update @@ -7,6 +8,8 @@ from core.errors.error import QuotaExceededError from extensions.ext_database import db from models import TenantCreditPool +logger = logging.getLogger(__name__) + class CreditPoolService: @classmethod @@ -72,8 +75,9 @@ class CreditPoolService: cls, tenant_id: str, credits_required: int, - ) -> bool: + ): """check and deduct credits""" + logger.info("check and deduct credits") pool = cls.get_pool(tenant_id) if not pool: raise QuotaExceededError("Credit pool not found") @@ -93,8 +97,6 @@ class CreditPoolService: stmt = update(TenantCreditPool).where(*where_conditions).values(**update_values) db.session.execute(stmt) - return True - @classmethod def check_deduct_credits(cls, tenant_id: str, credits_required: int) -> bool: """check and deduct credits""" From ab34cea714d32b49cb8278ed59c8e6e6e930f4cf Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 12:49:26 +0800 Subject: [PATCH 03/13] add paid credit --- .../feature/hosted_service/__init__.py | 40 ++++++++++ api/core/hosting_configuration.py | 18 ++++- api/core/provider_manager.py | 28 +++++-- api/core/workflow/nodes/llm/llm_utils.py | 51 +++++++----- .../update_provider_when_message_created.py | 12 ++- ...1520-58a70d22fdbd_add_table_credit_pool.py | 78 ++++++++++-------- api/services/credit_pool_service.py | 79 +++++-------------- api/services/workspace_service.py | 12 ++- 8 files changed, 192 insertions(+), 126 deletions(-) diff --git a/api/configs/feature/hosted_service/__init__.py b/api/configs/feature/hosted_service/__init__.py index 538c55d931..6415bd239d 100644 --- a/api/configs/feature/hosted_service/__init__.py +++ b/api/configs/feature/hosted_service/__init__.py @@ -128,6 +128,16 @@ class HostedGeminiConfig(BaseSettings): default="gemini-2.5-flash,gemini-2.0-flash,gemini-2.0-flash-lite,", ) + HOSTED_GEMINI_PAID_ENABLED: bool = Field( + description="Enable paid access to hosted gemini service", + default=False, + ) + + HOSTED_GEMINI_PAID_MODELS: str = Field( + description="Comma-separated list of available models for paid access", + default="gemini-2.5-flash,gemini-2.0-flash,gemini-2.0-flash-lite,", + ) + class HostedXAIConfig(BaseSettings): """ @@ -159,6 +169,16 @@ class HostedXAIConfig(BaseSettings): default="grok-3,grok-3-mini,grok-3-mini-fast", ) + HOSTED_XAI_PAID_ENABLED: bool = Field( + description="Enable paid access to hosted XAI service", + default=False, + ) + + HOSTED_XAI_PAID_MODELS: str = Field( + description="Comma-separated list of available models for paid access", + default="grok-3,grok-3-mini,grok-3-mini-fast", + ) + class HostedDeepseekConfig(BaseSettings): """ @@ -190,6 +210,16 @@ class HostedDeepseekConfig(BaseSettings): default="deepseek-chat,deepseek-reasoner", ) + HOSTED_DEEPSEEK_PAID_ENABLED: bool = Field( + description="Enable paid access to hosted XAI service", + default=False, + ) + + HOSTED_DEEPSEEK_PAID_MODELS: str = Field( + description="Comma-separated list of available models for paid access", + default="grok-3,grok-3-mini,grok-3-mini-fast", + ) + class HostedAzureOpenAiConfig(BaseSettings): """ @@ -252,6 +282,16 @@ class HostedAnthropicConfig(BaseSettings): "claude-3-7-sonnet-20250219," "claude-3-haiku-20240307", ) + HOSTED_ANTHROPIC_PAID_MODELS: str = Field( + description="Comma-separated list of available models for paid access", + default="claude-opus-4-20250514," + "claude-opus-4-20250514," + "claude-sonnet-4-20250514," + "claude-3-5-haiku-20241022," + "claude-3-opus-20240229," + "claude-3-7-sonnet-20250219," + "claude-3-haiku-20240307", + ) class HostedMinmaxConfig(BaseSettings): diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py index 7aafa4bc80..ed08ecf57b 100644 --- a/api/core/hosting_configuration.py +++ b/api/core/hosting_configuration.py @@ -169,6 +169,11 @@ class HostingConfiguration: trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models) quotas.append(trial_quota) + if dify_config.HOSTED_GEMINI_PAID_ENABLED: + paid_models = self.parse_restrict_models_from_env("HOSTED_GEMINI_PAID_MODELS") + paid_quota = PaidHostingQuota(restrict_models=paid_models) + quotas.append(paid_quota) + if len(quotas) > 0: credentials = { "google_api_key": dify_config.HOSTED_GEMINI_API_KEY, @@ -196,7 +201,8 @@ class HostingConfiguration: if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED: paid_quota = PaidHostingQuota() - quotas.append(paid_quota) + paid_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_PAID_MODELS") + quotas.append(paid_quota,restrict_models=paid_models) if len(quotas) > 0: credentials = { @@ -223,6 +229,11 @@ class HostingConfiguration: trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models) quotas.append(trial_quota) + if dify_config.HOSTED_XAI_PAID_ENABLED: + paid_models = self.parse_restrict_models_from_env("HOSTED_XAI_PAID_MODELS") + paid_quota = PaidHostingQuota(restrict_models=paid_models) + quotas.append(paid_quota) + if len(quotas) > 0: credentials = { "api_key": dify_config.HOSTED_XAI_API_KEY, @@ -248,6 +259,11 @@ class HostingConfiguration: trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trail_models) quotas.append(trial_quota) + if dify_config.HOSTED_DEEPSEEK_PAID_ENABLED: + paid_models = self.parse_restrict_models_from_env("HOSTED_DEEPSEEK_PAID_MODELS") + paid_quota = PaidHostingQuota(restrict_models=paid_models) + quotas.append(paid_quota) + if len(quotas) > 0: credentials = { "api_key": dify_config.HOSTED_DEEPSEEK_API_KEY, diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 1ac02d9b6a..2772048d26 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -916,11 +916,17 @@ class ProviderManager: if dify_config.EDITION == "CLOUD": from services.credit_pool_service import CreditPoolService - pool = CreditPoolService.get_or_create_pool( + trail_pool = CreditPoolService.get_pool( tenant_id=tenant_id, + pool_type="trial", + ) + paid_pool = CreditPoolService.get_pool( + tenant_id=tenant_id, + pool_type="paid", ) else: - pool = None + trail_pool = None + paid_pool = None for provider_quota in provider_hosting_configuration.quotas: if provider_quota.quota_type not in quota_type_to_provider_records_dict: @@ -942,13 +948,23 @@ class ProviderManager: raise ValueError("quota_used is None") if provider_record.quota_limit is None: raise ValueError("quota_limit is None") - if provider_quota.quota_type == ProviderQuotaType.TRIAL and pool is not None: + if provider_quota.quota_type == ProviderQuotaType.TRIAL and trail_pool is not None: quota_configuration = QuotaConfiguration( quota_type=provider_quota.quota_type, quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, - quota_used=pool.quota_used, - quota_limit=pool.quota_limit, - is_valid=pool.quota_limit > pool.quota_used or pool.quota_limit == -1, + quota_used=trail_pool.quota_used, + quota_limit=trail_pool.quota_limit, + is_valid=trail_pool.quota_limit > trail_pool.quota_used or trail_pool.quota_limit == -1, + restrict_models=provider_quota.restrict_models, + ) + + elif provider_quota.quota_type == ProviderQuotaType.PAID and paid_pool is not None: + quota_configuration = QuotaConfiguration( + quota_type=provider_quota.quota_type, + quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, + quota_used=paid_pool.quota_used, + quota_limit=paid_pool.quota_limit, + is_valid=paid_pool.quota_limit > paid_pool.quota_used or paid_pool.quota_limit == -1, restrict_models=provider_quota.restrict_models, ) diff --git a/api/core/workflow/nodes/llm/llm_utils.py b/api/core/workflow/nodes/llm/llm_utils.py index ad969cdad1..054fbe033d 100644 --- a/api/core/workflow/nodes/llm/llm_utils.py +++ b/api/core/workflow/nodes/llm/llm_utils.py @@ -23,7 +23,7 @@ from libs.datetime_utils import naive_utc_now from models.model import Conversation from models.provider import Provider, ProviderType from models.provider_ids import ModelProviderID - +from core.entities.provider_entities import ProviderQuotaType from .exc import InvalidVariableTypeError, LLMModeRequiredError, ModelNotExistError @@ -136,21 +136,36 @@ def deduct_llm_quota(tenant_id: str, model_instance: ModelInstance, usage: LLMUs used_quota = 1 if used_quota is not None and system_configuration.current_quota_type is not None: - with Session(db.engine) as session: - stmt = ( - update(Provider) - .where( - Provider.tenant_id == tenant_id, - # TODO: Use provider name with prefix after the data migration. - Provider.provider_name == ModelProviderID(model_instance.provider).provider_name, - Provider.provider_type == ProviderType.SYSTEM.value, - Provider.quota_type == system_configuration.current_quota_type.value, - Provider.quota_limit > Provider.quota_used, - ) - .values( - quota_used=Provider.quota_used + used_quota, - last_used=naive_utc_now(), - ) + + if system_configuration.current_quota_type == ProviderQuotaType.TRIAL: + from services.credit_pool_service import CreditPoolService + CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, + credits_required=used_quota, ) - session.execute(stmt) - session.commit() + elif system_configuration.current_quota_type == ProviderQuotaType.PAID: + from services.credit_pool_service import CreditPoolService + CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, + credits_required=used_quota, + pool_type="paid", + ) + else: + with Session(db.engine) as session: + stmt = ( + update(Provider) + .where( + Provider.tenant_id == tenant_id, + # TODO: Use provider name with prefix after the data migration. + Provider.provider_name == ModelProviderID(model_instance.provider).provider_name, + Provider.provider_type == ProviderType.SYSTEM.value, + Provider.quota_type == system_configuration.current_quota_type.value, + Provider.quota_limit > Provider.quota_used, + ) + .values( + quota_used=Provider.quota_used + used_quota, + last_used=naive_utc_now(), + ) + ) + session.execute(stmt) + session.commit() diff --git a/api/events/event_handlers/update_provider_when_message_created.py b/api/events/event_handlers/update_provider_when_message_created.py index 9efe9a79af..12e0961bcc 100644 --- a/api/events/event_handlers/update_provider_when_message_created.py +++ b/api/events/event_handlers/update_provider_when_message_created.py @@ -133,18 +133,24 @@ def handle(sender: Message, **kwargs): system_configuration=system_configuration, model_name=model_config.model, ) - logger.info("used_quota: %s", used_quota) if used_quota is not None: if provider_configuration.system_configuration.current_quota_type == ProviderQuotaType.TRIAL: - logger.info("deduct credits") from services.credit_pool_service import CreditPoolService CreditPoolService.check_and_deduct_credits( tenant_id=tenant_id, credits_required=used_quota, + pool_type="trial", + ) + elif provider_configuration.system_configuration.current_quota_type == ProviderQuotaType.PAID: + from services.credit_pool_service import CreditPoolService + + CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, + credits_required=used_quota, + pool_type="paid", ) else: - logger.info("update provider quota") quota_update = _ProviderUpdateOperation( filters=_ProviderUpdateFilters( tenant_id=tenant_id, diff --git a/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py b/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py index d298d885f4..b050008fc2 100644 --- a/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py +++ b/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py @@ -32,8 +32,8 @@ def upgrade(): with op.batch_alter_table('tenant_credit_pools', schema=None) as batch_op: batch_op.create_index('tenant_credit_pool_pool_type_idx', ['pool_type'], unique=False) batch_op.create_index('tenant_credit_pool_tenant_id_idx', ['tenant_id'], unique=False) - # Data migration: Move trial quota data from providers to tenant_credit_pools - migrate_trial_quota_data() + # Data migration: Move quota data from providers to tenant_credit_pools + migrate_quota_data() # ### end Alembic commands ### @@ -48,49 +48,57 @@ def downgrade(): # ### end Alembic commands ### -def migrate_trial_quota_data(): +def migrate_quota_data(): """ Migrate quota data from providers table to tenant_credit_pools table - for providers with quota_type='trial', provider_name='openai', provider_type='system' + for providers with quota_type='trial' or 'paid', provider_name='openai', provider_type='system' """ # Create connection bind = op.get_bind() - # Query providers that match the criteria - select_sql = sa.text(""" - SELECT tenant_id, quota_limit, quota_used - FROM providers - WHERE quota_type = 'trial' - AND provider_name = 'openai' - AND provider_type = 'system' - AND quota_limit IS NOT NULL - """) + # Define quota type mappings + quota_type_mappings = ['trial', 'paid'] - result = bind.execute(select_sql) - providers_data = result.fetchall() - - # Insert data into tenant_credit_pools - for provider_data in providers_data: - tenant_id, quota_limit, quota_used = provider_data - - # Check if credit pool already exists for this tenant - check_sql = sa.text(""" - SELECT COUNT(*) - FROM tenant_credit_pools - WHERE tenant_id = :tenant_id AND pool_type = 'trial' + for quota_type in quota_type_mappings: + # Query providers that match the criteria + select_sql = sa.text(""" + SELECT tenant_id, quota_limit, quota_used + FROM providers + WHERE quota_type = :quota_type + AND provider_name = 'openai' + AND provider_type = 'system' + AND quota_limit IS NOT NULL """) - existing_count = bind.execute(check_sql, {"tenant_id": tenant_id}).scalar() + result = bind.execute(select_sql, {"quota_type": quota_type}) + providers_data = result.fetchall() - if existing_count == 0: - # Insert new credit pool record - insert_sql = sa.text(""" - INSERT INTO tenant_credit_pools (tenant_id, pool_type, quota_limit, quota_used, created_at, updated_at) - VALUES (:tenant_id, 'trial', :quota_limit, :quota_used, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + # Insert data into tenant_credit_pools + for provider_data in providers_data: + tenant_id, quota_limit, quota_used = provider_data + + # Check if credit pool already exists for this tenant and pool type + check_sql = sa.text(""" + SELECT COUNT(*) + FROM tenant_credit_pools + WHERE tenant_id = :tenant_id AND pool_type = :pool_type """) - bind.execute(insert_sql, { + existing_count = bind.execute(check_sql, { "tenant_id": tenant_id, - "quota_limit": quota_limit or 0, - "quota_used": quota_used or 0 - }) + "pool_type": quota_type + }).scalar() + + if existing_count == 0: + # Insert new credit pool record + insert_sql = sa.text(""" + INSERT INTO tenant_credit_pools (tenant_id, pool_type, quota_limit, quota_used, created_at, updated_at) + VALUES (:tenant_id, :pool_type, :quota_limit, :quota_used, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + """) + + bind.execute(insert_sql, { + "tenant_id": tenant_id, + "pool_type": quota_type, + "quota_limit": quota_limit or 0, + "quota_used": quota_used or 0 + }) diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index f72686b9ff..fc2d875bd0 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -7,6 +7,7 @@ from configs import dify_config from core.errors.error import QuotaExceededError from extensions.ext_database import db from models import TenantCreditPool +from sqlalchemy.orm import Session logger = logging.getLogger(__name__) @@ -23,62 +24,27 @@ class CreditPoolService: return credit_pool @classmethod - def get_pool(cls, tenant_id: str) -> Optional[TenantCreditPool]: + def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> Optional[TenantCreditPool]: """get tenant credit pool""" return ( db.session.query(TenantCreditPool) .filter_by( tenant_id=tenant_id, + pool_type=pool_type, ) .first() ) - @classmethod - def get_or_create_pool(cls, tenant_id: str) -> TenantCreditPool: - """get or create credit pool""" - # First try to get existing pool - pool = cls.get_pool(tenant_id) - if pool: - return pool - - # Create new pool if not exists, handle race condition - try: - # Double-check in case another thread created it - pool = ( - db.session.query(TenantCreditPool) - .filter_by( - tenant_id=tenant_id, - ) - .first() - ) - if pool: - return pool - - # Create new pool - pool = TenantCreditPool( - tenant_id=tenant_id, quota_limit=dify_config.HOSTED_POOL_CREDITS, quota_used=0, pool_type="trial" - ) - db.session.add(pool) - db.session.commit() - - except Exception: - # If creation fails (e.g., due to race condition), rollback and try to get existing one - db.session.rollback() - pool = cls.get_pool(tenant_id) - if not pool: - raise - - return pool - @classmethod def check_and_deduct_credits( cls, tenant_id: str, credits_required: int, + pool_type: str = "trial", ): """check and deduct credits""" - logger.info("check and deduct credits") - pool = cls.get_pool(tenant_id) + + pool = cls.get_pool(tenant_id, pool_type) if not pool: raise QuotaExceededError("Credit pool not found") @@ -86,24 +52,17 @@ class CreditPoolService: raise QuotaExceededError( f"Insufficient credits. Required: {credits_required}, Available: {pool.remaining_credits}" ) + try: + with Session(db.engine) as session: + update_values = {"quota_used": pool.quota_used + credits_required} - with db.session.begin(): - update_values = {"quota_used": pool.quota_used + credits_required} - - where_conditions = [ - TenantCreditPool.tenant_id == tenant_id, - TenantCreditPool.quota_used + credits_required <= TenantCreditPool.quota_limit, - ] - stmt = update(TenantCreditPool).where(*where_conditions).values(**update_values) - db.session.execute(stmt) - - @classmethod - def check_deduct_credits(cls, tenant_id: str, credits_required: int) -> bool: - """check and deduct credits""" - pool = cls.get_pool(tenant_id) - if not pool: - return False - - if pool.remaining_credits < credits_required: - return False - return True + where_conditions = [ + TenantCreditPool.pool_type == pool_type, + TenantCreditPool.tenant_id == tenant_id, + TenantCreditPool.quota_used + credits_required <= TenantCreditPool.quota_limit, + ] + stmt = update(TenantCreditPool).where(*where_conditions).values(**update_values) + session.execute(stmt) + session.commit() + except Exception: + raise QuotaExceededError("Failed to deduct credits") diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index a21aac1984..c71a19636d 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -49,8 +49,14 @@ class WorkspaceService: if dify_config.EDITION == "CLOUD": from services.credit_pool_service import CreditPoolService - pool = CreditPoolService.get_or_create_pool(tenant_id=tenant.id) - tenant_info["trial_credits"] = pool.quota_limit - tenant_info["trial_credits_used"] = pool.quota_used + paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid") + if paid_pool: + tenant_info["trial_credits"] = paid_pool.quota_limit + tenant_info["trial_credits_used"] = paid_pool.quota_used + else: + trial_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="trial") + if trial_pool: + tenant_info["trial_credits"] = trial_pool.quota_limit + tenant_info["trial_credits_used"] = trial_pool.quota_used return tenant_info From c3e3a18ab465411e21bb54fe8bc8be56541976e1 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 12:49:35 +0800 Subject: [PATCH 04/13] add paid credit --- api/core/hosting_configuration.py | 2 +- api/core/provider_manager.py | 4 ++-- api/core/workflow/nodes/llm/llm_utils.py | 4 ++-- api/services/credit_pool_service.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py index ed08ecf57b..92babd5056 100644 --- a/api/core/hosting_configuration.py +++ b/api/core/hosting_configuration.py @@ -202,7 +202,7 @@ class HostingConfiguration: if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED: paid_quota = PaidHostingQuota() paid_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_PAID_MODELS") - quotas.append(paid_quota,restrict_models=paid_models) + quotas.append(paid_quota, restrict_models=paid_models) if len(quotas) > 0: credentials = { diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 2772048d26..a2c9dde63c 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -948,7 +948,7 @@ class ProviderManager: raise ValueError("quota_used is None") if provider_record.quota_limit is None: raise ValueError("quota_limit is None") - if provider_quota.quota_type == ProviderQuotaType.TRIAL and trail_pool is not None: + if provider_quota.quota_type == ProviderQuotaType.TRIAL and trail_pool is not None: quota_configuration = QuotaConfiguration( quota_type=provider_quota.quota_type, quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, @@ -958,7 +958,7 @@ class ProviderManager: restrict_models=provider_quota.restrict_models, ) - elif provider_quota.quota_type == ProviderQuotaType.PAID and paid_pool is not None: + elif provider_quota.quota_type == ProviderQuotaType.PAID and paid_pool is not None: quota_configuration = QuotaConfiguration( quota_type=provider_quota.quota_type, quota_unit=provider_hosting_configuration.quota_unit or QuotaUnit.TOKENS, diff --git a/api/core/workflow/nodes/llm/llm_utils.py b/api/core/workflow/nodes/llm/llm_utils.py index 054fbe033d..194ad43151 100644 --- a/api/core/workflow/nodes/llm/llm_utils.py +++ b/api/core/workflow/nodes/llm/llm_utils.py @@ -6,7 +6,7 @@ from sqlalchemy.orm import Session from configs import dify_config from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity -from core.entities.provider_entities import QuotaUnit +from core.entities.provider_entities import ProviderQuotaType, QuotaUnit from core.file.models import File from core.memory.token_buffer_memory import TokenBufferMemory from core.model_manager import ModelInstance, ModelManager @@ -23,7 +23,7 @@ from libs.datetime_utils import naive_utc_now from models.model import Conversation from models.provider import Provider, ProviderType from models.provider_ids import ModelProviderID -from core.entities.provider_entities import ProviderQuotaType + from .exc import InvalidVariableTypeError, LLMModeRequiredError, ModelNotExistError diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index fc2d875bd0..108ee05e45 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -2,12 +2,12 @@ import logging from typing import Optional from sqlalchemy import update +from sqlalchemy.orm import Session from configs import dify_config from core.errors.error import QuotaExceededError from extensions.ext_database import db from models import TenantCreditPool -from sqlalchemy.orm import Session logger = logging.getLogger(__name__) From da27d261b09830f969fe52bc80e28f7dc0af20d5 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 13:11:14 +0800 Subject: [PATCH 05/13] fix: add paid quota error for init_anthropic --- api/core/hosting_configuration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py index 92babd5056..29642d8c7a 100644 --- a/api/core/hosting_configuration.py +++ b/api/core/hosting_configuration.py @@ -200,9 +200,9 @@ class HostingConfiguration: quotas.append(trial_quota) if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED: - paid_quota = PaidHostingQuota() paid_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_PAID_MODELS") - quotas.append(paid_quota, restrict_models=paid_models) + paid_quota = PaidHostingQuota(restrict_models=paid_models) + quotas.append(paid_quota) if len(quotas) > 0: credentials = { From 560fe8a0f6b67c9397f6c9898454fdb46fa97a99 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 13:33:32 +0800 Subject: [PATCH 06/13] fix: format --- api/core/provider_manager.py | 2 +- api/core/workflow/nodes/llm/llm_utils.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index a2c9dde63c..489af29460 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -957,7 +957,7 @@ class ProviderManager: is_valid=trail_pool.quota_limit > trail_pool.quota_used or trail_pool.quota_limit == -1, restrict_models=provider_quota.restrict_models, ) - + elif provider_quota.quota_type == ProviderQuotaType.PAID and paid_pool is not None: quota_configuration = QuotaConfiguration( quota_type=provider_quota.quota_type, diff --git a/api/core/workflow/nodes/llm/llm_utils.py b/api/core/workflow/nodes/llm/llm_utils.py index 194ad43151..0af4024d3e 100644 --- a/api/core/workflow/nodes/llm/llm_utils.py +++ b/api/core/workflow/nodes/llm/llm_utils.py @@ -136,15 +136,16 @@ def deduct_llm_quota(tenant_id: str, model_instance: ModelInstance, usage: LLMUs used_quota = 1 if used_quota is not None and system_configuration.current_quota_type is not None: - if system_configuration.current_quota_type == ProviderQuotaType.TRIAL: from services.credit_pool_service import CreditPoolService + CreditPoolService.check_and_deduct_credits( tenant_id=tenant_id, credits_required=used_quota, ) elif system_configuration.current_quota_type == ProviderQuotaType.PAID: from services.credit_pool_service import CreditPoolService + CreditPoolService.check_and_deduct_credits( tenant_id=tenant_id, credits_required=used_quota, From 0360f0b33bd5f04abd63f6fd189485dcaad9402d Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 26 Sep 2025 14:31:26 +0800 Subject: [PATCH 07/13] fix: create paid provider auto --- api/core/provider_manager.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 489af29460..3ecd646862 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -618,9 +618,9 @@ class ProviderManager: ) for quota in configuration.quotas: - if quota.quota_type == ProviderQuotaType.TRIAL: + if quota.quota_type in (ProviderQuotaType.TRIAL, ProviderQuotaType.PAID): # Init trial provider records if not exists - if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict: + if quota.quota_type not in provider_quota_to_provider_record_dict: try: # FIXME ignore the type error, only TrialHostingQuota has limit need to change the logic new_provider_record = Provider( @@ -628,7 +628,7 @@ class ProviderManager: # TODO: Use provider name with prefix after the data migration. provider_name=ModelProviderID(provider_name).provider_name, provider_type=ProviderType.SYSTEM.value, - quota_type=ProviderQuotaType.TRIAL.value, + quota_type=quota.quota_type, quota_limit=0, # type: ignore quota_used=0, is_valid=True, @@ -642,7 +642,7 @@ class ProviderManager: Provider.tenant_id == tenant_id, Provider.provider_name == ModelProviderID(provider_name).provider_name, Provider.provider_type == ProviderType.SYSTEM.value, - Provider.quota_type == ProviderQuotaType.TRIAL.value, + Provider.quota_type == quota.quota_type, ) existed_provider_record = db.session.scalar(stmt) if not existed_provider_record: @@ -652,7 +652,7 @@ class ProviderManager: existed_provider_record.is_valid = True db.session.commit() - provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) + provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) return provider_name_to_provider_records_dict @@ -918,11 +918,11 @@ class ProviderManager: trail_pool = CreditPoolService.get_pool( tenant_id=tenant_id, - pool_type="trial", + pool_type=ProviderQuotaType.TRIAL.value, ) paid_pool = CreditPoolService.get_pool( tenant_id=tenant_id, - pool_type="paid", + pool_type=ProviderQuotaType.PAID.value, ) else: trail_pool = None From e1819fb7e574532c2dcd6a5ce7c572f8acb1e794 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 9 Oct 2025 02:50:01 +0000 Subject: [PATCH 08/13] [autofix.ci] apply automated fixes --- api/core/provider_manager.py | 4 ++-- api/services/credit_pool_service.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 3ecd646862..522dc6c372 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -620,7 +620,7 @@ class ProviderManager: for quota in configuration.quotas: if quota.quota_type in (ProviderQuotaType.TRIAL, ProviderQuotaType.PAID): # Init trial provider records if not exists - if quota.quota_type not in provider_quota_to_provider_record_dict: + if quota.quota_type not in provider_quota_to_provider_record_dict: try: # FIXME ignore the type error, only TrialHostingQuota has limit need to change the logic new_provider_record = Provider( @@ -652,7 +652,7 @@ class ProviderManager: existed_provider_record.is_valid = True db.session.commit() - provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) + provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) return provider_name_to_provider_records_dict diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index 108ee05e45..5f14735d1f 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -24,7 +24,7 @@ class CreditPoolService: return credit_pool @classmethod - def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> Optional[TenantCreditPool]: + def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | None: """get tenant credit pool""" return ( db.session.query(TenantCreditPool) From e056e0835aec0670ea073c804bdb580f3ac52d4e Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 9 Oct 2025 02:51:51 +0000 Subject: [PATCH 09/13] [autofix.ci] apply automated fixes (attempt 2/3) --- api/services/credit_pool_service.py | 1 - 1 file changed, 1 deletion(-) diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index 5f14735d1f..8ae409809a 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -1,5 +1,4 @@ import logging -from typing import Optional from sqlalchemy import update from sqlalchemy.orm import Session From 5b813970548cb66d3c3a69752c31a54cf9dadd0e Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Thu, 9 Oct 2025 11:08:24 +0800 Subject: [PATCH 10/13] fix test case --- api/tests/unit_tests/services/test_account_service.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index 737202f8de..aec8efd880 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -619,8 +619,13 @@ class TestTenantService: mock_tenant_instance.name = "Test User's Workspace" mock_tenant_class.return_value = mock_tenant_instance - # Execute test - TenantService.create_owner_tenant_if_not_exist(mock_account) + # Mock the db import in CreditPoolService to avoid database connection + with patch("services.credit_pool_service.db") as mock_credit_pool_db: + mock_credit_pool_db.session.add = MagicMock() + mock_credit_pool_db.session.commit = MagicMock() + + # Execute test + TenantService.create_owner_tenant_if_not_exist(mock_account) # Verify tenant was created with correct parameters mock_db_dependencies["db"].session.add.assert_called() From f71ad55d58bba2a7ef40268846239e3f2539f9c3 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Thu, 9 Oct 2025 11:08:52 +0800 Subject: [PATCH 11/13] fix test case --- api/core/provider_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 3ecd646862..522dc6c372 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -620,7 +620,7 @@ class ProviderManager: for quota in configuration.quotas: if quota.quota_type in (ProviderQuotaType.TRIAL, ProviderQuotaType.PAID): # Init trial provider records if not exists - if quota.quota_type not in provider_quota_to_provider_record_dict: + if quota.quota_type not in provider_quota_to_provider_record_dict: try: # FIXME ignore the type error, only TrialHostingQuota has limit need to change the logic new_provider_record = Provider( @@ -652,7 +652,7 @@ class ProviderManager: existed_provider_record.is_valid = True db.session.commit() - provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) + provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) return provider_name_to_provider_records_dict From c0ed353c105201161d27bdbfdc8ed646d2d14ed2 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Sat, 11 Oct 2025 11:45:18 +0800 Subject: [PATCH 12/13] add credit next_credit_reset_date --- api/services/feature_service.py | 4 ++++ api/services/workspace_service.py | 7 ++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/api/services/feature_service.py b/api/services/feature_service.py index 19d96cb972..b2b7df181a 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -131,6 +131,7 @@ class FeatureModel(BaseModel): # pydantic configs model_config = ConfigDict(protected_namespaces=()) knowledge_pipeline: KnowledgePipeline = KnowledgePipeline() + next_credit_reset_date: int = 0 class KnowledgeRateLimitModel(BaseModel): @@ -278,6 +279,9 @@ class FeatureService: if "knowledge_pipeline_publish_enabled" in billing_info: features.knowledge_pipeline.publish_enabled = billing_info["knowledge_pipeline_publish_enabled"] + + if "next_credit_reset_date" in billing_info: + features.next_credit_reset_date = billing_info["next_credit_reset_date"] @classmethod def _fulfill_params_from_enterprise(cls, features: SystemFeatureModel): diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index c71a19636d..f33489c7bc 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -1,3 +1,4 @@ +from dify.api.configs import feature from flask_login import current_user from configs import dify_config @@ -31,7 +32,8 @@ class WorkspaceService: assert tenant_account_join is not None, "TenantAccountJoin not found" tenant_info["role"] = tenant_account_join.role - can_replace_logo = FeatureService.get_features(tenant.id).can_replace_logo + feature = FeatureService.get_features(tenant.id) + can_replace_logo = feature.can_replace_logo if can_replace_logo and TenantService.has_roles(tenant, [TenantAccountRole.OWNER, TenantAccountRole.ADMIN]): base_url = dify_config.FILES_URL @@ -47,6 +49,9 @@ class WorkspaceService: "replace_webapp_logo": replace_webapp_logo, } if dify_config.EDITION == "CLOUD": + + tenant_info["next_credit_reset_date"] = feature.next_credit_reset_date + from services.credit_pool_service import CreditPoolService paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid") From 1d16528dff5588257cf8b8bedcd0e81348b3d23e Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Sat, 11 Oct 2025 11:45:24 +0800 Subject: [PATCH 13/13] add credit next_credit_reset_date --- api/services/workspace_service.py | 1 - 1 file changed, 1 deletion(-) diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index f33489c7bc..cf2c6cea1c 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -1,4 +1,3 @@ -from dify.api.configs import feature from flask_login import current_user from configs import dify_config