Underground-Digital/Workflow-Engine

1import logging2 3from flask_restful import reqparse4from werkzeug.exceptions import InternalServerError5 6from controllers.console import api7from controllers.console.app.error import (8    CompletionRequestError,9    ProviderModelCurrentlyNotSupportError,10    ProviderNotInitializeError,11    ProviderQuotaExceededError,12)13from controllers.console.explore.error import NotWorkflowAppError14from controllers.console.explore.wraps import InstalledAppResource15from core.app.apps.base_app_queue_manager import AppQueueManager16from core.app.entities.app_invoke_entities import InvokeFrom17from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError18from core.model_runtime.errors.invoke import InvokeError19from libs import helper20from libs.login import current_user21from models.model import AppMode, InstalledApp22from services.app_generate_service import AppGenerateService23 24logger = logging.getLogger(__name__)25 26 27class InstalledAppWorkflowRunApi(InstalledAppResource):28    def post(self, installed_app: InstalledApp):29        """30        Run workflow31        """32        app_model = installed_app.app33        app_mode = AppMode.value_of(app_model.mode)34        if app_mode != AppMode.WORKFLOW:35            raise NotWorkflowAppError()36 37        parser = reqparse.RequestParser()38        parser.add_argument("inputs", type=dict, required=True, nullable=False, location="json")39        parser.add_argument("files", type=list, required=False, location="json")40        args = parser.parse_args()41 42        try:43            response = AppGenerateService.generate(44                app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=True45            )46 47            return helper.compact_generate_response(response)48        except ProviderTokenNotInitError as ex:49            raise ProviderNotInitializeError(ex.description)50        except QuotaExceededError:51            raise ProviderQuotaExceededError()52        except ModelCurrentlyNotSupportError:53            raise ProviderModelCurrentlyNotSupportError()54        except InvokeError as e:55            raise CompletionRequestError(e.description)56        except ValueError as e:57            raise e58        except Exception as e:59            logging.exception("internal server error.")60            raise InternalServerError()61 62 63class InstalledAppWorkflowTaskStopApi(InstalledAppResource):64    def post(self, installed_app: InstalledApp, task_id: str):65        """66        Stop workflow task67        """68        app_model = installed_app.app69        app_mode = AppMode.value_of(app_model.mode)70        if app_mode != AppMode.WORKFLOW:71            raise NotWorkflowAppError()72 73        AppQueueManager.set_stop_flag(task_id, InvokeFrom.EXPLORE, current_user.id)74 75        return {"result": "success"}76 77 78api.add_resource(InstalledAppWorkflowRunApi, "/installed-apps/<uuid:installed_app_id>/workflows/run")79api.add_resource(80    InstalledAppWorkflowTaskStopApi, "/installed-apps/<uuid:installed_app_id>/workflows/tasks/<string:task_id>/stop"81)82