import contextvars import logging import threading import uuid from collections.abc import Generator, Mapping from typing import Any, Literal, overload from flask import Flask, current_app from pydantic import ValidationError from configs import dify_config from constants import UUID_NIL from core.app.app_config.easy_ui_based_app.model_config.converter import ModelConfigConverter from core.app.app_config.features.file_upload.manager import FileUploadConfigManager from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfigManager from core.app.apps.agent_chat.app_runner import AgentChatAppRunner from core.app.apps.agent_chat.generate_response_converter import AgentChatAppGenerateResponseConverter from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.message_based_app_generator import MessageBasedAppGenerator from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, InvokeFrom from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_database import db from factories import file_factory from graphon.model_runtime.errors.invoke import InvokeAuthorizationError from libs.flask_utils import preserve_flask_contexts from models import Account, App, EndUser from services.conversation_service import ConversationService logger = logging.getLogger(__name__) class AgentChatAppGenerator(MessageBasedAppGenerator): @overload def generate( self, *, app_model: App, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, streaming: Literal[False], ) -> Mapping[str, Any]: ... @overload def generate( self, *, app_model: App, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, streaming: Literal[True], ) -> Generator[Mapping | str, None, None]: ... @overload def generate( self, *, app_model: App, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, streaming: bool, ) -> Mapping | Generator[Mapping | str, None, None]: ... def generate( self, *, app_model: App, user: Account | EndUser, args: Mapping[str, Any], invoke_from: InvokeFrom, streaming: bool = True, ) -> Mapping | Generator[Mapping | str, None, None]: """ Generate App response. :param app_model: App :param user: account or end user :param args: request args :param invoke_from: invoke from source :param streaming: is stream """ if not streaming: raise ValueError("Agent Chat App does not support blocking mode") if not args.get("query"): raise ValueError("query is required") query = args["query"] if not isinstance(query, str): raise ValueError("query must be a string") query = query.replace("\x00", "") inputs = args["inputs"] extras = {"auto_generate_conversation_name": args.get("auto_generate_name", True)} # get conversation conversation = None conversation_id = args.get("conversation_id") if conversation_id: conversation = ConversationService.get_conversation( app_model=app_model, conversation_id=conversation_id, user=user ) # get app model config app_model_config = self._get_app_model_config(app_model=app_model, conversation=conversation) # validate override model config override_model_config_dict = None if args.get("model_config"): if invoke_from != InvokeFrom.DEBUGGER: raise ValueError("Only in App debug mode can override model config") # validate config override_model_config_dict = AgentChatAppConfigManager.config_validate( tenant_id=app_model.tenant_id, config=args["model_config"], ) # always enable retriever resource in debugger mode override_model_config_dict["retriever_resource"] = {"enabled": True} # parse files # TODO(QuantumGhost): Move file parsing logic to the API controller layer # for better separation of concerns. # # For implementation reference, see the `_parse_file` function and # `DraftWorkflowNodeRunApi` class which handle this properly. with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from): files = args.get("files") or [] file_extra_config = FileUploadConfigManager.convert( override_model_config_dict or app_model_config.to_dict() ) if file_extra_config: file_objs = file_factory.build_from_mappings( mappings=files, tenant_id=app_model.tenant_id, config=file_extra_config, access_controller=self._file_access_controller, ) else: file_objs = [] # convert to app config app_config = AgentChatAppConfigManager.get_app_config( app_model=app_model, app_model_config=app_model_config, conversation=conversation, override_config_dict=override_model_config_dict, ) # get tracing instance trace_manager = TraceQueueManager(app_model.id, user.id if isinstance(user, Account) else user.session_id) # init application generate entity application_generate_entity = AgentChatAppGenerateEntity( task_id=str(uuid.uuid4()), app_config=app_config, model_conf=ModelConfigConverter.convert(app_config), file_upload_config=file_extra_config, conversation_id=conversation.id if conversation else None, inputs=self._prepare_user_inputs( user_inputs=inputs, variables=app_config.variables, tenant_id=app_model.tenant_id ), query=query, files=list(file_objs), parent_message_id=( args.get("parent_message_id") if invoke_from not in {InvokeFrom.SERVICE_API, InvokeFrom.OPENAPI} else UUID_NIL ), user_id=user.id, stream=streaming, invoke_from=invoke_from, extras=extras, call_depth=0, trace_manager=trace_manager, ) # init generate records (conversation, message) = self._init_generate_records(application_generate_entity, conversation) # init queue manager queue_manager = MessageBasedAppQueueManager( task_id=application_generate_entity.task_id, user_id=application_generate_entity.user_id, invoke_from=application_generate_entity.invoke_from, conversation_id=conversation.id, app_mode=conversation.mode, message_id=message.id, ) # new thread with request context and contextvars context = contextvars.copy_context() worker_thread = threading.Thread( target=self._generate_worker, kwargs={ "flask_app": current_app._get_current_object(), # type: ignore "context": context, "application_generate_entity": application_generate_entity, "queue_manager": queue_manager, "conversation_id": conversation.id, "message_id": message.id, }, ) worker_thread.start() # return response or stream generator response = self._handle_response( application_generate_entity=application_generate_entity, queue_manager=queue_manager, conversation=conversation, message=message, user=user, stream=streaming, ) return AgentChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) def _generate_worker( self, flask_app: Flask, context: contextvars.Context, application_generate_entity: AgentChatAppGenerateEntity, queue_manager: AppQueueManager, conversation_id: str, message_id: str, ): """ Generate worker in a new thread. :param flask_app: Flask app :param application_generate_entity: application generate entity :param queue_manager: queue manager :param conversation_id: conversation ID :param message_id: message ID :return: """ with preserve_flask_contexts(flask_app, context_vars=context): try: # get conversation and message conversation = self._get_conversation(conversation_id) message = self._get_message(message_id) # chatbot app runner = AgentChatAppRunner() runner.run( application_generate_entity=application_generate_entity, queue_manager=queue_manager, conversation=conversation, message=message, ) except GenerateTaskStoppedError: pass except InvokeAuthorizationError: queue_manager.publish_error( InvokeAuthorizationError("Incorrect API key provided"), PublishFrom.APPLICATION_MANAGER ) except ValidationError as e: logger.exception("Validation Error when generating") queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER) except ValueError as e: if dify_config.DEBUG: logger.exception("Error when generating") queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER) except Exception as e: logger.exception("Unknown Error when generating") queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER) finally: db.session.close()