Underground-Digital/Workflow-Engine

1import logging2from datetime import datetime, timezone3 4from flask_login import current_user5from flask_restful import reqparse6from werkzeug.exceptions import InternalServerError, NotFound7 8import services9from controllers.console import api10from controllers.console.app.error import (11    AppUnavailableError,12    CompletionRequestError,13    ConversationCompletedError,14    ProviderModelCurrentlyNotSupportError,15    ProviderNotInitializeError,16    ProviderQuotaExceededError,17)18from controllers.console.explore.error import NotChatAppError, NotCompletionAppError19from controllers.console.explore.wraps import InstalledAppResource20from core.app.apps.base_app_queue_manager import AppQueueManager21from core.app.entities.app_invoke_entities import InvokeFrom22from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError23from core.model_runtime.errors.invoke import InvokeError24from extensions.ext_database import db25from libs import helper26from libs.helper import uuid_value27from models.model import AppMode28from services.app_generate_service import AppGenerateService29 30 31# define completion api for user32class CompletionApi(InstalledAppResource):33    def post(self, installed_app):34        app_model = installed_app.app35        if app_model.mode != "completion":36            raise NotCompletionAppError()37 38        parser = reqparse.RequestParser()39        parser.add_argument("inputs", type=dict, required=True, location="json")40        parser.add_argument("query", type=str, location="json", default="")41        parser.add_argument("files", type=list, required=False, location="json")42        parser.add_argument("response_mode", type=str, choices=["blocking", "streaming"], location="json")43        parser.add_argument("retriever_from", type=str, required=False, default="explore_app", location="json")44        args = parser.parse_args()45 46        streaming = args["response_mode"] == "streaming"47        args["auto_generate_name"] = False48 49        installed_app.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None)50        db.session.commit()51 52        try:53            response = AppGenerateService.generate(54                app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=streaming55            )56 57            return helper.compact_generate_response(response)58        except services.errors.conversation.ConversationNotExistsError:59            raise NotFound("Conversation Not Exists.")60        except services.errors.conversation.ConversationCompletedError:61            raise ConversationCompletedError()62        except services.errors.app_model_config.AppModelConfigBrokenError:63            logging.exception("App model config broken.")64            raise AppUnavailableError()65        except ProviderTokenNotInitError as ex:66            raise ProviderNotInitializeError(ex.description)67        except QuotaExceededError:68            raise ProviderQuotaExceededError()69        except ModelCurrentlyNotSupportError:70            raise ProviderModelCurrentlyNotSupportError()71        except InvokeError as e:72            raise CompletionRequestError(e.description)73        except ValueError as e:74            raise e75        except Exception as e:76            logging.exception("internal server error.")77            raise InternalServerError()78 79 80class CompletionStopApi(InstalledAppResource):81    def post(self, installed_app, task_id):82        app_model = installed_app.app83        if app_model.mode != "completion":84            raise NotCompletionAppError()85 86        AppQueueManager.set_stop_flag(task_id, InvokeFrom.EXPLORE, current_user.id)87 88        return {"result": "success"}, 20089 90 91class ChatApi(InstalledAppResource):92    def post(self, installed_app):93        app_model = installed_app.app94        app_mode = AppMode.value_of(app_model.mode)95        if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}:96            raise NotChatAppError()97 98        parser = reqparse.RequestParser()99        parser.add_argument("inputs", type=dict, required=True, location="json")100        parser.add_argument("query", type=str, required=True, location="json")101        parser.add_argument("files", type=list, required=False, location="json")102        parser.add_argument("conversation_id", type=uuid_value, location="json")103        parser.add_argument("parent_message_id", type=uuid_value, required=False, location="json")104        parser.add_argument("retriever_from", type=str, required=False, default="explore_app", location="json")105        args = parser.parse_args()106 107        args["auto_generate_name"] = False108 109        installed_app.last_used_at = datetime.now(timezone.utc).replace(tzinfo=None)110        db.session.commit()111 112        try:113            response = AppGenerateService.generate(114                app_model=app_model, user=current_user, args=args, invoke_from=InvokeFrom.EXPLORE, streaming=True115            )116 117            return helper.compact_generate_response(response)118        except services.errors.conversation.ConversationNotExistsError:119            raise NotFound("Conversation Not Exists.")120        except services.errors.conversation.ConversationCompletedError:121            raise ConversationCompletedError()122        except services.errors.app_model_config.AppModelConfigBrokenError:123            logging.exception("App model config broken.")124            raise AppUnavailableError()125        except ProviderTokenNotInitError as ex:126            raise ProviderNotInitializeError(ex.description)127        except QuotaExceededError:128            raise ProviderQuotaExceededError()129        except ModelCurrentlyNotSupportError:130            raise ProviderModelCurrentlyNotSupportError()131        except InvokeError as e:132            raise CompletionRequestError(e.description)133        except ValueError as e:134            raise e135        except Exception as e:136            logging.exception("internal server error.")137            raise InternalServerError()138 139 140class ChatStopApi(InstalledAppResource):141    def post(self, installed_app, task_id):142        app_model = installed_app.app143        app_mode = AppMode.value_of(app_model.mode)144        if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}:145            raise NotChatAppError()146 147        AppQueueManager.set_stop_flag(task_id, InvokeFrom.EXPLORE, current_user.id)148 149        return {"result": "success"}, 200150 151 152api.add_resource(153    CompletionApi, "/installed-apps/<uuid:installed_app_id>/completion-messages", endpoint="installed_app_completion"154)155api.add_resource(156    CompletionStopApi,157    "/installed-apps/<uuid:installed_app_id>/completion-messages/<string:task_id>/stop",158    endpoint="installed_app_stop_completion",159)160api.add_resource(161    ChatApi, "/installed-apps/<uuid:installed_app_id>/chat-messages", endpoint="installed_app_chat_completion"162)163api.add_resource(164    ChatStopApi,165    "/installed-apps/<uuid:installed_app_id>/chat-messages/<string:task_id>/stop",166    endpoint="installed_app_stop_chat_completion",167)168