mirror of https://github.com/langgenius/dify.git
fix migration and add tongyi
This commit is contained in:
Yansong Zhang
parent
ba0a59f998
commit
5c51d61049
|
|
@ -351,6 +351,42 @@ class HostedAnthropicConfig(BaseSettings):
|
|||
)
|
||||
|
||||
|
||||
class HostedTongyiConfig(BaseSettings):
|
||||
"""
|
||||
Configuration for hosted Tongyi service
|
||||
"""
|
||||
|
||||
HOSTED_TONGYI_API_KEY: str | None = Field(
|
||||
description="API key for hosted Tongyi service",
|
||||
default=None,
|
||||
)
|
||||
|
||||
HOSTED_TONGYI_USE_INTERNATIONAL_ENDPOINT: bool = Field(
|
||||
description="Use international endpoint for hosted Tongyi service",
|
||||
default=False,
|
||||
)
|
||||
|
||||
HOSTED_TONGYI_TRIAL_ENABLED: bool = Field(
|
||||
description="Enable trial access to hosted Anthropic service",
|
||||
default=False,
|
||||
)
|
||||
|
||||
HOSTED_TONGYI_PAID_ENABLED: bool = Field(
|
||||
description="Enable paid access to hosted Anthropic service",
|
||||
default=False,
|
||||
)
|
||||
|
||||
HOSTED_TONGYI_TRIAL_MODELS: str = Field(
|
||||
description="Comma-separated list of available models for trial access",
|
||||
default="",
|
||||
)
|
||||
|
||||
HOSTED_TONGYI_PAID_MODELS: str = Field(
|
||||
description="Comma-separated list of available models for paid access",
|
||||
default="",
|
||||
)
|
||||
|
||||
|
||||
class HostedMinmaxConfig(BaseSettings):
|
||||
"""
|
||||
Configuration for hosted Minmax service
|
||||
|
|
@ -442,6 +478,7 @@ class HostedServiceConfig(
|
|||
HostedOpenAiConfig,
|
||||
HostedSparkConfig,
|
||||
HostedZhipuAIConfig,
|
||||
HostedTongyiConfig,
|
||||
# moderation
|
||||
HostedModerationConfig,
|
||||
# credit config
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ class HostingConfiguration:
|
|||
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.provider_map[f"{DEFAULT_PLUGIN_ID}/tongyi/tongyi"] = self.init_tongyi()
|
||||
|
||||
self.moderation_config = self.init_moderation_config()
|
||||
|
||||
|
|
@ -219,6 +220,34 @@ class HostingConfiguration:
|
|||
quota_unit=quota_unit,
|
||||
)
|
||||
|
||||
def init_tongyi(self) -> HostingProvider:
|
||||
quota_unit = QuotaUnit.CREDITS
|
||||
quotas: list[HostingQuota] = []
|
||||
|
||||
if dify_config.HOSTED_TONGYI_TRIAL_ENABLED:
|
||||
hosted_quota_limit = 0
|
||||
trail_models = self.parse_restrict_models_from_env("HOSTED_TONGYI_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:
|
||||
paid_models = self.parse_restrict_models_from_env("HOSTED_ANTHROPIC_PAID_MODELS")
|
||||
paid_quota = PaidHostingQuota(restrict_models=paid_models)
|
||||
quotas.append(paid_quota)
|
||||
|
||||
if len(quotas) > 0:
|
||||
credentials = {
|
||||
"dashscope_api_key": dify_config.HOSTED_TONGYI_API_KEY,
|
||||
"use_international_endpoint": dify_config.HOSTED_TONGYI_USE_INTERNATIONAL_ENDPOINT,
|
||||
}
|
||||
|
||||
return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)
|
||||
|
||||
return HostingProvider(
|
||||
enabled=False,
|
||||
quota_unit=quota_unit,
|
||||
)
|
||||
|
||||
def init_xai(self) -> HostingProvider:
|
||||
quota_unit = QuotaUnit.CREDITS
|
||||
quotas: list[HostingQuota] = []
|
||||
|
|
|
|||
|
|
@ -1,104 +0,0 @@
|
|||
"""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 quota data from providers to tenant_credit_pools
|
||||
migrate_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_quota_data():
|
||||
"""
|
||||
Migrate quota data from providers table to tenant_credit_pools table
|
||||
for providers with quota_type='trial' or 'paid', provider_name='openai', provider_type='system'
|
||||
"""
|
||||
# Create connection
|
||||
bind = op.get_bind()
|
||||
|
||||
# Define quota type mappings
|
||||
quota_type_mappings = ['trial', 'paid']
|
||||
|
||||
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
|
||||
""")
|
||||
|
||||
result = bind.execute(select_sql, {"quota_type": quota_type})
|
||||
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 and pool type
|
||||
check_sql = sa.text("""
|
||||
SELECT COUNT(*)
|
||||
FROM tenant_credit_pools
|
||||
WHERE tenant_id = :tenant_id AND pool_type = :pool_type
|
||||
""")
|
||||
|
||||
existing_count = bind.execute(check_sql, {
|
||||
"tenant_id": tenant_id,
|
||||
"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
|
||||
})
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
"""add credit pool
|
||||
|
||||
Revision ID: 7df29de0f6be
|
||||
Revises: 03ea244985ce
|
||||
Create Date: 2025-12-25 10:39:15.139304
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import models as models
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '7df29de0f6be'
|
||||
down_revision = '03ea244985ce'
|
||||
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), server_default='trial', 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)
|
||||
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade():
|
||||
|
||||
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 ###
|
||||
Loading…
Reference in New Issue