Underground-Digital/Workflow-Engine
0
1from typing import Optional2 3from flask import Flask4from pydantic import BaseModel5 6from configs import dify_config7from core.entities.provider_entities import QuotaUnit, RestrictModel8from core.model_runtime.entities.model_entities import ModelType9from models.provider import ProviderQuotaType10 11 12class HostingQuota(BaseModel):13 quota_type: ProviderQuotaType14 restrict_models: list[RestrictModel] = []15 16 17class TrialHostingQuota(HostingQuota):18 quota_type: ProviderQuotaType = ProviderQuotaType.TRIAL19 quota_limit: int = 020 """Quota limit for the hosting provider models. -1 means unlimited."""21 22 23class PaidHostingQuota(HostingQuota):24 quota_type: ProviderQuotaType = ProviderQuotaType.PAID25 26 27class FreeHostingQuota(HostingQuota):28 quota_type: ProviderQuotaType = ProviderQuotaType.FREE29 30 31class HostingProvider(BaseModel):32 enabled: bool = False33 credentials: Optional[dict] = None34 quota_unit: Optional[QuotaUnit] = None35 quotas: list[HostingQuota] = []36 37 38class HostedModerationConfig(BaseModel):39 enabled: bool = False40 providers: list[str] = []41 42 43class HostingConfiguration:44 provider_map: dict[str, HostingProvider] = {}45 moderation_config: HostedModerationConfig = None46 47 def init_app(self, app: Flask) -> None:48 if dify_config.EDITION != "CLOUD":49 return50 51 self.provider_map["azure_openai"] = self.init_azure_openai()52 self.provider_map["openai"] = self.init_openai()53 self.provider_map["anthropic"] = self.init_anthropic()54 self.provider_map["minimax"] = self.init_minimax()55 self.provider_map["spark"] = self.init_spark()56 self.provider_map["zhipuai"] = self.init_zhipuai()57 58 self.moderation_config = self.init_moderation_config()59 60 @staticmethod61 def init_azure_openai() -> HostingProvider:62 quota_unit = QuotaUnit.TIMES63 if dify_config.HOSTED_AZURE_OPENAI_ENABLED:64 credentials = {65 "openai_api_key": dify_config.HOSTED_AZURE_OPENAI_API_KEY,66 "openai_api_base": dify_config.HOSTED_AZURE_OPENAI_API_BASE,67 "base_model_name": "gpt-35-turbo",68 }69 70 quotas = []71 hosted_quota_limit = dify_config.HOSTED_AZURE_OPENAI_QUOTA_LIMIT72 trial_quota = TrialHostingQuota(73 quota_limit=hosted_quota_limit,74 restrict_models=[75 RestrictModel(model="gpt-4", base_model_name="gpt-4", model_type=ModelType.LLM),76 RestrictModel(model="gpt-4o", base_model_name="gpt-4o", model_type=ModelType.LLM),77 RestrictModel(model="gpt-4o-mini", base_model_name="gpt-4o-mini", model_type=ModelType.LLM),78 RestrictModel(model="gpt-4-32k", base_model_name="gpt-4-32k", model_type=ModelType.LLM),79 RestrictModel(80 model="gpt-4-1106-preview", base_model_name="gpt-4-1106-preview", model_type=ModelType.LLM81 ),82 RestrictModel(83 model="gpt-4-vision-preview", base_model_name="gpt-4-vision-preview", model_type=ModelType.LLM84 ),85 RestrictModel(model="gpt-35-turbo", base_model_name="gpt-35-turbo", model_type=ModelType.LLM),86 RestrictModel(87 model="gpt-35-turbo-1106", base_model_name="gpt-35-turbo-1106", model_type=ModelType.LLM88 ),89 RestrictModel(90 model="gpt-35-turbo-instruct", base_model_name="gpt-35-turbo-instruct", model_type=ModelType.LLM91 ),92 RestrictModel(93 model="gpt-35-turbo-16k", base_model_name="gpt-35-turbo-16k", model_type=ModelType.LLM94 ),95 RestrictModel(96 model="text-davinci-003", base_model_name="text-davinci-003", model_type=ModelType.LLM97 ),98 RestrictModel(99 model="text-embedding-ada-002",100 base_model_name="text-embedding-ada-002",101 model_type=ModelType.TEXT_EMBEDDING,102 ),103 RestrictModel(104 model="text-embedding-3-small",105 base_model_name="text-embedding-3-small",106 model_type=ModelType.TEXT_EMBEDDING,107 ),108 RestrictModel(109 model="text-embedding-3-large",110 base_model_name="text-embedding-3-large",111 model_type=ModelType.TEXT_EMBEDDING,112 ),113 ],114 )115 quotas.append(trial_quota)116 117 return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)118 119 return HostingProvider(120 enabled=False,121 quota_unit=quota_unit,122 )123 124 def init_openai(self) -> HostingProvider:125 quota_unit = QuotaUnit.CREDITS126 quotas = []127 128 if dify_config.HOSTED_OPENAI_TRIAL_ENABLED:129 hosted_quota_limit = dify_config.HOSTED_OPENAI_QUOTA_LIMIT130 trial_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_TRIAL_MODELS")131 trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit, restrict_models=trial_models)132 quotas.append(trial_quota)133 134 if dify_config.HOSTED_OPENAI_PAID_ENABLED:135 paid_models = self.parse_restrict_models_from_env("HOSTED_OPENAI_PAID_MODELS")136 paid_quota = PaidHostingQuota(restrict_models=paid_models)137 quotas.append(paid_quota)138 139 if len(quotas) > 0:140 credentials = {141 "openai_api_key": dify_config.HOSTED_OPENAI_API_KEY,142 }143 144 if dify_config.HOSTED_OPENAI_API_BASE:145 credentials["openai_api_base"] = dify_config.HOSTED_OPENAI_API_BASE146 147 if dify_config.HOSTED_OPENAI_API_ORGANIZATION:148 credentials["openai_organization"] = dify_config.HOSTED_OPENAI_API_ORGANIZATION149 150 return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)151 152 return HostingProvider(153 enabled=False,154 quota_unit=quota_unit,155 )156 157 @staticmethod158 def init_anthropic() -> HostingProvider:159 quota_unit = QuotaUnit.TOKENS160 quotas = []161 162 if dify_config.HOSTED_ANTHROPIC_TRIAL_ENABLED:163 hosted_quota_limit = dify_config.HOSTED_ANTHROPIC_QUOTA_LIMIT164 trial_quota = TrialHostingQuota(quota_limit=hosted_quota_limit)165 quotas.append(trial_quota)166 167 if dify_config.HOSTED_ANTHROPIC_PAID_ENABLED:168 paid_quota = PaidHostingQuota()169 quotas.append(paid_quota)170 171 if len(quotas) > 0:172 credentials = {173 "anthropic_api_key": dify_config.HOSTED_ANTHROPIC_API_KEY,174 }175 176 if dify_config.HOSTED_ANTHROPIC_API_BASE:177 credentials["anthropic_api_url"] = dify_config.HOSTED_ANTHROPIC_API_BASE178 179 return HostingProvider(enabled=True, credentials=credentials, quota_unit=quota_unit, quotas=quotas)180 181 return HostingProvider(182 enabled=False,183 quota_unit=quota_unit,184 )185 186 @staticmethod187 def init_minimax() -> HostingProvider:188 quota_unit = QuotaUnit.TOKENS189 if dify_config.HOSTED_MINIMAX_ENABLED:190 quotas = [FreeHostingQuota()]191 192 return HostingProvider(193 enabled=True,194 credentials=None, # use credentials from the provider195 quota_unit=quota_unit,196 quotas=quotas,197 )198 199 return HostingProvider(200 enabled=False,201 quota_unit=quota_unit,202 )203 204 @staticmethod205 def init_spark() -> HostingProvider:206 quota_unit = QuotaUnit.TOKENS207 if dify_config.HOSTED_SPARK_ENABLED:208 quotas = [FreeHostingQuota()]209 210 return HostingProvider(211 enabled=True,212 credentials=None, # use credentials from the provider213 quota_unit=quota_unit,214 quotas=quotas,215 )216 217 return HostingProvider(218 enabled=False,219 quota_unit=quota_unit,220 )221 222 @staticmethod223 def init_zhipuai() -> HostingProvider:224 quota_unit = QuotaUnit.TOKENS225 if dify_config.HOSTED_ZHIPUAI_ENABLED:226 quotas = [FreeHostingQuota()]227 228 return HostingProvider(229 enabled=True,230 credentials=None, # use credentials from the provider231 quota_unit=quota_unit,232 quotas=quotas,233 )234 235 return HostingProvider(236 enabled=False,237 quota_unit=quota_unit,238 )239 240 @staticmethod241 def init_moderation_config() -> HostedModerationConfig:242 if dify_config.HOSTED_MODERATION_ENABLED and dify_config.HOSTED_MODERATION_PROVIDERS:243 return HostedModerationConfig(enabled=True, providers=dify_config.HOSTED_MODERATION_PROVIDERS.split(","))244 245 return HostedModerationConfig(enabled=False)246 247 @staticmethod248 def parse_restrict_models_from_env(env_var: str) -> list[RestrictModel]:249 models_str = dify_config.model_dump().get(env_var)250 models_list = models_str.split(",") if models_str else []251 return [252 RestrictModel(model=model_name.strip(), model_type=ModelType.LLM)253 for model_name in models_list254 if model_name.strip()255 ]256