Underground-Digital/Workflow-Engine

1import io2 3from flask import send_file4from flask_login import current_user5from flask_restful import Resource, reqparse6from werkzeug.exceptions import Forbidden7 8from controllers.console import api9from controllers.console.wraps import account_initialization_required, setup_required10from core.model_runtime.entities.model_entities import ModelType11from core.model_runtime.errors.validate import CredentialsValidateFailedError12from core.model_runtime.utils.encoders import jsonable_encoder13from libs.login import login_required14from services.billing_service import BillingService15from services.model_provider_service import ModelProviderService16 17 18class ModelProviderListApi(Resource):19    @setup_required20    @login_required21    @account_initialization_required22    def get(self):23        tenant_id = current_user.current_tenant_id24 25        parser = reqparse.RequestParser()26        parser.add_argument(27            "model_type",28            type=str,29            required=False,30            nullable=True,31            choices=[mt.value for mt in ModelType],32            location="args",33        )34        args = parser.parse_args()35 36        model_provider_service = ModelProviderService()37        provider_list = model_provider_service.get_provider_list(tenant_id=tenant_id, model_type=args.get("model_type"))38 39        return jsonable_encoder({"data": provider_list})40 41 42class ModelProviderCredentialApi(Resource):43    @setup_required44    @login_required45    @account_initialization_required46    def get(self, provider: str):47        tenant_id = current_user.current_tenant_id48 49        model_provider_service = ModelProviderService()50        credentials = model_provider_service.get_provider_credentials(tenant_id=tenant_id, provider=provider)51 52        return {"credentials": credentials}53 54 55class ModelProviderValidateApi(Resource):56    @setup_required57    @login_required58    @account_initialization_required59    def post(self, provider: str):60        parser = reqparse.RequestParser()61        parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json")62        args = parser.parse_args()63 64        tenant_id = current_user.current_tenant_id65 66        model_provider_service = ModelProviderService()67 68        result = True69        error = None70 71        try:72            model_provider_service.provider_credentials_validate(73                tenant_id=tenant_id, provider=provider, credentials=args["credentials"]74            )75        except CredentialsValidateFailedError as ex:76            result = False77            error = str(ex)78 79        response = {"result": "success" if result else "error"}80 81        if not result:82            response["error"] = error83 84        return response85 86 87class ModelProviderApi(Resource):88    @setup_required89    @login_required90    @account_initialization_required91    def post(self, provider: str):92        if not current_user.is_admin_or_owner:93            raise Forbidden()94 95        parser = reqparse.RequestParser()96        parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json")97        args = parser.parse_args()98 99        model_provider_service = ModelProviderService()100 101        try:102            model_provider_service.save_provider_credentials(103                tenant_id=current_user.current_tenant_id, provider=provider, credentials=args["credentials"]104            )105        except CredentialsValidateFailedError as ex:106            raise ValueError(str(ex))107 108        return {"result": "success"}, 201109 110    @setup_required111    @login_required112    @account_initialization_required113    def delete(self, provider: str):114        if not current_user.is_admin_or_owner:115            raise Forbidden()116 117        model_provider_service = ModelProviderService()118        model_provider_service.remove_provider_credentials(tenant_id=current_user.current_tenant_id, provider=provider)119 120        return {"result": "success"}, 204121 122 123class ModelProviderIconApi(Resource):124    """125    Get model provider icon126    """127 128    def get(self, provider: str, icon_type: str, lang: str):129        model_provider_service = ModelProviderService()130        icon, mimetype = model_provider_service.get_model_provider_icon(131            provider=provider,132            icon_type=icon_type,133            lang=lang,134        )135 136        return send_file(io.BytesIO(icon), mimetype=mimetype)137 138 139class PreferredProviderTypeUpdateApi(Resource):140    @setup_required141    @login_required142    @account_initialization_required143    def post(self, provider: str):144        if not current_user.is_admin_or_owner:145            raise Forbidden()146 147        tenant_id = current_user.current_tenant_id148 149        parser = reqparse.RequestParser()150        parser.add_argument(151            "preferred_provider_type",152            type=str,153            required=True,154            nullable=False,155            choices=["system", "custom"],156            location="json",157        )158        args = parser.parse_args()159 160        model_provider_service = ModelProviderService()161        model_provider_service.switch_preferred_provider(162            tenant_id=tenant_id, provider=provider, preferred_provider_type=args["preferred_provider_type"]163        )164 165        return {"result": "success"}166 167 168class ModelProviderPaymentCheckoutUrlApi(Resource):169    @setup_required170    @login_required171    @account_initialization_required172    def get(self, provider: str):173        if provider != "anthropic":174            raise ValueError(f"provider name {provider} is invalid")175        BillingService.is_tenant_owner_or_admin(current_user)176        data = BillingService.get_model_provider_payment_link(177            provider_name=provider,178            tenant_id=current_user.current_tenant_id,179            account_id=current_user.id,180            prefilled_email=current_user.email,181        )182        return data183 184 185class ModelProviderFreeQuotaSubmitApi(Resource):186    @setup_required187    @login_required188    @account_initialization_required189    def post(self, provider: str):190        model_provider_service = ModelProviderService()191        result = model_provider_service.free_quota_submit(tenant_id=current_user.current_tenant_id, provider=provider)192 193        return result194 195 196class ModelProviderFreeQuotaQualificationVerifyApi(Resource):197    @setup_required198    @login_required199    @account_initialization_required200    def get(self, provider: str):201        parser = reqparse.RequestParser()202        parser.add_argument("token", type=str, required=False, nullable=True, location="args")203        args = parser.parse_args()204 205        model_provider_service = ModelProviderService()206        result = model_provider_service.free_quota_qualification_verify(207            tenant_id=current_user.current_tenant_id, provider=provider, token=args["token"]208        )209 210        return result211 212 213api.add_resource(ModelProviderListApi, "/workspaces/current/model-providers")214 215api.add_resource(ModelProviderCredentialApi, "/workspaces/current/model-providers/<string:provider>/credentials")216api.add_resource(ModelProviderValidateApi, "/workspaces/current/model-providers/<string:provider>/credentials/validate")217api.add_resource(ModelProviderApi, "/workspaces/current/model-providers/<string:provider>")218api.add_resource(219    ModelProviderIconApi, "/workspaces/current/model-providers/<string:provider>/<string:icon_type>/<string:lang>"220)221 222api.add_resource(223    PreferredProviderTypeUpdateApi, "/workspaces/current/model-providers/<string:provider>/preferred-provider-type"224)225api.add_resource(226    ModelProviderPaymentCheckoutUrlApi, "/workspaces/current/model-providers/<string:provider>/checkout-url"227)228api.add_resource(229    ModelProviderFreeQuotaSubmitApi, "/workspaces/current/model-providers/<string:provider>/free-quota-submit"230)231api.add_resource(232    ModelProviderFreeQuotaQualificationVerifyApi,233    "/workspaces/current/model-providers/<string:provider>/free-quota-qualification-verify",234)235