Underground-Digital/Workflow-Engine

1import os2 3from flask_login import current_user4from flask_restful import Resource, reqparse5 6from controllers.console import api7from controllers.console.app.error import (8    CompletionRequestError,9    ProviderModelCurrentlyNotSupportError,10    ProviderNotInitializeError,11    ProviderQuotaExceededError,12)13from controllers.console.wraps import account_initialization_required, setup_required14from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError15from core.llm_generator.llm_generator import LLMGenerator16from core.model_runtime.errors.invoke import InvokeError17from libs.login import login_required18 19 20class RuleGenerateApi(Resource):21    @setup_required22    @login_required23    @account_initialization_required24    def post(self):25        parser = reqparse.RequestParser()26        parser.add_argument("instruction", type=str, required=True, nullable=False, location="json")27        parser.add_argument("model_config", type=dict, required=True, nullable=False, location="json")28        parser.add_argument("no_variable", type=bool, required=True, default=False, location="json")29        args = parser.parse_args()30 31        account = current_user32        PROMPT_GENERATION_MAX_TOKENS = int(os.getenv("PROMPT_GENERATION_MAX_TOKENS", "512"))33 34        try:35            rules = LLMGenerator.generate_rule_config(36                tenant_id=account.current_tenant_id,37                instruction=args["instruction"],38                model_config=args["model_config"],39                no_variable=args["no_variable"],40                rule_config_max_tokens=PROMPT_GENERATION_MAX_TOKENS,41            )42        except ProviderTokenNotInitError as ex:43            raise ProviderNotInitializeError(ex.description)44        except QuotaExceededError:45            raise ProviderQuotaExceededError()46        except ModelCurrentlyNotSupportError:47            raise ProviderModelCurrentlyNotSupportError()48        except InvokeError as e:49            raise CompletionRequestError(e.description)50 51        return rules52 53 54class RuleCodeGenerateApi(Resource):55    @setup_required56    @login_required57    @account_initialization_required58    def post(self):59        parser = reqparse.RequestParser()60        parser.add_argument("instruction", type=str, required=True, nullable=False, location="json")61        parser.add_argument("model_config", type=dict, required=True, nullable=False, location="json")62        parser.add_argument("no_variable", type=bool, required=True, default=False, location="json")63        parser.add_argument("code_language", type=str, required=False, default="javascript", location="json")64        args = parser.parse_args()65 66        account = current_user67        CODE_GENERATION_MAX_TOKENS = int(os.getenv("CODE_GENERATION_MAX_TOKENS", "1024"))68        try:69            code_result = LLMGenerator.generate_code(70                tenant_id=account.current_tenant_id,71                instruction=args["instruction"],72                model_config=args["model_config"],73                code_language=args["code_language"],74                max_tokens=CODE_GENERATION_MAX_TOKENS,75            )76        except ProviderTokenNotInitError as ex:77            raise ProviderNotInitializeError(ex.description)78        except QuotaExceededError:79            raise ProviderQuotaExceededError()80        except ModelCurrentlyNotSupportError:81            raise ProviderModelCurrentlyNotSupportError()82        except InvokeError as e:83            raise CompletionRequestError(e.description)84 85        return code_result86 87 88api.add_resource(RuleGenerateApi, "/rule-generate")89api.add_resource(RuleCodeGenerateApi, "/rule-code-generate")90