Underground-Digital/Workflow-Engine

1from flask_restful import Resource, reqparse2from werkzeug.exceptions import Forbidden3 4from controllers.console import api5from controllers.console.wraps import account_initialization_required, setup_required6from core.model_runtime.entities.model_entities import ModelType7from core.model_runtime.errors.validate import CredentialsValidateFailedError8from libs.login import current_user, login_required9from models.account import TenantAccountRole10from services.model_load_balancing_service import ModelLoadBalancingService11 12 13class LoadBalancingCredentialsValidateApi(Resource):14    @setup_required15    @login_required16    @account_initialization_required17    def post(self, provider: str):18        if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role):19            raise Forbidden()20 21        tenant_id = current_user.current_tenant_id22 23        parser = reqparse.RequestParser()24        parser.add_argument("model", type=str, required=True, nullable=False, location="json")25        parser.add_argument(26            "model_type",27            type=str,28            required=True,29            nullable=False,30            choices=[mt.value for mt in ModelType],31            location="json",32        )33        parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json")34        args = parser.parse_args()35 36        # validate model load balancing credentials37        model_load_balancing_service = ModelLoadBalancingService()38 39        result = True40        error = None41 42        try:43            model_load_balancing_service.validate_load_balancing_credentials(44                tenant_id=tenant_id,45                provider=provider,46                model=args["model"],47                model_type=args["model_type"],48                credentials=args["credentials"],49            )50        except CredentialsValidateFailedError as ex:51            result = False52            error = str(ex)53 54        response = {"result": "success" if result else "error"}55 56        if not result:57            response["error"] = error58 59        return response60 61 62class LoadBalancingConfigCredentialsValidateApi(Resource):63    @setup_required64    @login_required65    @account_initialization_required66    def post(self, provider: str, config_id: str):67        if not TenantAccountRole.is_privileged_role(current_user.current_tenant.current_role):68            raise Forbidden()69 70        tenant_id = current_user.current_tenant_id71 72        parser = reqparse.RequestParser()73        parser.add_argument("model", type=str, required=True, nullable=False, location="json")74        parser.add_argument(75            "model_type",76            type=str,77            required=True,78            nullable=False,79            choices=[mt.value for mt in ModelType],80            location="json",81        )82        parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json")83        args = parser.parse_args()84 85        # validate model load balancing config credentials86        model_load_balancing_service = ModelLoadBalancingService()87 88        result = True89        error = None90 91        try:92            model_load_balancing_service.validate_load_balancing_credentials(93                tenant_id=tenant_id,94                provider=provider,95                model=args["model"],96                model_type=args["model_type"],97                credentials=args["credentials"],98                config_id=config_id,99            )100        except CredentialsValidateFailedError as ex:101            result = False102            error = str(ex)103 104        response = {"result": "success" if result else "error"}105 106        if not result:107            response["error"] = error108 109        return response110 111 112# Load Balancing Config113api.add_resource(114    LoadBalancingCredentialsValidateApi,115    "/workspaces/current/model-providers/<string:provider>/models/load-balancing-configs/credentials-validate",116)117 118api.add_resource(119    LoadBalancingConfigCredentialsValidateApi,120    "/workspaces/current/model-providers/<string:provider>/models/load-balancing-configs/<string:config_id>/credentials-validate",121)122