diff --git a/.github/workflows/knowledge-fs-ci.yml b/.github/workflows/knowledge-fs-ci.yml index b4fc2ce288fe34..2239a0a2e092f9 100644 --- a/.github/workflows/knowledge-fs-ci.yml +++ b/.github/workflows/knowledge-fs-ci.yml @@ -62,13 +62,13 @@ jobs: - 'api/controllers/console/workspace/rbac/**' - 'api/controllers/inner_api/__init__.py' - 'api/controllers/service_api/__init__.py' - - 'api/core/agent/base_agent_runner.py' - - 'api/core/app/apps/agent_app/runtime_request_builder.py' + - 'api/services/agent/chat/base_runner.py' + - 'api/services/app/generation/adapters/agent_request_builder.py' - 'api/core/rbac/entities.py' - 'api/core/tools/__base/tool_runtime.py' - 'api/core/tools/builtin_tool/_position.yaml' - - 'api/core/workflow/node_runtime.py' - - 'api/core/workflow/nodes/agent_v2/runtime_request_builder.py' + - 'api/services/workflow/execution/adapters/node_runtime.py' + - 'api/services/workflow/execution/adapters/agent_v2/runtime_request_builder.py' - 'api/extensions/ext_celery.py' - 'api/extensions/ext_commands.py' - 'api/models/__init__.py' @@ -76,7 +76,8 @@ jobs: - 'api/models/model.py' - 'api/repositories/workspace/workspace_repository.py' - 'api/services/workspace/member_service.py' - - 'api/services/agent_tool_inner_service.py' + - 'api/services/agent/tool_invocation_service.py' + - 'api/services/tools/agent_invocation_gateway.py' - 'api/services/enterprise/rbac_service.py' - 'api/services/entities/agent_tool_inner.py' - 'api/services/vector_space_admission_service.py' @@ -211,12 +212,12 @@ jobs: api/controllers/console/workspace/rbac/schemas.py api/controllers/inner_api/__init__.py api/controllers/service_api/__init__.py - api/core/agent/base_agent_runner.py - api/core/app/apps/agent_app/runtime_request_builder.py + api/services/agent/chat/base_runner.py + api/services/app/generation/adapters/agent_request_builder.py api/core/rbac/entities.py api/core/tools/__base/tool_runtime.py - api/core/workflow/node_runtime.py - api/core/workflow/nodes/agent_v2/runtime_request_builder.py + api/services/workflow/execution/adapters/node_runtime.py + api/services/workflow/execution/adapters/agent_v2/runtime_request_builder.py api/extensions/ext_celery.py api/extensions/ext_commands.py api/models/__init__.py @@ -224,7 +225,8 @@ jobs: api/models/model.py api/repositories/workspace/workspace_repository.py api/services/workspace/member_service.py - api/services/agent_tool_inner_service.py + api/services/agent/tool_invocation_service.py + api/services/tools/agent_invocation_gateway.py api/services/enterprise/rbac_service.py api/services/entities/agent_tool_inner.py api/services/vector_space_admission_service.py diff --git a/api/.importlinter b/api/.importlinter index c158d3f3d49180..55a332cdb36f8b 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -32,40 +32,9 @@ layers = # error makes this list shrink whenever a legacy dependency is removed. unmatched_ignore_imports_alerting = error ignore_imports = - # The legacy base generator delegates file defaults to the application owner until cleanup. - core.app.apps.base_app_generator -> services.app.file_input_defaults core.app.app_config.easy_ui_based_app.dataset.manager -> services.knowledge.dataset_service - # Thin legacy generator modules re-export canonical service adapters during path migration. - core.app.apps.advanced_chat.app_generator -> services.workflow.execution.adapters.chatflow.app_generator - core.app.apps.advanced_chat.app_runner -> services.workflow.execution.adapters.chatflow.app_runner - core.app.apps.agent_app.app_generator -> services.agent.composer_service - core.app.apps.agent_app.app_generator -> services.agent.workspace_service - core.app.apps.agent_app.app_generator -> services.conversation_service - core.app.apps.agent_app.app_generator -> services.workflow.execution.adapters.agent_v2.dify_tools_builder - core.app.apps.agent_app.runtime_request_builder -> services.agent.prompt_mentions - core.app.apps.agent_app.runtime_request_builder -> services.workflow.execution.adapters.agent_v2.dify_tools_builder - core.app.apps.agent_app.runtime_request_builder -> services.workflow.execution.adapters.agent_v2.runtime_request_builder - # KnowledgeFS agent glue lives in services; agent runners only project its config and citations. - core.app.apps.agent_app.app_runner -> services.agent.knowledge_citations - core.app.apps.agent_app.runtime_request_builder -> services.agent.knowledge_runtime_config core.app.apps.agent_app.session_store -> services.agent.workspace_service - core.app.apps.agent_chat.app_generator -> services.conversation_service - core.app.apps.base_app_generator -> services.workflow_draft_variable_service - core.app.apps.chat.app_generator -> services.conversation_service - core.app.apps.common.workflow_response_converter -> services.variable_truncator - core.app.apps.common.workflow_response_converter -> services.workflow.execution.adapters.workflow_entry - core.app.apps.message_based_app_generator -> services.errors.app_model_config - core.app.apps.message_based_app_generator -> services.errors.conversation - core.app.apps.message_based_app_generator -> services.errors.message - core.app.apps.pipeline.pipeline_generator -> services.workflow.execution.adapters.pipeline.pipeline_generator - core.app.apps.pipeline.pipeline_runner -> services.workflow.execution.adapters.pipeline.pipeline_runner - core.app.apps.workflow.app_generator -> services.workflow.execution.adapters.workflow.app_generator - core.app.apps.workflow.app_runner -> services.workflow.execution.adapters.workflow.app_runner - # The legacy runners share contracts until their service replacements are activated. - core.app.apps.workflow_app_runner -> services.workflow.execution.contracts - core.app.features.annotation_reply.annotation_reply -> services.annotation_service core.app.llm.quota -> services.credit_pool_service - core.app.task_pipeline.message_cycle_manager -> services.annotation_service core.helper.credential_utils -> services.enterprise.plugin_manager_service core.helper.credential_utils -> services.system_feature_service core.mcp.auth_client -> services.tools.mcp_tools_manage_service @@ -86,11 +55,7 @@ ignore_imports = core.rag.summary_index.summary_index -> tasks.generate_summary_index_task core.repositories.sqlalchemy_workflow_node_execution_repository -> services.file_service core.repositories.sqlalchemy_workflow_node_execution_repository -> services.variable_truncator - core.tools.tool_manager -> services.tools.legacy_mcp_tools_manage_service - core.tools.tool_manager -> services.tools.legacy_tools_transform_service core.workflow.nodes.agent_v2.runtime_feature_manifest -> services.agent.knowledge_datasets - # Node construction delegates datasource dependency assembly to the composition root. - # The KnowledgeFS retrieval node is owned by services.knowledge_fs and only wired in here. core.workflow.nodes.agent_v2.session_store -> services.agent.workspace_service core.workflow.nodes.agent_v2.workspace_retirement_layer -> tasks.collect_agent_resources_task libs.email_i18n -> services.entities.feature_entities @@ -323,7 +288,7 @@ forbidden_modules = name = Tool runtime credential persistence is reached only through composition type = forbidden source_modules = - core.tools.tool_manager + services.tools.tool_manager forbidden_modules = repositories services.credentials.tool_runtime_adapters @@ -345,8 +310,8 @@ allow_indirect_imports = True name = Pipeline core depends on its document store port type = forbidden source_modules = - core.app.apps.pipeline.pipeline_generator - core.app.apps.pipeline.pipeline_runner + services.workflow.execution.adapters.pipeline.pipeline_generator + services.workflow.execution.adapters.pipeline.pipeline_runner forbidden_modules = repositories allow_indirect_imports = True diff --git a/api/controllers/console/human_input_form.py b/api/controllers/console/human_input_form.py index 9d0ed77955063c..85952c1672e400 100644 --- a/api/controllers/console/human_input_form.py +++ b/api/controllers/console/human_input_form.py @@ -28,6 +28,7 @@ with_current_user, ) from core.workflow.human_input_policy import HumanInputSurface, is_recipient_type_allowed_for_surface +from extensions.ext_application_services import application_services from extensions.ext_database import db from libs.external_api import error_responses from libs.login import login_required @@ -246,6 +247,7 @@ def _generate_stream_events(): tenant_id=workflow_run.tenant_id, app_id=workflow_run.app_id, session_maker=session_maker, + tool_providers=application_services().tools.tool_providers, human_input_surface=HumanInputSurface.CONSOLE, include_node_details=include_node_details, close_on_pause=not continue_on_pause, diff --git a/api/controllers/openapi/workflow_events.py b/api/controllers/openapi/workflow_events.py index a14b6155a397e4..13ad836bd6c465 100644 --- a/api/controllers/openapi/workflow_events.py +++ b/api/controllers/openapi/workflow_events.py @@ -35,6 +35,7 @@ from core.app.entities.task_entities import StreamEvent from core.db.session_factory import session_factory from core.workflow.human_input_policy import HumanInputSurface +from extensions.ext_application_services import application_services from models.model import AppMode from repositories.factory import DifyAPIRepositoryFactory from services.app.generation.response import convert_to_event_stream @@ -138,6 +139,7 @@ def _generate_stream_events(): tenant_id=tenant_id, app_id=owning_app_id, session_maker=session_maker, + tool_providers=application_services().tools.tool_providers, human_input_surface=HumanInputSurface.OPENAPI, close_on_pause=not continue_on_pause, ) diff --git a/api/controllers/service_api/app/workflow_events.py b/api/controllers/service_api/app/workflow_events.py index 90b3f3d0e7ea7b..de239eca34377e 100644 --- a/api/controllers/service_api/app/workflow_events.py +++ b/api/controllers/service_api/app/workflow_events.py @@ -19,6 +19,7 @@ from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token from core.app.entities.task_entities import StreamEvent from core.workflow.human_input_policy import HumanInputSurface +from extensions.ext_application_services import application_services from extensions.ext_database import db from libs.external_api import error_responses from models.enums import CreatorUserRole @@ -150,6 +151,7 @@ def _generate_stream_events(): tenant_id=app_model.tenant_id, app_id=app_model.id, session_maker=session_maker, + tool_providers=application_services().tools.tool_providers, human_input_surface=HumanInputSurface.SERVICE_API, close_on_pause=not continue_on_pause, ) diff --git a/api/controllers/web/workflow_events.py b/api/controllers/web/workflow_events.py index b6c045377c6b5c..42204effb1974a 100644 --- a/api/controllers/web/workflow_events.py +++ b/api/controllers/web/workflow_events.py @@ -13,6 +13,7 @@ from controllers.common.schema import register_response_schema_model from controllers.web import api, web_ns from controllers.web.wraps import WebApiResource +from extensions.ext_application_services import application_services from extensions.ext_database import db from libs.external_api import error_responses from models.enums import CreatorUserRole @@ -91,6 +92,7 @@ def _generate_stream_events(): tenant_id=app_model.tenant_id, app_id=app_model.id, session_maker=session_maker, + tool_providers=application_services().tools.tool_providers, close_on_pause=not continue_on_pause, ) ) diff --git a/api/core/agent/base_agent_runner.py b/api/core/agent/base_agent_runner.py deleted file mode 100644 index 28bbe1ffc35c69..00000000000000 --- a/api/core/agent/base_agent_runner.py +++ /dev/null @@ -1,506 +0,0 @@ -import json -import logging -import uuid -from decimal import Decimal -from typing import TYPE_CHECKING, cast - -from sqlalchemy import func, select -from sqlalchemy.orm import Session - -from core.agent.entities import AgentEntity, AgentToolEntity -from core.app.app_config.features.file_upload.manager import FileUploadConfigManager -from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig -from core.app.apps.base_app_queue_manager import AppQueueManager -from core.app.apps.base_app_runner import AppRunner -from core.app.entities.app_invoke_entities import ( - AgentChatAppGenerateEntity, - DifyRunContext, - InvokeFrom, - ModelConfigWithCredentialsEntity, - UserFrom, -) -from core.app.file_access import DatabaseFileAccessController -from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.memory.token_buffer_memory import TokenBufferMemory -from core.model_manager import ModelInstance -from core.prompt.utils.extract_thread_messages import extract_thread_messages -from core.tools.__base.tool import Tool -from core.tools.tool_manager import ToolManager -from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool -from extensions.ext_database import db -from factories import file_factory -from graphon.file import file_manager -from graphon.model_runtime.entities import ( - AssistantPromptMessage, - LLMUsage, - PromptMessage, - PromptMessageTool, - SystemPromptMessage, - TextPromptMessageContent, - ToolPromptMessage, - UserPromptMessage, -) -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes -from graphon.model_runtime.entities.model_entities import ModelFeature -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from models.enums import CreatorUserRole -from models.model import Conversation, Message, MessageAgentThought, MessageFile, load_annotation_reply_config - -if TYPE_CHECKING: - from core.app.apps.workflow_app_runner import WorkflowRunDriver - -logger = logging.getLogger(__name__) -_file_access_controller = DatabaseFileAccessController() - - -class BaseAgentRunner(AppRunner): - def __init__( - self, - *, - execution_driver: "WorkflowRunDriver", - session: Session, - tenant_id: str, - application_generate_entity: AgentChatAppGenerateEntity, - conversation: Conversation, - app_config: AgentChatAppConfig, - model_config: ModelConfigWithCredentialsEntity, - config: AgentEntity, - queue_manager: AppQueueManager, - message: Message, - user_id: str, - model_instance: ModelInstance, - memory: TokenBufferMemory | None = None, - prompt_messages: list[PromptMessage] | None = None, - ): - self._execution_driver = execution_driver - self.tenant_id = tenant_id - self.application_generate_entity = application_generate_entity - self.conversation = conversation - self.app_config = app_config - self.model_config = model_config - self.config = config - self.queue_manager = queue_manager - self.message = message - self.user_id = user_id - self.memory = memory - self.history_prompt_messages = self.organize_agent_history( - session=session, prompt_messages=prompt_messages or [] - ) - self.model_instance = model_instance - - # init callback - self.agent_callback = DifyAgentCallbackHandler() - # init dataset tools - hit_callback = DatasetIndexToolCallbackHandler( - queue_manager=queue_manager, - app_id=self.app_config.app_id, - message_id=message.id, - user_id=user_id, - invoke_from=self.application_generate_entity.invoke_from, - ) - self.dataset_tools = DatasetRetrieverTool.get_dataset_tools( - session=session, - tenant_id=tenant_id, - dataset_ids=app_config.dataset.dataset_ids if app_config.dataset else [], - retrieve_config=app_config.dataset.retrieve_config if app_config.dataset else None, - return_resource=( - app_config.additional_features.show_retrieve_source if app_config.additional_features else False - ), - invoke_from=application_generate_entity.invoke_from, - hit_callback=hit_callback, - user_id=user_id, - inputs=cast(dict, application_generate_entity.inputs), - ) - # get how many agent thoughts have been created - self.agent_thought_count = ( - session.scalar( - select(func.count()) - .select_from(MessageAgentThought) - .where( - MessageAgentThought.message_id == self.message.id, - ) - ) - or 0 - ) - session.close() - - # check if model supports stream tool call - llm_model = cast(LargeLanguageModel, model_instance.model_type_instance) - model_schema = llm_model.get_model_schema(model_instance.model_name, model_instance.credentials) - features = model_schema.features if model_schema and model_schema.features else [] - self.stream_tool_call = ModelFeature.STREAM_TOOL_CALL in features - self.vision_enabled = ModelFeature.VISION in features - self.files = application_generate_entity.files if self.vision_enabled else [] - self.query: str = "" - self._current_thoughts: list[PromptMessage] = [] - - def _repack_app_generate_entity( - self, app_generate_entity: AgentChatAppGenerateEntity - ) -> AgentChatAppGenerateEntity: - """ - Repack app generate entity - """ - if app_generate_entity.app_config.prompt_template.simple_prompt_template is None: - app_generate_entity.app_config.prompt_template.simple_prompt_template = "" - - return app_generate_entity - - def _convert_tool_to_prompt_message_tool(self, tool: AgentToolEntity) -> tuple[PromptMessageTool, Tool]: - """ - convert tool to prompt message tool - """ - tool_entity = ToolManager.get_agent_tool_runtime( - execution_driver=self._execution_driver, - tenant_id=self.tenant_id, - app_id=self.app_config.app_id, - agent_tool=tool, - user_id=self.user_id, - invoke_from=self.application_generate_entity.invoke_from, - ) - invoke_from = self.application_generate_entity.invoke_from - if isinstance(invoke_from, InvokeFrom): - tool_entity.runtime.dify_run_context = DifyRunContext( - tenant_id=self.tenant_id, - app_id=self.app_config.app_id, - user_id=self.user_id, - user_from=(UserFrom.ACCOUNT if invoke_from.runs_as_account() else UserFrom.END_USER), - invoke_from=invoke_from, - ) - assert tool_entity.entity.description - message_tool = PromptMessageTool( - name=tool.tool_name, - description=tool_entity.entity.description.llm, - parameters=tool_entity.get_llm_parameters_json_schema(), - ) - - return message_tool, tool_entity - - def _convert_dataset_retriever_tool_to_prompt_message_tool(self, tool: DatasetRetrieverTool) -> PromptMessageTool: - """ - convert dataset retriever tool to prompt message tool - """ - assert tool.entity.description - - prompt_tool = PromptMessageTool( - name=tool.entity.identity.name, - description=tool.entity.description.llm, - parameters={ - "type": "object", - "properties": {}, - "required": [], - }, - ) - - for parameter in tool.get_runtime_parameters(): - parameter_type = "string" - - prompt_tool.parameters["properties"][parameter.name] = { - "type": parameter_type, - "description": parameter.llm_description or "", - } - - if parameter.required: - if parameter.name not in prompt_tool.parameters["required"]: - prompt_tool.parameters["required"].append(parameter.name) - - return prompt_tool - - def _init_prompt_tools(self) -> tuple[dict[str, Tool], list[PromptMessageTool]]: - """ - Init tools - """ - tool_instances = {} - prompt_messages_tools = [] - - for tool in self.app_config.agent.tools or [] if self.app_config.agent else []: - try: - prompt_tool, tool_entity = self._convert_tool_to_prompt_message_tool(tool) - except Exception: - # api tool may be deleted - continue - # save tool entity - tool_instances[tool.tool_name] = tool_entity - # save prompt tool - prompt_messages_tools.append(prompt_tool) - - # convert dataset tools into ModelRuntime Tool format - for dataset_tool in self.dataset_tools: - prompt_tool = self._convert_dataset_retriever_tool_to_prompt_message_tool(dataset_tool) - # save prompt tool - prompt_messages_tools.append(prompt_tool) - # save tool entity - tool_instances[dataset_tool.entity.identity.name] = dataset_tool - - return tool_instances, prompt_messages_tools - - def update_prompt_message_tool(self, tool: Tool, prompt_tool: PromptMessageTool) -> PromptMessageTool: - """ - update prompt message tool - """ - prompt_tool.parameters = tool.get_llm_parameters_json_schema() - return prompt_tool - - def create_agent_thought( - self, message_id: str, message: str, tool_name: str, tool_input: str, messages_ids: list[str] - ) -> str: - """ - Create agent thought - """ - thought = MessageAgentThought( - message_id=message_id, - message_chain_id=None, - tool_process_data=None, - thought="", - tool=tool_name, - tool_labels_str="{}", - tool_meta_str="{}", - tool_input=tool_input, - message=message, - message_token=0, - message_unit_price=Decimal(0), - message_price_unit=Decimal("0.001"), - message_files=json.dumps(messages_ids) if messages_ids else "", - answer="", - observation="", - answer_token=0, - answer_unit_price=Decimal(0), - answer_price_unit=Decimal("0.001"), - tokens=0, - total_price=Decimal(0), - position=self.agent_thought_count + 1, - currency="USD", - latency=0, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=self.user_id, - ) - - db.session.add(thought) - db.session.commit() - agent_thought_id = str(thought.id) - self.agent_thought_count += 1 - db.session.close() - - return agent_thought_id - - def save_agent_thought( - self, - agent_thought_id: str, - tool_name: str | None, - tool_input: str | dict | None, - thought: str | None, - observation: str | dict | None, - tool_invoke_meta: str | dict | None, - answer: str | None, - messages_ids: list[str], - llm_usage: LLMUsage | None = None, - ): - """ - Save agent thought - """ - stmt = select(MessageAgentThought).where(MessageAgentThought.id == agent_thought_id) - agent_thought = db.session.scalar(stmt) - if not agent_thought: - raise ValueError("agent thought not found") - - if thought: - existing_thought = agent_thought.thought or "" - agent_thought.thought = f"{existing_thought}{thought}" - - if tool_name: - agent_thought.tool = tool_name - - if tool_input: - if isinstance(tool_input, dict): - try: - tool_input = json.dumps(tool_input, ensure_ascii=False) - except Exception: - tool_input = json.dumps(tool_input) - - agent_thought.tool_input = tool_input - - if observation: - if isinstance(observation, dict): - try: - observation = json.dumps(observation, ensure_ascii=False) - except Exception: - observation = json.dumps(observation) - - agent_thought.observation = observation - - if answer: - agent_thought.answer = answer - - if messages_ids is not None and len(messages_ids) > 0: - agent_thought.message_files = json.dumps(messages_ids) - - if llm_usage: - agent_thought.message_token = llm_usage.prompt_tokens - agent_thought.message_price_unit = llm_usage.prompt_price_unit - agent_thought.message_unit_price = llm_usage.prompt_unit_price - agent_thought.answer_token = llm_usage.completion_tokens - agent_thought.answer_price_unit = llm_usage.completion_price_unit - agent_thought.answer_unit_price = llm_usage.completion_unit_price - agent_thought.tokens = llm_usage.total_tokens - agent_thought.total_price = llm_usage.total_price - - # check if tool labels is not empty - labels = agent_thought.tool_labels or {} - tools = agent_thought.tool.split(";") if agent_thought.tool else [] - for tool in tools: - if not tool: - continue - if tool not in labels: - tool_label = ToolManager.get_tool_label(tool) - if tool_label: - labels[tool] = tool_label.to_dict() - else: - labels[tool] = {"en_US": tool, "zh_Hans": tool} - - agent_thought.tool_labels_str = json.dumps(labels) - - if tool_invoke_meta is not None: - if isinstance(tool_invoke_meta, dict): - try: - tool_invoke_meta = json.dumps(tool_invoke_meta, ensure_ascii=False) - except Exception: - tool_invoke_meta = json.dumps(tool_invoke_meta) - - agent_thought.tool_meta_str = tool_invoke_meta - - db.session.commit() - db.session.close() - - def organize_agent_history(self, prompt_messages: list[PromptMessage], *, session: Session) -> list[PromptMessage]: - """ - Organize agent history - """ - result: list[PromptMessage] = [] - # check if there is a system message in the beginning of the conversation - for prompt_message in prompt_messages: - if isinstance(prompt_message, SystemPromptMessage): - result.append(prompt_message) - - messages = ( - ( - session.execute( - select(Message) - .where(Message.conversation_id == self.message.conversation_id) - .order_by(Message.created_at.desc()) - ) - ) - .scalars() - .all() - ) - - messages = list(reversed(extract_thread_messages(messages))) - - for message in messages: - if message.id == self.message.id: - continue - - result.append(self.organize_agent_user_prompt(message, session=session)) - agent_thoughts = message.agent_thoughts_with_session(session=session) - if agent_thoughts: - for agent_thought in agent_thoughts: - tool_names_raw = agent_thought.tool - if tool_names_raw: - tool_names = tool_names_raw.split(";") - tool_calls: list[AssistantPromptMessage.ToolCall] = [] - tool_call_response: list[ToolPromptMessage] = [] - tool_input_payload = agent_thought.tool_input - if tool_input_payload: - try: - tool_inputs = json.loads(tool_input_payload) - except Exception: - tool_inputs = {tool: {} for tool in tool_names} - else: - tool_inputs = {tool: {} for tool in tool_names} - - observation_payload = agent_thought.observation - if observation_payload: - try: - tool_responses = json.loads(observation_payload) - except Exception: - tool_responses = dict.fromkeys(tool_names, observation_payload) - else: - tool_responses = dict.fromkeys(tool_names, observation_payload) - - for index, tool in enumerate(tool_names): - # Keep IDs stable across turns so replayed history can use prompt caching. - tool_call_id = str(uuid.uuid5(uuid.NAMESPACE_OID, f"{agent_thought.id}:{index}:{tool}")) - tool_calls.append( - AssistantPromptMessage.ToolCall( - id=tool_call_id, - type="function", - function=AssistantPromptMessage.ToolCall.ToolCallFunction( - name=tool, - arguments=json.dumps(tool_inputs.get(tool, {})), - ), - ) - ) - tool_call_response.append( - ToolPromptMessage( - content=tool_responses.get(tool, agent_thought.observation), - name=tool, - tool_call_id=tool_call_id, - ) - ) - - result.extend( - [ - AssistantPromptMessage( - content=agent_thought.thought, - tool_calls=tool_calls, - ), - *tool_call_response, - ] - ) - if not tool_names_raw: - result.append(AssistantPromptMessage(content=agent_thought.thought)) - else: - if message.answer: - result.append(AssistantPromptMessage(content=message.answer)) - - session.close() - - return result - - def organize_agent_user_prompt(self, message: Message, *, session: Session) -> UserPromptMessage: - stmt = select(MessageFile).where(MessageFile.message_id == message.id) - files = session.scalars(stmt).all() - if not files: - return UserPromptMessage(content=message.query) - app_model_config = message.app_model_config_with_session(session=session) - if app_model_config: - annotation_reply = load_annotation_reply_config(session, app_model_config.app_id) - file_extra_config = FileUploadConfigManager.convert( - app_model_config.to_dict(annotation_reply=annotation_reply) - ) - else: - file_extra_config = None - - if not file_extra_config: - return UserPromptMessage(content=message.query) - - image_detail_config = file_extra_config.image_config.detail if file_extra_config.image_config else None - image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW - - file_objs = file_factory.build_from_message_files( - message_files=files, - tenant_id=self.tenant_id, - access_controller=_file_access_controller, - ) - if not file_objs: - return UserPromptMessage(content=message.query) - prompt_message_contents: list[PromptMessageContentUnionTypes] = [] - for file in file_objs: - prompt_message_contents.append( - file_manager.to_prompt_message_content( - file, - image_detail_config=image_detail_config, - ) - ) - prompt_message_contents.append(TextPromptMessageContent(data=message.query)) - - return UserPromptMessage(content=prompt_message_contents) diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py deleted file mode 100644 index b508cc70497382..00000000000000 --- a/api/core/agent/cot_agent_runner.py +++ /dev/null @@ -1,471 +0,0 @@ -import contextlib -import json -import logging -from abc import ABC, abstractmethod -from collections.abc import Generator, Mapping, Sequence -from typing import Any, TypedDict - -from sqlalchemy.orm import Session - -from core.agent.base_agent_runner import BaseAgentRunner -from core.agent.entities import AgentScratchpadUnit -from core.agent.errors import AgentMaxIterationError -from core.agent.output_parser.cot_output_parser import CotAgentOutputParser -from core.app.apps.base_app_queue_manager import PublishFrom -from core.app.entities.queue_entities import QueueAgentThoughtEvent, QueueMessageEndEvent, QueueMessageFileEvent -from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy -from core.model_context import use_credit_usage_metadata -from core.ops.ops_trace_manager import TraceQueueManager -from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform -from core.tools.__base.tool import Tool -from core.tools.entities.tool_entities import ToolInvokeMeta -from core.tools.tool_engine import ToolEngine -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage -from graphon.model_runtime.entities.message_entities import ( - AssistantPromptMessage, - PromptMessage, - PromptMessageTool, - ToolPromptMessage, - UserPromptMessage, -) -from models.model import Message - -logger = logging.getLogger(__name__) - - -class ActionDict(TypedDict): - """Shape produced by AgentScratchpadUnit.Action.to_dict().""" - - action: str - action_input: dict[str, Any] | str - - -class CotAgentRunner(BaseAgentRunner, ABC): - _is_first_iteration = True - _ignore_observation_providers = ["wenxin"] - _historic_prompt_messages: list[PromptMessage] - _agent_scratchpad: list[AgentScratchpadUnit] - _instruction: str - _query: str - _prompt_messages_tools: Sequence[PromptMessageTool] - - def run( - self, - session: Session, - message: Message, - query: str, - inputs: Mapping[str, str], - ) -> Generator: - """ - Run Cot agent application - """ - - app_generate_entity = self.application_generate_entity - self._repack_app_generate_entity(app_generate_entity) - self._init_react_state(query) - - trace_manager = app_generate_entity.trace_manager - - # check model mode - if "Observation" not in app_generate_entity.model_conf.stop: - if app_generate_entity.model_conf.provider not in self._ignore_observation_providers: - app_generate_entity.model_conf.stop.append("Observation") - - app_config = self.app_config - assert app_config.agent - - # init instruction - inputs = inputs or {} - instruction = app_config.prompt_template.simple_prompt_template or "" - self._instruction = self._fill_in_inputs_from_external_data_tools(instruction, inputs) - - iteration_step = 1 - max_iteration_steps = min(app_config.agent.max_iteration, 99) + 1 - - # convert tools into ModelRuntime Tool format - tool_instances, prompt_messages_tools = self._init_prompt_tools() - self._prompt_messages_tools = prompt_messages_tools - - function_call_state = True - llm_usage: dict[str, LLMUsage | None] = {"usage": None} - final_answer = "" - prompt_messages: list = [] # Initialize prompt_messages - agent_thought_id = "" # Initialize agent_thought_id - - def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMUsage): - if not final_llm_usage_dict["usage"]: - final_llm_usage_dict["usage"] = usage - else: - llm_usage = final_llm_usage_dict["usage"] - llm_usage.prompt_tokens += usage.prompt_tokens - llm_usage.completion_tokens += usage.completion_tokens - llm_usage.total_tokens += usage.total_tokens - llm_usage.prompt_price += usage.prompt_price - llm_usage.completion_price += usage.completion_price - llm_usage.total_price += usage.total_price - - model_instance = self.model_instance - - while function_call_state and iteration_step <= max_iteration_steps: - # continue to run until there is not any tool call - function_call_state = False - - if iteration_step == max_iteration_steps: - # the last iteration, remove all tools - self._prompt_messages_tools = [] - - message_file_ids: list[str] = [] - - agent_thought_id = self.create_agent_thought( - message_id=message.id, - message="", - tool_name="", - tool_input="", - messages_ids=message_file_ids, - ) - - if iteration_step > 1: - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - # recalc llm max tokens - prompt_messages = self._organize_prompt_messages() - self.recalc_llm_max_tokens(self.model_config, prompt_messages) - - # Release any setup/tool transaction before waiting on the provider stream. - session.commit() - session.close() - - # invoke model - request_metadata: dict[str, object] = { - "app_id": self.app_config.app_id, - "app_type": CreditUsageAppType.AGENT, - "created_by": CreditUsageCreatedBy.APP, - } - - chunks = model_instance.invoke_llm( - prompt_messages=prompt_messages, - model_parameters=app_generate_entity.model_conf.parameters, - tools=[], - stop=app_generate_entity.model_conf.stop, - stream=True, - callbacks=[], - request_metadata=request_metadata, - ) - - usage_dict: dict[str, LLMUsage | None] = {} - react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks, usage_dict) - scratchpad = AgentScratchpadUnit( - agent_response="", - thought="", - action_str="", - observation="", - action=None, - ) - - # publish agent thought if it's first iteration - if iteration_step == 1: - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - for chunk in react_chunks: - if isinstance(chunk, AgentScratchpadUnit.Action): - action = chunk - # detect action - assert scratchpad.agent_response is not None - scratchpad.agent_response += json.dumps(chunk.model_dump()) - scratchpad.action_str = json.dumps(chunk.model_dump()) - scratchpad.action = action - else: - assert scratchpad.agent_response is not None - scratchpad.agent_response += chunk - assert scratchpad.thought is not None - scratchpad.thought += chunk - yield LLMResultChunk( - model=self.model_config.model, - prompt_messages=prompt_messages, - system_fingerprint="", - delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=chunk), usage=None), - ) - - assert scratchpad.thought is not None - scratchpad.thought = scratchpad.thought.strip() or "I am thinking about how to help you" - self._agent_scratchpad.append(scratchpad) - - # Check if max iteration is reached and model still wants to call tools - if iteration_step == max_iteration_steps and scratchpad.action: - if scratchpad.action.action_name.lower() != "final answer": - raise AgentMaxIterationError(app_config.agent.max_iteration) - - # get llm usage - if "usage" in usage_dict: - if usage_dict["usage"] is not None: - increase_usage(llm_usage, usage_dict["usage"]) - else: - usage_dict["usage"] = LLMUsage.empty_usage() - - self.save_agent_thought( - agent_thought_id=agent_thought_id, - tool_name=(scratchpad.action.action_name if scratchpad.action and not scratchpad.is_final() else ""), - tool_input={scratchpad.action.action_name: scratchpad.action.action_input} if scratchpad.action else {}, - tool_invoke_meta={}, - thought=scratchpad.thought or "", - observation="", - answer=scratchpad.agent_response or "", - messages_ids=[], - llm_usage=usage_dict["usage"], - ) - - if not scratchpad.is_final(): - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - if not scratchpad.action: - # failed to extract action, return final answer directly - final_answer = "" - else: - if scratchpad.action.action_name.lower() == "final answer": - # action is final answer, return final answer directly - try: - match scratchpad.action.action_input: - case dict(): - final_answer = json.dumps(scratchpad.action.action_input, ensure_ascii=False) - case str(): - final_answer = scratchpad.action.action_input - case _: - final_answer = f"{scratchpad.action.action_input}" - except TypeError: - final_answer = f"{scratchpad.action.action_input}" - else: - function_call_state = True - # action is tool call, invoke tool - tool_invoke_response, tool_invoke_meta = self._handle_invoke_action( - session=session, - action=scratchpad.action, - tool_instances=tool_instances, - message_file_ids=message_file_ids, - trace_manager=trace_manager, - ) - scratchpad.observation = tool_invoke_response - scratchpad.agent_response = tool_invoke_response - - self.save_agent_thought( - agent_thought_id=agent_thought_id, - tool_name=scratchpad.action.action_name, - tool_input={scratchpad.action.action_name: scratchpad.action.action_input}, - thought=scratchpad.thought or "", - observation={scratchpad.action.action_name: tool_invoke_response}, - tool_invoke_meta={scratchpad.action.action_name: tool_invoke_meta.to_dict()}, - answer=scratchpad.agent_response, - messages_ids=message_file_ids, - llm_usage=usage_dict["usage"], - ) - - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - # update prompt tool message - for prompt_tool in self._prompt_messages_tools: - self.update_prompt_message_tool(tool_instances[prompt_tool.name], prompt_tool) - - iteration_step += 1 - - yield LLMResultChunk( - model=model_instance.model_name, - prompt_messages=prompt_messages, - delta=LLMResultChunkDelta( - index=0, message=AssistantPromptMessage(content=final_answer), usage=llm_usage["usage"] - ), - system_fingerprint="", - ) - - # save agent thought - self.save_agent_thought( - agent_thought_id=agent_thought_id, - tool_name="", - tool_input={}, - tool_invoke_meta={}, - thought=final_answer, - observation={}, - answer=final_answer, - messages_ids=[], - ) - # publish end event - self.queue_manager.publish( - QueueMessageEndEvent( - llm_result=LLMResult( - model=model_instance.model_name, - prompt_messages=prompt_messages, - message=AssistantPromptMessage(content=final_answer), - usage=llm_usage["usage"] or LLMUsage.empty_usage(), - system_fingerprint="", - ) - ), - PublishFrom.APPLICATION_MANAGER, - ) - - def _handle_invoke_action( - self, - session: Session, - action: AgentScratchpadUnit.Action, - tool_instances: Mapping[str, Tool], - message_file_ids: list[str], - trace_manager: TraceQueueManager | None = None, - ) -> tuple[str, ToolInvokeMeta]: - """ - handle invoke action - :param action: action - :param tool_instances: tool instances - :param message_file_ids: message file ids - :param trace_manager: trace manager - :return: observation, meta - """ - # action is tool call, invoke tool - tool_call_name = action.action_name - tool_call_args = action.action_input - tool_instance = tool_instances.get(tool_call_name) - - if not tool_instance: - answer = f"there is not a tool named {tool_call_name}" - return answer, ToolInvokeMeta.error_instance(answer) - - if isinstance(tool_call_args, str): - with contextlib.suppress(json.JSONDecodeError): - tool_call_args = json.loads(tool_call_args) - - # invoke tool - with use_credit_usage_metadata({"app_type": CreditUsageAppType.AGENT}): - tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke( - session=session, - tool=tool_instance, - tool_parameters=tool_call_args, - user_id=self.user_id, - tenant_id=self.tenant_id, - message=self.message, - invoke_from=self.application_generate_entity.invoke_from, - agent_tool_callback=self.agent_callback, - trace_manager=trace_manager, - ) - session.commit() - session.close() - - # publish files - for message_file_id in message_files: - # publish message file - self.queue_manager.publish( - QueueMessageFileEvent(message_file_id=message_file_id), PublishFrom.APPLICATION_MANAGER - ) - # add message file ids - message_file_ids.append(message_file_id) - - return tool_invoke_response, tool_invoke_meta - - def _convert_dict_to_action(self, action: ActionDict) -> AgentScratchpadUnit.Action: - """ - convert dict to action - """ - return AgentScratchpadUnit.Action(action_name=action["action"], action_input=action["action_input"]) - - def _fill_in_inputs_from_external_data_tools(self, instruction: str, inputs: Mapping[str, Any]) -> str: - """ - fill in inputs from external data tools - """ - for key, value in inputs.items(): - try: - instruction = instruction.replace(f"{{{{{key}}}}}", str(value)) - except Exception: - continue - - return instruction - - def _init_react_state(self, query): - """ - init agent scratchpad - """ - self._query = query - self._agent_scratchpad = [] - self._historic_prompt_messages = self._organize_historic_prompt_messages() - - @abstractmethod - def _organize_prompt_messages(self) -> list[PromptMessage]: - """ - organize prompt messages - """ - - def _format_assistant_message(self, agent_scratchpad: list[AgentScratchpadUnit]) -> str: - """ - format assistant message - """ - message = "" - for scratchpad in agent_scratchpad: - if scratchpad.is_final(): - message += f"Final Answer: {scratchpad.agent_response}" - else: - message += f"Thought: {scratchpad.thought}\n\n" - if scratchpad.action_str: - message += f"Action: {scratchpad.action_str}\n\n" - if scratchpad.observation: - message += f"Observation: {scratchpad.observation}\n\n" - - return message - - def _organize_historic_prompt_messages( - self, current_session_messages: list[PromptMessage] | None = None - ) -> list[PromptMessage]: - """ - organize historic prompt messages - """ - result: list[PromptMessage] = [] - scratchpads: list[AgentScratchpadUnit] = [] - current_scratchpad: AgentScratchpadUnit | None = None - - for message in self.history_prompt_messages: - match message: - case AssistantPromptMessage(): - if not current_scratchpad: - assert isinstance(message.content, str) - current_scratchpad = AgentScratchpadUnit( - agent_response=message.content, - thought=message.content or "I am thinking about how to help you", - action_str="", - action=None, - observation=None, - ) - scratchpads.append(current_scratchpad) - if message.tool_calls: - try: - current_scratchpad.action = AgentScratchpadUnit.Action( - action_name=message.tool_calls[0].function.name, - action_input=json.loads(message.tool_calls[0].function.arguments), - ) - current_scratchpad.action_str = json.dumps(current_scratchpad.action.to_dict()) - except Exception: - logger.exception("Failed to parse tool call from assistant message") - case ToolPromptMessage(): - if current_scratchpad: - assert isinstance(message.content, str) - current_scratchpad.observation = message.content - else: - raise NotImplementedError("expected str type") - case UserPromptMessage(): - if scratchpads: - result.append(AssistantPromptMessage(content=self._format_assistant_message(scratchpads))) - scratchpads = [] - current_scratchpad = None - - result.append(message) - - if scratchpads: - result.append(AssistantPromptMessage(content=self._format_assistant_message(scratchpads))) - - historic_prompts = AgentHistoryPromptTransform( - model_config=self.model_config, - prompt_messages=current_session_messages or [], - history_messages=result, - memory=self.memory, - ).get_prompt() - return historic_prompts diff --git a/api/core/agent/cot_chat_agent_runner.py b/api/core/agent/cot_chat_agent_runner.py deleted file mode 100644 index 8d4a5931b821f7..00000000000000 --- a/api/core/agent/cot_chat_agent_runner.py +++ /dev/null @@ -1,117 +0,0 @@ -import json -from typing import override - -from core.agent.cot_agent_runner import CotAgentRunner -from graphon.file import file_manager -from graphon.model_runtime.entities import ( - AssistantPromptMessage, - PromptMessage, - SystemPromptMessage, - TextPromptMessageContent, - UserPromptMessage, -) -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes -from graphon.model_runtime.utils.encoders import jsonable_encoder - - -class CotChatAgentRunner(CotAgentRunner): - def _organize_system_prompt(self) -> SystemPromptMessage: - """ - Organize system prompt - """ - assert self.app_config.agent - assert self.app_config.agent.prompt - - prompt_entity = self.app_config.agent.prompt - if not prompt_entity: - raise ValueError("Agent prompt configuration is not set") - first_prompt = prompt_entity.first_prompt - - system_prompt = ( - first_prompt.replace("{{instruction}}", self._instruction) - .replace("{{tools}}", json.dumps(jsonable_encoder(self._prompt_messages_tools))) - .replace("{{tool_names}}", ", ".join([tool.name for tool in self._prompt_messages_tools])) - ) - - return SystemPromptMessage(content=system_prompt) - - def _organize_user_query(self, query, prompt_messages: list[PromptMessage]) -> list[PromptMessage]: - """ - Organize user query - """ - if self.files: - # get image detail config - image_detail_config = ( - self.application_generate_entity.file_upload_config.image_config.detail - if ( - self.application_generate_entity.file_upload_config - and self.application_generate_entity.file_upload_config.image_config - ) - else None - ) - image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW - - prompt_message_contents: list[PromptMessageContentUnionTypes] = [] - for file in self.files: - prompt_message_contents.append( - file_manager.to_prompt_message_content( - file, - image_detail_config=image_detail_config, - ) - ) - prompt_message_contents.append(TextPromptMessageContent(data=query)) - - prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) - else: - prompt_messages.append(UserPromptMessage(content=query)) - - return prompt_messages - - @override - def _organize_prompt_messages(self) -> list[PromptMessage]: - """ - Organize - """ - # organize system prompt - system_message = self._organize_system_prompt() - - # organize current assistant messages - agent_scratchpad = self._agent_scratchpad - if not agent_scratchpad: - assistant_messages = [] - else: - content = "" - for unit in agent_scratchpad: - if unit.is_final(): - content += f"Final Answer: {unit.agent_response}" - else: - content += f"Thought: {unit.thought}\n\n" - if unit.action_str: - content += f"Action: {unit.action_str}\n\n" - if unit.observation: - content += f"Observation: {unit.observation}\n\n" - - assistant_messages = [AssistantPromptMessage(content=content)] - - # query messages - query_messages = self._organize_user_query(self._query, []) - - if assistant_messages: - # organize historic prompt messages - historic_messages = self._organize_historic_prompt_messages( - [system_message, *query_messages, *assistant_messages, UserPromptMessage(content="continue")] - ) - messages = [ - system_message, - *historic_messages, - *query_messages, - *assistant_messages, - UserPromptMessage(content="continue"), - ] - else: - # organize historic prompt messages - historic_messages = self._organize_historic_prompt_messages([system_message, *query_messages]) - messages = [system_message, *historic_messages, *query_messages] - - # join all messages - return messages diff --git a/api/core/agent/cot_completion_agent_runner.py b/api/core/agent/cot_completion_agent_runner.py deleted file mode 100644 index 72d7831eb731b3..00000000000000 --- a/api/core/agent/cot_completion_agent_runner.py +++ /dev/null @@ -1,91 +0,0 @@ -import json -from typing import override - -from core.agent.cot_agent_runner import CotAgentRunner -from graphon.model_runtime.entities.message_entities import ( - AssistantPromptMessage, - PromptMessage, - TextPromptMessageContent, - UserPromptMessage, -) -from graphon.model_runtime.utils.encoders import jsonable_encoder - - -class CotCompletionAgentRunner(CotAgentRunner): - def _organize_instruction_prompt(self) -> str: - """ - Organize instruction prompt - """ - if self.app_config.agent is None: - raise ValueError("Agent configuration is not set") - prompt_entity = self.app_config.agent.prompt - if prompt_entity is None: - raise ValueError("prompt entity is not set") - first_prompt = prompt_entity.first_prompt - - system_prompt = ( - first_prompt.replace("{{instruction}}", self._instruction) - .replace("{{tools}}", json.dumps(jsonable_encoder(self._prompt_messages_tools))) - .replace("{{tool_names}}", ", ".join([tool.name for tool in self._prompt_messages_tools])) - ) - - return system_prompt - - def _organize_historic_prompt(self, current_session_messages: list[PromptMessage] | None = None) -> str: - """ - Organize historic prompt - """ - historic_prompt_messages = self._organize_historic_prompt_messages(current_session_messages) - historic_prompt = "" - - for message in historic_prompt_messages: - match message: - case UserPromptMessage(): - historic_prompt += f"Question: {message.content}\n\n" - case AssistantPromptMessage(): - match message.content: - case str(): - historic_prompt += message.content + "\n\n" - case list(): - for content in message.content: - if not isinstance(content, TextPromptMessageContent): - continue - historic_prompt += content.data - - return historic_prompt - - @override - def _organize_prompt_messages(self) -> list[PromptMessage]: - """ - Organize prompt messages - """ - # organize system prompt - system_prompt = self._organize_instruction_prompt() - - # organize historic prompt messages - historic_prompt = self._organize_historic_prompt() - - # organize current assistant messages - agent_scratchpad = self._agent_scratchpad - assistant_prompt = "" - for unit in agent_scratchpad or []: - if unit.is_final(): - assistant_prompt += f"Final Answer: {unit.agent_response}" - else: - assistant_prompt += f"Thought: {unit.thought}\n\n" - if unit.action_str: - assistant_prompt += f"Action: {unit.action_str}\n\n" - if unit.observation: - assistant_prompt += f"Observation: {unit.observation}\n\n" - - # query messages - query_prompt = f"Question: {self._query}" - - # join all messages - prompt = ( - system_prompt.replace("{{historic_messages}}", historic_prompt) - .replace("{{agent_scratchpad}}", assistant_prompt) - .replace("{{query}}", query_prompt) - ) - - return [UserPromptMessage(content=prompt)] diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py deleted file mode 100644 index 1b204608632436..00000000000000 --- a/api/core/agent/fc_agent_runner.py +++ /dev/null @@ -1,571 +0,0 @@ -import json -import logging -import re -from collections.abc import Generator -from copy import deepcopy -from typing import Any - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.agent.base_agent_runner import BaseAgentRunner -from core.agent.errors import AgentMaxIterationError -from core.app.apps.base_app_queue_manager import PublishFrom -from core.app.entities.queue_entities import QueueAgentThoughtEvent, QueueMessageEndEvent, QueueMessageFileEvent -from core.app.file_access import grant_upload_file_access -from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy -from core.model_context import use_credit_usage_metadata -from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform -from core.tools.entities.tool_entities import ToolInvokeMeta -from core.tools.signature import sign_upload_file_preview_url -from core.tools.tool_engine import ToolEngine -from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool -from core.workflow.file_reference import build_file_reference -from graphon.file import File, FileTransferMethod, FileType, file_manager -from graphon.model_runtime.entities import ( - AssistantPromptMessage, - LLMResult, - LLMResultChunk, - LLMResultChunkDelta, - LLMUsage, - PromptMessage, - PromptMessageContentType, - SystemPromptMessage, - TextPromptMessageContent, - ToolPromptMessage, - UserPromptMessage, -) -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes -from models import UploadFile -from models.model import Message - -logger = logging.getLogger(__name__) - -_FILE_PREVIEW_ID_PATTERN = re.compile(r"/files/([a-fA-F0-9-]{36})/file-preview") -_KNOWLEDGE_RETRIEVAL_PROMPT_NAME = "knowledge_retrieval" - - -class FunctionCallAgentRunner(BaseAgentRunner): - def _build_dataset_tool_image_contents( - self, session: Session, tool_response: str, tool_instance: Any - ) -> list[PromptMessageContentUnionTypes]: - if not self.vision_enabled or not isinstance(tool_instance, DatasetRetrieverTool): - return [] - - upload_file_ids = list(dict.fromkeys(_FILE_PREVIEW_ID_PATTERN.findall(tool_response))) - if not upload_file_ids: - return [] - - upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(upload_file_ids))).all() - upload_file_map = {str(upload_file.id): upload_file for upload_file in upload_files} - ordered_upload_files = [ - upload_file_map[upload_file_id] for upload_file_id in upload_file_ids if upload_file_id in upload_file_map - ] - image_upload_files = [ - upload_file for upload_file in ordered_upload_files if (upload_file.mime_type or "").startswith("image/") - ] - if not image_upload_files: - return [] - - grant_upload_file_access(str(upload_file.id) for upload_file in image_upload_files) - - image_detail_config = ( - self.application_generate_entity.file_upload_config.image_config.detail - if ( - self.application_generate_entity.file_upload_config - and self.application_generate_entity.file_upload_config.image_config - ) - else None - ) - image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW - - prompt_message_contents: list[PromptMessageContentUnionTypes] = [] - for upload_file in image_upload_files: - prompt_file = File( - file_id=upload_file.id, - filename=upload_file.name, - extension="." + upload_file.extension, - mime_type=upload_file.mime_type, - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.LOCAL_FILE, - remote_url=upload_file.source_url, - reference=build_file_reference(record_id=str(upload_file.id)), - size=upload_file.size, - storage_key=upload_file.key, - url=sign_upload_file_preview_url(upload_file.id, upload_file.extension), - ) - prompt_message_contents.append( - file_manager.to_prompt_message_content(prompt_file, image_detail_config=image_detail_config) - ) - - return prompt_message_contents - - def run(self, session: Session, message: Message, query: str, **kwargs: Any) -> Generator[LLMResultChunk]: - """ - Run FunctionCall agent application - """ - self.query = query - app_generate_entity = self.application_generate_entity - - app_config = self.app_config - assert app_config is not None, "app_config is required" - assert app_config.agent is not None, "app_config.agent is required" - - # convert tools into ModelRuntime Tool format - tool_instances, prompt_messages_tools = self._init_prompt_tools() - - assert app_config.agent - - iteration_step = 1 - max_iteration_steps = min(app_config.agent.max_iteration, 99) + 1 - - # continue to run until there is not any tool call - function_call_state = True - llm_usage: dict[str, LLMUsage | None] = {"usage": None} - final_answer = "" - prompt_messages: list = [] # Initialize prompt_messages - - # get tracing instance - trace_manager = app_generate_entity.trace_manager - - def increase_usage(final_llm_usage_dict: dict[str, LLMUsage | None], usage: LLMUsage): - if not final_llm_usage_dict["usage"]: - final_llm_usage_dict["usage"] = usage - else: - llm_usage = final_llm_usage_dict["usage"] - llm_usage.prompt_tokens += usage.prompt_tokens - llm_usage.completion_tokens += usage.completion_tokens - llm_usage.total_tokens += usage.total_tokens - llm_usage.prompt_price += usage.prompt_price - llm_usage.completion_price += usage.completion_price - llm_usage.total_price += usage.total_price - - model_instance = self.model_instance - - while function_call_state and iteration_step <= max_iteration_steps: - function_call_state = False - - if iteration_step == max_iteration_steps: - # the last iteration, remove all tools - prompt_messages_tools = [] - - message_file_ids: list[str] = [] - agent_thought_id = self.create_agent_thought( - message_id=message.id, - message="", - tool_name="", - tool_input="", - messages_ids=message_file_ids, - ) - - # recalc llm max tokens - prompt_messages = self._organize_prompt_messages() - self.recalc_llm_max_tokens(self.model_config, prompt_messages) - - # Release any setup/tool transaction before waiting on the provider stream. - session.commit() - session.close() - - # invoke model - request_metadata: dict[str, object] = { - "app_id": self.app_config.app_id, - "app_type": CreditUsageAppType.AGENT, - "created_by": CreditUsageCreatedBy.APP, - } - - chunks: Generator[LLMResultChunk] | LLMResult = model_instance.invoke_llm( - prompt_messages=prompt_messages, - model_parameters=app_generate_entity.model_conf.parameters, - tools=prompt_messages_tools, - stop=app_generate_entity.model_conf.stop, - stream=self.stream_tool_call, - callbacks=[], - request_metadata=request_metadata, - ) - - tool_calls: list[tuple[str, str, dict[str, Any]]] = [] - - # save full response - response = "" - - # save tool call names and inputs - tool_call_names = "" - tool_call_inputs = "" - - current_llm_usage = None - - if isinstance(chunks, Generator): - is_first_chunk = True - for chunk in chunks: - if is_first_chunk: - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - is_first_chunk = False - # check if there is any tool call - if self.check_tool_calls(chunk): - function_call_state = True - tool_calls.extend(self.extract_tool_calls(chunk) or []) - tool_call_names = ";".join([tool_call[1] for tool_call in tool_calls]) - try: - tool_call_inputs = json.dumps( - {tool_call[1]: tool_call[2] for tool_call in tool_calls}, ensure_ascii=False - ) - except TypeError: - # fallback: force ASCII to handle non-serializable objects - tool_call_inputs = json.dumps({tool_call[1]: tool_call[2] for tool_call in tool_calls}) - - if chunk.delta.message and chunk.delta.message.content: - if isinstance(chunk.delta.message.content, list): - for content in chunk.delta.message.content: - response += content.data - else: - response += str(chunk.delta.message.content) - - if chunk.delta.usage: - increase_usage(llm_usage, chunk.delta.usage) - current_llm_usage = chunk.delta.usage - - yield chunk - else: - result = chunks - # check if there is any tool call - if self.check_blocking_tool_calls(result): - function_call_state = True - tool_calls.extend(self.extract_blocking_tool_calls(result) or []) - tool_call_names = ";".join([tool_call[1] for tool_call in tool_calls]) - try: - tool_call_inputs = json.dumps( - {tool_call[1]: tool_call[2] for tool_call in tool_calls}, ensure_ascii=False - ) - except TypeError: - # fallback: force ASCII to handle non-serializable objects - tool_call_inputs = json.dumps({tool_call[1]: tool_call[2] for tool_call in tool_calls}) - - if result.usage: - increase_usage(llm_usage, result.usage) - current_llm_usage = result.usage - - if result.message and result.message.content: - if isinstance(result.message.content, list): - for content in result.message.content: - response += content.data - else: - response += result.message.content - - if not result.message.content: - result.message.content = "" - - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - yield LLMResultChunk( - model=model_instance.model_name, - prompt_messages=result.prompt_messages, - system_fingerprint=result.system_fingerprint, - delta=LLMResultChunkDelta( - index=0, - message=result.message, - usage=result.usage, - ), - ) - - assistant_message = AssistantPromptMessage(content=response, tool_calls=[]) - if tool_calls: - assistant_message.tool_calls = [ - AssistantPromptMessage.ToolCall( - id=tool_call[0], - type="function", - function=AssistantPromptMessage.ToolCall.ToolCallFunction( - name=tool_call[1], arguments=json.dumps(tool_call[2], ensure_ascii=False) - ), - ) - for tool_call in tool_calls - ] - - self._current_thoughts.append(assistant_message) - - # save thought - self.save_agent_thought( - agent_thought_id=agent_thought_id, - tool_name=tool_call_names, - tool_input=tool_call_inputs, - thought=response, - tool_invoke_meta=None, - observation=None, - answer=response, - messages_ids=[], - llm_usage=current_llm_usage, - ) - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - final_answer += response + "\n" - - # Check if max iteration is reached and model still wants to call tools - if iteration_step == max_iteration_steps and tool_calls: - raise AgentMaxIterationError(app_config.agent.max_iteration) - - # call tools - tool_responses = [] - for tool_call_id, tool_call_name, tool_call_args in tool_calls: - tool_instance = tool_instances.get(tool_call_name) - if not tool_instance: - tool_response = { - "tool_call_id": tool_call_id, - "tool_call_name": tool_call_name, - "tool_response": f"there is not a tool named {tool_call_name}", - "meta": ToolInvokeMeta.error_instance(f"there is not a tool named {tool_call_name}").to_dict(), - } - else: - # invoke tool - with use_credit_usage_metadata({"app_type": CreditUsageAppType.AGENT}): - tool_invoke_response, message_files, tool_invoke_meta = ToolEngine.agent_invoke( - session=session, - tool=tool_instance, - tool_parameters=tool_call_args, - user_id=self.user_id, - tenant_id=self.tenant_id, - message=self.message, - invoke_from=self.application_generate_entity.invoke_from, - agent_tool_callback=self.agent_callback, - trace_manager=trace_manager, - app_id=self.application_generate_entity.app_config.app_id, - message_id=self.message.id, - conversation_id=self.conversation.id, - ) - session.commit() - session.close() - # publish files - for message_file_id in message_files: - # publish message file - self.queue_manager.publish( - QueueMessageFileEvent(message_file_id=message_file_id), PublishFrom.APPLICATION_MANAGER - ) - # add message file ids - message_file_ids.append(message_file_id) - - tool_response = { - "tool_call_id": tool_call_id, - "tool_call_name": tool_call_name, - "tool_response": tool_invoke_response, - "meta": tool_invoke_meta.to_dict(), - } - - tool_responses.append(tool_response) - if tool_response["tool_response"] is not None: - tool_response_text = str(tool_response["tool_response"]) - dataset_image_contents = self._build_dataset_tool_image_contents( - session=session, - tool_response=tool_response_text, - tool_instance=tool_instance, - ) - self._current_thoughts.append( - ToolPromptMessage( - content=tool_response_text, - tool_call_id=tool_call_id, - name=tool_call_name, - ) - ) - if dataset_image_contents: - self._current_thoughts.append( - UserPromptMessage( - name=_KNOWLEDGE_RETRIEVAL_PROMPT_NAME, - content=[ - *dataset_image_contents, - TextPromptMessageContent(data=self.query or tool_response_text), - ], - ) - ) - - if len(tool_responses) > 0: - # save agent thought - self.save_agent_thought( - agent_thought_id=agent_thought_id, - tool_name="", - tool_input="", - thought="", - tool_invoke_meta={ - tool_response["tool_call_name"]: tool_response["meta"] for tool_response in tool_responses - }, - observation={ - tool_response["tool_call_name"]: tool_response["tool_response"] - for tool_response in tool_responses - }, - answer="", - messages_ids=message_file_ids, - ) - self.queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=agent_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - # update prompt tool - for prompt_tool in prompt_messages_tools: - tool_instance = tool_instances.get(prompt_tool.name) - if tool_instance: - self.update_prompt_message_tool(tool_instance, prompt_tool) - - iteration_step += 1 - - # publish end event - self.queue_manager.publish( - QueueMessageEndEvent( - llm_result=LLMResult( - model=model_instance.model_name, - prompt_messages=prompt_messages, - message=AssistantPromptMessage(content=final_answer), - usage=llm_usage["usage"] or LLMUsage.empty_usage(), - system_fingerprint="", - ) - ), - PublishFrom.APPLICATION_MANAGER, - ) - - def check_tool_calls(self, llm_result_chunk: LLMResultChunk) -> bool: - """ - Check if there is any tool call in llm result chunk - """ - return bool(llm_result_chunk.delta.message.tool_calls) - - def check_blocking_tool_calls(self, llm_result: LLMResult) -> bool: - """ - Check if there is any blocking tool call in llm result - """ - return bool(llm_result.message.tool_calls) - - def extract_tool_calls(self, llm_result_chunk: LLMResultChunk) -> list[tuple[str, str, dict[str, Any]]]: - """ - Extract tool calls from llm result chunk - - Returns: - List[Tuple[str, str, Dict[str, Any]]]: [(tool_call_id, tool_call_name, tool_call_args)] - """ - tool_calls = [] - for prompt_message in llm_result_chunk.delta.message.tool_calls: - args = {} - if prompt_message.function.arguments != "": - args = json.loads(prompt_message.function.arguments) - - tool_calls.append( - ( - prompt_message.id, - prompt_message.function.name, - args, - ) - ) - - return tool_calls - - def extract_blocking_tool_calls(self, llm_result: LLMResult) -> list[tuple[str, str, dict[str, Any]]]: - """ - Extract blocking tool calls from llm result - - Returns: - List[Tuple[str, str, Dict[str, Any]]]: [(tool_call_id, tool_call_name, tool_call_args)] - """ - tool_calls = [] - for prompt_message in llm_result.message.tool_calls: - args = {} - if prompt_message.function.arguments != "": - args = json.loads(prompt_message.function.arguments) - - tool_calls.append( - ( - prompt_message.id, - prompt_message.function.name, - args, - ) - ) - - return tool_calls - - def _init_system_message(self, prompt_template: str, prompt_messages: list[PromptMessage]) -> list[PromptMessage]: - """ - Initialize system message - """ - if not prompt_messages and prompt_template: - return [ - SystemPromptMessage(content=prompt_template), - ] - - if prompt_messages and not isinstance(prompt_messages[0], SystemPromptMessage) and prompt_template: - prompt_messages.insert(0, SystemPromptMessage(content=prompt_template)) - - return prompt_messages or [] - - def _organize_user_query(self, query: str, prompt_messages: list[PromptMessage]) -> list[PromptMessage]: - """ - Organize user query - """ - if self.files: - # get image detail config - image_detail_config = ( - self.application_generate_entity.file_upload_config.image_config.detail - if ( - self.application_generate_entity.file_upload_config - and self.application_generate_entity.file_upload_config.image_config - ) - else None - ) - image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW - - prompt_message_contents: list[PromptMessageContentUnionTypes] = [] - for file in self.files: - prompt_message_contents.append( - file_manager.to_prompt_message_content( - file, - image_detail_config=image_detail_config, - ) - ) - prompt_message_contents.append(TextPromptMessageContent(data=query)) - - prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) - else: - prompt_messages.append(UserPromptMessage(content=query)) - - return prompt_messages - - def _clear_user_prompt_image_messages(self, prompt_messages: list[PromptMessage]) -> list[PromptMessage]: - """ - As for now, gpt supports both fc and vision at the first iteration. - We need to remove the image messages from the prompt messages at the first iteration. - """ - prompt_messages = deepcopy(prompt_messages) - - for prompt_message in prompt_messages: - if isinstance(prompt_message, UserPromptMessage): - if prompt_message.name == _KNOWLEDGE_RETRIEVAL_PROMPT_NAME: - continue - if isinstance(prompt_message.content, list): - prompt_message.content = "\n".join( - [ - content.data - if content.type == PromptMessageContentType.TEXT - else "[image]" - if content.type == PromptMessageContentType.IMAGE - else "[file]" - for content in prompt_message.content - ] - ) - - return prompt_messages - - def _organize_prompt_messages(self): - prompt_template = self.app_config.prompt_template.simple_prompt_template or "" - self.history_prompt_messages = self._init_system_message(prompt_template, self.history_prompt_messages) - query_prompt_messages = self._organize_user_query(self.query or "", []) - - self.history_prompt_messages = AgentHistoryPromptTransform( - model_config=self.model_config, - prompt_messages=[*query_prompt_messages, *self._current_thoughts], - history_messages=self.history_prompt_messages, - memory=self.memory, - ).get_prompt() - - prompt_messages = [*self.history_prompt_messages, *query_prompt_messages, *self._current_thoughts] - if len(self._current_thoughts) != 0: - # clear messages after the first iteration - prompt_messages = self._clear_user_prompt_image_messages(prompt_messages) - return prompt_messages diff --git a/api/core/agent/tool_configuration.py b/api/core/agent/tool_configuration.py deleted file mode 100644 index 936f5cfca8c67f..00000000000000 --- a/api/core/agent/tool_configuration.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Build agent tool configuration for responses without changing stored config.""" - -import logging -from collections.abc import Mapping -from copy import deepcopy -from typing import cast - -from pydantic import JsonValue, ValidationError - -from core.agent.entities import AgentToolEntity -from core.tools.tool_manager import ToolManager -from core.tools.utils.configuration import ToolParameterConfigurationManager - -logger = logging.getLogger(__name__) - - -def mask_agent_tool_parameters( - *, agent_mode: Mapping[str, JsonValue], app_id: str, tenant_id: str, user_id: str -) -> dict[str, JsonValue]: - """Return a response copy, exposing parameters only after masking succeeds.""" - result = deepcopy(dict(agent_mode)) - for tool in cast(list[JsonValue], result.get("tools") or []): - if not isinstance(tool, dict): - continue - unmasked_tool = dict(tool) - if "tool_parameters" in tool: - tool["tool_parameters"] = {} - if len(unmasked_tool) <= 3: - continue - try: - agent_tool = AgentToolEntity.model_validate(unmasked_tool) - except ValidationError: - continue - try: - tool_runtime = ToolManager.get_agent_tool_runtime( - tenant_id=tenant_id, app_id=app_id, agent_tool=agent_tool, user_id=user_id - ) - manager = ToolParameterConfigurationManager( - tenant_id=tenant_id, - tool_runtime=tool_runtime, - provider_name=agent_tool.provider_id, - provider_type=agent_tool.provider_type, - identity_id=f"AGENT.{app_id}", - ) - if agent_tool.tool_parameters: - parameters = manager.decrypt_tool_parameters(agent_tool.tool_parameters) - tool["tool_parameters"] = manager.mask_tool_parameters(parameters or {}) - else: - tool["tool_parameters"] = {} - except Exception: - logger.exception("Failed to mask agent tool parameters for tool %s", agent_tool.tool_name) - return result diff --git a/api/core/app/apps/advanced_chat/__init__.py b/api/core/app/apps/advanced_chat/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/core/app/apps/advanced_chat/app_config_manager.py b/api/core/app/apps/advanced_chat/app_config_manager.py deleted file mode 100644 index c21c494efe6347..00000000000000 --- a/api/core/app/apps/advanced_chat/app_config_manager.py +++ /dev/null @@ -1,90 +0,0 @@ -from core.app.app_config.base_app_config_manager import BaseAppConfigManager -from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.app_config.features.file_upload.manager import FileUploadConfigManager -from core.app.app_config.features.opening_statement.manager import OpeningStatementConfigManager -from core.app.app_config.features.retrieval_resource.manager import RetrievalResourceConfigManager -from core.app.app_config.features.speech_to_text.manager import SpeechToTextConfigManager -from core.app.app_config.features.suggested_questions_after_answer.manager import ( - SuggestedQuestionsAfterAnswerConfigManager, -) -from core.app.app_config.features.text_to_speech.manager import TextToSpeechConfigManager -from core.app.app_config.workflow_ui_based_app.variables.manager import WorkflowVariablesConfigManager -from models.model import App, AppMode -from models.workflow import Workflow - - -class AdvancedChatAppConfig(WorkflowUIBasedAppConfig): - """ - Advanced Chatbot App Config Entity. - """ - - pass - - -class AdvancedChatAppConfigManager(BaseAppConfigManager): - @classmethod - def get_app_config(cls, app_model: App, workflow: Workflow) -> AdvancedChatAppConfig: - features_dict = workflow.features_dict - app_mode = AppMode.value_of(app_model.mode) - app_config = AdvancedChatAppConfig( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - app_mode=app_mode, - workflow_id=workflow.id, - sensitive_word_avoidance=SensitiveWordAvoidanceConfigManager.convert(config=features_dict), - variables=WorkflowVariablesConfigManager.convert(workflow=workflow), - additional_features=cls.convert_features(features_dict, app_mode), - ) - - return app_config - - @classmethod - def config_validate(cls, tenant_id: str, config: dict, only_structure_validate: bool = False): - """ - Validate for advanced chat app model config - - :param tenant_id: tenant id - :param config: app model config args - :param only_structure_validate: if True, only structure validation will be performed - """ - related_config_keys = [] - - # file upload validation - config, current_related_config_keys = FileUploadConfigManager.validate_and_set_defaults(config=config) - related_config_keys.extend(current_related_config_keys) - - # opening_statement - config, current_related_config_keys = OpeningStatementConfigManager.validate_and_set_defaults(config) - related_config_keys.extend(current_related_config_keys) - - # suggested_questions_after_answer - config, current_related_config_keys = SuggestedQuestionsAfterAnswerConfigManager.validate_and_set_defaults( - config - ) - related_config_keys.extend(current_related_config_keys) - - # speech_to_text - config, current_related_config_keys = SpeechToTextConfigManager.validate_and_set_defaults(config) - related_config_keys.extend(current_related_config_keys) - - # text_to_speech - config, current_related_config_keys = TextToSpeechConfigManager.validate_and_set_defaults(config) - related_config_keys.extend(current_related_config_keys) - - # return retriever resource - config, current_related_config_keys = RetrievalResourceConfigManager.validate_and_set_defaults(config) - related_config_keys.extend(current_related_config_keys) - - # moderation validation - config, current_related_config_keys = SensitiveWordAvoidanceConfigManager.validate_and_set_defaults( - tenant_id=tenant_id, config=config, only_structure_validate=only_structure_validate - ) - related_config_keys.extend(current_related_config_keys) - - related_config_keys = list(set(related_config_keys)) - - # Filter out extra parameters - filtered_config = {key: config.get(key) for key in related_config_keys} - - return filtered_config diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py deleted file mode 100644 index 65e1b4f6d77016..00000000000000 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility import for the chatflow generation service adapter.""" - -from services.workflow.execution.adapters.chatflow.app_generator import AdvancedChatAppGenerator - -__all__ = ["AdvancedChatAppGenerator"] diff --git a/api/core/app/apps/advanced_chat/app_runner.py b/api/core/app/apps/advanced_chat/app_runner.py deleted file mode 100644 index 617762dab5e11a..00000000000000 --- a/api/core/app/apps/advanced_chat/app_runner.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility import for the workflow execution service runner.""" - -from services.workflow.execution.adapters.chatflow.app_runner import AdvancedChatAppRunner - -__all__ = ["AdvancedChatAppRunner"] diff --git a/api/core/app/apps/advanced_chat/generate_response_converter.py b/api/core/app/apps/advanced_chat/generate_response_converter.py deleted file mode 100644 index c24be7bd03374d..00000000000000 --- a/api/core/app/apps/advanced_chat/generate_response_converter.py +++ /dev/null @@ -1,151 +0,0 @@ -from collections.abc import Generator -from typing import Any, cast, override - -from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter -from core.app.entities.task_entities import ( - AdvancedChatPausedBlockingResponse, - AppStreamResponse, - ChatbotAppBlockingResponse, - ChatbotAppStreamResponse, - ErrorStreamResponse, - MessageEndStreamResponse, - NodeFinishStreamResponse, - NodeStartStreamResponse, - PingStreamResponse, - StreamEvent, -) - - -class AdvancedChatAppGenerateResponseConverter( - AppGenerateResponseConverter[ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse] -): - @classmethod - @override - def convert_blocking_full_response( - cls, blocking_response: ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse - ) -> dict[str, Any]: - """ - Convert blocking full response. - :param blocking_response: blocking response - :return: - """ - if isinstance(blocking_response, AdvancedChatPausedBlockingResponse): - paused_data = blocking_response.data.model_dump(mode="json") - return { - "event": StreamEvent.WORKFLOW_PAUSED.value, - "task_id": blocking_response.task_id, - "id": blocking_response.data.id, - "message_id": blocking_response.data.message_id, - "conversation_id": blocking_response.data.conversation_id, - "mode": blocking_response.data.mode, - "answer": blocking_response.data.answer, - "metadata": blocking_response.data.metadata, - "created_at": blocking_response.data.created_at, - "workflow_run_id": blocking_response.data.workflow_run_id, - "data": paused_data, - } - - response = { - "event": StreamEvent.MESSAGE.value, - "task_id": blocking_response.task_id, - "id": blocking_response.data.id, - "message_id": blocking_response.data.message_id, - "conversation_id": blocking_response.data.conversation_id, - "mode": blocking_response.data.mode, - "answer": blocking_response.data.answer, - "metadata": blocking_response.data.metadata, - "created_at": blocking_response.data.created_at, - } - - return response - - @classmethod - @override - def convert_blocking_simple_response( - cls, blocking_response: ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse - ) -> dict[str, Any]: - """ - Convert blocking simple response. - :param blocking_response: blocking response - :return: - """ - response = cls.convert_blocking_full_response(blocking_response) - - metadata = response.get("metadata", {}) - if isinstance(metadata, dict): - response["metadata"] = cls._get_simple_metadata(metadata) - - return response - - @classmethod - @override - def convert_stream_full_response( - cls, stream_response: Generator[AppStreamResponse] - ) -> Generator[dict[str, Any] | str, Any]: - """ - Convert stream full response. - :param stream_response: stream response - :return: - """ - for chunk in stream_response: - chunk = cast(ChatbotAppStreamResponse, chunk) - sub_stream_response = chunk.stream_response - - if isinstance(sub_stream_response, PingStreamResponse): - yield "ping" - continue - - response_chunk: dict[str, Any] = { - "event": sub_stream_response.event.value, - "conversation_id": chunk.conversation_id, - "message_id": chunk.message_id, - "created_at": chunk.created_at, - } - - if isinstance(sub_stream_response, ErrorStreamResponse): - data = cls._error_to_stream_response(sub_stream_response.err) - response_chunk.update(data) - else: - response_chunk.update(sub_stream_response.model_dump(mode="json")) - yield response_chunk - - @classmethod - @override - def convert_stream_simple_response( - cls, stream_response: Generator[AppStreamResponse] - ) -> Generator[dict[str, Any] | str, Any]: - """ - Convert stream simple response. - :param stream_response: stream response - :return: - """ - for chunk in stream_response: - chunk = cast(ChatbotAppStreamResponse, chunk) - sub_stream_response = chunk.stream_response - - if isinstance(sub_stream_response, PingStreamResponse): - yield "ping" - continue - - response_chunk: dict[str, Any] = { - "event": sub_stream_response.event.value, - "conversation_id": chunk.conversation_id, - "message_id": chunk.message_id, - "created_at": chunk.created_at, - } - - match sub_stream_response: - case MessageEndStreamResponse(): - sub_stream_response_dict = sub_stream_response.model_dump(mode="json") - metadata = sub_stream_response_dict.get("metadata", {}) - sub_stream_response_dict["metadata"] = cls._get_simple_metadata(metadata) - response_chunk.update(sub_stream_response_dict) - case ErrorStreamResponse(): - data = cls._error_to_stream_response(sub_stream_response.err) - response_chunk.update(data) - case NodeStartStreamResponse() | NodeFinishStreamResponse(): - response_chunk.update(sub_stream_response.to_ignore_detail_dict()) - case _: - response_chunk.update(sub_stream_response.model_dump(mode="json")) - - yield response_chunk diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py deleted file mode 100644 index 3d9581f06e705d..00000000000000 --- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py +++ /dev/null @@ -1,1160 +0,0 @@ -import json -import logging -import time -from collections.abc import Callable, Generator, Mapping -from contextlib import contextmanager -from dataclasses import dataclass -from datetime import datetime -from threading import Thread -from typing import Any - -from sqlalchemy import select, update -from sqlalchemy.orm import Session - -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.common.graph_runtime_state_support import GraphRuntimeStateSupport -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.apps.draft_variable_saver import DraftVariableSaverFactory -from core.app.entities.app_invoke_entities import ( - AdvancedChatAppGenerateEntity, - InvokeFrom, - get_credit_usage_app_type, -) -from core.app.entities.queue_entities import ( - MessageQueueMessage, - QueueAdvancedChatMessageEndEvent, - QueueAgentLogEvent, - QueueAnnotationReplyEvent, - QueueErrorEvent, - QueueHumanInputFormFilledEvent, - QueueHumanInputFormTimeoutEvent, - QueueIterationCompletedEvent, - QueueIterationNextEvent, - QueueIterationStartEvent, - QueueLoopCompletedEvent, - QueueLoopNextEvent, - QueueLoopStartEvent, - QueueMessageReplaceEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueueNodeRetryEvent, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, - QueuePingEvent, - QueueReasoningChunkEvent, - QueueRetrieverResourcesEvent, - QueueStopEvent, - QueueTextChunkEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPartialSuccessEvent, - QueueWorkflowPausedEvent, - QueueWorkflowStartedEvent, - QueueWorkflowSucceededEvent, - WorkflowQueueMessage, -) -from core.app.entities.task_entities import ( - AdvancedChatPausedBlockingResponse, - ChatbotAppBlockingResponse, - ChatbotAppStreamResponse, - ErrorStreamResponse, - HumanInputRequiredPauseReasonPayload, - HumanInputRequiredResponse, - MessageAudioEndStreamResponse, - MessageAudioStreamResponse, - MessageEndStreamResponse, - PingStreamResponse, - ReasoningChunkStreamResponse, - StreamResponse, - WorkflowPauseStreamResponse, - WorkflowTaskState, -) -from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline -from core.app.task_pipeline.message_cycle_manager import MessageCycleManager -from core.base.tts import AppGeneratorTTSPublisher -from core.db.session_factory import session_factory -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.workflow.file_reference import resolve_file_record_id -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from core.workflow.system_variables import build_system_variables -from graphon.enums import WorkflowExecutionStatus -from graphon.model_runtime.entities.llm_entities import LLMUsage -from graphon.model_runtime.utils.encoders import jsonable_encoder -from graphon.nodes import BuiltinNodeTypes -from graphon.runtime import RuntimeState -from libs.datetime_utils import naive_utc_now -from models import Account, Conversation, EndUser, Message, MessageFile -from models.enums import CreatorUserRole, MessageFileBelongsTo, MessageStatus -from models.execution_extra_content import HumanInputContent -from models.model import AppMode -from models.workflow import Workflow, WorkflowRun - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True, slots=True) -class WorkflowSnapshot: - id: str - tenant_id: str - features_dict: Mapping[str, Any] - - @classmethod - def from_workflow(cls, workflow: Workflow) -> "WorkflowSnapshot": - return cls( - id=workflow.id, - tenant_id=workflow.tenant_id, - features_dict=dict(workflow.features_dict), - ) - - -@dataclass(frozen=True, slots=True) -class ConversationSnapshot: - id: str - mode: AppMode - - @classmethod - def from_conversation(cls, conversation: Conversation) -> "ConversationSnapshot": - return cls( - id=conversation.id, - mode=conversation.mode, - ) - - -@dataclass(frozen=True, slots=True) -class MessageSnapshot: - id: str - query: str - created_at: datetime - status: MessageStatus - answer: str - - @classmethod - def from_message(cls, message: Message) -> "MessageSnapshot": - return cls( - id=message.id, - query=message.query, - created_at=message.created_at, - status=message.status, - answer=message.answer, - ) - - -class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport): - """ - AdvancedChatAppGenerateTaskPipeline is a class that generate stream output and state management for Application. - """ - - def __init__( - self, - application_generate_entity: AdvancedChatAppGenerateEntity, - workflow: WorkflowSnapshot, - queue_manager: AppQueueManager, - conversation: ConversationSnapshot, - message: MessageSnapshot, - user: Account | EndUser, - stream: bool, - dialogue_count: int, - draft_var_saver_factory: DraftVariableSaverFactory, - ): - self._base_task_pipeline = BasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - stream=stream, - ) - - match user: - case EndUser(): - self._user_id = user.id - user_session_id = user.session_id - self._created_by_role = CreatorUserRole.END_USER - case Account(): - self._user_id = user.id - user_session_id = user.id - self._created_by_role = CreatorUserRole.ACCOUNT - case _: - raise NotImplementedError(f"User type not supported: {type(user)}") - - self._workflow_system_variables = build_system_variables( - query=message.query, - files=application_generate_entity.files, - conversation_id=conversation.id, - user_id=user_session_id, - dialogue_count=dialogue_count, - app_id=application_generate_entity.app_config.app_id, - workflow_id=workflow.id, - workflow_execution_id=application_generate_entity.workflow_run_id, - ) - self._workflow_response_converter = WorkflowResponseConverter( - application_generate_entity=application_generate_entity, - user=user, - system_variables=self._workflow_system_variables, - ) - - self._task_state = WorkflowTaskState() - self._seed_task_state_from_message(message) - self._message_cycle_manager = MessageCycleManager( - application_generate_entity=application_generate_entity, task_state=self._task_state - ) - - self._application_generate_entity = application_generate_entity - self._workflow_id = workflow.id - self._workflow_features_dict = workflow.features_dict - self._workflow_tenant_id = workflow.tenant_id - self._conversation_id = conversation.id - self._conversation_mode = conversation.mode - self._message_id = message.id - self._message_created_at = int(message.created_at.timestamp()) - self._conversation_name_generate_thread: Thread | None = None - self._recorded_files: list[Mapping[str, Any]] = [] - self._workflow_run_id: str = "" - self._draft_var_saver_factory = draft_var_saver_factory - self._graph_runtime_state: RuntimeState | None = None - self._message_saved_on_pause = False - self._seed_graph_runtime_state_from_queue_manager() - - def _seed_task_state_from_message(self, message: MessageSnapshot) -> None: - if message.status == MessageStatus.PAUSED and message.answer: - self._task_state.answer = message.answer - - def process( - self, - ) -> ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse | Generator[ChatbotAppStreamResponse]: - """ - Process generate task pipeline. - :return: - """ - self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name( - conversation_id=self._conversation_id, - query=self._application_generate_entity.query, - message_id=self._message_id, - ) - - generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) - - if self._base_task_pipeline.stream: - return self._to_stream_response(generator) - else: - return self._to_blocking_response(generator) - - def _to_blocking_response( - self, generator: Generator[StreamResponse] - ) -> ChatbotAppBlockingResponse | AdvancedChatPausedBlockingResponse: - """ - Process blocking response. - :return: - """ - human_input_responses: list[HumanInputRequiredResponse] = [] - for stream_response in generator: - match stream_response: - case ErrorStreamResponse(): - raise stream_response.err - case HumanInputRequiredResponse(): - human_input_responses.append(stream_response) - case WorkflowPauseStreamResponse(): - return AdvancedChatPausedBlockingResponse( - task_id=stream_response.task_id, - data=AdvancedChatPausedBlockingResponse.Data( - id=self._message_id, - mode=self._conversation_mode, - conversation_id=self._conversation_id, - message_id=self._message_id, - workflow_run_id=stream_response.data.workflow_run_id, - answer=self._task_state.answer, - metadata=self._message_end_to_stream_response().metadata, - created_at=self._message_created_at, - paused_nodes=stream_response.data.paused_nodes, - reasons=stream_response.data.reasons, - status=stream_response.data.status, - elapsed_time=stream_response.data.elapsed_time, - total_tokens=stream_response.data.total_tokens, - total_steps=stream_response.data.total_steps, - ), - ) - case MessageEndStreamResponse(): - extras = {} - if stream_response.metadata: - extras["metadata"] = stream_response.metadata - - return ChatbotAppBlockingResponse( - task_id=stream_response.task_id, - data=ChatbotAppBlockingResponse.Data( - id=self._message_id, - mode=self._conversation_mode, - conversation_id=self._conversation_id, - message_id=self._message_id, - answer=self._task_state.answer, - created_at=self._message_created_at, - **extras, - ), - ) - case _: - continue - - if human_input_responses: - return self._build_paused_blocking_response_from_human_input(human_input_responses) - - raise ValueError("queue listening stopped unexpectedly.") - - def _build_paused_blocking_response_from_human_input( - self, human_input_responses: list[HumanInputRequiredResponse] - ) -> AdvancedChatPausedBlockingResponse: - runtime_state = self._resolve_graph_runtime_state() - paused_nodes = list(dict.fromkeys(response.data.node_id for response in human_input_responses)) - reasons = [ - HumanInputRequiredPauseReasonPayload.from_response_data(response.data).model_dump(mode="json") - for response in human_input_responses - ] - - return AdvancedChatPausedBlockingResponse( - task_id=self._application_generate_entity.task_id, - data=AdvancedChatPausedBlockingResponse.Data( - id=self._message_id, - mode=self._conversation_mode, - conversation_id=self._conversation_id, - message_id=self._message_id, - workflow_run_id=human_input_responses[-1].workflow_run_id, - answer=self._task_state.answer, - metadata=self._message_end_to_stream_response().metadata, - created_at=self._message_created_at, - paused_nodes=paused_nodes, - reasons=reasons, - status=WorkflowExecutionStatus.PAUSED, - elapsed_time=self._application_generate_entity.workflow_execution_timing.finish(), - total_tokens=runtime_state.total_tokens, - total_steps=runtime_state.node_run_steps, - ), - ) - - def _to_stream_response(self, generator: Generator[StreamResponse]) -> Generator[ChatbotAppStreamResponse, Any]: - """ - To stream response. - :return: - """ - for stream_response in generator: - yield ChatbotAppStreamResponse( - conversation_id=self._conversation_id, - message_id=self._message_id, - created_at=self._message_created_at, - stream_response=stream_response, - ) - - def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: str): - if not publisher: - return None - audio_msg = publisher.check_and_get_audio() - if audio_msg is None: - return None - if audio_msg.status == "responding": - return MessageAudioStreamResponse(audio=audio_msg.audio, audio_type=audio_msg.audio_type, task_id=task_id) - if audio_msg.status in {"finish", "error"}: - return None - raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}") - - def _wrapper_process_stream_response( - self, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: - tts_publisher = None - task_id = self._application_generate_entity.task_id - tenant_id = self._application_generate_entity.app_config.tenant_id - features_dict = self._workflow_features_dict - - if ( - self._base_task_pipeline.stream - and features_dict.get("text_to_speech") - and features_dict["text_to_speech"].get("enabled") - and features_dict["text_to_speech"].get("autoPlay") == "enabled" - ): - tts_publisher = AppGeneratorTTSPublisher( - tenant_id, - features_dict["text_to_speech"].get("voice"), - features_dict["text_to_speech"].get("language"), - get_credit_usage_app_type(self._application_generate_entity.app_config.app_mode), - ) - - try: - for response in self._process_stream_response(tts_publisher=tts_publisher, trace_manager=trace_manager): - while audio_response := self._listen_audio_msg(publisher=tts_publisher, task_id=task_id): - yield audio_response - if tts_publisher and isinstance(response, ErrorStreamResponse): - tts_publisher.cancel() - yield MessageAudioEndStreamResponse(audio="", task_id=task_id) - yield response - return - yield response - - if tts_publisher is None: - return - - tts_publisher.publish(None) - while True: - audio_trunk = tts_publisher.check_and_get_audio(block=True) - assert audio_trunk is not None - if audio_trunk.status == "responding": - yield MessageAudioStreamResponse( - audio=audio_trunk.audio, audio_type=audio_trunk.audio_type, task_id=task_id - ) - continue - if audio_trunk.status not in {"finish", "error"}: - raise RuntimeError(f"TTS publisher returned an unknown status: {audio_trunk.status}") - - yield MessageAudioEndStreamResponse(audio="", task_id=task_id) - if audio_trunk.status == "error": - if audio_trunk.error is None: - raise RuntimeError("TTS publisher returned an error terminal without an exception") - yield ErrorStreamResponse(err=audio_trunk.error, task_id=task_id) - return - finally: - if tts_publisher: - tts_publisher.cancel() - - @contextmanager - def _database_session(self): - """Context manager for database sessions.""" - with session_factory.create_session() as session: - try: - yield session - session.commit() - except Exception: - session.rollback() - raise - - def _ensure_workflow_initialized(self): - """Fluent validation for workflow state.""" - if not self._workflow_run_id: - raise ValueError("workflow run not initialized.") - - def _handle_ping_event(self, event: QueuePingEvent, **kwargs) -> Generator[PingStreamResponse]: - """Handle ping events.""" - yield self._base_task_pipeline.ping_stream_response() - - def _handle_error_event(self, event: QueueErrorEvent, **kwargs) -> Generator[ErrorStreamResponse]: - """Handle error events.""" - with self._database_session() as session: - err = self._base_task_pipeline.handle_error(event=event, session=session, message_id=self._message_id) - yield self._base_task_pipeline.error_to_stream_response(err) - - def _handle_workflow_started_event( - self, - event: QueueWorkflowStartedEvent, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow started events.""" - runtime_state = self._resolve_graph_runtime_state() - run_id = self._extract_workflow_run_id(runtime_state) - self._workflow_run_id = run_id - - with self._database_session() as session: - session.execute(update(Message).where(Message.id == self._message_id).values(workflow_run_id=run_id)) - - workflow_start_resp = self._workflow_response_converter.workflow_start_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_run_id=run_id, - workflow_id=self._workflow_id, - reason=event.reason, - node_execution_snapshots=event.node_execution_snapshots, - ) - - yield workflow_start_resp - - def _handle_node_retry_event(self, event: QueueNodeRetryEvent, **kwargs) -> Generator[StreamResponse]: - """Handle node retry events.""" - self._ensure_workflow_initialized() - - node_retry_resp = self._workflow_response_converter.workflow_node_retry_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - if node_retry_resp: - yield node_retry_resp - - def _handle_node_started_event(self, event: QueueNodeStartedEvent, **kwargs) -> Generator[StreamResponse]: - """Handle node started events.""" - self._ensure_workflow_initialized() - - node_start_resp = self._workflow_response_converter.workflow_node_start_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - if node_start_resp: - yield node_start_resp - - def _handle_node_succeeded_event(self, event: QueueNodeSucceededEvent, **kwargs) -> Generator[StreamResponse]: - """Handle node succeeded events.""" - # Record files if it's an answer node or end node - if event.node_type in [BuiltinNodeTypes.ANSWER, BuiltinNodeTypes.END, BuiltinNodeTypes.LLM]: - self._recorded_files.extend( - self._workflow_response_converter.fetch_files_from_node_outputs(event.outputs or {}) - ) - - # Collect terminal reasoning (separated mode) per LLM node id for persistence. This is the - # authoritative source (outputs.reasoning_content), decoupled from the live delta stream. - # Accumulate across iteration/loop passes (same node_id) to match the live stream, which - # appends every pass under the same key — overwriting would keep only the last pass. - if event.node_type == BuiltinNodeTypes.LLM: - reasoning_content = (event.outputs or {}).get("reasoning_content") - if isinstance(reasoning_content, str) and reasoning_content: - self._task_state.metadata.reasoning[event.node_id] = ( - self._task_state.metadata.reasoning.get(event.node_id, "") + reasoning_content - ) - - node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - self._save_output_for_event(event, event.node_execution_id) - - if node_finish_resp: - yield node_finish_resp - - def _handle_node_failed_events( - self, - event: QueueNodeFailedEvent | QueueNodeExceptionEvent, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle various node failure events.""" - node_finish_resp = self._workflow_response_converter.workflow_node_finish_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - if isinstance(event, QueueNodeExceptionEvent): - self._save_output_for_event(event, event.node_execution_id) - - if node_finish_resp: - yield node_finish_resp - - def _handle_text_chunk_event( - self, - event: QueueTextChunkEvent, - *, - tts_publisher: AppGeneratorTTSPublisher | None = None, - queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle text chunk events.""" - delta_text = event.text - if delta_text is None: - return - - # Handle output moderation chunk - should_direct_answer = self._handle_output_moderation_chunk(delta_text) - if should_direct_answer: - return - - current_time = time.perf_counter() - if self._task_state.first_token_time is None and delta_text.strip(): - self._task_state.first_token_time = current_time - self._task_state.is_streaming_response = True - - if delta_text.strip(): - self._task_state.last_token_time = current_time - - # Only publish tts message at text chunk streaming - if tts_publisher and queue_message: - tts_publisher.publish(queue_message) - - self._task_state.answer += delta_text - yield self._message_cycle_manager.message_to_stream_response( - answer=delta_text, message_id=self._message_id, from_variable_selector=event.from_variable_selector - ) - - def _handle_reasoning_chunk_event(self, event: QueueReasoningChunkEvent, **kwargs) -> Generator[StreamResponse]: - """Handle out-of-band reasoning chunk events. - - Pure emit: reasoning is streamed on its own channel and never written to the - answer. The terminal marker (is_final) may carry an empty reasoning string, in - which case it is still forwarded as the "thinking finished" signal. - """ - if not event.reasoning and not event.is_final: - return - yield ReasoningChunkStreamResponse( - task_id=self._application_generate_entity.task_id, - data=ReasoningChunkStreamResponse.Data( - message_id=self._message_id, - reasoning=event.reasoning, - node_id=event.from_node_id, - is_final=event.is_final, - ), - ) - - def _handle_iteration_start_event(self, event: QueueIterationStartEvent, **kwargs) -> Generator[StreamResponse]: - """Handle iteration start events.""" - self._ensure_workflow_initialized() - - iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_run_id, - event=event, - ) - yield iter_start_resp - - def _handle_iteration_next_event(self, event: QueueIterationNextEvent, **kwargs) -> Generator[StreamResponse]: - """Handle iteration next events.""" - self._ensure_workflow_initialized() - - iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_run_id, - event=event, - ) - yield iter_next_resp - - def _handle_iteration_completed_event( - self, event: QueueIterationCompletedEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle iteration completed events.""" - self._ensure_workflow_initialized() - - iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_run_id, - event=event, - ) - yield iter_finish_resp - - def _handle_loop_start_event(self, event: QueueLoopStartEvent, **kwargs) -> Generator[StreamResponse]: - """Handle loop start events.""" - self._ensure_workflow_initialized() - - loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_run_id, - event=event, - ) - yield loop_start_resp - - def _handle_loop_next_event(self, event: QueueLoopNextEvent, **kwargs) -> Generator[StreamResponse]: - """Handle loop next events.""" - self._ensure_workflow_initialized() - - loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_run_id, - event=event, - ) - yield loop_next_resp - - def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) -> Generator[StreamResponse]: - """Handle loop completed events.""" - self._ensure_workflow_initialized() - - loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_run_id, - event=event, - ) - yield loop_finish_resp - - def _handle_workflow_succeeded_event( - self, - event: QueueWorkflowSucceededEvent, - *, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow succeeded events.""" - _ = trace_manager - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow_id, - status=WorkflowExecutionStatus.SUCCEEDED, - graph_runtime_state=validated_state, - ) - - yield from self._handle_advanced_chat_message_end_event( - QueueAdvancedChatMessageEndEvent(), graph_runtime_state=validated_state - ) - yield workflow_finish_resp - - def _handle_workflow_partial_success_event( - self, - event: QueueWorkflowPartialSuccessEvent, - *, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow partial success events.""" - _ = trace_manager - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow_id, - status=WorkflowExecutionStatus.PARTIAL_SUCCEEDED, - graph_runtime_state=validated_state, - exceptions_count=event.exceptions_count, - ) - - yield from self._handle_advanced_chat_message_end_event( - QueueAdvancedChatMessageEndEvent(), graph_runtime_state=validated_state - ) - yield workflow_finish_resp - - def _handle_workflow_paused_event( - self, - event: QueueWorkflowPausedEvent, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow paused events.""" - validated_state = self._ensure_graph_runtime_initialized() - responses = self._workflow_response_converter.workflow_pause_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - graph_runtime_state=validated_state, - ) - for reason in event.reasons: - if isinstance(reason, HumanInputRequired): - self._persist_human_input_extra_content(form_id=reason.form_id) - yield from responses - resolved_state: RuntimeState | None = None - try: - resolved_state = self._ensure_graph_runtime_initialized() - except ValueError: - resolved_state = None - - with self._database_session() as session: - # Serialize the message's pause state with durable stop on the run. - workflow_run_status = session.scalar( - select(WorkflowRun.status) - .where( - WorkflowRun.id == self._workflow_run_id, - WorkflowRun.tenant_id == self._workflow_tenant_id, - WorkflowRun.app_id == self._application_generate_entity.app_config.app_id, - ) - .with_for_update() - ) - self._save_message(session=session, graph_runtime_state=resolved_state) - message = self._get_message(session=session) - if message is not None and workflow_run_status == WorkflowExecutionStatus.PAUSED: - message.status = MessageStatus.PAUSED - self._message_saved_on_pause = True - self._base_task_pipeline.queue_manager.publish(QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE) - - def _handle_workflow_failed_event( - self, - event: QueueWorkflowFailedEvent, - *, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow failed events.""" - _ = trace_manager - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow_id, - status=WorkflowExecutionStatus.FAILED, - graph_runtime_state=validated_state, - error=event.error, - exceptions_count=event.exceptions_count, - ) - - with self._database_session() as session: - err_event = QueueErrorEvent(error=ValueError(f"Run failed: {event.error}")) - err = self._base_task_pipeline.handle_error(event=err_event, session=session, message_id=self._message_id) - - yield workflow_finish_resp - yield self._base_task_pipeline.error_to_stream_response(err) - - def _handle_stop_event( - self, - event: QueueStopEvent, - *, - graph_runtime_state: RuntimeState | None = None, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle stop events.""" - _ = trace_manager - resolved_state = None - if self._workflow_run_id: - resolved_state = self._resolve_graph_runtime_state(graph_runtime_state) - - if self._workflow_run_id and resolved_state: - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow_id, - status=WorkflowExecutionStatus.STOPPED, - graph_runtime_state=resolved_state, - error=event.get_stop_reason(), - ) - - with self._database_session() as session: - # Save message - self._save_message(session=session, graph_runtime_state=resolved_state) - self._emit_message_trace() - - yield workflow_finish_resp - elif event.stopped_by in ( - QueueStopEvent.StopBy.INPUT_MODERATION, - QueueStopEvent.StopBy.ANNOTATION_REPLY, - ): - # When hitting input-moderation or annotation-reply, the workflow will not start - with self._database_session() as session: - # Save message - self._save_message(session=session) - self._emit_message_trace() - - yield self._message_end_to_stream_response() - - def _handle_advanced_chat_message_end_event( - self, - event: QueueAdvancedChatMessageEndEvent, - *, - graph_runtime_state: RuntimeState | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle advanced chat message end events.""" - resolved_state = self._ensure_graph_runtime_initialized(graph_runtime_state) - - output_moderation_answer = self._base_task_pipeline.handle_output_moderation_when_task_finished( - self._task_state.answer - ) - if output_moderation_answer: - self._task_state.answer = output_moderation_answer - yield self._message_cycle_manager.message_replace_to_stream_response( - answer=output_moderation_answer, - reason=QueueMessageReplaceEvent.MessageReplaceReason.OUTPUT_MODERATION, - ) - - # Save message unless it has already been persisted on pause. - if not self._message_saved_on_pause: - with self._database_session() as session: - self._save_message(session=session, graph_runtime_state=resolved_state) - self._emit_message_trace() - - yield self._message_end_to_stream_response() - - def _handle_retriever_resources_event( - self, event: QueueRetrieverResourcesEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle retriever resources events.""" - self._message_cycle_manager.handle_retriever_resources(event) - yield from () - - def _handle_annotation_reply_event(self, event: QueueAnnotationReplyEvent, **kwargs) -> Generator[StreamResponse]: - """Handle annotation reply events.""" - with self._database_session() as session: - self._message_cycle_manager.handle_annotation_reply(event, session) - yield from () - - def _handle_message_replace_event(self, event: QueueMessageReplaceEvent, **kwargs) -> Generator[StreamResponse]: - """Handle message replace events.""" - yield self._message_cycle_manager.message_replace_to_stream_response(answer=event.text, reason=event.reason) - - def _handle_human_input_form_filled_event( - self, event: QueueHumanInputFormFilledEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle human input form filled events.""" - self._persist_human_input_extra_content(form_id=event.form_id) - yield self._workflow_response_converter.human_input_form_filled_to_stream_response( - event=event, task_id=self._application_generate_entity.task_id - ) - - def _handle_human_input_form_timeout_event( - self, event: QueueHumanInputFormTimeoutEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle human input form timeout events.""" - yield self._workflow_response_converter.human_input_form_timeout_to_stream_response( - event=event, task_id=self._application_generate_entity.task_id - ) - - def _persist_human_input_extra_content(self, *, form_id: str) -> None: - if not self._workflow_run_id or not self._message_id: - return - - with self._database_session() as session: - exists_stmt = select(HumanInputContent).where( - HumanInputContent.workflow_run_id == self._workflow_run_id, - HumanInputContent.message_id == self._message_id, - HumanInputContent.form_id == form_id, - ) - if session.scalar(exists_stmt) is not None: - return - - content = HumanInputContent( - workflow_run_id=self._workflow_run_id, - message_id=self._message_id, - form_id=form_id, - ) - session.add(content) - - def _handle_agent_log_event(self, event: QueueAgentLogEvent, **kwargs) -> Generator[StreamResponse]: - """Handle agent log events.""" - yield self._workflow_response_converter.handle_agent_log( - task_id=self._application_generate_entity.task_id, event=event - ) - - def _get_event_handlers(self) -> dict[type, Callable]: - """Get mapping of event types to their handlers using fluent pattern.""" - return { - # Basic events - QueuePingEvent: self._handle_ping_event, - QueueErrorEvent: self._handle_error_event, - QueueTextChunkEvent: self._handle_text_chunk_event, - QueueReasoningChunkEvent: self._handle_reasoning_chunk_event, - # Workflow events - QueueWorkflowStartedEvent: self._handle_workflow_started_event, - QueueWorkflowSucceededEvent: self._handle_workflow_succeeded_event, - QueueWorkflowPartialSuccessEvent: self._handle_workflow_partial_success_event, - QueueWorkflowPausedEvent: self._handle_workflow_paused_event, - QueueWorkflowFailedEvent: self._handle_workflow_failed_event, - # Node events - QueueNodeRetryEvent: self._handle_node_retry_event, - QueueNodeStartedEvent: self._handle_node_started_event, - QueueNodeSucceededEvent: self._handle_node_succeeded_event, - # Iteration events - QueueIterationStartEvent: self._handle_iteration_start_event, - QueueIterationNextEvent: self._handle_iteration_next_event, - QueueIterationCompletedEvent: self._handle_iteration_completed_event, - # Loop events - QueueLoopStartEvent: self._handle_loop_start_event, - QueueLoopNextEvent: self._handle_loop_next_event, - QueueLoopCompletedEvent: self._handle_loop_completed_event, - # Control events - QueueStopEvent: self._handle_stop_event, - # Message events - QueueRetrieverResourcesEvent: self._handle_retriever_resources_event, - QueueAnnotationReplyEvent: self._handle_annotation_reply_event, - QueueMessageReplaceEvent: self._handle_message_replace_event, - QueueAdvancedChatMessageEndEvent: self._handle_advanced_chat_message_end_event, - QueueAgentLogEvent: self._handle_agent_log_event, - QueueHumanInputFormFilledEvent: self._handle_human_input_form_filled_event, - QueueHumanInputFormTimeoutEvent: self._handle_human_input_form_timeout_event, - } - - def _dispatch_event( - self, - event: Any, - *, - tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, - queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None, - ) -> Generator[StreamResponse]: - """Dispatch events using elegant pattern matching.""" - handlers = self._get_event_handlers() - event_type = type(event) - - # Direct handler lookup - if handler := handlers.get(event_type): - yield from handler( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - return - - # Handle node failure events with isinstance check - if isinstance( - event, - ( - QueueNodeFailedEvent, - QueueNodeExceptionEvent, - ), - ): - yield from self._handle_node_failed_events( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - return - - # For unhandled events, we continue (original behavior) - return - - def _process_stream_response( - self, - tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, - ) -> Generator[StreamResponse]: - """ - Process stream response using elegant Fluent Python patterns. - Maintains exact same functionality as original 57-if-statement version. - """ - for queue_message in self._base_task_pipeline.queue_manager.listen(): - event = queue_message.event - - match event: - case QueueWorkflowStartedEvent(): - self._resolve_graph_runtime_state() - yield from self._handle_workflow_started_event(event) - - case QueueErrorEvent(): - yield from self._handle_error_event(event) - break - - case QueueWorkflowFailedEvent(): - yield from self._handle_workflow_failed_event(event, trace_manager=trace_manager) - break - case QueueWorkflowPausedEvent(): - yield from self._handle_workflow_paused_event(event) - break - - case QueueWorkflowSucceededEvent(): - yield from self._handle_workflow_succeeded_event(event, trace_manager=trace_manager) - break - - case QueueWorkflowPartialSuccessEvent(): - yield from self._handle_workflow_partial_success_event(event, trace_manager=trace_manager) - break - - case QueueStopEvent(): - yield from self._handle_stop_event(event, graph_runtime_state=None, trace_manager=trace_manager) - break - - # Handle all other events through elegant dispatch - case _: - if responses := list( - self._dispatch_event( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - ): - yield from responses - - if self._conversation_name_generate_thread: - logger.debug("Conversation name generation running as daemon thread") - - def _save_message(self, *, session: Session, graph_runtime_state: RuntimeState | None = None): - message = self._get_message(session=session) - if message is None: - return - - if message.status == MessageStatus.PAUSED: - message.status = MessageStatus.NORMAL - - answer_text = self._task_state.answer - - message.answer = answer_text - message.updated_at = naive_utc_now() - message.provider_response_latency = time.perf_counter() - self._base_task_pipeline.start_at - - # Set usage first before dumping metadata - if graph_runtime_state and graph_runtime_state.llm_usage: - usage = graph_runtime_state.llm_usage - message.message_tokens = usage.prompt_tokens - message.message_unit_price = usage.prompt_unit_price - message.message_price_unit = usage.prompt_price_unit - message.answer_tokens = usage.completion_tokens - message.answer_unit_price = usage.completion_unit_price - message.answer_price_unit = usage.completion_price_unit - message.total_price = usage.total_price - message.currency = usage.currency - self._task_state.metadata.usage = usage - else: - usage = LLMUsage.empty_usage() - self._task_state.metadata.usage = usage - - # Add streaming metrics to usage if available - if self._task_state.is_streaming_response and self._task_state.first_token_time: - start_time = self._base_task_pipeline.start_at - first_token_time = self._task_state.first_token_time - last_token_time = self._task_state.last_token_time or first_token_time - usage.time_to_first_token = round(first_token_time - start_time, 3) - usage.time_to_generate = round(last_token_time - first_token_time, 3) - - metadata = self._task_state.metadata.model_dump() - message.message_metadata = json.dumps(jsonable_encoder(metadata)) - message_files: list[MessageFile] = [] - for file in self._recorded_files: - reference = file.get("reference") or file.get("related_id") - message_files.append( - MessageFile( - message_id=message.id, - type=file["type"], - transfer_method=file["transfer_method"], - url=file["remote_url"], - belongs_to=MessageFileBelongsTo.ASSISTANT, - upload_file_id=resolve_file_record_id(reference if isinstance(reference, str) else None), - created_by_role=CreatorUserRole.ACCOUNT - if message.invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} - else CreatorUserRole.END_USER, - created_by=message.from_account_id or message.from_end_user_id or "", - ) - ) - session.add_all(message_files) - - def _emit_message_trace(self) -> None: - trace_manager = self._application_generate_entity.trace_manager - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MESSAGE_TRACE, - conversation_id=self._conversation_id, - message_id=self._message_id, - trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), - ) - ) - - def _seed_graph_runtime_state_from_queue_manager(self) -> None: - """Bootstrap the cached runtime state from the queue manager when present.""" - candidate = self._base_task_pipeline.queue_manager.graph_runtime_state - if candidate is not None: - self._graph_runtime_state = candidate - - def _message_end_to_stream_response(self) -> MessageEndStreamResponse: - """ - Message end to stream response. - :return: - """ - extras = self._task_state.metadata.model_dump() - - if self._task_state.metadata.annotation_reply: - del extras["annotation_reply"] - - return MessageEndStreamResponse( - task_id=self._application_generate_entity.task_id, - id=self._message_id, - files=self._recorded_files, - metadata=extras, - ) - - def _handle_output_moderation_chunk(self, text: str) -> bool: - """ - Handle output moderation chunk. - :param text: text - :return: True if output moderation should direct output, otherwise False - """ - if self._base_task_pipeline.output_moderation_handler: - if self._base_task_pipeline.output_moderation_handler.should_direct_output(): - self._task_state.answer = self._base_task_pipeline.output_moderation_handler.get_final_output() - self._base_task_pipeline.queue_manager.publish( - QueueTextChunkEvent(text=self._task_state.answer), PublishFrom.TASK_PIPELINE - ) - - self._base_task_pipeline.queue_manager.publish( - QueueStopEvent(stopped_by=QueueStopEvent.StopBy.OUTPUT_MODERATION), PublishFrom.TASK_PIPELINE - ) - return True - else: - self._base_task_pipeline.output_moderation_handler.append_new_token(text) - - return False - - def _get_message(self, *, session: Session): - stmt = select(Message).where(Message.id == self._message_id) - message = session.scalar(stmt) - if not message: - raise ValueError(f"Message not found: {self._message_id}") - return message - - def _save_output_for_event(self, event: QueueNodeSucceededEvent | QueueNodeExceptionEvent, node_execution_id: str): - saver = self._draft_var_saver_factory( - app_id=self._application_generate_entity.app_config.app_id, - node_id=event.node_id, - node_type=event.node_type, - node_execution_id=node_execution_id, - enclosing_node_id=event.in_loop_id or event.in_iteration_id, - ) - saver.save(event.process_data, event.outputs) diff --git a/api/core/app/apps/agent_app/app_generator.py b/api/core/app/apps/agent_app/app_generator.py deleted file mode 100644 index f00fa37616d266..00000000000000 --- a/api/core/app/apps/agent_app/app_generator.py +++ /dev/null @@ -1,606 +0,0 @@ -"""Agent App generator: orchestrate Agent App chat and finalize executions. - -Agent App turns mirror the agent_chat generator (conversation + message + -queue + streamed response over the EasyUI chat pipeline), but the backing -config comes from the bound Agent Soul and the answer is produced by -``AgentAppRunner`` calling the dify-agent backend rather than an in-process -LLM/ReAct loop. Build-chat finalization uses this same streamed path and only -changes the runtime exit policy carried to the backend. -""" - -from __future__ import annotations - -import contextvars -import logging -import threading -import uuid -from collections.abc import Generator, Mapping -from typing import Any, Literal - -from flask import Flask, current_app -from sqlalchemy import and_, or_, select -from sqlalchemy.orm import Session, sessionmaker - -from clients.agent_backend import AgentBackendRunEventAdapter -from clients.agent_backend.factory import create_agent_backend_run_client -from configs import dify_config -from constants import UUID_NIL -from core.agent.publish_visibility import agent_has_workflow_callable_active_snapshot -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_app.app_config_manager import AgentAppConfigManager -from core.app.apps.agent_app.app_runner import AgentAppRunner -from core.app.apps.agent_app.errors import ( - AgentAppGeneratorError, - AgentAppNotPublishedError, -) -from core.app.apps.agent_app.generate_response_converter import AgentAppGenerateResponseConverter -from core.app.apps.agent_app.runtime_request_builder import AgentAppRuntimeRequestBuilder -from core.app.apps.agent_app.session_store import AgentAppWorkspaceStore -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 ( - AgentAppGenerateEntity, - DifyRunContext, - InvokeFrom, - UserFrom, -) -from core.credit_usage import CreditUsageAppType -from core.db.session_factory import session_factory -from core.ops.ops_trace_manager import TraceQueueManager -from extensions.ext_database import db -from factories import file_factory -from models import Account, App, AppModelConfig, Conversation, EndUser, Message, MessageAnnotation -from models.agent import ( - APP_BACKED_AGENT_SOURCES, - Agent, - AgentConfigDraft, - AgentConfigDraftType, - AgentConfigSnapshot, - AgentConfigVersionKind, - AgentScope, - AgentStatus, - AgentWorkspaceBinding, - AgentWorkspaceOwnerType, -) -from models.agent_config_entities import AgentSoulConfig -from models.model import load_annotation_reply_config -from repositories.tools.provider_repository import ToolProviderRepository -from repositories.tools.workflow_repository import WorkflowToolRepository -from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope -from services.conversation_service import ConversationService -from services.workflow.execution.adapters.agent_v2.dify_tools_builder import WorkflowAgentDifyToolsBuilder - -logger = logging.getLogger(__name__) - - -class AgentAppGenerator(MessageBasedAppGenerator): - def generate( - self, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - session: Session, - streaming: bool = True, - ) -> Mapping[str, Any] | Generator[Mapping | str]: - if not streaming: - raise AgentAppGeneratorError("Agent App only supports streaming mode") - - query = self._require_query(args) - inputs = args["inputs"] - raw_files = args.get("files") or [] - - 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, session=session - ) - - # New conversations use the current Agent generation. Existing - # conversations use the immutable generation named by their Binding. - agent, agent_config_id, agent_config_version_kind, agent_soul = self._resolve_agent( - app_model, - invoke_from=invoke_from, - draft_type=args.get("draft_type"), - user=user, - session=session, - conversation=conversation, - ) - session_scope_config_version_id = self._session_scope_config_version_id( - invoke_from=invoke_from, - config_version_id=agent_config_id, - ) - - # Build the EasyUI-shaped config from the Agent Soul so the chat pipeline - # can persist usage; the answer itself comes from the agent backend. - app_model_config = ( - session.get(AppModelConfig, app_model.app_model_config_id) if app_model.app_model_config_id else None - ) - annotation_reply = load_annotation_reply_config(session, app_model.id) if app_model_config else None - app_config = AgentAppConfigManager.get_app_config( - app_model=app_model, - agent_soul=agent_soul, - annotation_reply=annotation_reply, - app_model_config=app_model_config, - conversation=conversation, - ) - model_conf = ModelConfigConverter.convert(app_config) - with self._bind_file_access_scope(tenant_id=app_model.tenant_id, user=user, invoke_from=invoke_from): - file_upload_config = ( - FileUploadConfigManager.convert(app_config.app_model_config_dict, is_vision=True) if raw_files else None - ) - file_objs = ( - file_factory.build_from_mappings( - mappings=raw_files, - tenant_id=app_model.tenant_id, - config=file_upload_config, - access_controller=self._file_access_controller, - ) - if raw_files and file_upload_config is not None - else [] - ) - - trace_manager = TraceQueueManager(app_model.id, user.id if isinstance(user, Account) else user.session_id) - application_generate_entity = AgentAppGenerateEntity( - task_id=str(uuid.uuid4()), - app_config=app_config, - model_conf=model_conf, - file_upload_config=file_upload_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={ - "auto_generate_conversation_name": args.get("auto_generate_name", True), - }, - call_depth=0, - trace_manager=trace_manager, - agent_id=agent.id, - agent_config_snapshot_id=agent_config_id, - agent_config_version_kind=agent_config_version_kind, - agent_session_scope_config_version_id=session_scope_config_version_id, - agent_llm_gateway_enabled=True, - ) - - conversation, message = self._init_generate_records( - application_generate_entity, - conversation, - session=session, - ) - - 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, - ) - - 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, - "user_from": UserFrom.ACCOUNT if isinstance(user, Account) else UserFrom.END_USER, - }, - ) - worker_thread.start() - - response = self._handle_response( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation=conversation, - message=message, - stream=streaming, - ) - return AgentAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) - - def _generate_worker( - self, - *, - flask_app: Flask, - context: contextvars.Context, - application_generate_entity: AgentAppGenerateEntity, - queue_manager: AppQueueManager, - conversation_id: str, - message_id: str, - user_from: UserFrom, - ) -> None: - from libs.flask_utils import preserve_flask_contexts - - with preserve_flask_contexts(flask_app, context_vars=context): - try: - conversation = self._get_conversation(conversation_id) - message = self._get_message(message_id) - app_config = application_generate_entity.app_config - - # Apply app-level input guards (content moderation + annotation - # reply) before reaching the Agent backend, mirroring the EasyUI - # chat / agent-chat runners. These can short-circuit the turn. - with session_factory.get_session_maker().begin() as session: - app_model = session.get(App, app_config.app_id) - if app_model is None: - raise AgentAppGeneratorError("App not found") - handled, query, annotation_reply = self._run_input_guards( - session=session, - application_generate_entity=application_generate_entity, - app_model=app_model, - message=message, - queue_manager=queue_manager, - ) - if annotation_reply: - from core.app.apps.agent_app.app_runner import publish_text_answer - from core.app.entities.queue_entities import QueueAnnotationReplyEvent - - queue_manager.publish( - QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id), - PublishFrom.APPLICATION_MANAGER, - ) - publish_text_answer( - queue_manager=queue_manager, - model_name=application_generate_entity.model_conf.model, - answer=annotation_reply.content, - user_query=query, - ) - if handled: - return - - dify_context = DifyRunContext( - tenant_id=app_config.tenant_id, - app_id=app_config.app_id, - user_id=application_generate_entity.user_id, - user_from=user_from, - invoke_from=application_generate_entity.invoke_from, - app_type=CreditUsageAppType.AGENT_V2, - ) - with session_factory.create_session() as session: - agent, config_version, agent_soul = self._resolve_agent_by_id( - tenant_id=app_config.tenant_id, - agent_id=application_generate_entity.agent_id, - snapshot_id=application_generate_entity.agent_config_snapshot_id, - session=session, - ) - - runner = self._build_runner() - image_detail_config = ( - application_generate_entity.file_upload_config.image_config.detail - if ( - application_generate_entity.file_upload_config - and application_generate_entity.file_upload_config.image_config - ) - else None - ) - runner.run( - dify_context=dify_context, - agent_id=application_generate_entity.agent_id, - agent_config_snapshot_id=application_generate_entity.agent_config_snapshot_id, - agent_config_version_kind=application_generate_entity.agent_config_version_kind, - agent_soul=agent_soul, - home_snapshot_id=config_version.home_snapshot_id, - conversation_id=conversation.id, - query=query, - message_id=message.id, - model_name=application_generate_entity.model_conf.model, - queue_manager=queue_manager, - files=list(application_generate_entity.files), - image_detail_config=image_detail_config, - session_scope_snapshot_id=application_generate_entity.agent_session_scope_config_version_id, - build_draft_id=( - application_generate_entity.agent_config_snapshot_id - if application_generate_entity.agent_config_version_kind == AgentConfigVersionKind.BUILD_DRAFT - else None - ), - ) - except GenerateTaskStoppedError: - # The runner hands off only a runtime terminal it observed. - # A local stop or failed cancellation wait is still just intent. - pass - except Exception as e: - from core.ops.base_trace_instance import application_tracing_hooks - - tracing_hooks = application_tracing_hooks() - if tracing_hooks is not None: - tracing_hooks.capture_agent_app_terminal( - tenant_id=application_generate_entity.app_config.tenant_id, - app_id=application_generate_entity.app_config.app_id, - message_id=message_id, - outcome="failed", - output=None, - ) - logger.exception("Unknown Error in Agent App generate worker") - queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER) - finally: - db.session.close() - - @staticmethod - def _require_query(args: Mapping[str, Any]) -> str: - query = args.get("query") - if not isinstance(query, str) or not query.strip(): - raise AgentAppGeneratorError("query is required") - return query.replace("\x00", "") - - @staticmethod - def _build_runner() -> AgentAppRunner: - database_client = sessionmaker(bind=db.engine, expire_on_commit=False) - return AgentAppRunner( - request_builder=AgentAppRuntimeRequestBuilder( - dify_tools_builder=WorkflowAgentDifyToolsBuilder( - tool_providers=ToolProviderRepository(database_client), - workflow_queries=WorkflowToolRepository(database_client), - ) - ), - agent_backend_client=create_agent_backend_run_client( - base_url=dify_config.AGENT_BACKEND_BASE_URL, - api_token=dify_config.AGENT_BACKEND_API_TOKEN, - use_fake=dify_config.AGENT_BACKEND_USE_FAKE, - fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO, - stream_read_timeout_seconds=dify_config.AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS, - stream_max_reconnects=dify_config.AGENT_BACKEND_STREAM_MAX_RECONNECTS, - ), - event_adapter=AgentBackendRunEventAdapter(), - session_store=AgentAppWorkspaceStore(), - text_delta_debounce_seconds=dify_config.AGENT_APP_TEXT_DELTA_DEBOUNCE_SECONDS, - ) - - def _run_input_guards( - self, - *, - session: Session, - application_generate_entity: AgentAppGenerateEntity, - app_model: App, - message: Message, - queue_manager: AppQueueManager, - ) -> tuple[bool, str, MessageAnnotation | None]: - """Apply input moderation + annotation reply before the backend call. - - Returns ``(handled, query, annotation_reply)``. Annotation output is - published by the caller only after this transaction commits. - """ - from core.app.apps.agent_app.app_runner import publish_text_answer - from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature - from core.moderation.base import ModerationError - from core.moderation.input_moderation import InputModeration - - app_config = application_generate_entity.app_config - model_name = application_generate_entity.model_conf.model - query = application_generate_entity.query or "" - - # content moderation (sensitive_word_avoidance); a blocked input yields a - # preset answer, an "overridden" action returns a sanitized query. - try: - _, _, query = InputModeration().check( - app_id=app_config.app_id, - tenant_id=app_config.tenant_id, - app_config=app_config, - inputs=dict(application_generate_entity.inputs), - query=query or "", - message_id=message.id, - trace_manager=application_generate_entity.trace_manager, - ) - except ModerationError as e: - publish_text_answer(queue_manager=queue_manager, model_name=model_name, answer=str(e), user_query=query) - return True, query, None - - # annotation reply: a matching annotation answers the turn deterministically. - if query: - annotation_reply = AnnotationReplyFeature().query( - app_record=app_model, - message=message, - query=query, - user_id=application_generate_entity.user_id, - invoke_from=application_generate_entity.invoke_from, - session=session, - ) - if annotation_reply: - return True, query, annotation_reply - - return False, query, None - - def _resolve_agent( - self, - app_model: App, - *, - invoke_from: InvokeFrom, - draft_type: Any, - draft_id: str | None = None, - user: Account | EndUser, - session: Session, - conversation: Conversation | None = None, - ) -> tuple[Agent, str, Literal["snapshot", "draft", "build_draft"], AgentSoulConfig]: - agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == app_model.tenant_id, - Agent.status == AgentStatus.ACTIVE, - or_( - and_( - Agent.app_id == app_model.id, - Agent.scope == AgentScope.ROSTER, - Agent.source.in_(APP_BACKED_AGENT_SOURCES), - ), - Agent.backing_app_id == app_model.id, - ), - ) - .order_by(Agent.created_at.desc()) - .limit(1) - ) - if agent is None: - raise AgentAppGeneratorError("Agent App has no bound Agent") - if invoke_from == InvokeFrom.DEBUGGER: - draft = self._resolve_debug_draft( - tenant_id=app_model.tenant_id, - agent=agent, - draft_type=draft_type, - draft_id=draft_id, - account_id=user.id if isinstance(user, Account) else None, - session=session, - ) - agent_soul = AgentSoulConfig.model_validate(draft.config_snapshot_dict) - config_version_kind: Literal["snapshot", "draft", "build_draft"] = ( - "build_draft" if draft.draft_type == AgentConfigDraftType.DEBUG_BUILD else "draft" - ) - return agent, draft.id, config_version_kind, agent_soul - # Dirty drafts do not revoke a published snapshot, while the seeded - # create/import snapshot must never become public runtime configuration. - if not agent_has_workflow_callable_active_snapshot(session=session, agent=agent): - raise AgentAppNotPublishedError("Agent has not been published") - conversation_binding = self._resolve_conversation_binding( - session=session, - tenant_id=app_model.tenant_id, - app_id=app_model.id, - agent_id=agent.id, - conversation=conversation, - ) - snapshot_id = ( - conversation_binding.agent_config_version_id - if conversation_binding is not None - else agent.active_config_snapshot_id - ) - _, snapshot, agent_soul = self._resolve_agent_by_id( - tenant_id=app_model.tenant_id, - agent_id=agent.id, - snapshot_id=snapshot_id, - session=session, - ) - if conversation_binding is not None: - AgentWorkspaceService.validate_binding_generation( - conversation_binding, - base_home_snapshot_id=snapshot.home_snapshot_id, - agent_config_version_id=snapshot.id, - agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, - ) - return agent, snapshot.id, "snapshot", agent_soul - - @staticmethod - def _resolve_conversation_binding( - *, - session: Session, - tenant_id: str, - app_id: str, - agent_id: str, - conversation: Conversation | None, - ) -> AgentWorkspaceBinding | None: - """Resolve the exact participant generation owned by an existing conversation.""" - - if conversation is None or conversation.agent_workspace_binding_id is None: - return None - binding = AgentWorkspaceService.get_active_binding( - session=session, - tenant_id=tenant_id, - binding_id=conversation.agent_workspace_binding_id, - expected_owner_scope=WorkspaceOwnerScope( - tenant_id=tenant_id, - app_id=app_id, - owner_type=AgentWorkspaceOwnerType.CONVERSATION, - owner_id=conversation.id, - ), - ) - if binding is None or binding.agent_id != agent_id: - raise AgentAppGeneratorError("Conversation participant Binding is unavailable") - return binding - - @staticmethod - def _session_scope_config_version_id(*, invoke_from: InvokeFrom, config_version_id: str) -> str | None: - """Return the config version id that scopes Agent App session reuse. - - Console preview/debug chat uses a stable Agent draft row id; build mode - uses the current user's build-draft row id. Published/web/API runs use - immutable published snapshot ids. This keeps Workspace Binding continuity - inside one editable surface without mixing draft/build/published state. - """ - del invoke_from - return config_version_id - - @staticmethod - def _resolve_debug_draft( - *, - tenant_id: str, - agent: Agent, - draft_type: Any, - account_id: str | None, - session: Session, - draft_id: str | None = None, - ) -> AgentConfigDraft: - effective_draft_type = ( - AgentConfigDraftType.DEBUG_BUILD - if draft_type == AgentConfigDraftType.DEBUG_BUILD.value - else AgentConfigDraftType.DRAFT - ) - if effective_draft_type == AgentConfigDraftType.DRAFT: - from services.agent.composer_service import AgentComposerService - - return AgentComposerService.get_or_create_normal_agent_draft( - session=session, - tenant_id=tenant_id, - agent=agent, - created_by=agent.updated_by or agent.created_by, - ) - if not account_id: - raise AgentAppGeneratorError("Build draft requires an account user") - stmt = select(AgentConfigDraft).where( - AgentConfigDraft.tenant_id == tenant_id, - AgentConfigDraft.agent_id == agent.id, - AgentConfigDraft.draft_type == AgentConfigDraftType.DEBUG_BUILD, - AgentConfigDraft.account_id == account_id, - ) - if draft_id is not None: - stmt = stmt.where(AgentConfigDraft.id == draft_id) - draft = session.scalar(stmt.order_by(AgentConfigDraft.updated_at.desc()).limit(1)) - if draft is not None: - return draft - raise AgentAppGeneratorError("Agent build draft not found") - - @staticmethod - def _resolve_agent_by_id( - *, tenant_id: str, agent_id: str, snapshot_id: str | None, session: Session - ) -> tuple[Agent, AgentConfigSnapshot | AgentConfigDraft, AgentSoulConfig]: - agent = session.scalar(select(Agent).where(Agent.id == agent_id, Agent.tenant_id == tenant_id)) - if agent is None: - raise AgentAppGeneratorError("Agent not found") - if not snapshot_id: - raise AgentAppGeneratorError("Agent has no published version") - snapshot = session.scalar( - select(AgentConfigSnapshot).where( - AgentConfigSnapshot.tenant_id == tenant_id, - AgentConfigSnapshot.agent_id == agent_id, - AgentConfigSnapshot.id == snapshot_id, - ) - ) - if snapshot is not None: - agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) - return agent, snapshot, agent_soul - draft = session.scalar( - select(AgentConfigDraft).where( - AgentConfigDraft.tenant_id == tenant_id, - AgentConfigDraft.agent_id == agent_id, - AgentConfigDraft.id == snapshot_id, - ) - ) - if draft is None: - raise AgentAppGeneratorError("Agent published version not found") - agent_soul = AgentSoulConfig.model_validate(draft.config_snapshot_dict) - return agent, draft, agent_soul - - -__all__ = ["AgentAppGenerator", "AgentAppGeneratorError", "AgentAppNotPublishedError"] diff --git a/api/core/app/apps/agent_app/app_runner.py b/api/core/app/apps/agent_app/app_runner.py deleted file mode 100644 index b2e7f04ae62c05..00000000000000 --- a/api/core/app/apps/agent_app/app_runner.py +++ /dev/null @@ -1,1187 +0,0 @@ -"""Agent App runner: drive Agent backend turns for chat and finalization flows. - -Unlike the legacy ``AgentChatAppRunner`` (which runs an in-process ReAct loop), -this runner delegates to the Agent backend, consumes the streamed event flow, -republishes the assistant answer through the existing EasyUI chat task -pipeline, and saves the latest module state snapshot on the persistent Binding. -""" - -from __future__ import annotations - -import json -import logging -import time -from collections.abc import Mapping -from decimal import Decimal -from typing import Any, Literal - -from pydantic import JsonValue - -from clients.agent_backend import ( - AgentBackendAgentMessageDeltaInternalEvent, - AgentBackendError, - AgentBackendInternalEventType, - AgentBackendRunCancelledInternalEvent, - AgentBackendRunClient, - AgentBackendRunEventAdapter, - AgentBackendRunFailedError, - AgentBackendRunFailedInternalEvent, - AgentBackendRunSucceededInternalEvent, - AgentBackendStreamInternalEvent, -) -from core.app.apps.agent_app.runtime_request_builder import ( - AgentAppRuntimeBuildContext, - AgentAppRuntimeRequest, - AgentAppRuntimeRequestBuilder, -) -from core.app.apps.agent_app.session_store import ( - AgentAppSessionScope, - AgentAppWorkspaceStore, - StoredAgentAppSession, -) -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import DifyRunContext -from core.app.entities.queue_entities import ( - QueueAgentMessageEvent, - QueueAgentThoughtEvent, - QueueLLMChunkEvent, - QueueMessageEndEvent, - QueueRetrieverResourcesEvent, -) -from extensions.ext_database import db -from graphon.file import File -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage -from graphon.model_runtime.entities.message_entities import ( - AssistantPromptMessage, - ImagePromptMessageContent, - PromptMessage, - UserPromptMessage, -) -from graphon.model_runtime.errors.invoke import ( - InvokeAuthorizationError, - InvokeBadRequestError, - InvokeConnectionError, - InvokeError, - InvokeRateLimitError, - InvokeServerUnavailableError, -) -from models.agent import AgentConfigVersionKind -from models.agent_config_entities import AgentSoulConfig -from models.enums import CreatorUserRole -from models.model import Message, MessageAgentThought -from services.agent.knowledge_citations import knowledge_sources_from_tool_part - -logger = logging.getLogger(__name__) - - -class _DefaultSessionScopeSnapshotId: - pass - - -_DEFAULT_SESSION_SCOPE_SNAPSHOT_ID = _DefaultSessionScopeSnapshotId() - -_AGENT_BACKEND_INVOKE_ERROR_BY_REASON: Mapping[str, type[InvokeError]] = { - "InvokeAuthorizationError": InvokeAuthorizationError, - "InvokeBadRequestError": InvokeBadRequestError, - "CredentialsValidateFailedError": InvokeBadRequestError, - "InvokeConnectionError": InvokeConnectionError, - "InvokeRateLimitError": InvokeRateLimitError, - "InvokeServerUnavailableError": InvokeServerUnavailableError, -} - - -def _agent_backend_failure_to_exception(event: AgentBackendRunFailedInternalEvent) -> Exception: - if event.error_type is None: - err_cls = _AGENT_BACKEND_INVOKE_ERROR_BY_REASON.get(event.reason or "") - if err_cls is not None: - return err_cls(event.error) - message = event.error or "Agent backend run did not complete successfully." - return AgentBackendRunFailedError( - event.run_id, - { - "error": event.error, - "reason": event.reason, - "source_event_id": event.source_event_id, - }, - message=message, - error_type=event.error_type, - reason=event.reason, - source_event_id=event.source_event_id, - ) - - -def _prompt_messages_from_query(user_query: str | None) -> list[PromptMessage]: - if not user_query: - return [] - return [UserPromptMessage(content=user_query)] - - -def _llm_usage_from_agent_backend(usage: Mapping[str, Any] | None) -> LLMUsage | None: - if usage is None: - return None - try: - return LLMUsage.from_metadata(usage) - except (TypeError, ValueError): - logger.warning("Failed to parse Agent backend usage metadata: %s", usage, exc_info=True) - return None - - -def publish_text_answer( - *, - queue_manager: AppQueueManager, - model_name: str, - answer: str, - user_query: str | None = None, - usage: LLMUsage | None = None, -) -> None: - """Publish a complete assistant answer as one chunk + message-end. - - The EasyUI chat task pipeline consumes a QueueLLMChunkEvent stream followed - by a QueueMessageEndEvent; emitting the whole answer as a single chunk lets - both the backend-produced answer and short-circuited answers (moderation / - annotation reply) share the exact same persistence + SSE path. - """ - publish_text_delta( - queue_manager=queue_manager, - model_name=model_name, - delta=answer, - user_query=user_query, - ) - publish_message_end( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - user_query=user_query, - usage=usage, - ) - - -def publish_text_delta( - *, - queue_manager: AppQueueManager, - model_name: str, - delta: str, - user_query: str | None = None, -) -> None: - """Publish one assistant text delta through the EasyUI chat pipeline.""" - if not delta: - return - prompt_messages = _prompt_messages_from_query(user_query) - chunk = LLMResultChunk( - model=model_name, - prompt_messages=prompt_messages, - delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=delta)), - ) - queue_manager.publish(QueueLLMChunkEvent(chunk=chunk), PublishFrom.APPLICATION_MANAGER) - - -def publish_agent_message_delta( - *, - queue_manager: AppQueueManager, - model_name: str, - delta: str, - user_query: str | None = None, -) -> None: - """Publish one agent-process text delta through the EasyUI chat pipeline.""" - if not delta: - return - prompt_messages = _prompt_messages_from_query(user_query) - chunk = LLMResultChunk( - model=model_name, - prompt_messages=prompt_messages, - delta=LLMResultChunkDelta(index=0, message=AssistantPromptMessage(content=delta)), - ) - queue_manager.publish(QueueAgentMessageEvent(chunk=chunk), PublishFrom.APPLICATION_MANAGER) - - -def publish_message_end( - *, - queue_manager: AppQueueManager, - model_name: str, - answer: str, - user_query: str | None = None, - usage: LLMUsage | None = None, -) -> None: - """Publish the terminal assistant result without emitting another delta.""" - prompt_messages = _prompt_messages_from_query(user_query) - queue_manager.publish( - QueueMessageEndEvent( - llm_result=LLMResult( - model=model_name, - prompt_messages=prompt_messages, - message=AssistantPromptMessage(content=answer), - usage=usage or LLMUsage.empty_usage(), - ), - ), - PublishFrom.APPLICATION_MANAGER, - ) - - -class _TextDeltaDebouncer: - """Batch independent model text deltas before agent-message SSE output.""" - - def __init__(self, *, debounce_seconds: float) -> None: - self._debounce_seconds = debounce_seconds - self._parts: list[str] = [] - self._first_pending_at: float | None = None - - def push(self, delta: str) -> str | None: - if not delta: - return None - if self._debounce_seconds <= 0: - return delta - - now = time.monotonic() - if not self._parts: - self._first_pending_at = now - self._parts.append(delta) - - if self._first_pending_at is not None and now - self._first_pending_at >= self._debounce_seconds: - return self.flush() - return None - - def flush(self) -> str | None: - if not self._parts: - return None - - text = "".join(self._parts) - self._parts = [] - self._first_pending_at = None - return text - - -class _AgentProcessRecorder: - """Persist Agent v2 process streams through the legacy thought model. - - Thinking and answer rows expose snapshot updates for contiguous model-text - segments. Tool events close currently open text segments so later model text - starts a fresh row instead of replaying content that was already streamed - before the tool. - """ - - def __init__( - self, - *, - dify_context: DifyRunContext, - message_id: str, - queue_manager: AppQueueManager, - ) -> None: - self._dify_context = dify_context - self._message_id = message_id - self._queue_manager = queue_manager - self._next_position = 1 - self._thinking_by_index: dict[int, str] = {} - self._answer_thought_id: str | None = None - self._tool_by_index: dict[int, str] = {} - self._tool_by_call_id: dict[str, str] = {} - self._open_tool_by_name: dict[str, set[str]] = {} - self._knowledge_receipts: set[str] = set() - - def handle_stream_event(self, event: AgentBackendStreamInternalEvent) -> None: - data = event.data - if not isinstance(data, dict): - return - - event_kind = data.get("event_kind") - if event_kind == "part_delta": - self._handle_part_delta(data) - elif event_kind == "part_start": - self._handle_part(data) - elif event_kind in {"function_tool_call", "output_tool_call"}: - self._handle_tool_call_event(data) - elif event_kind in {"function_tool_result", "output_tool_result"}: - self._handle_tool_result_event(data) - - def _handle_part_delta(self, data: dict[str, Any]) -> None: - delta = data.get("delta") - if not isinstance(delta, dict): - return - - index = _event_index(data) - delta_kind = delta.get("part_delta_kind") - if delta_kind == "thinking": - content_delta = delta.get("content_delta") - if isinstance(content_delta, str) and content_delta: - self._append_thinking(index, content_delta) - return - - if delta_kind == "tool_call": - self._record_tool_call_delta(index, delta) - - def _handle_part(self, data: dict[str, Any]) -> None: - part = data.get("part") - if not isinstance(part, dict): - return - - index = _event_index(data) - part_kind = part.get("part_kind") - if part_kind == "thinking": - content = part.get("content") - if isinstance(content, str) and content: - self._append_thinking(index, content) - return - - if part_kind in {"tool-call", "builtin-tool-call"}: - self._record_tool_call_part(index, part) - return - - if part_kind in {"tool-return", "builtin-tool-return"}: - self._record_tool_return_part(part) - - def append_answer_text(self, content_delta: str) -> None: - if not content_delta: - return - - self._thinking_by_index.clear() - if self._answer_thought_id is None: - self._answer_thought_id = self._create_thought(answer=content_delta) - return - self._update_thought(self._answer_thought_id, answer_delta=content_delta) - - def trim_answer_suffix(self, final_answer: str) -> None: - if not final_answer or self._answer_thought_id is None: - return - - row = db.session.get(MessageAgentThought, self._answer_thought_id) - if row is None: - return - - answer = row.answer or "" - overlap = _suffix_prefix_overlap_length(answer, final_answer) - if overlap == 0: - return - - row.answer = answer[:-overlap] - if _is_empty_answer_only_thought(row): - db.session.delete(row) - self._answer_thought_id = None - db.session.commit() - return - - db.session.commit() - self._queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=self._answer_thought_id), PublishFrom.APPLICATION_MANAGER - ) - - def _handle_tool_call_event(self, data: dict[str, Any]) -> None: - part = data.get("part") - if isinstance(part, dict): - self._record_tool_call_part(_event_index(data), part) - - def _handle_tool_result_event(self, data: dict[str, Any]) -> None: - part = data.get("part") or data.get("result") - if isinstance(part, dict): - self._record_tool_return_part(part) - return - - content = data.get("content") - if content is not None: - self._record_tool_observation( - tool_call_id=_string_or_none(data.get("tool_call_id")), - tool_name=_string_or_none(data.get("tool_name")), - observation=content, - ) - - def _append_thinking(self, index: int, content_delta: str) -> None: - self._answer_thought_id = None - thought_id = self._thinking_by_index.get(index) - if thought_id is None: - thought_id = self._create_thought(thought=content_delta) - self._thinking_by_index[index] = thought_id - return - self._update_thought(thought_id, thought_delta=content_delta) - - def _record_tool_call_delta(self, index: int, delta: dict[str, Any]) -> None: - self._close_thinking_segments() - tool_call_id = _string_or_none(delta.get("tool_call_id")) - tool_name = _string_or_none(delta.get("tool_name_delta")) - args_delta = delta.get("args_delta") - thought_id = self._lookup_tool_thought(index=index, tool_call_id=tool_call_id) - if thought_id is None: - thought_id = self._create_thought(tool=tool_name, tool_input=_json_or_text(args_delta)) - self._remember_tool_thought( - index=index, tool_call_id=tool_call_id, tool_name=tool_name, thought_id=thought_id - ) - return - - self._update_thought( - thought_id, - tool=tool_name, - tool_input_delta=_json_or_text(args_delta), - ) - self._remember_tool_thought(index=index, tool_call_id=tool_call_id, tool_name=tool_name, thought_id=thought_id) - - def _record_tool_call_part(self, index: int, part: dict[str, Any]) -> None: - self._close_thinking_segments() - tool_call_id = _string_or_none(part.get("tool_call_id")) - tool_name = _string_or_none(part.get("tool_name")) - thought_id = self._lookup_tool_thought(index=index, tool_call_id=tool_call_id) - if thought_id is None: - thought_id = self._create_thought(tool=tool_name, tool_input=_json_or_text(part.get("args"))) - self._remember_tool_thought( - index=index, tool_call_id=tool_call_id, tool_name=tool_name, thought_id=thought_id - ) - return - - self._update_thought( - thought_id, - tool=tool_name, - tool_input=_json_or_text(part.get("args")), - ) - self._remember_tool_thought(index=index, tool_call_id=tool_call_id, tool_name=tool_name, thought_id=thought_id) - - def _record_tool_return_part(self, part: dict[str, Any]) -> None: - self._close_thinking_segments() - sources = [] - for source in knowledge_sources_from_tool_part(part): - citation = source.knowledge_fs_citation - if citation and citation.id not in self._knowledge_receipts and len(self._knowledge_receipts) < 200: - self._knowledge_receipts.add(citation.id) - sources.append(source) - if sources: - self._queue_manager.publish( - QueueRetrieverResourcesEvent(retriever_resources=sources), PublishFrom.APPLICATION_MANAGER - ) - tool_call_id = _string_or_none(part.get("tool_call_id")) - tool_name = _string_or_none(part.get("tool_name")) - content = part.get("content") - if content is None: - content = part - self._record_tool_observation(tool_call_id=tool_call_id, tool_name=tool_name, observation=content) - - def _record_tool_observation(self, *, tool_call_id: str | None, tool_name: str | None, observation: Any) -> None: - self._close_thinking_segments() - thought_id = self._lookup_observation_thought(tool_call_id=tool_call_id, tool_name=tool_name) - if thought_id is None: - thought_id = self._create_thought(tool=tool_name) - else: - self._mark_tool_observed(thought_id) - self._update_thought(thought_id, observation=_json_or_text(observation)) - - def _lookup_tool_thought(self, *, index: int, tool_call_id: str | None) -> str | None: - if tool_call_id and tool_call_id in self._tool_by_call_id: - return self._tool_by_call_id[tool_call_id] - if index < 0: - return None - return self._tool_by_index.get(index) - - def _remember_tool_thought( - self, *, index: int, tool_call_id: str | None, tool_name: str | None, thought_id: str - ) -> None: - if index >= 0: - self._tool_by_index[index] = thought_id - if tool_call_id: - self._tool_by_call_id[tool_call_id] = thought_id - if tool_name: - self._open_tool_by_name.setdefault(tool_name, set()).add(thought_id) - - def _lookup_observation_thought(self, *, tool_call_id: str | None, tool_name: str | None) -> str | None: - if tool_call_id: - return self._tool_by_call_id.get(tool_call_id) - if tool_name: - open_thought_ids = self._open_tool_by_name.get(tool_name, set()) - if len(open_thought_ids) == 1: - return next(iter(open_thought_ids)) - return None - - def _mark_tool_observed(self, thought_id: str) -> None: - self._tool_by_index = {index: value for index, value in self._tool_by_index.items() if value != thought_id} - self._tool_by_call_id = { - tool_call_id: value for tool_call_id, value in self._tool_by_call_id.items() if value != thought_id - } - for open_thought_ids in self._open_tool_by_name.values(): - open_thought_ids.discard(thought_id) - - def _close_thinking_segments(self) -> None: - self._thinking_by_index.clear() - self._answer_thought_id = None - - def _create_thought( - self, - *, - thought: str | None = None, - answer: str | None = None, - tool: str | None = None, - tool_input: str | None = None, - ) -> str: - row = MessageAgentThought( - message_id=self._message_id, - message_chain_id=None, - thought=thought or "", - tool=tool or "", - tool_labels_str=_tool_labels(tool), - tool_meta_str="{}", - tool_input=tool_input or "", - observation="", - tool_process_data=None, - message="", - message_token=0, - message_unit_price=Decimal(0), - message_price_unit=Decimal("0.001"), - message_files="", - answer=answer or "", - answer_token=0, - answer_unit_price=Decimal(0), - answer_price_unit=Decimal("0.001"), - tokens=0, - total_price=Decimal(0), - position=self._next_position, - currency="USD", - latency=0, - created_by_role=self._created_by_role(), - created_by=self._dify_context.user_id, - ) - self._next_position += 1 - db.session.add(row) - db.session.commit() - thought_id = str(row.id) - self._queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=thought_id), PublishFrom.APPLICATION_MANAGER - ) - return thought_id - - def _update_thought( - self, - thought_id: str, - *, - thought_delta: str | None = None, - tool: str | None = None, - tool_input: str | None = None, - tool_input_delta: str | None = None, - observation: str | None = None, - answer_delta: str | None = None, - ) -> None: - row = db.session.get(MessageAgentThought, thought_id) - if row is None: - return - - if thought_delta: - row.thought = f"{row.thought or ''}{thought_delta}" - if tool: - row.tool = tool - row.tool_labels_str = _tool_labels(tool) - if tool_input is not None: - row.tool_input = tool_input - if tool_input_delta: - row.tool_input = f"{row.tool_input or ''}{tool_input_delta}" - if observation is not None: - row.observation = observation - if answer_delta: - row.answer = f"{row.answer or ''}{answer_delta}" - - db.session.commit() - self._queue_manager.publish( - QueueAgentThoughtEvent(agent_thought_id=thought_id), PublishFrom.APPLICATION_MANAGER - ) - - def _created_by_role(self) -> CreatorUserRole: - if self._dify_context.invoke_from.runs_as_account(): - return CreatorUserRole.ACCOUNT - return CreatorUserRole.END_USER - - -def _event_index(data: dict[str, Any]) -> int: - index = data.get("index") - return index if isinstance(index, int) else -1 - - -def _string_or_none(value: Any) -> str | None: - if not isinstance(value, str): - return None - normalized = value.strip() - if not normalized or normalized.lower() in {"none", "null"}: - return None - return normalized - - -def _json_or_text(value: Any) -> str | None: - if value is None: - return None - if isinstance(value, str): - return value - try: - return json.dumps(value, ensure_ascii=False) - except Exception: - return str(value) - - -def _tool_labels(tool: str | None) -> str: - if not tool: - return "{}" - return json.dumps({tool: {"en_US": tool, "zh_Hans": tool}}, ensure_ascii=False) - - -def _suffix_prefix_overlap_length(text: str, prefix_source: str) -> int: - max_length = min(len(text), len(prefix_source)) - for length in range(max_length, 0, -1): - if text.endswith(prefix_source[:length]): - return length - return 0 - - -def _is_empty_answer_only_thought(row: MessageAgentThought) -> bool: - return not any((row.thought, row.answer, row.tool, row.tool_input, row.observation)) - - -class AgentAppRunner: - """Runs one Agent App conversation turn against the Agent backend.""" - - def __init__( - self, - *, - request_builder: AgentAppRuntimeRequestBuilder, - agent_backend_client: AgentBackendRunClient, - event_adapter: AgentBackendRunEventAdapter, - session_store: AgentAppWorkspaceStore, - text_delta_debounce_seconds: float, - ) -> None: - self._request_builder = request_builder - self._agent_backend_client = agent_backend_client - self._event_adapter = event_adapter - self._session_store = session_store - self._text_delta_debounce_seconds = text_delta_debounce_seconds - - def run( - self, - *, - dify_context: DifyRunContext, - agent_id: str, - agent_config_snapshot_id: str, - agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot", - agent_soul: AgentSoulConfig, - home_snapshot_id: str | None, - conversation_id: str, - query: str, - message_id: str, - model_name: str, - queue_manager: AppQueueManager, - files: list[File] | None = None, - image_detail_config: ImagePromptMessageContent.DETAIL | None = None, - session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId = _DEFAULT_SESSION_SCOPE_SNAPSHOT_ID, - build_draft_id: str | None = None, - ) -> None: - scope = self._build_session_scope( - dify_context=dify_context, - agent_id=agent_id, - agent_config_snapshot_id=agent_config_snapshot_id, - home_snapshot_id=home_snapshot_id, - conversation_id=conversation_id, - session_scope_snapshot_id=session_scope_snapshot_id, - agent_config_version_kind=AgentConfigVersionKind(agent_config_version_kind), - build_draft_id=build_draft_id, - ) - stored = self._session_store.load_or_create(scope) - runtime = self._build_runtime( - dify_context=dify_context, - agent_id=agent_id, - agent_config_snapshot_id=agent_config_snapshot_id, - agent_config_version_kind=agent_config_version_kind, - agent_soul=agent_soul, - binding_id=stored.binding_id, - backend_binding_ref=stored.backend_binding_ref, - conversation_id=conversation_id, - query=query, - files=tuple(files or ()), - image_detail_config=image_detail_config, - idempotency_key=message_id, - stored=stored, - ) - - create_response = self._agent_backend_client.create_run(runtime.request) - terminal, process_recorder = self._consume_stream( - create_response.run_id, - dify_context=dify_context, - message_id=message_id, - queue_manager=queue_manager, - model_name=model_name, - query=query, - session_scope=scope, - binding_id=runtime.binding_id, - ) - - terminal_usage = None - if isinstance( - terminal, - AgentBackendRunSucceededInternalEvent - | AgentBackendRunFailedInternalEvent - | AgentBackendRunCancelledInternalEvent, - ): - terminal_usage = _llm_usage_from_agent_backend(terminal.usage) - self._persist_message_usage( - message_id=message_id, - usage=terminal_usage, - ) - - if isinstance(terminal, AgentBackendRunFailedInternalEvent | AgentBackendRunCancelledInternalEvent): - # None means no post-exit snapshot was produced; leave the previously stored session snapshot untouched. - if terminal.session_snapshot is not None: - self._save_session( - scope=scope, - binding_id=runtime.binding_id, - snapshot=terminal.session_snapshot, - ) - - if not isinstance(terminal, AgentBackendRunSucceededInternalEvent): - if isinstance(terminal, AgentBackendRunFailedInternalEvent): - reason = terminal.reason - if terminal.error_type is None and reason == "binding_lost": - raise AgentBackendError("The retained agent working environment is no longer available.") - raise _agent_backend_failure_to_exception(terminal) - raise AgentBackendError("Agent backend run did not complete successfully.") - - answer = self._terminal_output_to_answer(terminal.output) - try: - process_recorder.trim_answer_suffix(answer) - except Exception: - db.session.rollback() - logger.warning( - "Failed to trim Agent App answer text: run_id=%s message_id=%s", - terminal.run_id, - message_id, - exc_info=True, - ) - try: - self._publish_terminal_answer( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - query=query, - usage=terminal_usage, - ) - except GenerateTaskStoppedError: - # Stop can become visible after stream consumption returned its - # success but before the final chunk/end reaches the consumer. - self._finish_stopped_terminal( - terminal, - dify_context=dify_context, - session_scope=scope, - binding_id=runtime.binding_id, - message_id=message_id, - ) - raise - self._save_session( - scope=scope, - binding_id=runtime.binding_id, - snapshot=terminal.session_snapshot, - ) - from core.ops.base_trace_instance import application_tracing_hooks - - tracing_hooks = application_tracing_hooks() - if tracing_hooks is not None: - tracing_hooks.capture_agent_app_terminal( - tenant_id=dify_context.tenant_id, - app_id=dify_context.app_id, - message_id=message_id, - outcome="succeeded", - output=answer, - ) - - def _build_session_scope( - self, - *, - dify_context: DifyRunContext, - agent_id: str, - agent_config_snapshot_id: str, - home_snapshot_id: str | None, - conversation_id: str, - session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId, - agent_config_version_kind: AgentConfigVersionKind, - build_draft_id: str | None = None, - ) -> AgentAppSessionScope: - if isinstance(session_scope_snapshot_id, _DefaultSessionScopeSnapshotId): - effective_session_scope_snapshot_id: str | None = agent_config_snapshot_id - else: - effective_session_scope_snapshot_id = session_scope_snapshot_id - return AgentAppSessionScope( - tenant_id=dify_context.tenant_id, - app_id=dify_context.app_id, - conversation_id=conversation_id, - agent_id=agent_id, - agent_config_snapshot_id=effective_session_scope_snapshot_id or agent_config_snapshot_id, - home_snapshot_id=home_snapshot_id, - agent_config_version_kind=agent_config_version_kind, - build_draft_id=build_draft_id, - ) - - def _build_runtime( - self, - *, - dify_context: DifyRunContext, - agent_id: str, - agent_config_snapshot_id: str, - agent_config_version_kind: Literal["snapshot", "draft", "build_draft"], - agent_soul: AgentSoulConfig, - binding_id: str, - backend_binding_ref: str, - conversation_id: str, - query: str, - files: tuple[File, ...], - image_detail_config: ImagePromptMessageContent.DETAIL | None, - idempotency_key: str, - stored: StoredAgentAppSession, - ) -> AgentAppRuntimeRequest: - session_snapshot = stored.session_snapshot - return self._request_builder.build( - AgentAppRuntimeBuildContext( - dify_context=dify_context, - agent_id=agent_id, - agent_config_snapshot_id=agent_config_snapshot_id, - agent_config_version_kind=agent_config_version_kind, - agent_soul=agent_soul, - conversation_id=conversation_id, - user_query=query, - files=files, - image_detail_config=image_detail_config, - idempotency_key=idempotency_key, - binding_id=binding_id, - backend_binding_ref=backend_binding_ref, - session_snapshot=session_snapshot, - ) - ) - - def _consume_stream( - self, - run_id: str, - *, - dify_context: DifyRunContext, - message_id: str, - queue_manager: AppQueueManager, - model_name: str, - query: str | None, - session_scope: AgentAppSessionScope, - binding_id: str, - ): - """Consume backend events while preserving raw recorder granularity.""" - terminal = None - process_recorder = _AgentProcessRecorder( - dify_context=dify_context, - message_id=message_id, - queue_manager=queue_manager, - ) - text_delta_debouncer = _TextDeltaDebouncer(debounce_seconds=self._text_delta_debounce_seconds) - last_event_id: str | None = None - - def persist_answer_text(content_delta: str) -> None: - try: - process_recorder.append_answer_text(content_delta) - except Exception: - db.session.rollback() - logger.warning( - "Failed to persist Agent App answer text: run_id=%s message_id=%s", - run_id, - message_id, - exc_info=True, - ) - publish_agent_message_delta( - queue_manager=queue_manager, - model_name=model_name, - delta=content_delta, - user_query=query, - ) - - def flush_pending_agent_message_text() -> None: - pending_text = text_delta_debouncer.flush() - if pending_text: - persist_answer_text(pending_text) - - try: - public_events = self._agent_backend_client.stream_events( - run_id, - should_stop=queue_manager.is_stopped, - ) - for public_event in public_events: - if public_event.id is not None: - last_event_id = public_event.id - internal_events = self._event_adapter.adapt(public_event) - for internal_event in internal_events: - if isinstance( - internal_event, - AgentBackendRunSucceededInternalEvent - | AgentBackendRunFailedInternalEvent - | AgentBackendRunCancelledInternalEvent, - ): - # A terminal frame may already have arrived when stop - # becomes visible. Keep the observed runtime result. - terminal = internal_event - break - if terminal is not None: - break - if queue_manager.is_stopped(): - flush_pending_agent_message_text() - self._cancel_run( - run_id, - after=last_event_id, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - dify_context=dify_context, - queue_manager=queue_manager, - ) - raise GenerateTaskStoppedError() - for internal_event in internal_events: - if queue_manager.is_stopped(): - flush_pending_agent_message_text() - self._cancel_run( - run_id, - after=last_event_id, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - dify_context=dify_context, - queue_manager=queue_manager, - ) - raise GenerateTaskStoppedError() - if internal_event.type in ( - AgentBackendInternalEventType.RUN_STARTED, - AgentBackendInternalEventType.STREAM_EVENT, - AgentBackendInternalEventType.AGENT_MESSAGE_DELTA, - ): - if isinstance(internal_event, AgentBackendAgentMessageDeltaInternalEvent): - debounced_delta = text_delta_debouncer.push(internal_event.delta) - if debounced_delta: - persist_answer_text(debounced_delta) - continue - - if isinstance(internal_event, AgentBackendStreamInternalEvent): - flush_pending_agent_message_text() - try: - process_recorder.handle_stream_event(internal_event) - except Exception: - db.session.rollback() - logger.warning( - "Failed to persist Agent App process event: run_id=%s message_id=%s event_kind=%s", - run_id, - message_id, - internal_event.event_kind, - exc_info=True, - ) - continue - continue - except GenerateTaskStoppedError: - raise - except Exception as error: - flush_pending_agent_message_text() - self._cancel_run( - run_id, - after=last_event_id, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - dify_context=dify_context, - queue_manager=queue_manager, - ) - if queue_manager.is_stopped(): - raise GenerateTaskStoppedError() from error - raise - try: - flush_pending_agent_message_text() - except GenerateTaskStoppedError: - if terminal is not None: - self._finish_stopped_terminal( - terminal, - dify_context=dify_context, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - ) - raise - if queue_manager.is_stopped(): - if terminal is not None: - self._finish_stopped_terminal( - terminal, - dify_context=dify_context, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - ) - else: - self._cancel_run( - run_id, - after=last_event_id, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - dify_context=dify_context, - queue_manager=queue_manager, - ) - raise GenerateTaskStoppedError() - return terminal, process_recorder - - def _cancel_run( - self, - run_id: str, - *, - after: str | None, - session_scope: AgentAppSessionScope, - binding_id: str, - message_id: str, - dify_context: DifyRunContext, - queue_manager: AppQueueManager, - ) -> None: - try: - public_event = self._agent_backend_client.cancel_run_and_wait(run_id, after=after) - for internal_event in self._event_adapter.adapt(public_event): - if isinstance(internal_event, AgentBackendRunCancelledInternalEvent): - self._finish_stopped_terminal( - internal_event, - dify_context=dify_context, - session_scope=session_scope, - binding_id=binding_id, - message_id=message_id, - queue_manager=queue_manager, - ) - except Exception: - logger.warning( - "Failed to finish cancelling stopped Agent App backend run: run_id=%s", - run_id, - exc_info=True, - ) - - def _finish_stopped_terminal( - self, - terminal: AgentBackendRunSucceededInternalEvent - | AgentBackendRunFailedInternalEvent - | AgentBackendRunCancelledInternalEvent, - *, - dify_context: DifyRunContext, - session_scope: AgentAppSessionScope, - binding_id: str, - message_id: str, - queue_manager: AppQueueManager | None = None, - ) -> None: - """Keep the observed runtime result when the client has requested stop.""" - if terminal.session_snapshot is not None: - self._save_session(scope=session_scope, binding_id=binding_id, snapshot=terminal.session_snapshot) - self._persist_message_usage(message_id=message_id, usage=_llm_usage_from_agent_backend(terminal.usage)) - if queue_manager is not None and not queue_manager.is_stopped(): - # A stream error can cancel its runtime only for cleanup. Keep its - # snapshot/usage, but let the original business failure own the root. - # Check after cleanup so a Stop observed during its wait still wins. - return - from core.ops.base_trace_instance import application_tracing_hooks - - tracing_hooks = application_tracing_hooks() - if tracing_hooks is not None: - outcome: Literal["succeeded", "failed", "cancelled"] - if isinstance(terminal, AgentBackendRunSucceededInternalEvent): - outcome = "succeeded" - elif isinstance(terminal, AgentBackendRunFailedInternalEvent): - outcome = "failed" - else: - outcome = "cancelled" - tracing_hooks.capture_agent_app_terminal( - tenant_id=dify_context.tenant_id, - app_id=dify_context.app_id, - message_id=message_id, - outcome=outcome, - output=( - self._terminal_output_to_answer(terminal.output) - if isinstance(terminal, AgentBackendRunSucceededInternalEvent) - else None - ), - ) - - def _publish_terminal_answer( - self, - *, - queue_manager: AppQueueManager, - model_name: str, - answer: str, - query: str | None, - usage: LLMUsage | None, - ) -> None: - """Finish a successful turn from the backend terminal output.""" - publish_text_answer( - queue_manager=queue_manager, - model_name=model_name, - answer=answer, - usage=usage, - user_query=query, - ) - - @staticmethod - def _persist_message_usage(*, message_id: str, usage: LLMUsage | None) -> None: - """Persist terminal usage independently of the client-facing stream lifecycle.""" - if usage is None or (usage.total_tokens <= 0 and usage.total_price <= 0): - return - try: - # A stopped consumer may have committed its stop intent/partial - # answer through another session. Merge usage with that latest row - # under the same short lock used by its save, not an identity-cache - # snapshot that can erase the intent or newer reported usage. - message = db.session.get(Message, message_id, populate_existing=True, with_for_update=True) - if message is None: - logger.warning("Cannot persist Agent App usage: message not found: %s", message_id) - return - message.message_tokens = usage.prompt_tokens - message.message_unit_price = usage.prompt_unit_price - message.message_price_unit = usage.prompt_price_unit - message.answer_tokens = usage.completion_tokens - message.answer_unit_price = usage.completion_unit_price - message.answer_price_unit = usage.completion_price_unit - message.total_price = usage.total_price - message.currency = usage.currency - if usage.latency > 0: - message.provider_response_latency = usage.latency - - try: - metadata = json.loads(message.message_metadata) if message.message_metadata else {} - except (json.JSONDecodeError, TypeError): - metadata = {} - if not isinstance(metadata, dict): - metadata = {} - metadata["usage"] = usage.model_dump(mode="json") - message.message_metadata = json.dumps(metadata, ensure_ascii=False) - db.session.commit() - except Exception: - db.session.rollback() - logger.warning("Failed to persist partial Agent App usage: message_id=%s", message_id, exc_info=True) - - def _save_session( - self, - *, - scope: AgentAppSessionScope, - binding_id: str, - snapshot: Any, - ) -> bool: - try: - self._session_store.save_active_snapshot( - scope=scope, - binding_id=binding_id, - snapshot=snapshot, - ) - return True - except Exception: - logger.warning( - "Failed to persist Agent App conversation session snapshot: " - "tenant_id=%s app_id=%s conversation_id=%s agent_id=%s", - scope.tenant_id, - scope.app_id, - scope.conversation_id, - scope.agent_id, - exc_info=True, - ) - return False - - @staticmethod - def _terminal_output_to_answer(output: JsonValue) -> str: - """Normalize the backend's terminal output to assistant text. - - Free-text Agent Apps return a plain string; if a structured output is - configured the value is a JSON object, which we serialize so the chat - message always has a string body. - """ - if output is None: - return "" - if isinstance(output, str): - return output - if isinstance(output, dict): - text = output.get("text") - if isinstance(text, str): - return text - return json.dumps(output, ensure_ascii=False) - return json.dumps(output, ensure_ascii=False) - - -__all__ = ["AgentAppRunner", "publish_message_end", "publish_text_answer", "publish_text_delta"] diff --git a/api/core/app/apps/agent_app/errors.py b/api/core/app/apps/agent_app/errors.py deleted file mode 100644 index 51b4e77116a291..00000000000000 --- a/api/core/app/apps/agent_app/errors.py +++ /dev/null @@ -1,6 +0,0 @@ -class AgentAppGeneratorError(ValueError): - """Raised when an Agent App turn cannot be set up.""" - - -class AgentAppNotPublishedError(AgentAppGeneratorError): - """Raised when a public Agent App runtime is requested before publish.""" diff --git a/api/core/app/apps/agent_app/runtime_request_builder.py b/api/core/app/apps/agent_app/runtime_request_builder.py deleted file mode 100644 index 608194e8d07970..00000000000000 --- a/api/core/app/apps/agent_app/runtime_request_builder.py +++ /dev/null @@ -1,349 +0,0 @@ -"""Build dify-agent run requests for one Agent App conversation turn. - -Mirrors the workflow ``WorkflowAgentRuntimeRequestBuilder`` but for the Agent -App surface: the user prompt is the chat message (no workflow-node job / no -previous-node context), multi-turn continuity flows through the -conversation-keyed ``session_snapshot`` plus the history layer, and Agent Soul -knowledge config is mapped into the same fixed ``dify.knowledge_base`` layer -used by workflow runs. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Literal, cast - -from dify_agent.layers.execution_context import ( - DifyExecutionContextInvokeFrom, - DifyExecutionContextLayerConfig, - DifyExecutionContextUserFrom, -) -from dify_agent.layers.user_prompt import ( - DifyUserPromptDownloadConfig, - DifyUserPromptFileConfig, - DifyUserPromptFileType, - DifyUserPromptImageConfig, -) -from dify_agent.protocol import CreateRunRequest -from dify_agent.protocol.knowledge_fs import KnowledgeFsError -from dify_agent.protocol.snapshot import SessionSnapshot - -from clients.agent_backend import ( - AgentBackendAgentAppRunInput, - AgentBackendModelConfig, - AgentBackendRunRequestBuilder, - redact_for_agent_backend_log, -) -from configs import dify_config -from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom -from core.app.llm.model_access import resolve_model_context_window, resolve_model_supports_vision -from core.ops.base_trace_instance import application_tracing_hooks -from core.plugin.provider_identity import normalize_plugin_daemon_provider_identity -from core.workflow.file_reference import build_file_reference, is_canonical_file_reference -from graphon.file import File, FileTransferMethod, FileType, file_manager -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent -from models.agent_config_entities import AgentKnowledgeSpaceConfig, AgentSoulConfig, AgentSoulToolsConfig -from models.provider_ids import ModelProviderID -from services.agent.knowledge_runtime_config import build_knowledge_fs_layer_config -from services.agent.prompt_mentions import expand_prompt_mentions -from services.workflow.execution.adapters.agent_v2.dify_tools_builder import ( - WorkflowAgentDifyToolLayersBuilder, - WorkflowAgentDifyToolsBuildError, - WorkflowAgentToolLayers, -) -from services.workflow.execution.adapters.agent_v2.runtime_request_builder import ( - append_runtime_warnings, - build_config_aware_soul_mention_resolver, - build_config_layer_config, - build_knowledge_layer_config, - build_shell_layer_config, - load_runtime_agent_skill_configs, -) - - -class AgentAppRuntimeRequestBuildError(ValueError): - """Raised when Agent App state cannot be mapped to a valid run request.""" - - def __init__(self, error_code: str, message: str) -> None: - self.error_code = error_code - super().__init__(message) - - -type _ReferenceFileTransferMethod = Literal["local_file", "tool_file", "datasource_file"] - - -@dataclass(frozen=True, slots=True) -class AgentAppRuntimeBuildContext: - dify_context: DifyRunContext - agent_id: str - agent_config_snapshot_id: str - agent_soul: AgentSoulConfig - conversation_id: str - user_query: str - idempotency_key: str - binding_id: str - backend_binding_ref: str - files: tuple[File, ...] = () - image_detail_config: ImagePromptMessageContent.DETAIL | None = None - agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot" - session_snapshot: SessionSnapshot | None = None - - -@dataclass(frozen=True, slots=True) -class AgentAppRuntimeRequest: - request: CreateRunRequest - redacted_request: dict[str, Any] - metadata: dict[str, Any] - binding_id: str - - -class AgentAppRuntimeRequestBuilder: - """Build dify-agent run requests from Agent App conversation state.""" - - def __init__( - self, - *, - request_builder: AgentBackendRunRequestBuilder | None = None, - dify_tools_builder: WorkflowAgentDifyToolLayersBuilder, - ) -> None: - self._request_builder = request_builder or AgentBackendRunRequestBuilder() - self._dify_tools_builder = dify_tools_builder - - def build(self, context: AgentAppRuntimeBuildContext) -> AgentAppRuntimeRequest: - agent_soul = context.agent_soul - if agent_soul.model is None: - raise AgentAppRuntimeRequestBuildError( - "agent_model_not_configured", - "Agent App requires the Agent Soul model to be configured.", - ) - - metadata = self._build_metadata(context) - try: - tool_layers = self._build_tool_layers( - tenant_id=context.dify_context.tenant_id, - app_id=context.dify_context.app_id, - user_id=context.dify_context.user_id, - tools=agent_soul.tools, - invoke_from=context.dify_context.invoke_from, - knowledge_fs_spaces=tuple(agent_soul.knowledge.spaces), - ) - except WorkflowAgentDifyToolsBuildError as error: - raise AgentAppRuntimeRequestBuildError(error.error_code, str(error)) from error - if tool_layers.plugin_tools is not None or tool_layers.core_tools is not None or agent_soul.tools.cli_tools: - metadata["agent_tools"] = { - "dify_tool_count": len(tool_layers.exposed_tool_names()), - "dify_tool_names": tool_layers.exposed_tool_names(), - "cli_tool_count": len(agent_soul.tools.cli_tools), - } - - runtime_config_skills = load_runtime_agent_skill_configs( - tenant_id=context.dify_context.tenant_id, - agent_id=context.agent_id, - ) - config_layer_config, config_warnings = build_config_layer_config( - agent_soul, - agent_id=context.agent_id, - config_version_id=context.agent_config_snapshot_id, - config_version_kind=context.agent_config_version_kind, - runtime_config_skills=runtime_config_skills, - ) - append_runtime_warnings(metadata, config_warnings) - soul_prompt_resolver = build_config_aware_soul_mention_resolver( - agent_soul, - runtime_config_skills=runtime_config_skills, - ) - try: - knowledge_fs_config = build_knowledge_fs_layer_config(agent_soul, run_context=context.dify_context) - except KnowledgeFsError as exc: - raise AgentAppRuntimeRequestBuildError(exc.code, exc.message) from exc - context_window_tokens = resolve_model_context_window( - run_context=context.dify_context, - provider_name=agent_soul.model.model_provider, - model_name=agent_soul.model.model, - ) - model_plugin_id, model_provider = normalize_plugin_daemon_provider_identity( - ModelProviderID(agent_soul.model.model_provider), - agent_soul.model.plugin_id, - ) - user_files = self._build_user_files( - files=context.files, - run_context=context.dify_context, - provider_name=agent_soul.model.model_provider, - model_name=agent_soul.model.model, - image_detail_config=context.image_detail_config, - include_image_references=knowledge_fs_config is not None, - ) - - tracing_hooks = application_tracing_hooks() - request = self._request_builder.build_for_agent_app( - AgentBackendAgentAppRunInput( - model=AgentBackendModelConfig( - plugin_id=model_plugin_id, - model_provider=model_provider, - model=agent_soul.model.model, - model_settings=agent_soul.model.model_settings.model_dump(mode="json", exclude_none=True), - context_window_tokens=context_window_tokens, - ), - execution_context=DifyExecutionContextLayerConfig( - tenant_id=context.dify_context.tenant_id, - user_id=context.dify_context.user_id, - app_id=context.dify_context.app_id, - conversation_id=context.conversation_id, - agent_id=context.agent_id, - agent_config_version_id=context.agent_config_snapshot_id, - agent_config_version_kind=context.agent_config_version_kind, - # Agent Files §1.3: real Dify access context + agent run mode. - user_from=cast(DifyExecutionContextUserFrom, context.dify_context.user_from.value), - invoke_from=cast(DifyExecutionContextInvokeFrom, context.dify_context.invoke_from.value), - agent_mode="agent_app", - trace_id=context.dify_context.trace_session_id, - tracing=tracing_hooks.prepare_agent_app(context) if tracing_hooks is not None else None, - ), - backend_binding_ref=context.backend_binding_ref, - # ENG-616: expand slash-menu mention tokens to canonical names so - # no frontend-internal {{#…#}} marker ever reaches the model. - agent_soul_prompt=expand_prompt_mentions(agent_soul.prompt.system_prompt, soul_prompt_resolver).strip() - or None, - agent_config_version_kind=context.agent_config_version_kind, - user_prompt=context.user_query, - user_files=user_files, - tools=tool_layers.plugin_tools, - core_tools=tool_layers.core_tools, - knowledge=build_knowledge_layer_config(agent_soul), - knowledge_fs=knowledge_fs_config, - config_layer_config=config_layer_config, - include_shell=dify_config.AGENT_SHELL_ENABLED, - shell_config=build_shell_layer_config(agent_soul), - session_snapshot=context.session_snapshot, - idempotency_key=context.idempotency_key, - metadata=metadata, - ) - ) - redacted = cast(dict[str, Any], redact_for_agent_backend_log(request)) - return AgentAppRuntimeRequest( - request=request, - redacted_request=redacted, - metadata=metadata, - binding_id=context.binding_id, - ) - - @staticmethod - def _build_user_files( - *, - files: tuple[File, ...], - run_context: DifyRunContext, - provider_name: str, - model_name: str, - image_detail_config: ImagePromptMessageContent.DETAIL | None, - include_image_references: bool, - ) -> list[DifyUserPromptFileConfig]: - supports_vision = any(file.type == FileType.IMAGE for file in files) and resolve_model_supports_vision( - run_context=run_context, - provider_name=provider_name, - model_name=model_name, - ) - user_files: list[DifyUserPromptFileConfig] = [] - for file in files: - if supports_vision and file.type == FileType.IMAGE: - user_files.append(_build_user_image(file, image_detail_config=image_detail_config)) - # KnowledgeFS image search needs the upload reference even when - # the model also receives the image as native vision content. - if not include_image_references or file.transfer_method != FileTransferMethod.LOCAL_FILE: - continue - user_files.append(_build_user_download(file)) - return user_files - - def _build_tool_layers( - self, - *, - tenant_id: str, - app_id: str, - user_id: str | None, - tools: AgentSoulToolsConfig, - invoke_from: InvokeFrom, - knowledge_fs_spaces: tuple[AgentKnowledgeSpaceConfig, ...], - ) -> WorkflowAgentToolLayers: - # Production Agent App runs intentionally keep existing plugin configs - # on the direct `dify.plugin.tools` route. This builder emits plugin - # tools directly and non-plugin Dify tools through `dify.core.tools`. - return self._dify_tools_builder.build_layers( - tenant_id=tenant_id, - app_id=app_id, - user_id=user_id, - tools=tools, - invoke_from=invoke_from, - knowledge_fs_spaces=knowledge_fs_spaces, - ) - - @staticmethod - def _build_metadata(context: AgentAppRuntimeBuildContext) -> dict[str, Any]: - return { - "tenant_id": context.dify_context.tenant_id, - "app_id": context.dify_context.app_id, - "conversation_id": context.conversation_id, - "agent_id": context.agent_id, - "agent_config_snapshot_id": context.agent_config_snapshot_id, - } - - -def _build_user_image( - file: File, - *, - image_detail_config: ImagePromptMessageContent.DETAIL | None, -) -> DifyUserPromptImageConfig: - content = file_manager.to_prompt_message_content(file, image_detail_config=image_detail_config) - if not isinstance(content, ImagePromptMessageContent): - raise AgentAppRuntimeRequestBuildError( - "agent_user_file_unsupported", - f"Agent App cannot send file '{file.filename or 'image'}' as vision content.", - ) - detail = content.detail.value - return DifyUserPromptImageConfig( - filename=content.filename or file.filename or f"image.{content.format}", - mime_type=content.mime_type, - format=content.format, - url=content.url or None, - base64_data=content.base64_data or None, - detail=detail if detail in {"low", "high"} else None, - ) - - -def _build_user_download(file: File) -> DifyUserPromptDownloadConfig: - file_type = cast(DifyUserPromptFileType, file.type.value) - if file.transfer_method == FileTransferMethod.REMOTE_URL: - if file.remote_url is None: - raise AgentAppRuntimeRequestBuildError("agent_user_file_invalid", "Remote user file is missing its URL.") - return DifyUserPromptDownloadConfig( - type=file_type, filename=file.filename or None, transfer_method="remote_url", url=file.remote_url - ) - if file.reference is None: - raise AgentAppRuntimeRequestBuildError("agent_user_file_invalid", "User file is missing its reference.") - reference = file.reference - if not reference.startswith("dify-file-ref:"): - reference = build_file_reference(record_id=reference) - elif not is_canonical_file_reference(reference): - raise AgentAppRuntimeRequestBuildError("agent_user_file_invalid", "User file reference is invalid.") - transfer_method: _ReferenceFileTransferMethod - match file.transfer_method: - case FileTransferMethod.LOCAL_FILE: - transfer_method = "local_file" - case FileTransferMethod.TOOL_FILE: - transfer_method = "tool_file" - case FileTransferMethod.DATASOURCE_FILE: - transfer_method = "datasource_file" - case _: - raise AgentAppRuntimeRequestBuildError( - "agent_user_file_invalid", - f"User file transfer method '{file.transfer_method.value}' is unsupported.", - ) - return DifyUserPromptDownloadConfig( - type=file_type, filename=file.filename or None, transfer_method=transfer_method, reference=reference - ) - - -__all__ = [ - "AgentAppRuntimeBuildContext", - "AgentAppRuntimeRequest", - "AgentAppRuntimeRequestBuildError", - "AgentAppRuntimeRequestBuilder", -] diff --git a/api/core/app/apps/agent_chat/app_generator.py b/api/core/app/apps/agent_chat/app_generator.py deleted file mode 100644 index 2f661f066452a5..00000000000000 --- a/api/core/app/apps/agent_chat/app_generator.py +++ /dev/null @@ -1,304 +0,0 @@ -import contextvars -import logging -import threading -import uuid -from collections.abc import Generator, Mapping -from typing import TYPE_CHECKING, Any, Literal, overload - -from flask import Flask, current_app -from pydantic import ValidationError -from sqlalchemy.orm import Session - -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.db.session_factory import session_factory -from core.helper.trace_id_helper import extract_trace_session_id_from_args -from core.ops.ops_trace_manager import TraceQueueManager -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 models.model import load_annotation_reply_config -from services.conversation_service import ConversationService - -if TYPE_CHECKING: - from core.app.apps.workflow_app_runner import WorkflowRunDriver - -logger = logging.getLogger(__name__) - - -class AgentChatAppGenerator(MessageBasedAppGenerator): - def __init__(self, *, execution_driver: "WorkflowRunDriver") -> None: - self._execution_driver = execution_driver - - @overload - def generate( - self, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: Literal[False], - session: Session, - ) -> Mapping[str, Any]: ... - - @overload - def generate( - self, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: Literal[True], - session: Session, - ) -> Generator[Mapping | str]: ... - - @overload - def generate( - self, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: bool, - session: Session, - ) -> Mapping | Generator[Mapping | str]: ... - - def generate( - self, - *, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: bool = True, - session: Session, - ) -> Mapping | Generator[Mapping | str]: - """ - 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), - **extract_trace_session_id_from_args(args), - } - - # 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, session=session - ) - # get app model config - app_model_config = self._get_app_model_config( - app_model=app_model, - conversation=conversation, - session=session, - ) - - # 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"], - session=session, - ) - - # always enable retriever resource in debugger mode - override_model_config_dict["retriever_resource"] = {"enabled": True} - - if override_model_config_dict: - annotation_reply = None - effective_model_config_dict = override_model_config_dict - else: - annotation_reply = load_annotation_reply_config(session, app_model_config.app_id) - effective_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply) - - # 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(effective_model_config_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, - annotation_reply=annotation_reply, - ) - - # 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, - session=session, - ) - - # 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, - 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(execution_driver=self._execution_driver) - with session_factory.create_session() as session: - runner.run( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation=conversation, - message=message, - session=session, - ) - 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) diff --git a/api/core/app/apps/agent_chat/app_runner.py b/api/core/app/apps/agent_chat/app_runner.py deleted file mode 100644 index b67565f5d9e1ee..00000000000000 --- a/api/core/app/apps/agent_chat/app_runner.py +++ /dev/null @@ -1,269 +0,0 @@ -import logging -from typing import TYPE_CHECKING, cast - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.agent.cot_chat_agent_runner import CotChatAgentRunner -from core.agent.cot_completion_agent_runner import CotCompletionAgentRunner -from core.agent.entities import AgentEntity -from core.agent.fc_agent_runner import FunctionCallAgentRunner -from core.app.apps.agent_chat.app_config_manager import AgentChatAppConfig -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.base_app_runner import AppRunner -from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity -from core.app.entities.queue_entities import QueueAnnotationReplyEvent -from core.db.session_factory import create_session -from core.memory.token_buffer_memory import TokenBufferMemory -from core.model_manager import ModelInstance -from core.moderation.base import ModerationError -from graphon.model_runtime.entities.llm_entities import LLMMode -from graphon.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from models.model import App, Conversation, Message - -if TYPE_CHECKING: - from core.app.apps.workflow_app_runner import WorkflowRunDriver - -logger = logging.getLogger(__name__) - - -class AgentChatAppRunner(AppRunner): - """ - Agent Application Runner - """ - - def __init__(self, *, execution_driver: "WorkflowRunDriver") -> None: - self._execution_driver = execution_driver - - def run( - self, - application_generate_entity: AgentChatAppGenerateEntity, - queue_manager: AppQueueManager, - conversation: Conversation, - message: Message, - session: Session, - ): - """Run the assistant application with bounded explicit transactions. - - The setup session is committed and released before the multi-step agent - runner begins model or tool I/O. - - :param application_generate_entity: application generate entity - :param queue_manager: application queue manager - :param conversation: conversation - :param message: message - :return: - """ - app_config = application_generate_entity.app_config - app_config = cast(AgentChatAppConfig, app_config) - app_stmt = select(App).where(App.id == app_config.app_id) - with create_session() as read_session: - app_record = read_session.scalar(app_stmt) - if app_record: - read_session.expunge(app_record) - if not app_record: - raise ValueError("App not found") - - inputs = application_generate_entity.inputs - query = application_generate_entity.query - files = application_generate_entity.files - - memory = None - if application_generate_entity.conversation_id: - # get memory of conversation (read-only) - model_instance = ModelInstance( - provider_model_bundle=application_generate_entity.model_conf.provider_model_bundle, - model=application_generate_entity.model_conf.model, - ) - - memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance) - - # organize all inputs and template to prompt messages - # Include: prompt template, inputs, query(optional), files(optional) - # memory(optional) - prompt_messages, _ = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=dict(inputs), - files=list(files), - query=query, - memory=memory, - ) - - # moderation - try: - # process sensitive_word_avoidance - _, inputs, query = self.moderation_for_inputs( - app_id=app_record.id, - tenant_id=app_config.tenant_id, - app_generate_entity=application_generate_entity, - inputs=dict(inputs), - query=query or "", - message_id=message.id, - ) - except ModerationError as e: - self.direct_output( - queue_manager=queue_manager, - app_generate_entity=application_generate_entity, - prompt_messages=prompt_messages, - text=str(e), - stream=application_generate_entity.stream, - ) - return - - if query: - # annotation reply - annotation_reply = self.query_app_annotations_to_reply( - app_record=app_record, - message=message, - query=query, - user_id=application_generate_entity.user_id, - invoke_from=application_generate_entity.invoke_from, - session=session, - ) - session.commit() - session.close() - - if annotation_reply: - queue_manager.publish( - QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id), - PublishFrom.APPLICATION_MANAGER, - ) - - self.direct_output( - queue_manager=queue_manager, - app_generate_entity=application_generate_entity, - prompt_messages=prompt_messages, - text=annotation_reply.content, - stream=application_generate_entity.stream, - ) - return - - # fill in variable inputs from external data tools if exists - external_data_tools = app_config.external_data_variables - if external_data_tools: - inputs = self.fill_in_inputs_from_external_data_tools( - tenant_id=app_record.tenant_id, - app_id=app_record.id, - external_data_tools=external_data_tools, - inputs=inputs, - query=query, - ) - - # reorganize all inputs and template to prompt messages - # Include: prompt template, inputs, query(optional), files(optional) - # memory(optional), external data, dataset context(optional) - prompt_messages, _ = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=dict(inputs), - files=list(files), - query=query, - memory=memory, - ) - - # check hosting moderation - hosting_moderation_result = self.check_hosting_moderation( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - prompt_messages=prompt_messages, - ) - - if hosting_moderation_result: - return - - agent_entity = app_config.agent - assert agent_entity is not None - - # init model instance - model_instance = ModelInstance( - provider_model_bundle=application_generate_entity.model_conf.provider_model_bundle, - model=application_generate_entity.model_conf.model, - ) - prompt_message, _ = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=dict(inputs), - files=list(files), - query=query, - memory=memory, - ) - - # change function call strategy based on LLM model - llm_model = cast(LargeLanguageModel, model_instance.model_type_instance) - model_schema = llm_model.get_model_schema(model_instance.model_name, model_instance.credentials) - if not model_schema: - raise ValueError("Model schema not found") - - if {ModelFeature.MULTI_TOOL_CALL, ModelFeature.TOOL_CALL}.intersection(model_schema.features or []): - agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING - conversation_stmt = select(Conversation).where(Conversation.id == conversation.id) - msg_stmt = select(Message).where(Message.id == message.id) - with create_session() as read_session: - conversation_result = read_session.scalar(conversation_stmt) - if conversation_result is None: - raise ValueError("Conversation not found") - - message_result = read_session.scalar(msg_stmt) - if message_result is not None: - read_session.expunge(message_result) - read_session.expunge(conversation_result) - if message_result is None: - raise ValueError("Message not found") - - runner_cls: type[FunctionCallAgentRunner] | type[CotChatAgentRunner] | type[CotCompletionAgentRunner] - # start agent runner - if agent_entity.strategy == AgentEntity.Strategy.CHAIN_OF_THOUGHT: - # check LLM mode - if model_schema.model_properties.get(ModelPropertyKey.MODE) == LLMMode.CHAT: - runner_cls = CotChatAgentRunner - elif model_schema.model_properties.get(ModelPropertyKey.MODE) == LLMMode.COMPLETION: - runner_cls = CotCompletionAgentRunner - else: - raise ValueError(f"Invalid LLM mode: {model_schema.model_properties.get(ModelPropertyKey.MODE)}") - elif agent_entity.strategy == AgentEntity.Strategy.FUNCTION_CALLING: - runner_cls = FunctionCallAgentRunner - else: - raise ValueError(f"Invalid agent strategy: {agent_entity.strategy}") - - runner = runner_cls( - execution_driver=self._execution_driver, - session=session, - tenant_id=app_config.tenant_id, - application_generate_entity=application_generate_entity, - conversation=conversation_result, - app_config=app_config, - model_config=application_generate_entity.model_conf, - config=agent_entity, - queue_manager=queue_manager, - message=message_result, - user_id=application_generate_entity.user_id, - memory=memory, - prompt_messages=prompt_message, - model_instance=model_instance, - ) - - session.commit() - session.close() - - invoke_result = runner.run( - session=session, - message=message, - query=query, - inputs=inputs, - ) - - # handle invoke result - self._handle_invoke_result( - invoke_result=invoke_result, - queue_manager=queue_manager, - stream=application_generate_entity.stream, - agent=True, - message_id=message.id, - user_id=application_generate_entity.user_id, - tenant_id=app_config.tenant_id, - ) diff --git a/api/core/app/apps/base_app_generator.py b/api/core/app/apps/base_app_generator.py deleted file mode 100644 index 4a23bb71af4cbe..00000000000000 --- a/api/core/app/apps/base_app_generator.py +++ /dev/null @@ -1,376 +0,0 @@ -import logging -import threading -from collections.abc import Generator, Mapping, Sequence -from contextlib import AbstractContextManager, nullcontext -from typing import TYPE_CHECKING, Any, final - -from sqlalchemy.orm import Session -from sqlalchemy.orm.attributes import set_committed_value - -from core.app.apps.draft_variable_saver import ( - DraftVariableSaver, - DraftVariableSaverFactory, - NoopDraftVariableSaver, -) -from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.app.file_access import ( - DatabaseFileAccessController, - FileAccessScope, - bind_file_access_scope, -) -from extensions.ext_database import db -from factories import file_factory -from graphon.enums import NodeType -from graphon.file import File, FileUploadConfig -from graphon.variables.input_entities import VariableEntityType -from libs.orjson import orjson_dumps -from models import Account, EndUser, Workflow, WorkflowRun -from services.app.file_input_defaults import apply_file_input_defaults -from services.workflow_draft_variable_service import DraftVariableSaver as DraftVariableSaverImpl - -if TYPE_CHECKING: - from graphon.variables.input_entities import VariableEntity - -logger = logging.getLogger(__name__) - -_WORKER_THREAD_JOIN_TIMEOUT_SECONDS = 300 - - -@final -class _DebuggerDraftVariableSaver: - """Adapter that binds SQLAlchemy session setup outside the saver port.""" - - def __init__( - self, - *, - account: Account, - tenant_id: str, - app_id: str, - node_id: str, - node_type: NodeType, - node_execution_id: str, - enclosing_node_id: str | None = None, - ) -> None: - self._account = account - self._tenant_id = tenant_id - self._app_id = app_id - self._node_id = node_id - self._node_type = node_type - self._node_execution_id = node_execution_id - self._enclosing_node_id = enclosing_node_id - - def save(self, process_data: Mapping[str, Any] | None, outputs: Mapping[str, Any] | None) -> None: - with Session(db.engine) as session, session.begin(): - DraftVariableSaverImpl( - session=session, - tenant_id=self._tenant_id, - app_id=self._app_id, - node_id=self._node_id, - node_type=self._node_type, - node_execution_id=self._node_execution_id, - enclosing_node_id=self._enclosing_node_id, - user=self._account, - ).save(process_data, outputs) - - -class BaseAppGenerator: - _file_access_controller: DatabaseFileAccessController = DatabaseFileAccessController() - - @staticmethod - def _restore_workflow_run_graph(*, session: Session, workflow: Workflow, workflow_run_id: str | None) -> None: - if workflow_run_id is None: - raise ValueError("Workflow run id is required when resuming") - workflow_run = session.get(WorkflowRun, workflow_run_id) - if workflow_run is None or workflow_run.graph is None: - raise ValueError(f"Workflow run graph not found: {workflow_run_id}") - set_committed_value(workflow, "graph", workflow_run.graph) - - @staticmethod - def _join_worker_thread(worker_thread: threading.Thread) -> None: - # Bound the wait so a leaked app worker cannot occupy an execution slot indefinitely. - worker_thread.join(timeout=_WORKER_THREAD_JOIN_TIMEOUT_SECONDS) - if worker_thread.is_alive(): - logger.warning( - "Possible app worker thread leak: thread_name=%s timeout_seconds=%s; " - "continuing without waiting further to avoid occupying an execution slot indefinitely", - worker_thread.name, - _WORKER_THREAD_JOIN_TIMEOUT_SECONDS, - ) - - @staticmethod - def _wrap_stream_with_worker_thread_join[ResponseT]( - response_stream: Generator[ResponseT], - worker_thread: threading.Thread, - ) -> Generator[ResponseT]: - """Keep the producer owned by the response stream until both finish.""" - try: - yield from response_stream - finally: - BaseAppGenerator._join_worker_thread(worker_thread) - - @staticmethod - def _bind_file_access_scope( - *, - tenant_id: str, - user: Account | EndUser, - invoke_from: InvokeFrom, - ) -> AbstractContextManager[None]: - """Bind request-scoped file ownership markers for downstream file lookups.""" - - user_id = getattr(user, "id", None) - if not isinstance(user_id, str) or not user_id: - return nullcontext() - - user_from = UserFrom.ACCOUNT if isinstance(user, Account) else UserFrom.END_USER - return bind_file_access_scope( - FileAccessScope( - tenant_id=tenant_id, - user_id=user_id, - user_from=user_from, - invoke_from=invoke_from, - ) - ) - - def _prepare_user_inputs( - self, - *, - user_inputs: Mapping[str, Any] | None, - variables: Sequence["VariableEntity"], - tenant_id: str, - strict_type_validation: bool = False, - ) -> Mapping[str, Any]: - user_inputs = apply_file_input_defaults(user_inputs, variables) - # Filter input variables from form configuration, handle required fields, default values, and option values - user_inputs = { - var.variable: self._validate_inputs(value=user_inputs.get(var.variable), variable_entity=var) - for var in variables - } - user_inputs = {k: self._sanitize_value(v) for k, v in user_inputs.items()} - # Convert files in inputs to File - entity_dictionary = {item.variable: item for item in variables} - # Convert single file to File - files_inputs = { - k: file_factory.build_from_mapping( - mapping=v, - tenant_id=tenant_id, - config=FileUploadConfig( - allowed_file_types=entity_dictionary[k].allowed_file_types or [], - allowed_file_extensions=entity_dictionary[k].allowed_file_extensions or [], - allowed_file_upload_methods=entity_dictionary[k].allowed_file_upload_methods or [], - ), - strict_type_validation=strict_type_validation, - access_controller=self._file_access_controller, - ) - for k, v in user_inputs.items() - if isinstance(v, dict) and entity_dictionary[k].type == VariableEntityType.FILE - } - # Convert list of files to File - file_list_inputs = { - k: file_factory.build_from_mappings( - mappings=v, - tenant_id=tenant_id, - config=FileUploadConfig( - allowed_file_types=entity_dictionary[k].allowed_file_types or [], - allowed_file_extensions=entity_dictionary[k].allowed_file_extensions or [], - allowed_file_upload_methods=entity_dictionary[k].allowed_file_upload_methods or [], - ), - access_controller=self._file_access_controller, - ) - for k, v in user_inputs.items() - if isinstance(v, list) - # Ensure skip List - and all(isinstance(item, dict) for item in v) - and entity_dictionary[k].type == VariableEntityType.FILE_LIST - } - # Merge all inputs - user_inputs = {**user_inputs, **files_inputs, **file_list_inputs} - - # Check if all files are converted to File - invalid_dict_keys = [ - k - for k, v in user_inputs.items() - if isinstance(v, dict) - and entity_dictionary[k].type not in {VariableEntityType.FILE, VariableEntityType.JSON_OBJECT} - ] - if invalid_dict_keys: - raise ValueError(f"Invalid input type for {invalid_dict_keys}") - - invalid_list_dict_keys = [ - k - for k, v in user_inputs.items() - if isinstance(v, list) - and any(isinstance(item, dict) for item in v) - and entity_dictionary[k].type != VariableEntityType.FILE_LIST - ] - if invalid_list_dict_keys: - raise ValueError(f"Invalid input type for {invalid_list_dict_keys}") - - return user_inputs - - def _validate_inputs( - self, - *, - variable_entity: "VariableEntity", - value: Any, - ): - if value is None: - if variable_entity.required: - raise ValueError(f"{variable_entity.variable} is required in input form") - # Use default value and continue validation to ensure type conversion - value = variable_entity.default - # If default is also None, return None directly - if value is None: - return None - - # Treat empty placeholders for optional file inputs as unset - if ( - variable_entity.type in {VariableEntityType.FILE, VariableEntityType.FILE_LIST} - and not variable_entity.required - ): - # Treat empty string (frontend default) as unset - # For FILE_LIST, allow empty list [] to pass through - if isinstance(value, str) and not value: - return None - - if variable_entity.type in { - VariableEntityType.TEXT_INPUT, - VariableEntityType.SELECT, - VariableEntityType.PARAGRAPH, - } and not isinstance(value, str): - raise ValueError( - f"(type '{variable_entity.type}') {variable_entity.variable} in input form must be a string" - ) - - if variable_entity.type == VariableEntityType.NUMBER: - match value: - case int() | float(): - return value - case str(): - # handle empty string case - if not value.strip(): - return None - # may raise ValueError if user_input_value is not a valid number - try: - if "." in value: - return float(value) - else: - return int(value) - except ValueError as caught_error: - raise ValueError( - f"{variable_entity.variable} in input form must be a valid number" - ) from caught_error - case _: - raise TypeError(f"expected value type int, float or str, got {type(value)}, value: {value}") - - match variable_entity.type: - case VariableEntityType.SELECT: - if value not in variable_entity.options: - raise ValueError( - f"{variable_entity.variable} in input form must be one of the following: " - f"{variable_entity.options}" - ) - case VariableEntityType.TEXT_INPUT | VariableEntityType.PARAGRAPH: - if variable_entity.max_length and len(value) > variable_entity.max_length: - raise ValueError( - f"{variable_entity.variable} in input form must be less than {variable_entity.max_length} " - "characters" - ) - case VariableEntityType.FILE: - if not isinstance(value, dict) and not isinstance(value, File): - raise ValueError(f"{variable_entity.variable} in input form must be a file") - case VariableEntityType.FILE_LIST: - # if number of files exceeds the limit, raise ValueError - if not ( - isinstance(value, list) - and (all(isinstance(item, dict) for item in value) or all(isinstance(item, File) for item in value)) - ): - raise ValueError(f"{variable_entity.variable} in input form must be a list of files") - - if variable_entity.max_length and len(value) > variable_entity.max_length: - raise ValueError( - f"{variable_entity.variable} in input form must be less than {variable_entity.max_length} files" - ) - case VariableEntityType.CHECKBOX: - match value: - case str(): - normalized_value = value.strip().lower() - if normalized_value in {"true", "1", "yes", "on"}: - value = True - elif normalized_value in {"false", "0", "no", "off"}: - value = False - case int() | float(): - if value == 1: - value = True - elif value == 0: - value = False - case VariableEntityType.JSON_OBJECT: - if value and not isinstance(value, dict): - raise ValueError(f"{variable_entity.variable} in input form must be a dict") - case _: - raise AssertionError("this statement should be unreachable.") - - return value - - def _sanitize_value(self, value: Any): - if isinstance(value, str): - return value.replace("\x00", "") - return value - - @classmethod - def convert_to_event_stream(cls, generator: Mapping | Generator[Mapping | str]): - """ - Convert messages into event stream - """ - if isinstance(generator, dict): - return generator - else: - - def gen(): - for message in generator: - if isinstance(message, Mapping | dict): - yield f"data: {orjson_dumps(message)}\n\n" - else: - yield f"event: {message}\n\n" - - return gen() - - @final - @staticmethod - def _get_draft_var_saver_factory( - invoke_from: InvokeFrom, - account: Account | EndUser, - *, - tenant_id: str, - ) -> DraftVariableSaverFactory: - if invoke_from == InvokeFrom.DEBUGGER: - assert isinstance(account, Account) - - def draft_var_saver_factory( - app_id: str, - node_id: str, - node_type: NodeType, - node_execution_id: str, - enclosing_node_id: str | None = None, - ) -> DraftVariableSaver: - return _DebuggerDraftVariableSaver( - account=account, - tenant_id=tenant_id, - app_id=app_id, - node_id=node_id, - node_type=node_type, - node_execution_id=node_execution_id, - enclosing_node_id=enclosing_node_id, - ) - else: - - def draft_var_saver_factory( - app_id: str, - node_id: str, - node_type: NodeType, - node_execution_id: str, - enclosing_node_id: str | None = None, - ) -> DraftVariableSaver: - _ = app_id, node_id, node_type, node_execution_id, enclosing_node_id - return NoopDraftVariableSaver() - - return draft_var_saver_factory diff --git a/api/core/app/apps/base_app_runner.py b/api/core/app/apps/base_app_runner.py deleted file mode 100644 index 774cf255968bc7..00000000000000 --- a/api/core/app/apps/base_app_runner.py +++ /dev/null @@ -1,560 +0,0 @@ -import base64 -import logging -import time -from collections.abc import Generator, Mapping, Sequence -from mimetypes import guess_extension -from typing import TYPE_CHECKING, Any - -from sqlalchemy.orm import Session - -from core.app.app_config.entities import ExternalDataVariableEntity, PromptTemplateEntity -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.entities.app_invoke_entities import ( - AppGenerateEntity, - EasyUIBasedAppGenerateEntity, - InvokeFrom, - ModelConfigWithCredentialsEntity, -) -from core.app.entities.queue_entities import ( - QueueAgentMessageEvent, - QueueLLMChunkEvent, - QueueMessageEndEvent, - QueueMessageFileEvent, -) -from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature -from core.app.features.hosting_moderation.hosting_moderation import HostingModerationFeature -from core.db.session_factory import session_factory -from core.external_data_tool.external_data_fetch import ExternalDataFetch -from core.memory.token_buffer_memory import TokenBufferMemory -from core.model_manager import ModelInstance -from core.moderation.input_moderation import InputModeration -from core.prompt.advanced_prompt_transform import AdvancedPromptTransform -from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig -from core.prompt.simple_prompt_transform import ModelMode, SimplePromptTransform -from core.tools.tool_file_manager import ToolFileManager -from graphon.file import FileTransferMethod, FileType -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage -from graphon.model_runtime.entities.message_entities import ( - AssistantPromptMessage, - ImagePromptMessageContent, - PromptMessage, - TextPromptMessageContent, -) -from graphon.model_runtime.entities.model_entities import ModelPropertyKey -from graphon.model_runtime.errors.invoke import InvokeBadRequestError -from models.enums import CreatorUserRole, MessageFileBelongsTo -from models.model import App, AppMode, Message, MessageAnnotation, MessageFile - -if TYPE_CHECKING: - from graphon.file import File - -_logger = logging.getLogger(__name__) - - -class AppRunner: - def recalc_llm_max_tokens( - self, model_config: ModelConfigWithCredentialsEntity, prompt_messages: list[PromptMessage] - ): - # recalc max_tokens if sum(prompt_token + max_tokens) over model token limit - model_instance = ModelInstance( - provider_model_bundle=model_config.provider_model_bundle, model=model_config.model - ) - - model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE) - - max_tokens = 0 - for parameter_rule in model_config.model_schema.parameter_rules: - if parameter_rule.name == "max_tokens" or ( - parameter_rule.use_template and parameter_rule.use_template == "max_tokens" - ): - max_tokens = ( - model_config.parameters.get(parameter_rule.name) - or model_config.parameters.get(parameter_rule.use_template or "") - ) or 0 - - if model_context_tokens is None: - return -1 - - prompt_tokens = model_instance.get_llm_num_tokens(prompt_messages) - - if prompt_tokens + max_tokens > model_context_tokens: - max_tokens = max(model_context_tokens - prompt_tokens, 16) - - for parameter_rule in model_config.model_schema.parameter_rules: - if parameter_rule.name == "max_tokens" or ( - parameter_rule.use_template and parameter_rule.use_template == "max_tokens" - ): - model_config.parameters[parameter_rule.name] = max_tokens - - def organize_prompt_messages( - self, - app_record: App, - model_config: ModelConfigWithCredentialsEntity, - prompt_template_entity: PromptTemplateEntity, - inputs: Mapping[str, str], - files: Sequence["File"], - query: str = "", - context: str | None = None, - memory: TokenBufferMemory | None = None, - image_detail_config: ImagePromptMessageContent.DETAIL | None = None, - context_files: list["File"] | None = None, - ) -> tuple[list[PromptMessage], list[str] | None]: - """ - Organize prompt messages - :param context: - :param app_record: app record - :param model_config: model config entity - :param prompt_template_entity: prompt template entity - :param inputs: inputs - :param files: files - :param query: query - :param memory: memory - :param image_detail_config: the image quality config - :return: - """ - # get prompt without memory and context - if prompt_template_entity.prompt_type == PromptTemplateEntity.PromptType.SIMPLE: - prompt_transform: SimplePromptTransform | AdvancedPromptTransform - prompt_transform = SimplePromptTransform() - prompt_messages, stop = prompt_transform.get_prompt( - app_mode=AppMode.value_of(app_record.mode), - prompt_template_entity=prompt_template_entity, - inputs=inputs, - query=query, - files=files, - context=context, - memory=memory, - model_config=model_config, - image_detail_config=image_detail_config, - context_files=context_files, - ) - else: - memory_config = MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False)) - - model_mode = ModelMode(model_config.mode) - prompt_template: CompletionModelPromptTemplate | list[ChatModelMessage] - if model_mode == ModelMode.COMPLETION: - advanced_completion_prompt_template = prompt_template_entity.advanced_completion_prompt_template - if not advanced_completion_prompt_template: - raise InvokeBadRequestError("Advanced completion prompt template is required.") - prompt_template = CompletionModelPromptTemplate(text=advanced_completion_prompt_template.prompt) - - if advanced_completion_prompt_template.role_prefix: - memory_config.role_prefix = MemoryConfig.RolePrefix( - user=advanced_completion_prompt_template.role_prefix.user, - assistant=advanced_completion_prompt_template.role_prefix.assistant, - ) - else: - if not prompt_template_entity.advanced_chat_prompt_template: - raise InvokeBadRequestError("Advanced chat prompt template is required.") - prompt_template = [] - for message in prompt_template_entity.advanced_chat_prompt_template.messages: - prompt_template.append(ChatModelMessage(text=message.text, role=message.role)) - - prompt_transform = AdvancedPromptTransform() - prompt_messages = prompt_transform.get_prompt( - prompt_template=prompt_template, - inputs=inputs, - query=query or "", - files=files, - context=context, - memory_config=memory_config, - memory=memory, - model_config=model_config, - image_detail_config=image_detail_config, - ) - stop = model_config.stop - - return prompt_messages, stop - - def direct_output( - self, - queue_manager: AppQueueManager, - app_generate_entity: EasyUIBasedAppGenerateEntity, - prompt_messages: list, - text: str, - stream: bool, - usage: LLMUsage | None = None, - ): - """ - Direct output - :param queue_manager: application queue manager - :param app_generate_entity: app generate entity - :param prompt_messages: prompt messages - :param text: text - :param stream: stream - :param usage: usage - :return: - """ - if stream: - for index, token in enumerate(text): - chunk = LLMResultChunk( - model=app_generate_entity.model_conf.model, - prompt_messages=prompt_messages, - delta=LLMResultChunkDelta(index=index, message=AssistantPromptMessage(content=token)), - ) - - queue_manager.publish(QueueLLMChunkEvent(chunk=chunk), PublishFrom.APPLICATION_MANAGER) - time.sleep(0.01) - - queue_manager.publish( - QueueMessageEndEvent( - llm_result=LLMResult( - model=app_generate_entity.model_conf.model, - prompt_messages=prompt_messages, - message=AssistantPromptMessage(content=text), - usage=usage or LLMUsage.empty_usage(), - ), - ), - PublishFrom.APPLICATION_MANAGER, - ) - - def _handle_invoke_result( - self, - invoke_result: LLMResult | Generator[Any], - queue_manager: AppQueueManager, - stream: bool, - agent: bool = False, - message_id: str | None = None, - user_id: str | None = None, - tenant_id: str | None = None, - ): - """ - Handle invoke result - :param invoke_result: invoke result - :param queue_manager: application queue manager - :param stream: stream - :param agent: agent - :param message_id: message id for multimodal output - :param user_id: user id for multimodal output - :param tenant_id: tenant id for multimodal output - :return: - """ - match invoke_result: - case LLMResult() if not stream: - self._handle_invoke_result_direct( - invoke_result=invoke_result, - queue_manager=queue_manager, - ) - case _ if stream and isinstance(invoke_result, Generator): - self._handle_invoke_result_stream( - invoke_result=invoke_result, - queue_manager=queue_manager, - agent=agent, - message_id=message_id, - user_id=user_id, - tenant_id=tenant_id, - ) - case _: - raise NotImplementedError(f"unsupported invoke result type: {type(invoke_result)}") - - def _handle_invoke_result_direct( - self, - invoke_result: LLMResult, - queue_manager: AppQueueManager, - ): - """ - Handle invoke result direct - :param invoke_result: invoke result - :param queue_manager: application queue manager - :param agent: agent - :param message_id: message id for multimodal output - :param user_id: user id for multimodal output - :param tenant_id: tenant id for multimodal output - :return: - """ - queue_manager.publish( - QueueMessageEndEvent( - llm_result=invoke_result, - ), - PublishFrom.APPLICATION_MANAGER, - ) - - def _handle_invoke_result_stream( - self, - invoke_result: Generator[LLMResultChunk], - queue_manager: AppQueueManager, - agent: bool, - message_id: str | None = None, - user_id: str | None = None, - tenant_id: str | None = None, - ): - """ - Handle invoke result - :param invoke_result: invoke result - :param queue_manager: application queue manager - :param agent: agent - :param message_id: message id for multimodal output - :param user_id: user id for multimodal output - :param tenant_id: tenant id for multimodal output - :return: - """ - model: str = "" - prompt_messages: list[PromptMessage] = [] - text = "" - usage = None - try: - for result in invoke_result: - if not agent: - queue_manager.publish(QueueLLMChunkEvent(chunk=result), PublishFrom.APPLICATION_MANAGER) - else: - queue_manager.publish(QueueAgentMessageEvent(chunk=result), PublishFrom.APPLICATION_MANAGER) - - message = result.delta.message - match message.content: - case str(): - text += message.content - case list(): - for content in message.content: - match content: - case TextPromptMessageContent(): - text += content.data - case ImagePromptMessageContent(): - if message_id and user_id and tenant_id: - try: - with session_factory.create_session() as session: - message_file_id = self._handle_multimodal_image_content( - session=session, - content=content, - message_id=message_id, - user_id=user_id, - tenant_id=tenant_id, - queue_manager=queue_manager, - ) - session.commit() - if message_file_id: - queue_manager.publish( - QueueMessageFileEvent(message_file_id=message_file_id), - PublishFrom.APPLICATION_MANAGER, - ) - _logger.info( - "QueueMessageFileEvent published for message_file_id: %s", - message_file_id, - ) - except Exception: - _logger.exception("Failed to handle multimodal image output") - else: - _logger.warning("Received multimodal output but missing required parameters") - case _: - text += content.data if hasattr(content, "data") else str(content) - - if not model: - model = result.model - - if not prompt_messages: - prompt_messages = list(result.prompt_messages) - - if result.delta.usage: - usage = result.delta.usage - except GenerateTaskStoppedError: - # Explicitly close provider stream to stop in-flight token generation ASAP. - invoke_result.close() - raise - - if usage is None: - usage = LLMUsage.empty_usage() - - llm_result = LLMResult( - model=model, prompt_messages=prompt_messages, message=AssistantPromptMessage(content=text), usage=usage - ) - - queue_manager.publish( - QueueMessageEndEvent( - llm_result=llm_result, - ), - PublishFrom.APPLICATION_MANAGER, - ) - - def _handle_multimodal_image_content( - self, - content: ImagePromptMessageContent, - message_id: str, - user_id: str, - tenant_id: str, - queue_manager: AppQueueManager, - session: Session, - ) -> str | None: - """ - Handle multimodal image content from LLM response. - Save the image and create a MessageFile record. - - :param content: ImagePromptMessageContent instance - :param message_id: message id - :param user_id: user id - :param tenant_id: tenant id - :param queue_manager: queue manager - :return: - """ - _logger.info("Handling multimodal image content for message %s", message_id) - - image_url = content.url - base64_data = content.base64_data - - _logger.info("Image URL: %s, Base64 data present: %s", image_url, base64_data) - - if not image_url and not base64_data: - _logger.warning("Image content has neither URL nor base64 data") - return None - - tool_file_manager = ToolFileManager() - - # Save the image file - try: - if image_url: - # Download image from URL - _logger.info("Downloading image from URL: %s", image_url) - tool_file = tool_file_manager.create_file_by_url( - user_id=user_id, - tenant_id=tenant_id, - file_url=image_url, - conversation_id=None, - ) - _logger.info("Image saved successfully, tool_file_id: %s", tool_file.id) - elif base64_data: - if base64_data.startswith("data:"): - base64_data = base64_data.split(",", 1)[1] - - image_binary = base64.b64decode(base64_data) - mimetype = content.mime_type or "image/png" - extension = guess_extension(mimetype) or ".png" - - tool_file = tool_file_manager.create_file_by_raw( - user_id=user_id, - tenant_id=tenant_id, - conversation_id=None, - file_binary=image_binary, - mimetype=mimetype, - filename=f"generated_image{extension}", - ) - _logger.info("Image saved successfully, tool_file_id: %s", tool_file.id) - else: - return None - except Exception: - _logger.exception("Failed to save image file") - return None - - # Create MessageFile record. - # Use an independent session so this side-effect write does not - # commit or close the caller's request-scoped session. - message_file = MessageFile( - message_id=message_id, - type=FileType.IMAGE, - transfer_method=FileTransferMethod.TOOL_FILE, - belongs_to=MessageFileBelongsTo.ASSISTANT, - url=f"/files/tools/{tool_file.id}", - upload_file_id=tool_file.id, - created_by_role=( - CreatorUserRole.ACCOUNT if queue_manager.invoke_from.runs_as_account() else CreatorUserRole.END_USER - ), - created_by=user_id, - ) - - session.add(message_file) - session.flush() - return message_file.id - - def moderation_for_inputs( - self, - *, - app_id: str, - tenant_id: str, - app_generate_entity: AppGenerateEntity, - inputs: Mapping[str, Any], - query: str | None = None, - message_id: str, - ) -> tuple[bool, Mapping[str, Any], str]: - """ - Process sensitive_word_avoidance. - :param app_id: app id - :param tenant_id: tenant id - :param app_generate_entity: app generate entity - :param inputs: inputs - :param query: query - :param message_id: message id - :return: - """ - moderation_feature = InputModeration() - return moderation_feature.check( - app_id=app_id, - tenant_id=tenant_id, - app_config=app_generate_entity.app_config, - inputs=dict(inputs), - query=query or "", - message_id=message_id, - trace_manager=app_generate_entity.trace_manager, - ) - - def check_hosting_moderation( - self, - application_generate_entity: EasyUIBasedAppGenerateEntity, - queue_manager: AppQueueManager, - prompt_messages: list[PromptMessage], - ) -> bool: - """ - Check hosting moderation - :param application_generate_entity: application generate entity - :param queue_manager: queue manager - :param prompt_messages: prompt messages - :return: - """ - hosting_moderation_feature = HostingModerationFeature() - moderation_result = hosting_moderation_feature.check( - application_generate_entity=application_generate_entity, prompt_messages=prompt_messages - ) - - if moderation_result: - self.direct_output( - queue_manager=queue_manager, - app_generate_entity=application_generate_entity, - prompt_messages=prompt_messages, - text="I apologize for any confusion, but I'm an AI assistant to be helpful, harmless, and honest.", - stream=application_generate_entity.stream, - ) - - return moderation_result - - def fill_in_inputs_from_external_data_tools( - self, - tenant_id: str, - app_id: str, - external_data_tools: list[ExternalDataVariableEntity], - inputs: Mapping[str, Any], - query: str, - ) -> Mapping[str, Any]: - """ - Fill in variable inputs from external data tools if exists. - - :param tenant_id: workspace id - :param app_id: app id - :param external_data_tools: external data tools configs - :param inputs: the inputs - :param query: the query - :return: the filled inputs - """ - external_data_fetch_feature = ExternalDataFetch() - return external_data_fetch_feature.fetch( - tenant_id=tenant_id, app_id=app_id, external_data_tools=external_data_tools, inputs=inputs, query=query - ) - - def query_app_annotations_to_reply( - self, app_record: App, message: Message, query: str, user_id: str, invoke_from: InvokeFrom, session: Session - ) -> MessageAnnotation | None: - """ - Query app annotations to reply - :param app_record: app record - :param message: message - :param query: query - :param user_id: user id - :param invoke_from: invoke from - :return: - """ - annotation_reply_feature = AnnotationReplyFeature() - return annotation_reply_feature.query( - app_record=app_record, - message=message, - query=query, - user_id=user_id, - invoke_from=invoke_from, - session=session, - ) diff --git a/api/core/app/apps/chat/app_generator.py b/api/core/app/apps/chat/app_generator.py deleted file mode 100644 index d0be650e4871a2..00000000000000 --- a/api/core/app/apps/chat/app_generator.py +++ /dev/null @@ -1,296 +0,0 @@ -import contextvars -import logging -import threading -import uuid -from collections.abc import Generator, Mapping -from typing import Any, Literal, overload - -from flask import Flask, copy_current_request_context, current_app -from pydantic import ValidationError -from sqlalchemy.orm import Session - -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.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.chat.app_config_manager import ChatAppConfigManager -from core.app.apps.chat.app_runner import ChatAppRunner -from core.app.apps.chat.generate_response_converter import ChatAppGenerateResponseConverter -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 ChatAppGenerateEntity, InvokeFrom -from core.db.session_factory import session_factory -from core.helper.trace_id_helper import extract_trace_session_id_from_args -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 models import Account -from models.model import App, EndUser, load_annotation_reply_config -from services.conversation_service import ConversationService - -logger = logging.getLogger(__name__) - - -class ChatAppGenerator(MessageBasedAppGenerator): - @overload - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: Literal[True], - *, - session: Session, - ) -> Generator[Mapping | str]: ... - - @overload - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: Literal[False], - *, - session: Session, - ) -> Mapping[str, Any]: ... - - @overload - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: bool, - *, - session: Session, - ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str]: ... - - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: bool = True, - *, - session: Session, - ) -> Mapping[str, Any] | Generator[Mapping[str, Any] | str]: - """ - 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 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), - **extract_trace_session_id_from_args(args), - } - - # 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, session=session - ) - # get app model config - app_model_config = self._get_app_model_config( - app_model=app_model, - conversation=conversation, - session=session, - ) - - # 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 = ChatAppConfigManager.config_validate( - tenant_id=app_model.tenant_id, config=args.get("model_config", {}), session=session - ) - - # always enable retriever resource in debugger mode - override_model_config_dict["retriever_resource"] = {"enabled": True} - - if override_model_config_dict: - annotation_reply = None - effective_model_config_dict = override_model_config_dict - else: - annotation_reply = load_annotation_reply_config(session, app_model_config.app_id) - effective_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply) - - # 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["files"] if args.get("files") else [] - file_extra_config = FileUploadConfigManager.convert(effective_model_config_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 = ChatAppConfigManager.get_app_config( - app_model=app_model, - app_model_config=app_model_config, - conversation=conversation, - override_config_dict=override_model_config_dict, - annotation_reply=annotation_reply, - ) - - # get tracing instance - trace_manager = TraceQueueManager( - app_id=app_model.id, user_id=user.id if isinstance(user, Account) else user.session_id - ) - - # init application generate entity - application_generate_entity = ChatAppGenerateEntity( - 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, - invoke_from=invoke_from, - extras=extras, - trace_manager=trace_manager, - stream=streaming, - ) - - # init generate records - (conversation, message) = self._init_generate_records( - application_generate_entity, - conversation, - session=session, - ) - - # 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, - ) - - context = contextvars.copy_context() - - # new thread with request context - @copy_current_request_context - def worker_with_context(): - return context.run( - self._generate_worker, - flask_app=current_app._get_current_object(), # type: ignore - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation_id=conversation.id, - message_id=message.id, - ) - - worker_thread = threading.Thread(target=worker_with_context) - - 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, - stream=streaming, - ) - - return ChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) - - def _generate_worker( - self, - flask_app: Flask, - application_generate_entity: ChatAppGenerateEntity, - 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 flask_app.app_context(): - try: - # get conversation and message - conversation = self._get_conversation(conversation_id) - message = self._get_message(message_id) - - # chatbot app - runner = ChatAppRunner() - with session_factory.create_session() as session: - runner.run( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation=conversation, - message=message, - session=session, - ) - 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() diff --git a/api/core/app/apps/chat/app_runner.py b/api/core/app/apps/chat/app_runner.py deleted file mode 100644 index 65d60748d1e5e6..00000000000000 --- a/api/core/app/apps/chat/app_runner.py +++ /dev/null @@ -1,256 +0,0 @@ -import logging -from typing import cast - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.base_app_runner import AppRunner -from core.app.apps.chat.app_config_manager import ChatAppConfig -from core.app.entities.app_invoke_entities import ( - ChatAppGenerateEntity, - get_credit_usage_app_type, - get_credit_usage_created_by, -) -from core.app.entities.queue_entities import QueueAnnotationReplyEvent -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.db.session_factory import create_session -from core.memory.token_buffer_memory import TokenBufferMemory -from core.model_manager import ModelInstance -from core.moderation.base import ModerationError -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from graphon.file import File -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent -from models.model import App, Conversation, Message - -logger = logging.getLogger(__name__) - - -class ChatAppRunner(AppRunner): - """ - Chat Application Runner - """ - - def run( - self, - application_generate_entity: ChatAppGenerateEntity, - queue_manager: AppQueueManager, - conversation: Conversation, - message: Message, - session: Session, - ): - """Run the application without retaining ``session`` during model I/O. - - Database preparation is committed and the connection is released before - the provider response is requested or consumed. - - :param application_generate_entity: application generate entity - :param queue_manager: application queue manager - :param conversation: conversation - :param message: message - :return: - """ - app_config = application_generate_entity.app_config - app_config = cast(ChatAppConfig, app_config) - stmt = select(App).where(App.id == app_config.app_id) - with create_session() as read_session: - app_record = read_session.scalar(stmt) - if app_record: - read_session.expunge(app_record) - if not app_record: - raise ValueError("App not found") - - inputs = application_generate_entity.inputs - query = application_generate_entity.query - files = application_generate_entity.files - - image_detail_config = ( - application_generate_entity.file_upload_config.image_config.detail - if ( - application_generate_entity.file_upload_config - and application_generate_entity.file_upload_config.image_config - ) - else None - ) - image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW - - memory = None - if application_generate_entity.conversation_id: - # get memory of conversation (read-only) - model_instance = ModelInstance( - provider_model_bundle=application_generate_entity.model_conf.provider_model_bundle, - model=application_generate_entity.model_conf.model, - ) - - memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance) - - # organize all inputs and template to prompt messages - # Include: prompt template, inputs, query(optional), files(optional) - # memory(optional) - prompt_messages, stop = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=inputs, - files=files, - query=query, - memory=memory, - image_detail_config=image_detail_config, - ) - - # moderation - try: - # process sensitive_word_avoidance - _, inputs, query = self.moderation_for_inputs( - app_id=app_record.id, - tenant_id=app_config.tenant_id, - app_generate_entity=application_generate_entity, - inputs=inputs, - query=query, - message_id=message.id, - ) - except ModerationError as e: - self.direct_output( - queue_manager=queue_manager, - app_generate_entity=application_generate_entity, - prompt_messages=prompt_messages, - text=str(e), - stream=application_generate_entity.stream, - ) - return - - if query: - # annotation reply - annotation_reply = self.query_app_annotations_to_reply( - app_record=app_record, - message=message, - query=query, - user_id=application_generate_entity.user_id, - invoke_from=application_generate_entity.invoke_from, - session=session, - ) - session.commit() - session.close() - - if annotation_reply: - queue_manager.publish( - QueueAnnotationReplyEvent(message_annotation_id=annotation_reply.id), - PublishFrom.APPLICATION_MANAGER, - ) - - self.direct_output( - queue_manager=queue_manager, - app_generate_entity=application_generate_entity, - prompt_messages=prompt_messages, - text=annotation_reply.content, - stream=application_generate_entity.stream, - ) - return - - # fill in variable inputs from external data tools if exists - external_data_tools = app_config.external_data_variables - if external_data_tools: - inputs = self.fill_in_inputs_from_external_data_tools( - tenant_id=app_record.tenant_id, - app_id=app_record.id, - external_data_tools=external_data_tools, - inputs=inputs, - query=query, - ) - - # get context from datasets - context = None - context_files: list[File] = [] - if app_config.dataset and app_config.dataset.dataset_ids: - hit_callback = DatasetIndexToolCallbackHandler( - queue_manager, - app_record.id, - message.id, - application_generate_entity.user_id, - application_generate_entity.invoke_from, - ) - - dataset_retrieval = DatasetRetrieval(application_generate_entity) - context, retrieved_files = dataset_retrieval.retrieve( - session=session, - app_id=app_record.id, - user_id=application_generate_entity.user_id, - tenant_id=app_record.tenant_id, - model_config=application_generate_entity.model_conf, - config=app_config.dataset, - query=query, - invoke_from=application_generate_entity.invoke_from, - show_retrieve_source=( - app_config.additional_features.show_retrieve_source if app_config.additional_features else False - ), - hit_callback=hit_callback, - memory=memory, - message_id=message.id, - inputs=inputs, - vision_enabled=bool( - application_generate_entity.app_config.app_model_config_dict.get("file_upload", {}) - .get("image", {}) - .get("enabled", False) - ), - ) - context_files = retrieved_files or [] - - session.commit() - session.close() - - # reorganize all inputs and template to prompt messages - # Include: prompt template, inputs, query(optional), files(optional) - # memory(optional), external data, dataset context(optional) - prompt_messages, stop = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=inputs, - files=files, - query=query, - context=context, - memory=memory, - image_detail_config=image_detail_config, - context_files=context_files, - ) - - # check hosting moderation - hosting_moderation_result = self.check_hosting_moderation( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - prompt_messages=prompt_messages, - ) - - if hosting_moderation_result: - return - - # Re-calculate the max tokens if sum(prompt_token + max_tokens) over model token limit - self.recalc_llm_max_tokens(model_config=application_generate_entity.model_conf, prompt_messages=prompt_messages) - - # Invoke model - model_instance = ModelInstance( - provider_model_bundle=application_generate_entity.model_conf.provider_model_bundle, - model=application_generate_entity.model_conf.model, - ) - - request_metadata: dict[str, object] = {"app_id": app_config.app_id} - request_metadata["app_type"] = get_credit_usage_app_type(app_config.app_mode) - request_metadata["created_by"] = get_credit_usage_created_by(app_config.app_mode) - - invoke_result = model_instance.invoke_llm( - prompt_messages=prompt_messages, - model_parameters=application_generate_entity.model_conf.parameters, - stop=stop, - stream=application_generate_entity.stream, - request_metadata=request_metadata, - ) - - # handle invoke result - self._handle_invoke_result( - invoke_result=invoke_result, - queue_manager=queue_manager, - stream=application_generate_entity.stream, - message_id=message.id, - user_id=application_generate_entity.user_id, - tenant_id=app_config.tenant_id, - ) diff --git a/api/core/app/apps/common/workflow_response_converter.py b/api/core/app/apps/common/workflow_response_converter.py deleted file mode 100644 index 4c4e7bce03e96f..00000000000000 --- a/api/core/app/apps/common/workflow_response_converter.py +++ /dev/null @@ -1,1031 +0,0 @@ -import json -import logging -import time -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime -from typing import Any, NewType, TypedDict - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity -from core.app.entities.queue_entities import ( - NodeExecutionSnapshot, - QueueAgentLogEvent, - QueueHumanInputFormFilledEvent, - QueueHumanInputFormTimeoutEvent, - QueueIterationCompletedEvent, - QueueIterationNextEvent, - QueueIterationStartEvent, - QueueLoopCompletedEvent, - QueueLoopNextEvent, - QueueLoopStartEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueueNodeRetryEvent, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, - QueueWorkflowPausedEvent, -) -from core.app.entities.task_entities import ( - AgentLogStreamResponse, - HumanInputFormFilledResponse, - HumanInputFormTimeoutResponse, - HumanInputRequiredResponse, - IterationNodeCompletedStreamResponse, - IterationNodeNextStreamResponse, - IterationNodeStartStreamResponse, - LoopNodeCompletedStreamResponse, - LoopNodeNextStreamResponse, - LoopNodeStartStreamResponse, - NodeFinishStreamResponse, - NodeRetryStreamResponse, - NodeStartStreamResponse, - StreamResponse, - WorkflowFinishStreamResponse, - WorkflowPauseStreamResponse, - WorkflowStartStreamResponse, -) -from core.plugin.impl.datasource import PluginDatasourceManager -from core.tools.entities.tool_entities import ToolProviderType -from core.tools.tool_manager import ToolManager -from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE -from core.trigger.trigger_manager import TriggerManager -from core.workflow.human_input_forms import ( - load_form_dispositions_by_form_id, -) -from core.workflow.human_input_policy import ( - HumanInputSurface, - enrich_human_input_pause_reasons, - resolve_variable_select_input_options, -) -from core.workflow.nodes.human_input.pause_reason import HumanInputRequired -from core.workflow.nodes.human_input_v2.presentation import build_human_input_v2_pause_reason -from core.workflow.nodes.human_input_v2.runtime import PreparedForm -from core.workflow.system_variables import SystemVariableKey, system_variables_to_mapping -from extensions.ext_database import db -from graphon.entities import WorkflowStartReason -from graphon.enums import ( - BuiltinNodeTypes, - WorkflowExecutionStatus, - WorkflowNodeExecutionMetadataKey, - WorkflowNodeExecutionStatus, -) -from graphon.file import FILE_MODEL_IDENTITY, File -from graphon.runtime import RuntimeState -from graphon.runtime.runtime_state_protocol import ReadOnlyVariablePool -from graphon.variables.segments import ArrayFileSegment, FileSegment, Segment -from graphon.variables.variables import Variable -from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter -from libs.datetime_utils import naive_utc_now, to_utc_timestamp -from models import Account, EndUser -from models.human_input import HumanInputForm -from models.workflow import WorkflowRun -from services.variable_truncator import BaseTruncator, DummyVariableTruncator, VariableTruncator -from services.workflow.execution.adapters.workflow_entry import WorkflowEntry - -# Maps the entry surface a workflow was invoked from to the HITL surface that -# its resume tokens must be filtered for. Surfaces not in this map fall back to -# the general priority ordering (typically CONSOLE > BACKSTAGE). -_INVOKE_FROM_TO_HITL_SURFACE: Mapping[InvokeFrom, HumanInputSurface] = { - InvokeFrom.SERVICE_API: HumanInputSurface.SERVICE_API, - InvokeFrom.OPENAPI: HumanInputSurface.OPENAPI, -} - -NodeExecutionId = NewType("NodeExecutionId", str) -logger = logging.getLogger(__name__) - - -class AccountCreatedByDict(TypedDict): - id: str - name: str - email: str - - -class EndUserCreatedByDict(TypedDict): - id: str - user: str - - -CreatedByDict = AccountCreatedByDict | EndUserCreatedByDict - - -@dataclass(frozen=True, slots=True) -class _HumanInputPauseOutput: - """Keep a form event paired with its entry in the workflow pause response.""" - - response: HumanInputRequiredResponse - pause_reason: dict[str, Any] - - -class WorkflowResponseConverter: - @staticmethod - def human_input_v2_required_data(prepared: PreparedForm, *, node_id: str) -> HumanInputRequiredResponse.Data: - """Project Dify-owned form data into the existing HITL wire contract. - - Live execution supplies the node's snapshot; reconnect loads the form - from its run-owned v2 records. Neither path resolves a v1 token. - """ - reason = build_human_input_v2_pause_reason( - prepared.form.resolved_form, form_id=prepared.form.id, node_id=node_id - ) - return HumanInputRequiredResponse.Data( - form_version=reason.form_version, - form_id=reason.form_id, - node_id=node_id, - node_title=reason.node_title, - form_content=reason.form_content, - inputs=reason.inputs, - actions=reason.actions, - resolved_default_values=reason.resolved_default_values, - expiration_time=to_utc_timestamp(prepared.form.expiration_time), - display_in_ui=prepared.form_token is not None, - form_token=prepared.form_token.get_secret_value() if prepared.form_token is not None else None, - ) - - _truncator: BaseTruncator - - def __init__( - self, - *, - application_generate_entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, - user: Account | EndUser, - system_variables: Sequence[Variable], - ): - self._application_generate_entity = application_generate_entity - self._user = user - self._system_variables = system_variables_to_mapping(system_variables) - self._workflow_inputs = self._prepare_workflow_inputs() - - # Disable truncation for SERVICE_API calls to keep backward compatibility. - if application_generate_entity.invoke_from == InvokeFrom.SERVICE_API: - self._truncator = DummyVariableTruncator() - else: - self._truncator = VariableTruncator.default() - - self._node_snapshots: dict[NodeExecutionId, NodeExecutionSnapshot] = {} - self._workflow_execution_id: str | None = None - self._workflow_started_at: datetime | None = None - - # ------------------------------------------------------------------ - # Workflow lifecycle helpers - # ------------------------------------------------------------------ - def _prepare_workflow_inputs(self) -> Mapping[str, Any]: - inputs = dict(self._application_generate_entity.inputs) - for field_name, value in self._system_variables.items(): - # TODO(@future-refactor): store system variables separately from user inputs so we don't - # need to flatten `sys.*` entries into the input payload just for rerun/export tooling. - if field_name == SystemVariableKey.CONVERSATION_ID: - # Conversation IDs are session-scoped; omitting them keeps workflow inputs - # reusable without pinning new runs to a prior conversation. - continue - inputs[f"sys.{field_name}"] = value - handled = WorkflowEntry.handle_special_values(inputs) - return dict(handled or {}) - - def _ensure_workflow_run_id(self, workflow_run_id: str | None = None) -> str: - """Return the memoized workflow run id, optionally seeding it during start events.""" - if workflow_run_id is not None: - self._workflow_execution_id = workflow_run_id - if not self._workflow_execution_id: - raise ValueError("workflow_run_id missing before streaming workflow events") - return self._workflow_execution_id - - # ------------------------------------------------------------------ - # Node snapshot helpers - # ------------------------------------------------------------------ - def _store_snapshot(self, event: QueueNodeStartedEvent) -> NodeExecutionSnapshot: - # The engine may have persisted this start before resumption metadata was loaded. - snapshot = self._get_snapshot(event.node_execution_id) - if snapshot is not None: - return snapshot - - snapshot = NodeExecutionSnapshot( - execution_id=event.node_execution_id, - title=event.node_title, - index=event.node_run_index, - start_at=event.start_at, - iteration_id=event.in_iteration_id or "", - loop_id=event.in_loop_id or "", - ) - node_execution_id = NodeExecutionId(event.node_execution_id) - self._node_snapshots[node_execution_id] = snapshot - return snapshot - - def _get_snapshot(self, node_execution_id: str) -> NodeExecutionSnapshot | None: - return self._node_snapshots.get(NodeExecutionId(node_execution_id)) - - def _pop_snapshot(self, node_execution_id: str) -> NodeExecutionSnapshot | None: - return self._node_snapshots.pop(NodeExecutionId(node_execution_id), None) - - @staticmethod - def _merge_metadata( - base_metadata: Mapping[WorkflowNodeExecutionMetadataKey, Any] | None, - snapshot: NodeExecutionSnapshot | None, - ) -> Mapping[WorkflowNodeExecutionMetadataKey, Any] | None: - if not base_metadata and not snapshot: - return base_metadata - - merged: dict[WorkflowNodeExecutionMetadataKey, Any] = {} - if base_metadata: - merged.update(base_metadata) - - if snapshot: - if snapshot.iteration_id: - merged[WorkflowNodeExecutionMetadataKey.ITERATION_ID] = snapshot.iteration_id - if snapshot.loop_id: - merged[WorkflowNodeExecutionMetadataKey.LOOP_ID] = snapshot.loop_id - - return merged or None - - def _truncate_mapping( - self, - mapping: Mapping[str, Any] | None, - ) -> tuple[Mapping[str, Any] | None, bool]: - if mapping is None: - return None, False - if not mapping: - return {}, False - - normalized = WorkflowEntry.handle_special_values(dict(mapping)) - if normalized is None: - return None, False - - truncated, is_truncated = self._truncator.truncate_variable_mapping(dict(normalized)) - return truncated, is_truncated - - @staticmethod - def _encode_outputs(outputs: Mapping[str, Any] | None) -> Mapping[str, Any] | None: - if outputs is None: - return None - converter = WorkflowRuntimeTypeConverter() - return converter.to_json_encodable(outputs) - - def workflow_start_to_stream_response( - self, - *, - task_id: str, - workflow_run_id: str, - workflow_id: str, - reason: WorkflowStartReason, - node_execution_snapshots: Sequence[NodeExecutionSnapshot] = (), - ) -> WorkflowStartStreamResponse: - run_id = self._ensure_workflow_run_id(workflow_run_id) - started_at = naive_utc_now() - self._workflow_started_at = started_at - self._application_generate_entity.workflow_execution_timing.start() - if reason == WorkflowStartReason.RESUMPTION: - self._node_snapshots.update( - (NodeExecutionId(snapshot.execution_id), snapshot) for snapshot in node_execution_snapshots - ) - - return WorkflowStartStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=WorkflowStartStreamResponse.Data( - id=run_id, - workflow_id=workflow_id, - inputs=self._workflow_inputs, - created_at=int(started_at.timestamp()), - reason=reason, - ), - ) - - def workflow_finish_to_stream_response( - self, - *, - task_id: str, - workflow_id: str, - status: WorkflowExecutionStatus, - graph_runtime_state: RuntimeState, - error: str | None = None, - exceptions_count: int = 0, - ) -> WorkflowFinishStreamResponse: - run_id = self._ensure_workflow_run_id() - started_at = self._workflow_started_at - if started_at is None: - raise ValueError( - "workflow_finish_to_stream_response called before workflow_start_to_stream_response", - ) - - finished_at = naive_utc_now() - elapsed_time = self._application_generate_entity.workflow_execution_timing.finish() - - outputs_mapping = graph_runtime_state.outputs or {} - encoded_outputs = WorkflowRuntimeTypeConverter().to_json_encodable(outputs_mapping) - - created_by: CreatedByDict | dict[str, object] = {} - user = self._user - match user: - case Account(): - created_by = AccountCreatedByDict( - id=user.id, - name=user.name, - email=user.email, - ) - case EndUser(): - created_by = EndUserCreatedByDict( - id=user.id, - user=user.session_id, - ) - - return WorkflowFinishStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=WorkflowFinishStreamResponse.Data( - id=run_id, - workflow_id=workflow_id, - status=status, - outputs=encoded_outputs, - error=error, - elapsed_time=elapsed_time, - total_tokens=graph_runtime_state.total_tokens, - total_steps=graph_runtime_state.node_run_steps, - created_by=created_by, - created_at=int(started_at.timestamp()), - finished_at=int(finished_at.timestamp()), - files=self.fetch_files_from_node_outputs(outputs_mapping), - exceptions_count=exceptions_count, - ), - ) - - def workflow_pause_to_stream_response( - self, - *, - event: QueueWorkflowPausedEvent, - task_id: str, - graph_runtime_state: RuntimeState, - ) -> list[StreamResponse]: - run_id = self._ensure_workflow_run_id() - started_at = self._workflow_started_at - if started_at is None: - raise ValueError( - "workflow_pause_to_stream_response called before workflow_start_to_stream_response", - ) - elapsed_time = self._application_generate_entity.workflow_execution_timing.finish() - encoded_outputs = self._encode_outputs(event.outputs) or {} - if self._application_generate_entity.invoke_from == InvokeFrom.SERVICE_API: - encoded_outputs = {} - v1_outputs = iter( - self._human_input_v1_pause_responses( - reasons=[ - reason - for reason in event.reasons - if isinstance(reason, HumanInputRequired) and reason.form_version == "1" - ], - task_id=task_id, - variable_pool=graph_runtime_state.variable_pool, - ) - ) - v2_outputs = iter( - self._human_input_v2_pause_responses( - reasons=[ - reason - for reason in event.reasons - if isinstance(reason, HumanInputRequired) and reason.form_version == "2" - ], - forms=event.human_input_v2_forms, - task_id=task_id, - ) - ) - responses: list[StreamResponse] = [] - pause_reasons: list[dict[str, Any]] = [] - # Each version owns its complete conversion; merge in the original event - # order, including scheduling reasons interleaved with approval forms. - for reason in event.reasons: - if not isinstance(reason, HumanInputRequired): - pause_reasons.append(reason.model_dump(mode="json")) - continue - match reason.form_version: - case "1": - converted = next(v1_outputs) - case "2": - converted = next(v2_outputs) - case _: - raise ValueError(f"Unsupported Human Input form version: {reason.form_version}") - responses.append(converted.response) - pause_reasons.append(converted.pause_reason) - - responses.append( - WorkflowPauseStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=WorkflowPauseStreamResponse.Data( - workflow_run_id=run_id, - paused_nodes=list(event.paused_nodes), - outputs=encoded_outputs, - reasons=pause_reasons, - status=WorkflowExecutionStatus.PAUSED, - created_at=int(started_at.timestamp()), - elapsed_time=elapsed_time, - total_tokens=graph_runtime_state.total_tokens, - total_steps=graph_runtime_state.node_run_steps, - ), - ) - ) - - return responses - - def human_input_form_filled_to_stream_response( - self, *, event: QueueHumanInputFormFilledEvent, task_id: str - ) -> HumanInputFormFilledResponse: - run_id = self._ensure_workflow_run_id() - data = HumanInputFormFilledResponse.Data( - form_id=event.form_id, - node_id=event.node_id, - node_title=event.node_title, - rendered_content=event.rendered_content, - action_id=event.action_id, - action_text=event.action_text, - ) - if event.submitted_data is not None: - runtime_type_converter = WorkflowRuntimeTypeConverter() - - data.submitted_data = runtime_type_converter.value_to_json_encodable_recursive(event.submitted_data) - - return HumanInputFormFilledResponse(task_id=task_id, workflow_run_id=run_id, data=data) - - def human_input_form_timeout_to_stream_response( - self, *, event: QueueHumanInputFormTimeoutEvent, task_id: str - ) -> HumanInputFormTimeoutResponse: - run_id = self._ensure_workflow_run_id() - return HumanInputFormTimeoutResponse( - task_id=task_id, - workflow_run_id=run_id, - data=HumanInputFormTimeoutResponse.Data( - form_id=event.form_id, - node_id=event.node_id, - node_title=event.node_title, - expiration_time=to_utc_timestamp(event.expiration_time), - ), - ) - - @classmethod - def workflow_run_result_to_finish_response( - cls, - *, - task_id: str, - workflow_run: WorkflowRun, - creator_user: Account | EndUser, - ) -> WorkflowFinishStreamResponse: - run_id = workflow_run.id - elapsed_time = workflow_run.elapsed_time - - encoded_outputs = workflow_run.outputs_dict - finished_at = workflow_run.finished_at - assert finished_at is not None - - created_by: Mapping[str, object] - user = creator_user - match user: - case Account(): - created_by = { - "id": user.id, - "name": user.name, - "email": user.email, - } - case _: - created_by = { - "id": user.id, - "user": user.session_id, - } - - return WorkflowFinishStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=WorkflowFinishStreamResponse.Data( - id=run_id, - workflow_id=workflow_run.workflow_id, - status=workflow_run.status, - outputs=encoded_outputs, - error=workflow_run.error, - elapsed_time=elapsed_time, - total_tokens=workflow_run.total_tokens, - total_steps=workflow_run.total_steps, - created_by=created_by, - created_at=int(workflow_run.created_at.timestamp()), - finished_at=int(finished_at.timestamp()), - files=cls.fetch_files_from_node_outputs(encoded_outputs), - exceptions_count=workflow_run.exceptions_count, - ), - ) - - def workflow_node_start_to_stream_response( - self, - *, - event: QueueNodeStartedEvent, - task_id: str, - ) -> NodeStartStreamResponse | None: - snapshot = self._store_snapshot(event) - if event.node_type in {BuiltinNodeTypes.ITERATION, BuiltinNodeTypes.LOOP}: - return None - run_id = self._ensure_workflow_run_id() - - response = NodeStartStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=NodeStartStreamResponse.Data( - id=event.node_execution_id, - node_id=event.node_id, - node_type=event.node_type, - title=snapshot.title, - index=snapshot.index, - created_at=int(snapshot.start_at.timestamp()), - iteration_id=event.in_iteration_id, - loop_id=event.in_loop_id, - agent_strategy=event.agent_strategy, - ), - ) - - try: - if event.node_type == BuiltinNodeTypes.TOOL: - if event.provider_type == ToolProviderType.WORKFLOW: - response.data.extras["workflow_tool"] = True - response.data.extras["icon"] = ToolManager.get_tool_icon( - tenant_id=self._application_generate_entity.app_config.tenant_id, - provider_type=ToolProviderType(event.provider_type), - provider_id=event.provider_id, - ) - elif event.node_type == BuiltinNodeTypes.DATASOURCE: - manager = PluginDatasourceManager() - provider_entity = manager.fetch_datasource_provider( - self._application_generate_entity.app_config.tenant_id, - event.provider_id, - ) - response.data.extras["icon"] = provider_entity.declaration.identity.generate_datasource_icon_url( - self._application_generate_entity.app_config.tenant_id - ) - elif event.node_type == TRIGGER_PLUGIN_NODE_TYPE: - response.data.extras["icon"] = TriggerManager.get_trigger_plugin_icon( - self._application_generate_entity.app_config.tenant_id, - event.provider_id, - ) - except Exception: - # metadata fetch may fail, for example, the plugin daemon is down or plugin is uninstalled. - logger.warning("failed to fetch icon for %s", event.provider_id) - - return response - - def workflow_node_finish_to_stream_response( - self, - *, - event: QueueNodeSucceededEvent | QueueNodeFailedEvent | QueueNodeExceptionEvent, - task_id: str, - ) -> NodeFinishStreamResponse | None: - if event.node_type in {BuiltinNodeTypes.ITERATION, BuiltinNodeTypes.LOOP}: - return None - run_id = self._ensure_workflow_run_id() - snapshot = self._pop_snapshot(event.node_execution_id) - - start_at = snapshot.start_at if snapshot else event.start_at - finished_at = event.finished_at or naive_utc_now() - elapsed_time = (finished_at - start_at).total_seconds() - - inputs, inputs_truncated = self._truncate_mapping(event.inputs) - process_data, process_data_truncated = self._truncate_mapping(event.process_data) - encoded_outputs = self._encode_outputs(event.outputs) - outputs, outputs_truncated = self._truncate_mapping(encoded_outputs) - metadata = self._merge_metadata(event.execution_metadata, snapshot) - - match event: - case QueueNodeSucceededEvent(): - status = WorkflowNodeExecutionStatus.SUCCEEDED - error_message = event.error - case QueueNodeFailedEvent(): - status = WorkflowNodeExecutionStatus.FAILED - error_message = event.error - case _: - status = WorkflowNodeExecutionStatus.EXCEPTION - error_message = event.error - - return NodeFinishStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=NodeFinishStreamResponse.Data( - id=event.node_execution_id, - node_id=event.node_id, - node_type=event.node_type, - index=snapshot.index if snapshot else 0, - title=snapshot.title if snapshot else "", - inputs=inputs, - inputs_truncated=inputs_truncated, - process_data=process_data, - process_data_truncated=process_data_truncated, - outputs=outputs, - outputs_truncated=outputs_truncated, - status=status, - error=error_message, - elapsed_time=elapsed_time, - execution_metadata=metadata, - created_at=int(start_at.timestamp()), - finished_at=int(finished_at.timestamp()), - files=self.fetch_files_from_node_outputs(event.outputs or {}), - iteration_id=event.in_iteration_id, - loop_id=event.in_loop_id, - ), - ) - - def workflow_node_retry_to_stream_response( - self, - *, - event: QueueNodeRetryEvent, - task_id: str, - ) -> NodeRetryStreamResponse | None: - if event.node_type in {BuiltinNodeTypes.ITERATION, BuiltinNodeTypes.LOOP}: - return None - run_id = self._ensure_workflow_run_id() - - snapshot = self._get_snapshot(event.node_execution_id) - if snapshot is None: - raise AssertionError("node retry event arrived without a stored snapshot") - finished_at = naive_utc_now() - elapsed_time = (finished_at - event.start_at).total_seconds() - - inputs, inputs_truncated = self._truncate_mapping(event.inputs) - process_data, process_data_truncated = self._truncate_mapping(event.process_data) - encoded_outputs = self._encode_outputs(event.outputs) - outputs, outputs_truncated = self._truncate_mapping(encoded_outputs) - metadata = self._merge_metadata(event.execution_metadata, snapshot) - - return NodeRetryStreamResponse( - task_id=task_id, - workflow_run_id=run_id, - data=NodeRetryStreamResponse.Data( - id=event.node_execution_id, - node_id=event.node_id, - node_type=event.node_type, - index=snapshot.index, - title=snapshot.title, - inputs=inputs, - inputs_truncated=inputs_truncated, - process_data=process_data, - process_data_truncated=process_data_truncated, - outputs=outputs, - outputs_truncated=outputs_truncated, - status=WorkflowNodeExecutionStatus.RETRY, - error=event.error, - elapsed_time=elapsed_time, - execution_metadata=metadata, - created_at=int(snapshot.start_at.timestamp()), - finished_at=int(finished_at.timestamp()), - files=self.fetch_files_from_node_outputs(event.outputs or {}), - iteration_id=event.in_iteration_id, - loop_id=event.in_loop_id, - retry_index=event.retry_index, - ), - ) - - def workflow_iteration_start_to_stream_response( - self, - *, - task_id: str, - workflow_execution_id: str, - event: QueueIterationStartEvent, - ) -> IterationNodeStartStreamResponse: - new_inputs, truncated = self._truncator.truncate_variable_mapping(event.inputs or {}) - return IterationNodeStartStreamResponse( - task_id=task_id, - workflow_run_id=workflow_execution_id, - data=IterationNodeStartStreamResponse.Data( - id=event.node_id, - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - created_at=int(time.time()), - extras={}, - inputs=new_inputs, - inputs_truncated=truncated, - metadata=event.metadata or {}, - ), - ) - - def workflow_iteration_next_to_stream_response( - self, - *, - task_id: str, - workflow_execution_id: str, - event: QueueIterationNextEvent, - ) -> IterationNodeNextStreamResponse: - return IterationNodeNextStreamResponse( - task_id=task_id, - workflow_run_id=workflow_execution_id, - data=IterationNodeNextStreamResponse.Data( - id=event.node_id, - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - index=event.index, - created_at=int(time.time()), - extras={}, - ), - ) - - def workflow_iteration_completed_to_stream_response( - self, - *, - task_id: str, - workflow_execution_id: str, - event: QueueIterationCompletedEvent, - ) -> IterationNodeCompletedStreamResponse: - json_converter = WorkflowRuntimeTypeConverter() - - new_inputs, inputs_truncated = self._truncator.truncate_variable_mapping(event.inputs or {}) - new_outputs, outputs_truncated = self._truncator.truncate_variable_mapping( - json_converter.to_json_encodable(event.outputs) or {} - ) - return IterationNodeCompletedStreamResponse( - task_id=task_id, - workflow_run_id=workflow_execution_id, - data=IterationNodeCompletedStreamResponse.Data( - id=event.node_id, - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - outputs=new_outputs, - outputs_truncated=outputs_truncated, - created_at=int(time.time()), - extras={}, - inputs=new_inputs, - inputs_truncated=inputs_truncated, - status=WorkflowNodeExecutionStatus.SUCCEEDED - if event.error is None - else WorkflowNodeExecutionStatus.FAILED, - error=None, - elapsed_time=(naive_utc_now() - event.start_at).total_seconds(), - total_tokens=(lambda x: x if isinstance(x, int) else 0)(event.metadata.get("total_tokens", 0)), - execution_metadata=event.metadata, - finished_at=int(time.time()), - steps=event.steps, - ), - ) - - def workflow_loop_start_to_stream_response( - self, *, task_id: str, workflow_execution_id: str, event: QueueLoopStartEvent - ) -> LoopNodeStartStreamResponse: - new_inputs, truncated = self._truncator.truncate_variable_mapping(event.inputs or {}) - return LoopNodeStartStreamResponse( - task_id=task_id, - workflow_run_id=workflow_execution_id, - data=LoopNodeStartStreamResponse.Data( - id=event.node_id, - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - created_at=int(time.time()), - extras={}, - inputs=new_inputs, - inputs_truncated=truncated, - metadata=event.metadata or {}, - ), - ) - - def workflow_loop_next_to_stream_response( - self, - *, - task_id: str, - workflow_execution_id: str, - event: QueueLoopNextEvent, - ) -> LoopNodeNextStreamResponse: - return LoopNodeNextStreamResponse( - task_id=task_id, - workflow_run_id=workflow_execution_id, - data=LoopNodeNextStreamResponse.Data( - id=event.node_id, - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - index=event.index, - # The `pre_loop_output` field is not utilized by the frontend. - # Previously, it was assigned the value of `event.output`. - pre_loop_output={}, - created_at=int(time.time()), - extras={}, - ), - ) - - def workflow_loop_completed_to_stream_response( - self, - *, - task_id: str, - workflow_execution_id: str, - event: QueueLoopCompletedEvent, - ) -> LoopNodeCompletedStreamResponse: - json_converter = WorkflowRuntimeTypeConverter() - new_inputs, inputs_truncated = self._truncator.truncate_variable_mapping(event.inputs or {}) - new_outputs, outputs_truncated = self._truncator.truncate_variable_mapping( - json_converter.to_json_encodable(event.outputs) or {} - ) - return LoopNodeCompletedStreamResponse( - task_id=task_id, - workflow_run_id=workflow_execution_id, - data=LoopNodeCompletedStreamResponse.Data( - id=event.node_id, - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - outputs=new_outputs, - outputs_truncated=outputs_truncated, - created_at=int(time.time()), - extras={}, - inputs=new_inputs, - inputs_truncated=inputs_truncated, - status=WorkflowNodeExecutionStatus.SUCCEEDED - if event.error is None - else WorkflowNodeExecutionStatus.FAILED, - error=None, - elapsed_time=(naive_utc_now() - event.start_at).total_seconds(), - total_tokens=(lambda x: x if isinstance(x, int) else 0)(event.metadata.get("total_tokens", 0)), - execution_metadata=event.metadata, - finished_at=int(time.time()), - steps=event.steps, - ), - ) - - @classmethod - def fetch_files_from_node_outputs(cls, outputs_dict: Mapping[str, Any] | None) -> Sequence[Mapping[str, Any]]: - """ - Fetch files from node outputs - :param outputs_dict: node outputs dict - :return: - """ - if not outputs_dict: - return [] - - files = [cls._fetch_files_from_variable_value(output_value) for output_value in outputs_dict.values()] - # Remove None - files = [file for file in files if file] - # Flatten list - # Flatten the list of sequences into a single list of mappings - flattened_files = [file for sublist in files if sublist for file in sublist] - - # Convert to tuple to match Sequence type - return tuple(flattened_files) - - @classmethod - def _fetch_files_from_variable_value( - cls, value: dict | list | Segment | File | None - ) -> Sequence[Mapping[str, Any]]: - """ - Fetch files from variable value - :param value: variable value - :return: - """ - if not value: - return [] - - files: list[Mapping[str, Any]] = [] - match value: - case FileSegment(): - files.append(value.value.to_dict()) - case ArrayFileSegment(): - files.extend([i.to_dict() for i in value.value]) - case File(): - files.append(value.to_dict()) - case list(): - for item in value: - file = cls._get_file_var_from_value(item) - if file: - files.append(file) - case dict(): - file = cls._get_file_var_from_value(value) - if file: - files.append(file) - case _: - pass - - return files - - @classmethod - def _get_file_var_from_value(cls, value: object) -> Mapping[str, Any] | None: - """ - Get file var from value - :param value: variable value - :return: - """ - if not value: - return None - - match value: - case dict() if value.get("dify_model_identity") == FILE_MODEL_IDENTITY: - return value - case File(): - return value.to_dict() - - return None - - def handle_agent_log(self, task_id: str, event: QueueAgentLogEvent) -> AgentLogStreamResponse: - """ - Handle agent log - :param task_id: task id - :param event: agent log event - :return: - """ - return AgentLogStreamResponse( - task_id=task_id, - data=AgentLogStreamResponse.Data( - node_execution_id=event.node_execution_id, - id=event.id, - parent_id=event.parent_id, - label=event.label, - error=event.error, - status=event.status, - data=event.data, - metadata=event.metadata, - node_id=event.node_id, - ), - ) - - def _human_input_v1_pause_responses( - self, - *, - reasons: Sequence[HumanInputRequired], - task_id: str, - variable_pool: ReadOnlyVariablePool, - ) -> list[_HumanInputPauseOutput]: - """Resolve legacy inputs, access channels, and both V1 response payloads.""" - if not reasons: - return [] - run_id = self._ensure_workflow_run_id() - resolved_reasons = [ - reason.model_copy( - update={ - "inputs": resolve_variable_select_input_options(reason.inputs, variable_pool=variable_pool), - } - ) - for reason in reasons - ] - form_ids = [reason.form_id for reason in resolved_reasons] - expiration_times: dict[str, datetime] = {} - display_in_ui: dict[str, bool] = {} - stmt = select(HumanInputForm.id, HumanInputForm.expiration_time, HumanInputForm.form_definition).where( - HumanInputForm.id.in_(form_ids) - ) - surface = _INVOKE_FROM_TO_HITL_SURFACE.get(self._application_generate_entity.invoke_from) - with Session(bind=db.engine) as session: - for form_id, expiration_time, form_definition in session.execute(stmt): - expiration_times[str(form_id)] = expiration_time - try: - definition_payload = json.loads(form_definition) if form_definition else {} - except (TypeError, json.JSONDecodeError): - definition_payload = {} - display_in_ui[str(form_id)] = bool(definition_payload.get("display_in_ui")) - dispositions = load_form_dispositions_by_form_id(form_ids, session=session, surface=surface) - # Preserve the legacy pause-reason fields, which differ from the - # human_input_required data payload (notably display_in_ui). - pause_reasons = enrich_human_input_pause_reasons( - [reason.model_dump(mode="json") for reason in resolved_reasons], - dispositions_by_form_id=dispositions, - expiration_times_by_form_id={ - form_id: to_utc_timestamp(value) for form_id, value in expiration_times.items() - }, - ) - outputs: list[_HumanInputPauseOutput] = [] - for reason, payload in zip(resolved_reasons, pause_reasons, strict=True): - expiration_time = expiration_times.get(reason.form_id) - if expiration_time is None: - raise ValueError(f"HumanInputForm not found for pause reason, form_id={reason.form_id}") - disposition = dispositions.get(reason.form_id) - response = HumanInputRequiredResponse( - task_id=task_id, - workflow_run_id=run_id, - data=HumanInputRequiredResponse.Data( - form_version=reason.form_version, - form_id=reason.form_id, - node_id=reason.node_id, - node_title=reason.node_title, - form_content=reason.form_content, - inputs=reason.inputs, - actions=reason.actions, - display_in_ui=display_in_ui.get(reason.form_id, False), - form_token=disposition.form_token if disposition else None, - approval_channels=list(disposition.approval_channels) if disposition else [], - resolved_default_values=reason.resolved_default_values, - expiration_time=to_utc_timestamp(expiration_time), - ), - ) - outputs.append(_HumanInputPauseOutput(response=response, pause_reason=payload)) - return outputs - - def _human_input_v2_pause_responses( - self, - *, - reasons: Sequence[HumanInputRequired], - forms: Mapping[str, PreparedForm], - task_id: str, - ) -> list[_HumanInputPauseOutput]: - """Build both V2 response payloads from the snapshots emitted by nodes.""" - if not reasons: - return [] - run_id = self._ensure_workflow_run_id() - outputs: list[_HumanInputPauseOutput] = [] - for reason in reasons: - prepared = forms.get(reason.form_id) - if prepared is None: - raise ValueError(f"Human Input v2 node did not supply its form snapshot: {reason.form_id}") - data = self.human_input_v2_required_data(prepared, node_id=reason.node_id) - response = HumanInputRequiredResponse(task_id=task_id, workflow_run_id=run_id, data=data) - payload = reason.model_dump(mode="json") - payload.update(data.model_dump(mode="json")) - outputs.append(_HumanInputPauseOutput(response=response, pause_reason=payload)) - return outputs diff --git a/api/core/app/apps/completion/app_generator.py b/api/core/app/apps/completion/app_generator.py deleted file mode 100644 index 7312b079fe8e75..00000000000000 --- a/api/core/app/apps/completion/app_generator.py +++ /dev/null @@ -1,290 +0,0 @@ -import contextvars -import logging -import threading -import uuid -from collections.abc import Generator, Mapping -from typing import Any, Literal, overload - -from flask import Flask, copy_current_request_context, current_app -from pydantic import ValidationError -from sqlalchemy.orm import Session, sessionmaker - -from configs import dify_config -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.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.completion.app_config_manager import CompletionAppConfigManager -from core.app.apps.completion.app_runner import CompletionAppRunner -from core.app.apps.completion.generate_response_converter import CompletionAppGenerateResponseConverter -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 CompletionAppGenerateEntity, InvokeFrom -from core.app.entities.task_entities import AppBlockingResponse, AppStreamResponse -from core.db.session_factory import session_factory -from core.helper.trace_id_helper import extract_trace_session_id_from_args -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 models import Account, App, Conversation, EndUser, Message -from models.model import load_annotation_reply_config - -logger = logging.getLogger(__name__) - - -class CompletionAppGenerator(MessageBasedAppGenerator): - @overload - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: Literal[True], - *, - session: Session, - ) -> Generator[str | Mapping[str, Any]]: ... - - @overload - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: Literal[False], - *, - session: Session, - ) -> Mapping[str, Any]: ... - - @overload - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: bool = False, - *, - session: Session, - ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any]]: ... - - def generate( - self, - app_model: App, - user: Account | EndUser, - args: Mapping[str, Any], - invoke_from: InvokeFrom, - streaming: bool = True, - *, - session: Session, - ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any]]: - """ - 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 - """ - query = args["query"] - if not isinstance(query, str): - raise ValueError("query must be a string") - - query = query.replace("\x00", "") - inputs = args["inputs"] - - # get conversation - conversation = None - - # get app model config - app_model_config = self._get_app_model_config( - app_model=app_model, - conversation=conversation, - session=session, - ) - - # 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 = CompletionAppConfigManager.config_validate( - tenant_id=app_model.tenant_id, config=args.get("model_config", {}), session=session - ) - - if override_model_config_dict: - annotation_reply = None - effective_model_config_dict = override_model_config_dict - else: - annotation_reply = load_annotation_reply_config(session, app_model_config.app_id) - effective_model_config_dict = app_model_config.to_dict(annotation_reply=annotation_reply) - - # 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["files"] if args.get("files") else [] - file_extra_config = FileUploadConfigManager.convert(effective_model_config_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 = CompletionAppConfigManager.get_app_config( - app_model=app_model, - app_model_config=app_model_config, - override_config_dict=override_model_config_dict, - annotation_reply=annotation_reply, - ) - - # get tracing instance - trace_manager = TraceQueueManager( - app_id=app_model.id, user_id=user.id if isinstance(user, Account) else user.session_id - ) - - # init application generate entity - application_generate_entity = CompletionAppGenerateEntity( - task_id=str(uuid.uuid4()), - app_config=app_config, - model_conf=ModelConfigConverter.convert(app_config), - file_upload_config=file_extra_config, - inputs=self._prepare_user_inputs( - user_inputs=inputs, variables=app_config.variables, tenant_id=app_model.tenant_id - ), - query=query, - files=list(file_objs), - user_id=user.id, - stream=streaming, - invoke_from=invoke_from, - extras={ - **extract_trace_session_id_from_args(args), - }, - trace_manager=trace_manager, - ) - - # init generate records - (conversation, message) = self._init_generate_records( - application_generate_entity, - session=session, - ) - - response = self._start_generation(application_generate_entity, conversation, message) - return CompletionAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) - - def generate_from_entity( - self, - application_generate_entity: CompletionAppGenerateEntity, - *, - session_factory: sessionmaker[Session], - ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any]]: - """Persist prepared completion records before starting provider work. - - Preparation owns configuration, inputs and files. This path preserves - those values and closes its write session before starting the worker or - waiting for a blocking response. - """ - with session_factory(expire_on_commit=False) as session: - conversation, message = self._init_generate_records(application_generate_entity, session=session) - response = self._start_generation(application_generate_entity, conversation, message) - return CompletionAppGenerateResponseConverter.convert( - response=response, invoke_from=application_generate_entity.invoke_from - ) - - def _start_generation( - self, - application_generate_entity: CompletionAppGenerateEntity, - conversation: Conversation, - message: Message, - ) -> AppBlockingResponse | Generator[AppStreamResponse]: - 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, - ) - context = contextvars.copy_context() - - @copy_current_request_context - def worker_with_context() -> None: - context.run( - self._generate_worker, - flask_app=current_app._get_current_object(), # type: ignore - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - message_id=message.id, - ) - - worker_thread = threading.Thread(target=worker_with_context) - worker_thread.start() - - return self._handle_response( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation=conversation, - message=message, - stream=application_generate_entity.stream, - ) - - def _generate_worker( - self, - flask_app: Flask, - application_generate_entity: CompletionAppGenerateEntity, - queue_manager: AppQueueManager, - message_id: str, - ) -> None: - """ - Generate worker in a new thread. - :param flask_app: Flask app - :param application_generate_entity: application generate entity - :param queue_manager: queue manager - :param message_id: message ID - :return: - """ - with flask_app.app_context(): - try: - # get message - message = self._get_message(message_id) - - # chatbot app - runner = CompletionAppRunner() - with session_factory.create_session() as session: - runner.run( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - message=message, - session=session, - ) - 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() diff --git a/api/core/app/apps/completion/app_runner.py b/api/core/app/apps/completion/app_runner.py deleted file mode 100644 index 25828e5dd254aa..00000000000000 --- a/api/core/app/apps/completion/app_runner.py +++ /dev/null @@ -1,217 +0,0 @@ -import logging -from typing import cast - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.apps.base_app_queue_manager import AppQueueManager -from core.app.apps.base_app_runner import AppRunner -from core.app.apps.completion.app_config_manager import CompletionAppConfig -from core.app.entities.app_invoke_entities import ( - CompletionAppGenerateEntity, - get_credit_usage_app_type, - get_credit_usage_created_by, -) -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.db.session_factory import create_session -from core.model_manager import ModelManager -from core.moderation.base import ModerationError -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from graphon.file import File -from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent -from graphon.model_runtime.entities.model_entities import ModelType -from models.model import App, Message - -logger = logging.getLogger(__name__) - - -class CompletionAppRunner(AppRunner): - """ - Completion Application Runner - """ - - def run( - self, - application_generate_entity: CompletionAppGenerateEntity, - queue_manager: AppQueueManager, - message: Message, - session: Session, - ): - """Run the application without retaining ``session`` during model I/O. - - Database preparation is committed and the connection is released before - the provider response is requested or consumed. - - :param application_generate_entity: application generate entity - :param queue_manager: application queue manager - :param message: message - :return: - """ - app_config = application_generate_entity.app_config - app_config = cast(CompletionAppConfig, app_config) - stmt = select(App).where(App.id == app_config.app_id) - with create_session() as read_session: - app_record = read_session.scalar(stmt) - if app_record: - read_session.expunge(app_record) - if not app_record: - raise ValueError("App not found") - - inputs = application_generate_entity.inputs - query = application_generate_entity.query - files = application_generate_entity.files - - image_detail_config = ( - application_generate_entity.file_upload_config.image_config.detail - if ( - application_generate_entity.file_upload_config - and application_generate_entity.file_upload_config.image_config - ) - else None - ) - image_detail_config = image_detail_config or ImagePromptMessageContent.DETAIL.LOW - - # organize all inputs and template to prompt messages - # Include: prompt template, inputs, query(optional), files(optional) - prompt_messages, stop = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=inputs, - files=files, - query=query, - image_detail_config=image_detail_config, - ) - - # moderation - try: - # process sensitive_word_avoidance - _, inputs, query = self.moderation_for_inputs( - app_id=app_record.id, - tenant_id=app_config.tenant_id, - app_generate_entity=application_generate_entity, - inputs=inputs, - query=query or "", - message_id=message.id, - ) - except ModerationError as e: - self.direct_output( - queue_manager=queue_manager, - app_generate_entity=application_generate_entity, - prompt_messages=prompt_messages, - text=str(e), - stream=application_generate_entity.stream, - ) - return - - # fill in variable inputs from external data tools if exists - external_data_tools = app_config.external_data_variables - if external_data_tools: - inputs = self.fill_in_inputs_from_external_data_tools( - tenant_id=app_record.tenant_id, - app_id=app_record.id, - external_data_tools=external_data_tools, - inputs=inputs, - query=query, - ) - - # get context from datasets - context = None - context_files: list[File] = [] - if app_config.dataset and app_config.dataset.dataset_ids: - hit_callback = DatasetIndexToolCallbackHandler( - queue_manager, - app_record.id, - message.id, - application_generate_entity.user_id, - application_generate_entity.invoke_from, - ) - - dataset_config = app_config.dataset - if dataset_config and dataset_config.retrieve_config.query_variable: - query = inputs.get(dataset_config.retrieve_config.query_variable, "") - - dataset_retrieval = DatasetRetrieval(application_generate_entity) - context, retrieved_files = dataset_retrieval.retrieve( - session=session, - app_id=app_record.id, - user_id=application_generate_entity.user_id, - tenant_id=app_record.tenant_id, - model_config=application_generate_entity.model_conf, - config=dataset_config, - query=query or "", - invoke_from=application_generate_entity.invoke_from, - show_retrieve_source=app_config.additional_features.show_retrieve_source - if app_config.additional_features - else False, - hit_callback=hit_callback, - message_id=message.id, - inputs=inputs, - vision_enabled=bool( - application_generate_entity.app_config.app_model_config_dict.get("file_upload", {}) - .get("image", {}) - .get("enabled", False) - ), - ) - context_files = retrieved_files or [] - - session.commit() - session.close() - - # reorganize all inputs and template to prompt messages - # Include: prompt template, inputs, query(optional), files(optional) - # memory(optional), external data, dataset context(optional) - prompt_messages, stop = self.organize_prompt_messages( - app_record=app_record, - model_config=application_generate_entity.model_conf, - prompt_template_entity=app_config.prompt_template, - inputs=inputs, - files=files, - query=query, - context=context, - image_detail_config=image_detail_config, - context_files=context_files, - ) - - # check hosting moderation - hosting_moderation_result = self.check_hosting_moderation( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - prompt_messages=prompt_messages, - ) - - if hosting_moderation_result: - return - - # Re-calculate the max tokens if sum(prompt_token + max_tokens) over model token limit - self.recalc_llm_max_tokens(model_config=application_generate_entity.model_conf, prompt_messages=prompt_messages) - - # Invoke model - model_instance = ModelManager.for_tenant(tenant_id=app_config.tenant_id).get_model_instance( - tenant_id=app_config.tenant_id, - provider=application_generate_entity.model_conf.provider, - model_type=ModelType.LLM, - model=application_generate_entity.model_conf.model, - ) - - request_metadata: dict[str, object] = {"app_id": app_config.app_id} - request_metadata["app_type"] = get_credit_usage_app_type(app_config.app_mode) - request_metadata["created_by"] = get_credit_usage_created_by(app_config.app_mode) - - invoke_result = model_instance.invoke_llm( - prompt_messages=prompt_messages, - model_parameters=application_generate_entity.model_conf.parameters, - stop=stop, - stream=application_generate_entity.stream, - request_metadata=request_metadata, - ) - - # handle invoke result - self._handle_invoke_result( - invoke_result=invoke_result, - queue_manager=queue_manager, - stream=application_generate_entity.stream, - message_id=message.id, - user_id=application_generate_entity.user_id, - tenant_id=app_config.tenant_id, - ) diff --git a/api/core/app/apps/message_based_app_generator.py b/api/core/app/apps/message_based_app_generator.py deleted file mode 100644 index 7a6ffa334e14f7..00000000000000 --- a/api/core/app/apps/message_based_app_generator.py +++ /dev/null @@ -1,336 +0,0 @@ -import json -import logging -from collections.abc import Callable, Generator, Mapping -from decimal import Decimal -from typing import cast - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.app_config.entities import EasyUIBasedAppConfig, EasyUIBasedAppModelConfigFrom -from core.app.apps.base_app_generator import BaseAppGenerator -from core.app.apps.base_app_queue_manager import AppQueueManager -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.streaming_utils import stream_topic_events -from core.app.entities.app_invoke_entities import ( - AdvancedChatAppGenerateEntity, - AgentChatAppGenerateEntity, - AppGenerateEntity, - ChatAppGenerateEntity, - CompletionAppGenerateEntity, - ConversationAppGenerateEntity, - InvokeFrom, -) -from core.app.entities.task_entities import ( - ChatbotAppBlockingResponse, - ChatbotAppStreamResponse, - CompletionAppBlockingResponse, - CompletionAppStreamResponse, -) -from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline -from core.prompt.utils.prompt_template_parser import PromptTemplateParser -from core.workflow.file_reference import resolve_file_record_id -from extensions.ext_database import db -from extensions.ext_redis import get_pubsub_broadcast_channel -from libs.broadcast_channel.channel import SupportsPreparedSubscription, Topic -from libs.datetime_utils import naive_utc_now -from models.enums import ( - ConversationFromSource, - ConversationStatus, - CreatorUserRole, - MessageFileBelongsTo, -) -from models.enums import ( - InvokeFrom as ModelInvokeFrom, -) -from models.model import App, AppMode, AppModelConfig, Conversation, Message, MessageFile -from services.errors.app_model_config import AppModelConfigBrokenError -from services.errors.conversation import ConversationNotExistsError -from services.errors.message import MessageNotExistsError - -logger = logging.getLogger(__name__) - - -class MessageBasedAppGenerator(BaseAppGenerator): - def _handle_response( - self, - application_generate_entity: ChatAppGenerateEntity | CompletionAppGenerateEntity | AgentChatAppGenerateEntity, - queue_manager: AppQueueManager, - conversation: Conversation, - message: Message, - stream: bool = False, - ) -> ( - ChatbotAppBlockingResponse - | CompletionAppBlockingResponse - | Generator[ChatbotAppStreamResponse | CompletionAppStreamResponse] - ): - """ - Handle response. - :param application_generate_entity: application generate entity - :param queue_manager: queue manager - :param conversation: conversation - :param message: message - :param stream: is stream - :return: - """ - # init generate task pipeline - generate_task_pipeline = EasyUIBasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - conversation=conversation, - message=message, - stream=stream, - ) - - try: - return generate_task_pipeline.process() - except ValueError as e: - if len(e.args) > 0 and e.args[0] == "I/O operation on closed file.": # ignore this error - raise GenerateTaskStoppedError() from e - else: - logger.exception("Failed to handle response, conversation_id: %s", conversation.id) - raise e - - def _get_app_model_config( - self, - app_model: App, - conversation: Conversation | None = None, - *, - session: Session, - ) -> AppModelConfig: - if conversation: - stmt = select(AppModelConfig).where( - AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id - ) - app_model_config = session.scalar(stmt) - - if not app_model_config: - raise AppModelConfigBrokenError() - else: - if app_model.app_model_config_id is None: - raise AppModelConfigBrokenError() - - app_model_config = session.get(AppModelConfig, app_model.app_model_config_id) - - if not app_model_config: - raise AppModelConfigBrokenError() - - return app_model_config - - def _init_generate_records( - self, - application_generate_entity: ChatAppGenerateEntity - | CompletionAppGenerateEntity - | AgentChatAppGenerateEntity - | AdvancedChatAppGenerateEntity, - conversation: Conversation | None = None, - *, - session: Session, - ) -> tuple[Conversation, Message]: - """ - Initialize generate records - :param application_generate_entity: application generate entity - :conversation conversation - :return: - """ - app_config: EasyUIBasedAppConfig = cast(EasyUIBasedAppConfig, application_generate_entity.app_config) - - # get from source - end_user_id = None - account_id = None - if application_generate_entity.invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}: - from_source = ConversationFromSource.API - end_user_id = application_generate_entity.user_id - else: - from_source = ConversationFromSource.CONSOLE - account_id = application_generate_entity.user_id - - if isinstance(application_generate_entity, AdvancedChatAppGenerateEntity): - app_model_config_id = None - override_model_configs = None - model_provider = None - model_id = None - else: - app_model_config_id = app_config.app_model_config_id - model_provider = application_generate_entity.model_conf.provider - model_id = application_generate_entity.model_conf.model - override_model_configs = None - if app_config.app_model_config_from == EasyUIBasedAppModelConfigFrom.ARGS and app_config.app_mode in { - AppMode.AGENT_CHAT, - AppMode.CHAT, - AppMode.COMPLETION, - }: - override_model_configs = app_config.app_model_config_dict - - # get conversation introduction - introduction = self._get_conversation_introduction(application_generate_entity) - - # get conversation name - query = application_generate_entity.query or "New conversation" - conversation_name = (query[:20] + "…") if len(query) > 20 else query - - created_new_conversation = conversation is None - try: - if not conversation: - conversation = Conversation( - app_id=app_config.app_id, - app_model_config_id=app_model_config_id, - model_provider=model_provider, - model_id=model_id, - override_model_configs=json.dumps(override_model_configs) if override_model_configs else None, - mode=AppMode(app_config.app_mode.value), - name=conversation_name, - _inputs=dict(application_generate_entity.inputs), - introduction=introduction, - system_instruction="", - system_instruction_tokens=0, - status=ConversationStatus.NORMAL, - invoke_from=ModelInvokeFrom(application_generate_entity.invoke_from.value), - from_source=from_source, - from_end_user_id=end_user_id, - from_account_id=account_id, - ) - - session.add(conversation) - session.flush() - session.refresh(conversation) - else: - conversation.updated_at = naive_utc_now() - - message = Message( - app_id=app_config.app_id, - model_provider=model_provider, - model_id=model_id, - override_model_configs=json.dumps(override_model_configs) if override_model_configs else None, - conversation_id=conversation.id, - _inputs=dict(application_generate_entity.inputs), - query=application_generate_entity.query, - message="", - message_tokens=0, - message_unit_price=Decimal(0), - message_price_unit=Decimal(0), - answer="", - answer_tokens=0, - answer_unit_price=Decimal(0), - answer_price_unit=Decimal(0), - parent_message_id=getattr(application_generate_entity, "parent_message_id", None), - provider_response_latency=0, - total_price=Decimal(0), - currency="USD", - invoke_from=ModelInvokeFrom(application_generate_entity.invoke_from.value), - from_source=from_source, - from_end_user_id=end_user_id, - from_account_id=account_id, - app_mode=app_config.app_mode, - ) - - session.add(message) - session.flush() - session.refresh(message) - - message_files = [] - for file in application_generate_entity.files: - message_file = MessageFile( - message_id=message.id, - type=file.type, - transfer_method=file.transfer_method, - belongs_to=MessageFileBelongsTo.USER, - url=file.remote_url, - upload_file_id=resolve_file_record_id(file.reference), - created_by_role=(CreatorUserRole.ACCOUNT if account_id else CreatorUserRole.END_USER), - created_by=account_id or end_user_id or "", - ) - message_files.append(message_file) - - if message_files: - session.add_all(message_files) - - session.commit() - - if isinstance(application_generate_entity, ConversationAppGenerateEntity): - application_generate_entity.conversation_id = conversation.id - application_generate_entity.is_new_conversation = created_new_conversation - return conversation, message - except Exception: - session.rollback() - raise - - def _get_conversation_introduction(self, application_generate_entity: AppGenerateEntity) -> str: - """ - Get conversation introduction - :param application_generate_entity: application generate entity - :return: conversation introduction - """ - app_config = application_generate_entity.app_config - introduction = app_config.additional_features.opening_statement - - if introduction: - try: - inputs = application_generate_entity.inputs - prompt_template = PromptTemplateParser(template=introduction) - prompt_inputs = {k: inputs[k] for k in prompt_template.variable_keys if k in inputs} - introduction = prompt_template.format(prompt_inputs) - except KeyError: - pass - - return introduction or "" - - def _get_conversation(self, conversation_id: str) -> Conversation: - """ - Get conversation by conversation id - :param conversation_id: conversation id - :return: conversation - """ - with Session(db.engine, expire_on_commit=False) as session: - conversation = session.scalar(select(Conversation).where(Conversation.id == conversation_id)) - - if not conversation: - raise ConversationNotExistsError("Conversation not exists") - - return conversation - - def _get_message(self, message_id: str) -> Message: - """ - Get message by message id - :param message_id: message id - :return: message - """ - with Session(db.engine, expire_on_commit=False) as session: - message = session.scalar(select(Message).where(Message.id == message_id)) - - if message is None: - raise MessageNotExistsError("Message not exists") - - return message - - @staticmethod - def _make_channel_key(app_mode: AppMode, workflow_run_id: str): - return f"channel:{app_mode}:{workflow_run_id}" - - @classmethod - def get_response_topic(cls, app_mode: AppMode, workflow_run_id: str) -> Topic: - key = cls._make_channel_key(app_mode, workflow_run_id) - channel = get_pubsub_broadcast_channel() - topic = channel.topic(key) - return topic - - @classmethod - def retrieve_events( - cls, - app_mode: AppMode, - workflow_run_id: str, - idle_timeout=300, - on_subscribe: Callable[[], None] | None = None, - ) -> Generator[Mapping | str]: - topic = cls.get_response_topic(app_mode, workflow_run_id) - subscriber = topic.as_subscriber() - subscription = ( - subscriber.prepare_subscription() - if isinstance(subscriber, SupportsPreparedSubscription) - else subscriber.subscribe() - ) - return stream_topic_events( - subscription=subscription, - idle_timeout=idle_timeout, - on_subscribe=on_subscribe, - ) diff --git a/api/core/app/apps/message_generator.py b/api/core/app/apps/message_generator.py deleted file mode 100644 index 79ff9cad14bba8..00000000000000 --- a/api/core/app/apps/message_generator.py +++ /dev/null @@ -1,45 +0,0 @@ -from collections.abc import Callable, Generator, Iterable, Mapping - -from core.app.apps.streaming_utils import stream_topic_events -from core.app.entities.task_entities import StreamEvent -from extensions.ext_redis import get_pubsub_broadcast_channel -from libs.broadcast_channel.channel import SupportsPreparedSubscription, Topic -from models.model import AppMode - - -class MessageGenerator: - @staticmethod - def _make_channel_key(app_mode: AppMode, workflow_run_id: str): - return f"channel:{app_mode}:{str(workflow_run_id)}" - - @classmethod - def get_response_topic(cls, app_mode: AppMode, workflow_run_id: str) -> Topic: - key = cls._make_channel_key(app_mode, workflow_run_id) - channel = get_pubsub_broadcast_channel() - topic = channel.topic(key) - return topic - - @classmethod - def retrieve_events( - cls, - app_mode: AppMode, - workflow_run_id: str, - idle_timeout=300, - ping_interval: float = 10.0, - on_subscribe: Callable[[], None] | None = None, - terminal_events: Iterable[str | StreamEvent] | None = None, - ) -> Generator[Mapping | str]: - topic = cls.get_response_topic(app_mode, workflow_run_id) - subscriber = topic.as_subscriber() - subscription = ( - subscriber.prepare_subscription() - if isinstance(subscriber, SupportsPreparedSubscription) - else subscriber.subscribe() - ) - return stream_topic_events( - subscription=subscription, - idle_timeout=idle_timeout, - ping_interval=ping_interval, - on_subscribe=on_subscribe, - terminal_events=terminal_events, - ) diff --git a/api/core/app/apps/pipeline/__init__.py b/api/core/app/apps/pipeline/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/core/app/apps/pipeline/generate_response_converter.py b/api/core/app/apps/pipeline/generate_response_converter.py deleted file mode 100644 index 2b9391fed795f4..00000000000000 --- a/api/core/app/apps/pipeline/generate_response_converter.py +++ /dev/null @@ -1,109 +0,0 @@ -from collections.abc import Generator -from typing import Any, cast, override - -from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter -from core.app.entities.task_entities import ( - AppStreamResponse, - ErrorStreamResponse, - NodeFinishStreamResponse, - NodeStartStreamResponse, - PingStreamResponse, - WorkflowAppBlockingResponse, - WorkflowAppStreamResponse, -) - - -class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter[WorkflowAppBlockingResponse]): - @classmethod - @override - def convert_blocking_full_response(cls, blocking_response: WorkflowAppBlockingResponse) -> dict[str, object]: - """ - Convert blocking full response. - :param blocking_response: blocking response - :return: - """ - return dict(blocking_response.model_dump()) - - @classmethod - @override - def convert_blocking_simple_response(cls, blocking_response: WorkflowAppBlockingResponse) -> dict[str, object]: - """ - Convert blocking simple response. - :param blocking_response: blocking response - :return: - """ - return cls.convert_blocking_full_response(blocking_response) - - @classmethod - @override - def convert_stream_full_response( - cls, stream_response: Generator[AppStreamResponse] - ) -> Generator[dict[str, Any] | str]: - """ - Convert stream full response. - :param stream_response: stream response - :return: - """ - for chunk in stream_response: - chunk = cast(WorkflowAppStreamResponse, chunk) - sub_stream_response = chunk.stream_response - - match sub_stream_response: - case PingStreamResponse(): - yield "ping" - continue - case ErrorStreamResponse(): - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - data = cls._error_to_stream_response(sub_stream_response.err) - response_chunk.update(cast(dict, data)) - case _: - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - response_chunk.update(sub_stream_response.model_dump()) - - yield response_chunk - - @classmethod - @override - def convert_stream_simple_response( - cls, stream_response: Generator[AppStreamResponse] - ) -> Generator[dict[str, Any] | str]: - """ - Convert stream simple response. - :param stream_response: stream response - :return: - """ - for chunk in stream_response: - chunk = cast(WorkflowAppStreamResponse, chunk) - sub_stream_response = chunk.stream_response - - match sub_stream_response: - case PingStreamResponse(): - yield "ping" - continue - case ErrorStreamResponse(): - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - data = cls._error_to_stream_response(sub_stream_response.err) - response_chunk.update(cast(dict, data)) - case NodeStartStreamResponse() | NodeFinishStreamResponse(): - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - response_chunk.update(cast(dict, sub_stream_response.to_ignore_detail_dict())) - case _: - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - response_chunk.update(sub_stream_response.model_dump()) - - yield response_chunk diff --git a/api/core/app/apps/pipeline/pipeline_config_manager.py b/api/core/app/apps/pipeline/pipeline_config_manager.py deleted file mode 100644 index 53b85f9f27b3a2..00000000000000 --- a/api/core/app/apps/pipeline/pipeline_config_manager.py +++ /dev/null @@ -1,72 +0,0 @@ -from typing import Any - -from core.app.app_config.base_app_config_manager import BaseAppConfigManager -from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager -from core.app.app_config.entities import RagPipelineVariableEntity, WorkflowUIBasedAppConfig -from core.app.app_config.features.file_upload.manager import FileUploadConfigManager -from core.app.app_config.features.text_to_speech.manager import TextToSpeechConfigManager -from core.app.app_config.workflow_ui_based_app.variables.manager import WorkflowVariablesConfigManager -from models.dataset import Pipeline -from models.model import AppMode -from models.workflow import Workflow - - -class PipelineConfig(WorkflowUIBasedAppConfig): - """ - Pipeline Config Entity. - """ - - rag_pipeline_variables: list[RagPipelineVariableEntity] = [] - pass - - -class PipelineConfigManager(BaseAppConfigManager): - @classmethod - def get_pipeline_config(cls, pipeline: Pipeline, workflow: Workflow, start_node_id: str) -> PipelineConfig: - pipeline_config = PipelineConfig( - tenant_id=pipeline.tenant_id, - app_id=pipeline.id, - app_mode=AppMode.RAG_PIPELINE, - workflow_id=workflow.id, - rag_pipeline_variables=WorkflowVariablesConfigManager.convert_rag_pipeline_variable( - workflow=workflow, - start_node_id=start_node_id, - rag_pipeline_variables=workflow.rag_pipeline_variables, - ), - ) - - return pipeline_config - - @classmethod - def config_validate( - cls, tenant_id: str, config: dict[str, Any], only_structure_validate: bool = False - ) -> dict[str, Any]: - """ - Validate for pipeline config - - :param tenant_id: tenant id - :param config: app model config args - :param only_structure_validate: only validate the structure of the config - """ - related_config_keys = [] - - # file upload validation - config, current_related_config_keys = FileUploadConfigManager.validate_and_set_defaults(config=config) - related_config_keys.extend(current_related_config_keys) - - # text_to_speech - config, current_related_config_keys = TextToSpeechConfigManager.validate_and_set_defaults(config) - related_config_keys.extend(current_related_config_keys) - - # moderation validation - config, current_related_config_keys = SensitiveWordAvoidanceConfigManager.validate_and_set_defaults( - tenant_id=tenant_id, config=config, only_structure_validate=only_structure_validate - ) - related_config_keys.extend(current_related_config_keys) - - related_config_keys = list(set(related_config_keys)) - - # Filter out extra parameters - filtered_config = {key: config.get(key) for key in related_config_keys} - - return filtered_config diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py deleted file mode 100644 index 9e174d01ae028b..00000000000000 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility import for the pipeline generation service adapter.""" - -from services.workflow.execution.adapters.pipeline.pipeline_generator import PipelineGenerator - -__all__ = ["PipelineGenerator"] diff --git a/api/core/app/apps/pipeline/pipeline_queue_manager.py b/api/core/app/apps/pipeline/pipeline_queue_manager.py deleted file mode 100644 index 13459fb94d15e0..00000000000000 --- a/api/core/app/apps/pipeline/pipeline_queue_manager.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import override - -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.execution_coordinator import AppExecutionState -from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.entities.queue_entities import ( - AppQueueEvent, - QueueErrorEvent, - QueueMessageEndEvent, - QueueStopEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPartialSuccessEvent, - QueueWorkflowSucceededEvent, - WorkflowQueueMessage, -) - - -class PipelineQueueManager(AppQueueManager): - def __init__(self, task_id: str, user_id: str, invoke_from: InvokeFrom, app_mode: str) -> None: - super().__init__(task_id, user_id, invoke_from) - - self._app_mode = app_mode - - @override - def _publish(self, event: AppQueueEvent, pub_from: PublishFrom) -> None: - """ - Publish event to queue - :param event: - :param pub_from: - :return: - """ - message = WorkflowQueueMessage(task_id=self._task_id, app_mode=self._app_mode, event=event) - - self._q.put(message) - - if isinstance( - event, - QueueStopEvent - | QueueErrorEvent - | QueueMessageEndEvent - | QueueWorkflowSucceededEvent - | QueueWorkflowFailedEvent - | QueueWorkflowPartialSuccessEvent, - ): - self.stop_listen(execution_state=AppExecutionState.TERMINAL) - - if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped(): - raise GenerateTaskStoppedError() diff --git a/api/core/app/apps/pipeline/pipeline_runner.py b/api/core/app/apps/pipeline/pipeline_runner.py deleted file mode 100644 index 5bcacca845477d..00000000000000 --- a/api/core/app/apps/pipeline/pipeline_runner.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Compatibility imports for the pipeline execution service runner.""" - -from models.pipeline_execution import PipelineDocumentStore -from services.workflow.execution.adapters.pipeline.pipeline_runner import PipelineRunner - -__all__ = ["PipelineDocumentStore", "PipelineRunner"] diff --git a/api/core/app/apps/workflow/__init__.py b/api/core/app/apps/workflow/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/core/app/apps/workflow/active_workflow_tasks.py b/api/core/app/apps/workflow/active_workflow_tasks.py deleted file mode 100644 index ad2d8cfb5e029d..00000000000000 --- a/api/core/app/apps/workflow/active_workflow_tasks.py +++ /dev/null @@ -1,38 +0,0 @@ -"""In-process registry for workflow application task IDs.""" - -import threading -from collections.abc import Generator -from contextlib import contextmanager - -_active_task_ids: set[str] = set() -_active_task_ids_lock = threading.RLock() - - -@contextmanager -def active_workflow_task(task_id: str) -> Generator[None]: - """Register a workflow application task ID for the duration of a workflow run.""" - if not task_id: - raise ValueError("task_id must not be empty") - - with _active_task_ids_lock: - if task_id in _active_task_ids: - raise ValueError(f"Workflow task already active for task_id={task_id}") - _active_task_ids.add(task_id) - - try: - yield - finally: - with _active_task_ids_lock: - _active_task_ids.discard(task_id) - - -def get_active_workflow_task_count() -> int: - """Return the number of active workflow application task IDs in this process.""" - with _active_task_ids_lock: - return len(_active_task_ids) - - -def reset_active_workflow_tasks() -> None: - """Clear active workflow application task IDs for worker initialization and tests.""" - with _active_task_ids_lock: - _active_task_ids.clear() diff --git a/api/core/app/apps/workflow/app_config_manager.py b/api/core/app/apps/workflow/app_config_manager.py deleted file mode 100644 index e72da91c21f35b..00000000000000 --- a/api/core/app/apps/workflow/app_config_manager.py +++ /dev/null @@ -1,67 +0,0 @@ -from core.app.app_config.base_app_config_manager import BaseAppConfigManager -from core.app.app_config.common.sensitive_word_avoidance.manager import SensitiveWordAvoidanceConfigManager -from core.app.app_config.entities import WorkflowUIBasedAppConfig -from core.app.app_config.features.file_upload.manager import FileUploadConfigManager -from core.app.app_config.features.text_to_speech.manager import TextToSpeechConfigManager -from core.app.app_config.workflow_ui_based_app.variables.manager import WorkflowVariablesConfigManager -from models.model import App, AppMode -from models.workflow import Workflow - - -class WorkflowAppConfig(WorkflowUIBasedAppConfig): - """ - Workflow App Config Entity. - """ - - pass - - -class WorkflowAppConfigManager(BaseAppConfigManager): - @classmethod - def get_app_config(cls, app_model: App, workflow: Workflow) -> WorkflowAppConfig: - features_dict = workflow.features_dict - - app_mode = AppMode.value_of(app_model.mode) - app_config = WorkflowAppConfig( - tenant_id=app_model.tenant_id, - app_id=app_model.id, - app_mode=app_mode, - workflow_id=workflow.id, - sensitive_word_avoidance=SensitiveWordAvoidanceConfigManager.convert(config=features_dict), - variables=WorkflowVariablesConfigManager.convert(workflow=workflow), - additional_features=cls.convert_features(features_dict, app_mode), - ) - - return app_config - - @classmethod - def config_validate(cls, tenant_id: str, config: dict, only_structure_validate: bool = False): - """ - Validate for workflow app model config - - :param tenant_id: tenant id - :param config: app model config args - :param only_structure_validate: only validate the structure of the config - """ - related_config_keys = [] - - # file upload validation - config, current_related_config_keys = FileUploadConfigManager.validate_and_set_defaults(config=config) - related_config_keys.extend(current_related_config_keys) - - # text_to_speech - config, current_related_config_keys = TextToSpeechConfigManager.validate_and_set_defaults(config) - related_config_keys.extend(current_related_config_keys) - - # moderation validation - config, current_related_config_keys = SensitiveWordAvoidanceConfigManager.validate_and_set_defaults( - tenant_id=tenant_id, config=config, only_structure_validate=only_structure_validate - ) - related_config_keys.extend(current_related_config_keys) - - related_config_keys = list(set(related_config_keys)) - - # Filter out extra parameters - filtered_config = {key: config.get(key) for key in related_config_keys} - - return filtered_config diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py deleted file mode 100644 index d83867dcc45c82..00000000000000 --- a/api/core/app/apps/workflow/app_generator.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility import for the workflow generation service adapter.""" - -from services.workflow.execution.adapters.workflow.app_generator import WorkflowAppGenerator - -__all__ = ["WorkflowAppGenerator"] diff --git a/api/core/app/apps/workflow/app_queue_manager.py b/api/core/app/apps/workflow/app_queue_manager.py deleted file mode 100644 index f837805f7b56cd..00000000000000 --- a/api/core/app/apps/workflow/app_queue_manager.py +++ /dev/null @@ -1,48 +0,0 @@ -from typing import override - -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.apps.execution_coordinator import AppExecutionState -from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.entities.queue_entities import ( - AppQueueEvent, - QueueErrorEvent, - QueueMessageEndEvent, - QueueStopEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPartialSuccessEvent, - QueueWorkflowPausedEvent, - QueueWorkflowSucceededEvent, - WorkflowQueueMessage, -) - - -class WorkflowAppQueueManager(AppQueueManager): - def __init__(self, task_id: str, user_id: str, invoke_from: InvokeFrom, app_mode: str): - super().__init__(task_id, user_id, invoke_from) - - self._app_mode = app_mode - - @override - def _publish(self, event: AppQueueEvent, pub_from: PublishFrom): - """ - Publish event to queue - :param event: - :param pub_from: - :return: - """ - message = WorkflowQueueMessage(task_id=self._task_id, app_mode=self._app_mode, event=event) - - self._q.put(message) - - if isinstance(event, QueueWorkflowPausedEvent): - self.stop_listen(execution_state=AppExecutionState.PAUSED) - elif isinstance( - event, - QueueStopEvent - | QueueErrorEvent - | QueueMessageEndEvent - | QueueWorkflowSucceededEvent - | QueueWorkflowFailedEvent - | QueueWorkflowPartialSuccessEvent, - ): - self.stop_listen(execution_state=AppExecutionState.TERMINAL) diff --git a/api/core/app/apps/workflow/app_runner.py b/api/core/app/apps/workflow/app_runner.py deleted file mode 100644 index d94196e23dc10b..00000000000000 --- a/api/core/app/apps/workflow/app_runner.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility imports for the workflow execution service runner.""" - -from services.workflow.execution.adapters.workflow.app_runner import WorkflowAppRunner, WorkflowExecutionCancellation - -__all__ = ["WorkflowAppRunner", "WorkflowExecutionCancellation"] diff --git a/api/core/app/apps/workflow/command_channels.py b/api/core/app/apps/workflow/command_channels.py deleted file mode 100644 index 234a5e4307ab98..00000000000000 --- a/api/core/app/apps/workflow/command_channels.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Command channels used by Dify workflow runners.""" - -import logging -from collections.abc import Callable, Sequence -from typing import final, override - -from core.app.apps.execution_coordinator import is_app_task_stop_flag_set -from graphon.engine.command import AbortCommand, Command, CommandChannel - -logger = logging.getLogger(__name__) -ShutdownStateGetter = Callable[[], bool] - - -@final -class CombinedCommandChannel: - """Fetch commands from all sources and send outbound commands through the primary source.""" - - _command_channels: tuple[CommandChannel, ...] - - def __init__(self, command_channels: Sequence[CommandChannel]) -> None: - if not command_channels: - raise ValueError("command_channels must not be empty") - self._command_channels = tuple(command_channels) - - def fetch_commands(self) -> list[Command]: - commands: list[Command] = [] - for channel in self._command_channels: - try: - commands.extend(channel.fetch_commands()) - except Exception: - logger.exception("Failed to fetch Engine commands from %s", channel.__class__.__name__) - return commands - - def send_command(self, command: Command) -> None: - """Send commands through the first channel, which is the runner's primary command sink.""" - self._command_channels[0].send_command(command) - - -@final -class CelerySignalCommandChannel(CommandChannel): - """Translate process-local Celery shutdown state into one Engine abort command.""" - - _shutdown_state_getter: ShutdownStateGetter - _abort_reason: str - _abort_emitted: bool - - def __init__( - self, - *, - shutdown_state_getter: ShutdownStateGetter, - abort_reason: str, - ) -> None: - self._shutdown_state_getter = shutdown_state_getter - self._abort_reason = abort_reason - self._abort_emitted = False - - @override - def fetch_commands(self) -> list[Command]: - if self._abort_emitted or not self._shutdown_state_getter(): - return [] - - self._abort_emitted = True - return [AbortCommand(reason=self._abort_reason)] - - @override - def send_command(self, command: Command) -> None: - _ = command - - -@final -class StopFlagCommandChannel(CommandChannel): - """Translate the legacy Redis stop flag into one GraphEngine abort command.""" - - _task_id: str - _abort_reason: str - _abort_emitted: bool - - def __init__( - self, - *, - task_id: str, - abort_reason: str = "User requested stop", - ) -> None: - self._task_id = task_id - self._abort_reason = abort_reason - self._abort_emitted = False - - @override - def fetch_commands(self) -> list[Command]: - if self._abort_emitted or not is_app_task_stop_flag_set(self._task_id): - return [] - - self._abort_emitted = True - return [AbortCommand(reason=self._abort_reason)] - - @override - def send_command(self, command: Command) -> None: - _ = command diff --git a/api/core/app/apps/workflow/errors.py b/api/core/app/apps/workflow/errors.py deleted file mode 100644 index 16cd864209d51c..00000000000000 --- a/api/core/app/apps/workflow/errors.py +++ /dev/null @@ -1,7 +0,0 @@ -from libs.exception import BaseHTTPException - - -class WorkflowPausedInBlockingModeError(BaseHTTPException): - error_code = "workflow_paused_in_blocking_mode" - description = "Workflow execution paused for human input; blocking response mode is not supported." - code = 400 diff --git a/api/core/app/apps/workflow/generate_response_converter.py b/api/core/app/apps/workflow/generate_response_converter.py deleted file mode 100644 index 929f01157b3d13..00000000000000 --- a/api/core/app/apps/workflow/generate_response_converter.py +++ /dev/null @@ -1,116 +0,0 @@ -from collections.abc import Generator -from typing import Any, cast, override - -from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter -from core.app.entities.task_entities import ( - AppStreamResponse, - ErrorStreamResponse, - NodeFinishStreamResponse, - NodeStartStreamResponse, - PingStreamResponse, - WorkflowAppBlockingResponse, - WorkflowAppPausedBlockingResponse, - WorkflowAppStreamResponse, -) - - -class WorkflowAppGenerateResponseConverter( - AppGenerateResponseConverter[WorkflowAppBlockingResponse | WorkflowAppPausedBlockingResponse] -): - @classmethod - @override - def convert_blocking_full_response( - cls, blocking_response: WorkflowAppBlockingResponse | WorkflowAppPausedBlockingResponse - ) -> dict[str, Any]: - """ - Convert blocking full response. - :param blocking_response: blocking response - :return: - """ - return dict(blocking_response.model_dump()) - - @classmethod - @override - def convert_blocking_simple_response( - cls, blocking_response: WorkflowAppBlockingResponse | WorkflowAppPausedBlockingResponse - ) -> dict[str, Any]: - """ - Convert blocking simple response. - :param blocking_response: blocking response - :return: - """ - return cls.convert_blocking_full_response(blocking_response) - - @classmethod - @override - def convert_stream_full_response( - cls, stream_response: Generator[AppStreamResponse] - ) -> Generator[dict[str, Any] | str]: - """ - Convert stream full response. - :param stream_response: stream response - :return: - """ - for chunk in stream_response: - chunk = cast(WorkflowAppStreamResponse, chunk) - sub_stream_response = chunk.stream_response - - match sub_stream_response: - case PingStreamResponse(): - yield "ping" - continue - case ErrorStreamResponse(): - response_chunk: dict[str, object] = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - data = cls._error_to_stream_response(sub_stream_response.err) - response_chunk.update(data) - case _: - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - response_chunk.update(sub_stream_response.model_dump(mode="json")) - - yield response_chunk - - @classmethod - @override - def convert_stream_simple_response( - cls, stream_response: Generator[AppStreamResponse] - ) -> Generator[dict[str, Any] | str]: - """ - Convert stream simple response. - :param stream_response: stream response - :return: - """ - for chunk in stream_response: - chunk = cast(WorkflowAppStreamResponse, chunk) - sub_stream_response = chunk.stream_response - - match sub_stream_response: - case PingStreamResponse(): - yield "ping" - continue - case ErrorStreamResponse(): - response_chunk: dict[str, object] = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - data = cls._error_to_stream_response(sub_stream_response.err) - response_chunk.update(data) - case NodeStartStreamResponse() | NodeFinishStreamResponse(): - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - response_chunk.update(sub_stream_response.to_ignore_detail_dict()) - case _: - response_chunk = { - "event": sub_stream_response.event.value, - "workflow_run_id": chunk.workflow_run_id, - } - response_chunk.update(sub_stream_response.model_dump(mode="json")) - - yield response_chunk diff --git a/api/core/app/apps/workflow/generate_task_pipeline.py b/api/core/app/apps/workflow/generate_task_pipeline.py deleted file mode 100644 index 1ea0fb1127700e..00000000000000 --- a/api/core/app/apps/workflow/generate_task_pipeline.py +++ /dev/null @@ -1,792 +0,0 @@ -import logging -from collections.abc import Callable, Generator -from contextlib import contextmanager - -from sqlalchemy.orm import Session, sessionmaker - -from core.app.apps.base_app_queue_manager import AppQueueManager -from core.app.apps.common.graph_runtime_state_support import GraphRuntimeStateSupport -from core.app.apps.common.workflow_response_converter import WorkflowResponseConverter -from core.app.apps.draft_variable_saver import DraftVariableSaverFactory -from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity, get_credit_usage_app_type -from core.app.entities.queue_entities import ( - AppQueueEvent, - MessageQueueMessage, - QueueAgentLogEvent, - QueueErrorEvent, - QueueHumanInputFormFilledEvent, - QueueHumanInputFormTimeoutEvent, - QueueIterationCompletedEvent, - QueueIterationNextEvent, - QueueIterationStartEvent, - QueueLoopCompletedEvent, - QueueLoopNextEvent, - QueueLoopStartEvent, - QueueNodeExceptionEvent, - QueueNodeFailedEvent, - QueueNodeRetryEvent, - QueueNodeStartedEvent, - QueueNodeSucceededEvent, - QueuePingEvent, - QueueReasoningChunkEvent, - QueueStopEvent, - QueueTextChunkEvent, - QueueWorkflowFailedEvent, - QueueWorkflowPartialSuccessEvent, - QueueWorkflowPausedEvent, - QueueWorkflowStartedEvent, - QueueWorkflowSucceededEvent, - WorkflowQueueMessage, -) -from core.app.entities.task_entities import ( - ErrorStreamResponse, - HumanInputRequiredPauseReasonPayload, - HumanInputRequiredResponse, - MessageAudioEndStreamResponse, - MessageAudioStreamResponse, - PingStreamResponse, - ReasoningChunkStreamResponse, - StreamResponse, - TextChunkStreamResponse, - WorkflowAppBlockingResponse, - WorkflowAppPausedBlockingResponse, - WorkflowAppStreamResponse, - WorkflowFinishStreamResponse, - WorkflowPauseStreamResponse, - WorkflowStartStreamResponse, -) -from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline -from core.base.tts import AppGeneratorTTSPublisher -from core.ops.ops_trace_manager import TraceQueueManager -from core.workflow.system_variables import build_system_variables -from extensions.ext_database import db -from graphon.entities import WorkflowStartReason -from graphon.enums import WorkflowExecutionStatus -from graphon.runtime import RuntimeState -from models import Account -from models.enums import CreatorUserRole -from models.model import EndUser -from models.workflow import Workflow, WorkflowAppLog, WorkflowAppLogCreatedFrom - -logger = logging.getLogger(__name__) - - -class WorkflowAppGenerateTaskPipeline(GraphRuntimeStateSupport): - """ - WorkflowAppGenerateTaskPipeline is a class that generate stream output and state management for Application. - """ - - def __init__( - self, - application_generate_entity: WorkflowAppGenerateEntity, - workflow: Workflow, - queue_manager: AppQueueManager, - user: Account | EndUser, - stream: bool, - draft_var_saver_factory: DraftVariableSaverFactory, - ): - self._base_task_pipeline = BasedGenerateTaskPipeline( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - stream=stream, - ) - - if isinstance(user, EndUser): - self._user_id = user.id - user_session_id = user.session_id - self._created_by_role = CreatorUserRole.END_USER - else: - self._user_id = user.id - user_session_id = user.id - self._created_by_role = CreatorUserRole.ACCOUNT - - self._application_generate_entity = application_generate_entity - self._workflow_features_dict = workflow.features_dict - self._workflow_execution_id = "" - self._invoke_from = queue_manager.invoke_from - self._draft_var_saver_factory = draft_var_saver_factory - self._workflow = workflow - self._workflow_system_variables = build_system_variables( - files=application_generate_entity.files, - user_id=user_session_id, - app_id=application_generate_entity.app_config.app_id, - workflow_id=workflow.id, - workflow_execution_id=application_generate_entity.workflow_execution_id, - ) - self._workflow_response_converter = WorkflowResponseConverter( - application_generate_entity=application_generate_entity, - user=user, - system_variables=self._workflow_system_variables, - ) - self._graph_runtime_state: RuntimeState | None = self._base_task_pipeline.queue_manager.graph_runtime_state - - def process( - self, - ) -> WorkflowAppBlockingResponse | WorkflowAppPausedBlockingResponse | Generator[WorkflowAppStreamResponse]: - """ - Process generate task pipeline. - :return: - """ - generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) - if self._base_task_pipeline.stream: - return self._to_stream_response(generator) - else: - return self._to_blocking_response(generator) - - def _to_blocking_response( - self, generator: Generator[StreamResponse] - ) -> WorkflowAppBlockingResponse | WorkflowAppPausedBlockingResponse: - """ - To blocking response. - :return: - """ - human_input_responses: list[HumanInputRequiredResponse] = [] - for stream_response in generator: - match stream_response: - case ErrorStreamResponse(): - raise stream_response.err - case HumanInputRequiredResponse(): - human_input_responses.append(stream_response) - case WorkflowPauseStreamResponse(): - return WorkflowAppPausedBlockingResponse( - task_id=self._application_generate_entity.task_id, - workflow_run_id=stream_response.data.workflow_run_id, - data=WorkflowAppPausedBlockingResponse.Data( - id=stream_response.data.workflow_run_id, - workflow_id=self._workflow.id, - status=stream_response.data.status, - outputs=stream_response.data.outputs or {}, - error=None, - elapsed_time=stream_response.data.elapsed_time, - total_tokens=stream_response.data.total_tokens, - total_steps=stream_response.data.total_steps, - created_at=stream_response.data.created_at, - finished_at=None, - paused_nodes=stream_response.data.paused_nodes, - reasons=stream_response.data.reasons, - ), - ) - case WorkflowFinishStreamResponse(): - return WorkflowAppBlockingResponse( - task_id=self._application_generate_entity.task_id, - workflow_run_id=stream_response.data.id, - data=WorkflowAppBlockingResponse.Data( - id=stream_response.data.id, - workflow_id=stream_response.data.workflow_id, - status=stream_response.data.status, - outputs=stream_response.data.outputs, - error=stream_response.data.error, - elapsed_time=stream_response.data.elapsed_time, - total_tokens=stream_response.data.total_tokens, - total_steps=stream_response.data.total_steps, - created_at=int(stream_response.data.created_at), - finished_at=int(stream_response.data.finished_at) - if stream_response.data.finished_at - else None, - ), - ) - case _: - continue - - if human_input_responses: - return self._build_paused_blocking_response_from_human_input(human_input_responses) - - raise ValueError("queue listening stopped unexpectedly.") - - def _build_paused_blocking_response_from_human_input( - self, human_input_responses: list[HumanInputRequiredResponse] - ) -> WorkflowAppPausedBlockingResponse: - runtime_state = self._resolve_graph_runtime_state() - paused_nodes = list(dict.fromkeys(response.data.node_id for response in human_input_responses)) - created_at = int(runtime_state.start_at) - reasons = [ - HumanInputRequiredPauseReasonPayload.from_response_data(response.data).model_dump(mode="json") - for response in human_input_responses - ] - - return WorkflowAppPausedBlockingResponse( - task_id=self._application_generate_entity.task_id, - workflow_run_id=human_input_responses[-1].workflow_run_id, - data=WorkflowAppPausedBlockingResponse.Data( - id=human_input_responses[-1].workflow_run_id, - workflow_id=self._workflow.id, - status=WorkflowExecutionStatus.PAUSED, - outputs={}, - error=None, - elapsed_time=self._application_generate_entity.workflow_execution_timing.finish(), - total_tokens=runtime_state.total_tokens, - total_steps=runtime_state.node_run_steps, - created_at=created_at, - finished_at=None, - paused_nodes=paused_nodes, - reasons=reasons, - ), - ) - - def _to_stream_response(self, generator: Generator[StreamResponse]) -> Generator[WorkflowAppStreamResponse]: - """ - To stream response. - :return: - """ - workflow_run_id = None - for stream_response in generator: - if isinstance(stream_response, WorkflowStartStreamResponse): - workflow_run_id = stream_response.workflow_run_id - - yield WorkflowAppStreamResponse(workflow_run_id=workflow_run_id, stream_response=stream_response) - - def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: str): - if not publisher: - return None - audio_msg = publisher.check_and_get_audio() - if audio_msg is None: - return None - if audio_msg.status == "responding": - return MessageAudioStreamResponse(audio=audio_msg.audio, audio_type=audio_msg.audio_type, task_id=task_id) - if audio_msg.status in {"finish", "error"}: - return None - raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}") - - def _wrapper_process_stream_response( - self, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: - tts_publisher = None - task_id = self._application_generate_entity.task_id - tenant_id = self._application_generate_entity.app_config.tenant_id - features_dict = self._workflow_features_dict - - if ( - self._base_task_pipeline.stream - and features_dict.get("text_to_speech") - and features_dict["text_to_speech"].get("enabled") - and features_dict["text_to_speech"].get("autoPlay") == "enabled" - ): - tts_publisher = AppGeneratorTTSPublisher( - tenant_id, - features_dict["text_to_speech"].get("voice"), - features_dict["text_to_speech"].get("language"), - get_credit_usage_app_type(self._application_generate_entity.app_config.app_mode), - ) - - try: - for response in self._process_stream_response(tts_publisher=tts_publisher, trace_manager=trace_manager): - while audio_response := self._listen_audio_msg(publisher=tts_publisher, task_id=task_id): - yield audio_response - if tts_publisher and isinstance(response, ErrorStreamResponse): - tts_publisher.cancel() - yield MessageAudioEndStreamResponse(audio="", task_id=task_id) - yield response - return - yield response - - if tts_publisher is None: - return - - tts_publisher.publish(None) - while True: - audio_trunk = tts_publisher.check_and_get_audio(block=True) - assert audio_trunk is not None - if audio_trunk.status == "responding": - yield MessageAudioStreamResponse( - audio=audio_trunk.audio, audio_type=audio_trunk.audio_type, task_id=task_id - ) - continue - if audio_trunk.status not in {"finish", "error"}: - raise RuntimeError(f"TTS publisher returned an unknown status: {audio_trunk.status}") - - yield MessageAudioEndStreamResponse(audio="", task_id=task_id) - if audio_trunk.status == "error": - if audio_trunk.error is None: - raise RuntimeError("TTS publisher returned an error terminal without an exception") - yield ErrorStreamResponse(err=audio_trunk.error, task_id=task_id) - return - finally: - if tts_publisher: - tts_publisher.cancel() - - @contextmanager - def _database_session(self): - """Context manager for database sessions.""" - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - yield session - - def _ensure_workflow_initialized(self): - """Fluent validation for workflow state.""" - if not self._workflow_execution_id: - raise ValueError("workflow run not initialized.") - - def _handle_ping_event(self, event: QueuePingEvent, **kwargs) -> Generator[PingStreamResponse]: - """Handle ping events.""" - yield self._base_task_pipeline.ping_stream_response() - - def _handle_error_event(self, event: QueueErrorEvent, **kwargs) -> Generator[ErrorStreamResponse]: - """Handle error events.""" - err = self._base_task_pipeline.handle_error(event=event) - yield self._base_task_pipeline.error_to_stream_response(err) - - def _handle_workflow_started_event(self, event: QueueWorkflowStartedEvent, **kwargs) -> Generator[StreamResponse]: - """Handle workflow started events.""" - runtime_state = self._resolve_graph_runtime_state() - - run_id = self._extract_workflow_run_id(runtime_state) - self._workflow_execution_id = run_id - - if event.reason == WorkflowStartReason.INITIAL: - with self._database_session() as session: - self._save_workflow_app_log(session=session, workflow_run_id=self._workflow_execution_id) - - start_resp = self._workflow_response_converter.workflow_start_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_run_id=run_id, - workflow_id=self._workflow.id, - reason=event.reason, - node_execution_snapshots=event.node_execution_snapshots, - ) - yield start_resp - - def _handle_node_retry_event(self, event: QueueNodeRetryEvent, **kwargs) -> Generator[StreamResponse]: - """Handle node retry events.""" - self._ensure_workflow_initialized() - - response = self._workflow_response_converter.workflow_node_retry_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - if response: - yield response - - def _handle_node_started_event(self, event: QueueNodeStartedEvent, **kwargs) -> Generator[StreamResponse]: - """Handle node started events.""" - self._ensure_workflow_initialized() - - node_start_response = self._workflow_response_converter.workflow_node_start_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - if node_start_response: - yield node_start_response - - def _handle_node_succeeded_event(self, event: QueueNodeSucceededEvent, **kwargs) -> Generator[StreamResponse]: - """Handle node succeeded events.""" - node_success_response = self._workflow_response_converter.workflow_node_finish_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - self._save_output_for_event(event, event.node_execution_id) - - if node_success_response: - yield node_success_response - - def _handle_node_failed_events( - self, - event: QueueNodeFailedEvent | QueueNodeExceptionEvent, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle various node failure events.""" - node_failed_response = self._workflow_response_converter.workflow_node_finish_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - ) - - if isinstance(event, QueueNodeExceptionEvent): - self._save_output_for_event(event, event.node_execution_id) - - if node_failed_response: - yield node_failed_response - - def _handle_iteration_start_event(self, event: QueueIterationStartEvent, **kwargs) -> Generator[StreamResponse]: - """Handle iteration start events.""" - self._ensure_workflow_initialized() - - iter_start_resp = self._workflow_response_converter.workflow_iteration_start_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_execution_id, - event=event, - ) - yield iter_start_resp - - def _handle_iteration_next_event(self, event: QueueIterationNextEvent, **kwargs) -> Generator[StreamResponse]: - """Handle iteration next events.""" - self._ensure_workflow_initialized() - - iter_next_resp = self._workflow_response_converter.workflow_iteration_next_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_execution_id, - event=event, - ) - yield iter_next_resp - - def _handle_iteration_completed_event( - self, event: QueueIterationCompletedEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle iteration completed events.""" - self._ensure_workflow_initialized() - - iter_finish_resp = self._workflow_response_converter.workflow_iteration_completed_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_execution_id, - event=event, - ) - yield iter_finish_resp - - def _handle_loop_start_event(self, event: QueueLoopStartEvent, **kwargs) -> Generator[StreamResponse]: - """Handle loop start events.""" - self._ensure_workflow_initialized() - - loop_start_resp = self._workflow_response_converter.workflow_loop_start_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_execution_id, - event=event, - ) - yield loop_start_resp - - def _handle_loop_next_event(self, event: QueueLoopNextEvent, **kwargs) -> Generator[StreamResponse]: - """Handle loop next events.""" - self._ensure_workflow_initialized() - - loop_next_resp = self._workflow_response_converter.workflow_loop_next_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_execution_id, - event=event, - ) - yield loop_next_resp - - def _handle_loop_completed_event(self, event: QueueLoopCompletedEvent, **kwargs) -> Generator[StreamResponse]: - """Handle loop completed events.""" - self._ensure_workflow_initialized() - - loop_finish_resp = self._workflow_response_converter.workflow_loop_completed_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_execution_id=self._workflow_execution_id, - event=event, - ) - yield loop_finish_resp - - def _handle_workflow_succeeded_event( - self, - event: QueueWorkflowSucceededEvent, - *, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow succeeded events.""" - _ = trace_manager - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow.id, - status=WorkflowExecutionStatus.SUCCEEDED, - graph_runtime_state=validated_state, - ) - - yield workflow_finish_resp - - def _handle_workflow_partial_success_event( - self, - event: QueueWorkflowPartialSuccessEvent, - *, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow partial success events.""" - _ = trace_manager - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow.id, - status=WorkflowExecutionStatus.PARTIAL_SUCCEEDED, - graph_runtime_state=validated_state, - exceptions_count=event.exceptions_count, - ) - yield workflow_finish_resp - - def _handle_workflow_paused_event( - self, - event: QueueWorkflowPausedEvent, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow paused events.""" - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - responses = self._workflow_response_converter.workflow_pause_to_stream_response( - event=event, - task_id=self._application_generate_entity.task_id, - graph_runtime_state=validated_state, - ) - yield from responses - - def _handle_workflow_failed_and_stop_events( - self, - event: QueueWorkflowFailedEvent | QueueStopEvent, - *, - trace_manager: TraceQueueManager | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle workflow failed and stop events.""" - _ = trace_manager - self._ensure_workflow_initialized() - validated_state = self._ensure_graph_runtime_initialized() - - if isinstance(event, QueueWorkflowFailedEvent): - status = WorkflowExecutionStatus.FAILED - error = event.error - exceptions_count = event.exceptions_count - else: - status = WorkflowExecutionStatus.STOPPED - error = event.get_stop_reason() - exceptions_count = 0 - workflow_finish_resp = self._workflow_response_converter.workflow_finish_to_stream_response( - task_id=self._application_generate_entity.task_id, - workflow_id=self._workflow.id, - status=status, - graph_runtime_state=validated_state, - error=error, - exceptions_count=exceptions_count, - ) - yield workflow_finish_resp - - def _handle_text_chunk_event( - self, - event: QueueTextChunkEvent, - *, - tts_publisher: AppGeneratorTTSPublisher | None = None, - queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None, - **kwargs, - ) -> Generator[StreamResponse]: - """Handle text chunk events.""" - delta_text = event.text - if delta_text is None: - return - - # only publish tts message at text chunk streaming - if tts_publisher and queue_message: - tts_publisher.publish(queue_message) - - yield self._text_chunk_to_stream_response(delta_text, from_variable_selector=event.from_variable_selector) - - def _handle_reasoning_chunk_event(self, event: QueueReasoningChunkEvent, **kwargs) -> Generator[StreamResponse]: - """Handle reasoning chunk events.""" - # is_final with empty reasoning is still forwarded as the "thinking finished" signal - if not event.reasoning and not event.is_final: - return - yield ReasoningChunkStreamResponse( - task_id=self._application_generate_entity.task_id, - data=ReasoningChunkStreamResponse.Data( - reasoning=event.reasoning, - node_id=event.from_node_id, - is_final=event.is_final, - ), - ) - - def _handle_agent_log_event(self, event: QueueAgentLogEvent, **kwargs) -> Generator[StreamResponse]: - """Handle agent log events.""" - yield self._workflow_response_converter.handle_agent_log( - task_id=self._application_generate_entity.task_id, event=event - ) - - def _handle_human_input_form_filled_event( - self, event: QueueHumanInputFormFilledEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle human input form filled events.""" - yield self._workflow_response_converter.human_input_form_filled_to_stream_response( - event=event, task_id=self._application_generate_entity.task_id - ) - - def _handle_human_input_form_timeout_event( - self, event: QueueHumanInputFormTimeoutEvent, **kwargs - ) -> Generator[StreamResponse]: - """Handle human input form timeout events.""" - yield self._workflow_response_converter.human_input_form_timeout_to_stream_response( - event=event, task_id=self._application_generate_entity.task_id - ) - - def _get_event_handlers(self) -> dict[type, Callable]: - """Get mapping of event types to their handlers using fluent pattern.""" - return { - # Basic events - QueuePingEvent: self._handle_ping_event, - QueueErrorEvent: self._handle_error_event, - QueueTextChunkEvent: self._handle_text_chunk_event, - QueueReasoningChunkEvent: self._handle_reasoning_chunk_event, - # Workflow events - QueueWorkflowStartedEvent: self._handle_workflow_started_event, - QueueWorkflowSucceededEvent: self._handle_workflow_succeeded_event, - QueueWorkflowPartialSuccessEvent: self._handle_workflow_partial_success_event, - QueueWorkflowPausedEvent: self._handle_workflow_paused_event, - # Node events - QueueNodeRetryEvent: self._handle_node_retry_event, - QueueNodeStartedEvent: self._handle_node_started_event, - QueueNodeSucceededEvent: self._handle_node_succeeded_event, - # Iteration events - QueueIterationStartEvent: self._handle_iteration_start_event, - QueueIterationNextEvent: self._handle_iteration_next_event, - QueueIterationCompletedEvent: self._handle_iteration_completed_event, - # Loop events - QueueLoopStartEvent: self._handle_loop_start_event, - QueueLoopNextEvent: self._handle_loop_next_event, - QueueLoopCompletedEvent: self._handle_loop_completed_event, - # Agent events - QueueAgentLogEvent: self._handle_agent_log_event, - QueueHumanInputFormFilledEvent: self._handle_human_input_form_filled_event, - QueueHumanInputFormTimeoutEvent: self._handle_human_input_form_timeout_event, - } - - def _dispatch_event( - self, - event: AppQueueEvent, - *, - tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, - queue_message: WorkflowQueueMessage | MessageQueueMessage | None = None, - ) -> Generator[StreamResponse]: - """Dispatch events using elegant pattern matching.""" - handlers = self._get_event_handlers() - event_type = type(event) - - # Direct handler lookup - if handler := handlers.get(event_type): - yield from handler( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - return - - # Handle node failure events with isinstance check - if isinstance( - event, - ( - QueueNodeFailedEvent, - QueueNodeExceptionEvent, - ), - ): - yield from self._handle_node_failed_events( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - return - - # Handle workflow failed and stop events with isinstance check - if isinstance(event, (QueueWorkflowFailedEvent, QueueStopEvent)): - yield from self._handle_workflow_failed_and_stop_events( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - return - - # For unhandled events, we continue (original behavior) - return - - def _process_stream_response( - self, - tts_publisher: AppGeneratorTTSPublisher | None = None, - trace_manager: TraceQueueManager | None = None, - ) -> Generator[StreamResponse]: - """ - Process stream response using elegant Fluent Python patterns. - Maintains exact same functionality as original 44-if-statement version. - """ - for queue_message in self._base_task_pipeline.queue_manager.listen(): - event = queue_message.event - - match event: - case QueueWorkflowStartedEvent(): - self._resolve_graph_runtime_state() - yield from self._handle_workflow_started_event(event) - - case QueueTextChunkEvent(): - yield from self._handle_text_chunk_event( - event, tts_publisher=tts_publisher, queue_message=queue_message - ) - - case QueueErrorEvent(): - yield from self._handle_error_event(event) - break - - case QueueWorkflowFailedEvent(): - yield from self._handle_workflow_failed_and_stop_events(event) - break - case QueueWorkflowPausedEvent(): - yield from self._handle_workflow_paused_event(event) - break - - case QueueStopEvent(): - yield from self._handle_workflow_failed_and_stop_events(event) - break - - # Handle all other events through elegant dispatch - case _: - if responses := list( - self._dispatch_event( - event, - tts_publisher=tts_publisher, - trace_manager=trace_manager, - queue_message=queue_message, - ) - ): - yield from responses - - def _save_workflow_app_log(self, *, session: Session, workflow_run_id: str | None): - invoke_from = self._application_generate_entity.invoke_from - match invoke_from: - case InvokeFrom.SERVICE_API: - created_from = WorkflowAppLogCreatedFrom.SERVICE_API - case InvokeFrom.OPENAPI: - created_from = WorkflowAppLogCreatedFrom.OPENAPI - case InvokeFrom.EXPLORE: - created_from = WorkflowAppLogCreatedFrom.INSTALLED_APP - case InvokeFrom.WEB_APP: - created_from = WorkflowAppLogCreatedFrom.WEB_APP - case InvokeFrom.DEBUGGER | InvokeFrom.TRIGGER | InvokeFrom.PUBLISHED_PIPELINE | InvokeFrom.VALIDATION: - # not save log for debugging - return - - if not workflow_run_id: - return - - workflow_app_log = WorkflowAppLog( - tenant_id=self._application_generate_entity.app_config.tenant_id, - app_id=self._application_generate_entity.app_config.app_id, - workflow_id=self._workflow.id, - workflow_run_id=workflow_run_id, - created_from=created_from, - created_by_role=self._created_by_role, - created_by=self._user_id, - ) - - session.add(workflow_app_log) - - def _text_chunk_to_stream_response( - self, text: str, from_variable_selector: list[str] | None = None - ) -> TextChunkStreamResponse: - """ - Handle completed event. - :param text: text - :return: - """ - response = TextChunkStreamResponse( - task_id=self._application_generate_entity.task_id, - data=TextChunkStreamResponse.Data(text=text, from_variable_selector=from_variable_selector), - ) - - return response - - def _save_output_for_event(self, event: QueueNodeSucceededEvent | QueueNodeExceptionEvent, node_execution_id: str): - saver = self._draft_var_saver_factory( - app_id=self._application_generate_entity.app_config.app_id, - node_id=event.node_id, - node_type=event.node_type, - node_execution_id=node_execution_id, - enclosing_node_id=event.in_loop_id or event.in_iteration_id, - ) - saver.save(event.process_data, event.outputs) diff --git a/api/core/app/apps/workflow/stop_aware_ready_queue.py b/api/core/app/apps/workflow/stop_aware_ready_queue.py deleted file mode 100644 index ef5edb193acc48..00000000000000 --- a/api/core/app/apps/workflow/stop_aware_ready_queue.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Ready-queue wrapper that honors stop before the next node is scheduled.""" - -from __future__ import annotations - -from typing import final - -from core.app.apps.execution_coordinator import is_app_task_stop_flag_set -from graphon.engine.ready_queue import ReadyTask -from graphon.runtime import RuntimeState -from graphon.runtime.execution import GraphExecution -from graphon.runtime.ready_queue import ReadyQueue - - -@final -class StopAwareReadyQueue: - """Reject newly ready nodes once the run has been stopped. - - Graph execution still enqueues successors after abort. Drop those puts so - later nodes do not start; the in-flight node can finish. - """ - - def __init__( - self, - inner: ReadyQueue, - *, - task_id: str, - graph_execution: GraphExecution, - ) -> None: - self._inner = inner - self._task_id = task_id - self._graph_execution = graph_execution - - def _should_reject(self) -> bool: - return self._graph_execution.aborted or is_app_task_stop_flag_set(self._task_id) - - def put(self, item: ReadyTask) -> None: - if self._should_reject(): - return - self._inner.put(item) - - def get(self, timeout: float | None = None) -> ReadyTask: - return self._inner.get(timeout=timeout) - - def task_done(self) -> None: - self._inner.task_done() - - def qsize(self) -> int: - return self._inner.qsize() - - def take_all(self) -> list[ReadyTask]: - return self._inner.take_all() - - def dumps(self) -> str: - return self._inner.dumps() - - def loads(self, data: str) -> None: - self._inner.loads(data) - - -def attach_stop_aware_ready_queue(graph_runtime_state: RuntimeState, *, task_id: str) -> None: - """Install stop-aware enqueue policy on an existing runtime state.""" - current = graph_runtime_state.ready_queue - if isinstance(current, StopAwareReadyQueue): - return - graph_runtime_state._ready_queue = StopAwareReadyQueue( - current, - task_id=task_id, - graph_execution=graph_runtime_state.graph_execution, - ) diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py deleted file mode 100644 index 4712f3fd0f8b2d..00000000000000 --- a/api/core/app/apps/workflow_app_runner.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Compatibility aliases for workflow execution service contracts.""" - -from services.workflow.execution.contracts import PreparedWorkflowRun, WorkflowRunDriver, WorkflowRunRunner - -WorkflowBasedAppRunner = WorkflowRunRunner - -__all__ = [ - "PreparedWorkflowRun", - "WorkflowBasedAppRunner", - "WorkflowRunDriver", -] diff --git a/api/core/app/features/annotation_reply/__init__.py b/api/core/app/features/annotation_reply/__init__.py deleted file mode 100644 index e69de29bb2d1d6..00000000000000 diff --git a/api/core/app/features/annotation_reply/annotation_reply.py b/api/core/app/features/annotation_reply/annotation_reply.py deleted file mode 100644 index 4cf21bed4d6c81..00000000000000 --- a/api/core/app/features/annotation_reply/annotation_reply.py +++ /dev/null @@ -1,97 +0,0 @@ -import logging -from typing import cast - -from sqlalchemy.orm import Session - -from core.app.entities.app_invoke_entities import InvokeFrom -from core.rag.datasource.vdb.vector_factory import Vector -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from models.dataset import Dataset -from models.enums import CollectionBindingType, ConversationFromSource -from models.model import AnnotationReplyEnabledConfig, App, Message, MessageAnnotation, load_annotation_reply_config -from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository -from services.annotation_service import AppAnnotationService - -logger = logging.getLogger(__name__) - - -class AnnotationReplyFeature: - def query( - self, - app_record: App, - message: Message, - query: str, - user_id: str, - invoke_from: InvokeFrom, - *, - session: Session, - ) -> MessageAnnotation | None: - """Return the closest annotation reply and record a hit in ``session``. - - The setting lookup, vector access, annotation lookup, and hit-history - write share the caller-owned session. Vector-search failures are logged - and return ``None``; transaction cleanup remains the caller's responsibility. - """ - try: - annotation_reply_config = load_annotation_reply_config(session, app_record.id) - except ValueError: - return None - - if not annotation_reply_config["enabled"]: - return None - enabled_config = cast(AnnotationReplyEnabledConfig, annotation_reply_config) - - try: - score_threshold = enabled_config["score_threshold"] or 1 - embedding_provider_name = enabled_config["embedding_model"]["embedding_provider_name"] - embedding_model_name = enabled_config["embedding_model"]["embedding_model_name"] - - dataset_collection_binding = DatasetCollectionBindingRepository.get_dataset_collection_binding( - embedding_provider_name, embedding_model_name, session, CollectionBindingType.ANNOTATION - ) - - dataset = Dataset( - id=app_record.id, - tenant_id=app_record.tenant_id, - indexing_technique=IndexTechniqueType.HIGH_QUALITY, - embedding_model_provider=embedding_provider_name, - embedding_model=embedding_model_name, - collection_binding_id=dataset_collection_binding.id, - ) - - vector = Vector(dataset, attributes=["doc_id", "annotation_id", "app_id"], session=session) - - documents = vector.search_by_vector( - query=query, top_k=1, score_threshold=score_threshold, filter={"group_id": [dataset.id]} - ) - - if documents and documents[0].metadata: - annotation_id = documents[0].metadata["annotation_id"] - score = documents[0].metadata["score"] - annotation = AppAnnotationService.get_annotation_by_id(annotation_id, session=session) - if annotation: - if invoke_from in {InvokeFrom.SERVICE_API, InvokeFrom.WEB_APP}: - from_source = ConversationFromSource.API - else: - from_source = ConversationFromSource.CONSOLE - - # insert annotation history - AppAnnotationService.add_annotation_history( - annotation.id, - app_record.id, - annotation.question_text, - annotation.content, - query, - user_id, - message.id, - from_source, - score, - session=session, - ) - - return annotation - except Exception: - logger.warning("Query annotation failed.", exc_info=True) - return None - - return None diff --git a/api/core/app/layers/conversation_variable_persist_layer.py b/api/core/app/layers/conversation_variable_persist_layer.py deleted file mode 100644 index d9048d57cbe963..00000000000000 --- a/api/core/app/layers/conversation_variable_persist_layer.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Persist conversation-scoped variable updates emitted by the graph engine. - -The graph package emits generic variable update events and stays unaware of -conversation identity or storage concerns. This layer lives in the application -core, listens to those generic events, and persists only the `conversation.*` -scope updates that matter to chat applications. -""" - -import logging -from typing import override - -from core.workflow.system_variables import SystemVariableKey, get_system_text -from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID -from graphon.engine.layer import Layer -from graphon.engine_events import EngineEvent, NodeRunVariableUpdatedEvent -from services.conversation_variable_updater import ConversationVariableUpdater - -logger = logging.getLogger(__name__) - - -class ConversationVariablePersistenceLayer(Layer): - def __init__(self, conversation_variable_updater: ConversationVariableUpdater) -> None: - super().__init__() - self._conversation_variable_updater = conversation_variable_updater - - @override - def on_event(self, event: EngineEvent) -> None: - if not isinstance(event, NodeRunVariableUpdatedEvent): - return - - selector = event.variable.selector - if len(selector) < 2: - logger.warning("Conversation variable selector invalid. selector=%s", selector) - return - - conversation_id = get_system_text(self.runtime_state.variable_pool, SystemVariableKey.CONVERSATION_ID) - if conversation_id is None: - return - - if selector[0] != CONVERSATION_VARIABLE_NODE_ID: - return - - self._conversation_variable_updater.update(conversation_id=conversation_id, variable=event.variable) diff --git a/api/core/app/task_pipeline/based_generate_task_pipeline.py b/api/core/app/task_pipeline/based_generate_task_pipeline.py index 1ff89b92db82c7..44677e7ab40819 100644 --- a/api/core/app/task_pipeline/based_generate_task_pipeline.py +++ b/api/core/app/task_pipeline/based_generate_task_pipeline.py @@ -1,9 +1,6 @@ import logging import time -from sqlalchemy import select -from sqlalchemy.orm import Session - from clients.agent_backend.errors import AgentBackendError from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.entities.app_invoke_entities import ( @@ -19,8 +16,6 @@ from core.errors.error import QuotaExceededError from core.moderation.output_moderation import ModerationRule, OutputModeration from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError -from models.enums import MessageStatus -from models.model import Message logger = logging.getLogger(__name__) @@ -47,9 +42,7 @@ def __init__( self.output_moderation_handler = self._init_output_moderation() self.stream = stream - def handle_error( - self, *, event: QueueErrorEvent, session: Session | None = None, message_id: str = "" - ) -> Exception: + def handle_error(self, *, event: QueueErrorEvent) -> Exception: logger.debug("error: %s", event.error) e = event.error err: Exception @@ -63,15 +56,6 @@ def handle_error( description = getattr(e, "description", None) err = Exception(description if description is not None else str(e)) - if not message_id or session is None: - return err - - message = session.scalar(select(Message).where(Message.id == message_id)) - if message is None: - return err - - message.status = MessageStatus.ERROR - message.error = self._error_to_desc(err) return err def _error_to_desc(self, e: Exception) -> str: diff --git a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py b/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py deleted file mode 100644 index aa490fbf5462fa..00000000000000 --- a/api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py +++ /dev/null @@ -1,659 +0,0 @@ -import json -import logging -import time -from collections.abc import Generator, Mapping, Sequence -from threading import Thread -from typing import Any, cast - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.entities.app_invoke_entities import ( - AgentChatAppGenerateEntity, - ChatAppGenerateEntity, - CompletionAppGenerateEntity, - get_credit_usage_app_type, -) -from core.app.entities.queue_entities import ( - QueueAgentMessageEvent, - QueueAgentThoughtEvent, - QueueAnnotationReplyEvent, - QueueErrorEvent, - QueueLLMChunkEvent, - QueueMessageEndEvent, - QueueMessageFileEvent, - QueueMessageReplaceEvent, - QueuePingEvent, - QueueRetrieverResourcesEvent, - QueueStopEvent, -) -from core.app.entities.task_entities import ( - AgentMessageStreamResponse, - AgentThoughtStreamResponse, - ChatbotAppBlockingResponse, - ChatbotAppStreamResponse, - CompletionAppBlockingResponse, - CompletionAppStreamResponse, - EasyUITaskState, - ErrorStreamResponse, - MessageAudioEndStreamResponse, - MessageAudioStreamResponse, - MessageEndStreamResponse, - StreamEvent, - StreamResponse, -) -from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline -from core.app.task_pipeline.message_cycle_manager import MessageCycleManager -from core.app.task_pipeline.message_file_utils import prepare_file_dict -from core.base.tts import AppGeneratorTTSPublisher -from core.db.session_factory import session_factory -from core.model_manager import ModelInstance -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.prompt.utils.prompt_message_util import PromptMessageUtil -from core.prompt.utils.prompt_template_parser import PromptTemplateParser -from events.message_event import message_was_created -from graphon.file import FileTransferMethod -from graphon.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage -from graphon.model_runtime.entities.message_entities import ( - AssistantPromptMessage, - TextPromptMessageContent, -) -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from libs.datetime_utils import naive_utc_now -from models.model import AppMode, Conversation, Message, MessageAgentThought, MessageFile, UploadFile - -logger = logging.getLogger(__name__) - -type EasyUIAppGenerateEntity = ChatAppGenerateEntity | CompletionAppGenerateEntity | AgentChatAppGenerateEntity - - -class EasyUIBasedGenerateTaskPipeline(BasedGenerateTaskPipeline[EasyUIAppGenerateEntity]): - """ - EasyUIBasedGenerateTaskPipeline is a class that generate stream output and state management for Application. - """ - - _task_state: EasyUITaskState - _precomputed_event_type: StreamEvent | None = None - - def __init__( - self, - application_generate_entity: EasyUIAppGenerateEntity, - queue_manager: AppQueueManager, - conversation: Conversation, - message: Message, - stream: bool, - ): - super().__init__( - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - stream=stream, - ) - self._model_config = application_generate_entity.model_conf - self._app_config = application_generate_entity.app_config - - self._conversation_id = conversation.id - self._conversation_mode = conversation.mode - - self._message_id = message.id - self._message_created_at = int(message.created_at.timestamp()) - - self._task_state = EasyUITaskState( - llm_result=LLMResult( - model=self._model_config.model, - prompt_messages=[], - message=AssistantPromptMessage(content=""), - usage=LLMUsage.empty_usage(), - ) - ) - - self._message_cycle_manager = MessageCycleManager( - application_generate_entity=application_generate_entity, - task_state=self._task_state, - ) - - self._conversation_name_generate_thread: Thread | None = None - - def process( - self, - ) -> ( - ChatbotAppBlockingResponse - | CompletionAppBlockingResponse - | Generator[ChatbotAppStreamResponse | CompletionAppStreamResponse] - ): - if self._application_generate_entity.app_config.app_mode != AppMode.COMPLETION: - # start generate conversation name thread - self._conversation_name_generate_thread = self._message_cycle_manager.generate_conversation_name( - conversation_id=self._conversation_id, - query=self._application_generate_entity.query, - message_id=self._message_id, - ) - - generator = self._wrapper_process_stream_response(trace_manager=self._application_generate_entity.trace_manager) - if self.stream: - return self._to_stream_response(generator) - else: - return self._to_blocking_response(generator) - - def _to_blocking_response( - self, generator: Generator[StreamResponse] - ) -> ChatbotAppBlockingResponse | CompletionAppBlockingResponse: - """ - Process blocking response. - :return: - """ - for stream_response in generator: - match stream_response: - case ErrorStreamResponse(): - raise stream_response.err - case MessageEndStreamResponse(): - extras = {"usage": self._task_state.llm_result.usage.model_dump()} - if self._task_state.metadata: - extras["metadata"] = self._task_state.metadata.model_dump() - response: ChatbotAppBlockingResponse | CompletionAppBlockingResponse - if self._conversation_mode == AppMode.COMPLETION: - response = CompletionAppBlockingResponse( - task_id=self._application_generate_entity.task_id, - data=CompletionAppBlockingResponse.Data( - id=self._message_id, - mode=self._conversation_mode, - message_id=self._message_id, - answer=self._task_state.llm_result.message.get_text_content(), - created_at=self._message_created_at, - **extras, - ), - ) - else: - response = ChatbotAppBlockingResponse( - task_id=self._application_generate_entity.task_id, - data=ChatbotAppBlockingResponse.Data( - id=self._message_id, - mode=self._conversation_mode, - conversation_id=self._conversation_id, - message_id=self._message_id, - answer=self._task_state.llm_result.message.get_text_content(), - created_at=self._message_created_at, - **extras, - ), - ) - - return response - case _: - continue - - raise RuntimeError("queue listening stopped unexpectedly.") - - def _to_stream_response( - self, generator: Generator[StreamResponse] - ) -> Generator[ChatbotAppStreamResponse | CompletionAppStreamResponse]: - """ - To stream response. - :return: - """ - for stream_response in generator: - if isinstance(self._application_generate_entity, CompletionAppGenerateEntity): - yield CompletionAppStreamResponse( - message_id=self._message_id, - created_at=self._message_created_at, - stream_response=stream_response, - ) - else: - yield ChatbotAppStreamResponse( - conversation_id=self._conversation_id, - message_id=self._message_id, - created_at=self._message_created_at, - stream_response=stream_response, - ) - - def _listen_audio_msg(self, publisher: AppGeneratorTTSPublisher | None, task_id: str): - if publisher is None: - return None - audio_msg = publisher.check_and_get_audio() - if audio_msg is None: - return None - if audio_msg.status == "responding": - return MessageAudioStreamResponse(audio=audio_msg.audio, audio_type=audio_msg.audio_type, task_id=task_id) - if audio_msg.status in {"finish", "error"}: - return None - raise RuntimeError(f"TTS publisher returned an unknown status: {audio_msg.status}") - - def _wrapper_process_stream_response( - self, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: - tenant_id = self._application_generate_entity.app_config.tenant_id - task_id = self._application_generate_entity.task_id - publisher = None - text_to_speech_dict = cast(dict[str, Any], self._app_config.app_model_config_dict.get("text_to_speech")) - if ( - self.stream - and text_to_speech_dict - and text_to_speech_dict.get("autoPlay") == "enabled" - and text_to_speech_dict.get("enabled") - ): - publisher = AppGeneratorTTSPublisher( - tenant_id, - text_to_speech_dict.get("voice", ""), - text_to_speech_dict.get("language", None), - get_credit_usage_app_type(self._app_config.app_mode), - ) - try: - for response in self._process_stream_response(publisher=publisher, trace_manager=trace_manager): - while audio_response := self._listen_audio_msg(publisher, task_id): - yield audio_response - if publisher and isinstance(response, ErrorStreamResponse): - publisher.cancel() - yield MessageAudioEndStreamResponse(audio="", task_id=task_id) - yield response - return - yield response - - if publisher is None: - return - - publisher.publish(None) - while True: - audio = publisher.check_and_get_audio(block=True) - assert audio is not None - if audio.status == "responding": - yield MessageAudioStreamResponse(audio=audio.audio, audio_type=audio.audio_type, task_id=task_id) - continue - if audio.status not in {"finish", "error"}: - raise RuntimeError(f"TTS publisher returned an unknown status: {audio.status}") - - yield MessageAudioEndStreamResponse(audio="", task_id=task_id) - if audio.status == "error": - if audio.error is None: - raise RuntimeError("TTS publisher returned an error terminal without an exception") - yield ErrorStreamResponse(err=audio.error, task_id=task_id) - return - finally: - if publisher: - publisher.cancel() - - def _process_stream_response( - self, publisher: AppGeneratorTTSPublisher | None, trace_manager: TraceQueueManager | None = None - ) -> Generator[StreamResponse]: - """ - Process stream response. - :return: - """ - for message in self.queue_manager.listen(): - if publisher: - publisher.publish(message) - event = message.event - - match event: - case QueueErrorEvent(): - with session_factory.create_session() as session: - err = self.handle_error(event=event, session=session, message_id=self._message_id) - session.commit() - - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MESSAGE_TRACE, - conversation_id=self._conversation_id, - message_id=self._message_id, - trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), - ) - ) - - yield self.error_to_stream_response(err) - break - case QueueStopEvent() | QueueMessageEndEvent(): - if isinstance(event, QueueMessageEndEvent): - if event.llm_result: - self._task_state.llm_result = event.llm_result - else: - self._handle_stop(event) - - # handle output moderation - output_moderation_answer = self.handle_output_moderation_when_task_finished( - self._task_state.llm_result.message.get_text_content() - ) - if output_moderation_answer: - self._task_state.llm_result.message.content = output_moderation_answer - yield self._message_cycle_manager.message_replace_to_stream_response( - answer=output_moderation_answer - ) - - stopped_agent_message: Message | None = None - with session_factory.create_session() as session: - # A stopped Agent run may persist provider-reported usage after - # cancellation completes. Do not replace it with local token estimates. - if isinstance(event, QueueStopEvent): - saved_message = self._save_message( - session=session, - trace_manager=trace_manager, - preserve_existing_usage=True, - ) - if self._app_config.app_mode == AppMode.AGENT: - stopped_agent_message = saved_message - else: - self._save_message(session=session, trace_manager=trace_manager) - session.commit() - if stopped_agent_message is not None: - # Provider/billing subscribers may perform external I/O. - # Release the stop-save row lock before notifying them. - message_was_created.send( - stopped_agent_message, - application_generate_entity=self._application_generate_entity, - ) - message_end_resp = self._message_end_to_stream_response() - yield message_end_resp - case QueueRetrieverResourcesEvent(): - self._message_cycle_manager.handle_retriever_resources(event) - case QueueAnnotationReplyEvent(): - annotation_content = None - with session_factory.create_session() as session: - annotation = self._message_cycle_manager.handle_annotation_reply(event, session) - if annotation: - annotation_content = annotation.content - if annotation_content: - self._task_state.llm_result.message.content = annotation_content - case QueueAgentThoughtEvent(): - agent_thought_response = self._agent_thought_to_stream_response(event) - if agent_thought_response is not None: - yield agent_thought_response - case QueueMessageFileEvent(): - response = self._message_cycle_manager.message_file_to_stream_response(event) - if response: - yield response - case QueueAgentMessageEvent(): - chunk = event.chunk - delta_text = self._chunk_delta_text(chunk) - if delta_text is None: - continue - yield self._agent_message_to_stream_response( - answer=delta_text, - message_id=self._message_id, - ) - case QueueLLMChunkEvent(): - chunk = event.chunk - delta_text = self._chunk_delta_text(chunk) - if delta_text is None: - continue - - if not self._task_state.llm_result.prompt_messages: - self._task_state.llm_result.prompt_messages = chunk.prompt_messages - - # handle output moderation chunk - should_direct_answer = self._handle_output_moderation_chunk(delta_text) - if should_direct_answer: - continue - - current_content = cast(str, self._task_state.llm_result.message.content) - current_content += delta_text - self._task_state.llm_result.message.content = current_content - - # Determine the event type once, on first LLM chunk, and reuse for subsequent chunks - if not hasattr(self, "_precomputed_event_type") or self._precomputed_event_type is None: - self._precomputed_event_type = self._message_cycle_manager.get_message_event_type( - message_id=self._message_id - ) - yield self._message_cycle_manager.message_to_stream_response( - answer=delta_text, - message_id=self._message_id, - event_type=self._precomputed_event_type, - ) - case QueueMessageReplaceEvent(): - yield self._message_cycle_manager.message_replace_to_stream_response(answer=event.text) - case QueuePingEvent(): - yield self.ping_stream_response() - case _: - continue - if self._conversation_name_generate_thread: - logger.debug("Conversation name generation running as daemon thread") - - @staticmethod - def _chunk_delta_text(chunk: LLMResultChunk) -> str | None: - delta_content = chunk.delta.message.content - if delta_content is None: - return None - if not isinstance(delta_content, list): - return delta_content - - delta_text = "" - # EasyUI streams text only; structured multimodal chunks contribute their text parts. - for content in cast(list[object], delta_content): - logger.debug("The content type %s in LLM chunk delta message content.: %r", type(content), content) - match content: - case TextPromptMessageContent(): - delta_text += content.data - case str(): - delta_text += content - case _: - logger.warning( - "Unsupported content type %s in LLM chunk delta message content.: %r", - type(content), - content, - ) - continue - return delta_text - - def _save_message( - self, - *, - session: Session, - trace_manager: TraceQueueManager | None = None, - preserve_existing_usage: bool = False, - ) -> Message: - """ - Save message. - :return: - """ - llm_result = self._task_state.llm_result - usage = llm_result.usage - - agent_stop_requested = preserve_existing_usage and self._app_config.app_mode == AppMode.AGENT - message_stmt = select(Message).where(Message.id == self._message_id) - if agent_stop_requested: - message_stmt = message_stmt.with_for_update().execution_options(populate_existing=True) - message = session.scalar(message_stmt) - if not message: - raise ValueError(f"message {self._message_id} not found") - conversation_stmt = select(Conversation).where(Conversation.id == self._conversation_id) - conversation = session.scalar(conversation_stmt) - if not conversation: - raise ValueError(f"Conversation {self._conversation_id} not found") - - saved_prompt = PromptMessageUtil.prompt_messages_to_prompt_for_saving( - self._model_config.mode, self._task_state.llm_result.prompt_messages - ) - object.__setattr__(message, "message", saved_prompt) - try: - existing_metadata = json.loads(message.message_metadata) if message.message_metadata else {} - except (json.JSONDecodeError, TypeError): - existing_metadata = {} - if not isinstance(existing_metadata, dict): - existing_metadata = {} - has_persisted_usage = preserve_existing_usage and ( - int(message.message_tokens or 0) + int(message.answer_tokens or 0) > 0 or bool(message.total_price) - ) - if not has_persisted_usage: - message.message_tokens = usage.prompt_tokens - message.message_unit_price = usage.prompt_unit_price - message.message_price_unit = usage.prompt_price_unit - message.answer = ( - PromptTemplateParser.remove_template_variables(llm_result.message.get_text_content().strip()) - if llm_result.message.content - else "" - ) - message.updated_at = naive_utc_now() - if not has_persisted_usage: - message.answer_tokens = usage.completion_tokens - message.answer_unit_price = usage.completion_unit_price - message.answer_price_unit = usage.completion_price_unit - message.provider_response_latency = time.perf_counter() - self.start_at - message.total_price = usage.total_price - message.currency = usage.currency - self._task_state.llm_result.usage.latency = message.provider_response_latency - self._task_state.metadata.usage = self._task_state.llm_result.usage - - metadata = self._task_state.metadata.model_dump(mode="json") - if has_persisted_usage and "usage" in existing_metadata: - metadata["usage"] = existing_metadata["usage"] - if agent_stop_requested: - # Closing the consumer records a request, not a runtime outcome. - # The Agent worker owns the terminal after cancellation/cleanup. - metadata["agent_stop_requested"] = True - message.message_metadata = json.dumps(metadata, ensure_ascii=False) - - if trace_manager and not agent_stop_requested: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.MESSAGE_TRACE, - conversation_id=self._conversation_id, - message_id=self._message_id, - trace_session_id=self._application_generate_entity.extras.get("trace_session_id"), - ) - ) - - if not agent_stop_requested: - message_was_created.send( - message, - application_generate_entity=self._application_generate_entity, - ) - return message - - def _handle_stop(self, event: QueueStopEvent): - """ - Handle stop. - :return: - """ - model_config = self._model_config - model = model_config.model - - model_instance = ModelInstance( - provider_model_bundle=model_config.provider_model_bundle, model=model_config.model - ) - - # calculate num tokens - prompt_tokens = 0 - if event.stopped_by != QueueStopEvent.StopBy.ANNOTATION_REPLY: - prompt_tokens = model_instance.get_llm_num_tokens(self._task_state.llm_result.prompt_messages) - - completion_tokens = 0 - if event.stopped_by == QueueStopEvent.StopBy.USER_MANUAL: - completion_tokens = model_instance.get_llm_num_tokens([self._task_state.llm_result.message]) - - credentials = model_config.credentials - - # transform usage - model_type_instance = model_config.provider_model_bundle.model_type_instance - model_type_instance = cast(LargeLanguageModel, model_type_instance) - self._task_state.llm_result.usage = model_type_instance.calc_response_usage( - model, credentials, prompt_tokens, completion_tokens - ) - - def _message_end_to_stream_response(self) -> MessageEndStreamResponse: - """ - Message end to stream response. - :return: - """ - self._task_state.metadata.usage = self._task_state.llm_result.usage - metadata_dict = self._task_state.metadata.model_dump(exclude_none=True) - - # Fetch files associated with this message - files: Sequence[Mapping[str, Any]] = [] - with session_factory.create_session() as session: - message_files = session.scalars(select(MessageFile).where(MessageFile.message_id == self._message_id)).all() - - if message_files: - # Fetch all required UploadFile objects in a single query to avoid N+1 problem - upload_file_ids = list( - dict.fromkeys( - mf.upload_file_id - for mf in message_files - if mf.transfer_method == FileTransferMethod.LOCAL_FILE and mf.upload_file_id - ) - ) - upload_files_map = {} - if upload_file_ids: - upload_files = session.scalars(select(UploadFile).where(UploadFile.id.in_(upload_file_ids))).all() - upload_files_map = {uf.id: uf for uf in upload_files} - - files_list = [] - for message_file in message_files: - file_dict = prepare_file_dict(message_file, upload_files_map) - files_list.append(file_dict) - - files = cast(Sequence[Mapping[str, Any]], files_list) - - return MessageEndStreamResponse( - task_id=self._application_generate_entity.task_id, - id=self._message_id, - metadata=metadata_dict, - files=files, - ) - - def _agent_message_to_stream_response(self, answer: str, message_id: str) -> AgentMessageStreamResponse: - """ - Agent message to stream response. - :param answer: answer - :param message_id: message id - :return: - """ - return AgentMessageStreamResponse( - task_id=self._application_generate_entity.task_id, id=message_id, answer=answer - ) - - def _agent_thought_to_stream_response(self, event: QueueAgentThoughtEvent) -> AgentThoughtStreamResponse | None: - """ - Agent thought to stream response. - :param event: agent thought event - :return: - """ - with session_factory.create_session() as session: - agent_thought: MessageAgentThought | None = session.scalar( - select(MessageAgentThought).where(MessageAgentThought.id == event.agent_thought_id).limit(1) - ) - - if agent_thought: - return AgentThoughtStreamResponse( - task_id=self._application_generate_entity.task_id, - id=agent_thought.id, - position=agent_thought.position, - thought=agent_thought.thought or "", - observation=agent_thought.observation or "", - tool=agent_thought.tool or "", - tool_labels=agent_thought.tool_labels, - tool_input=agent_thought.tool_input or "", - message_files=agent_thought.files, - ) - - return None - - def _handle_output_moderation_chunk(self, text: str) -> bool: - """ - Handle output moderation chunk. - :param text: text - :return: True if output moderation should direct output, otherwise False - """ - if self.output_moderation_handler: - if self.output_moderation_handler.should_direct_output(): - # stop subscribe new token when output moderation should direct output - self._task_state.llm_result.message.content = self.output_moderation_handler.get_final_output() - self.queue_manager.publish( - QueueLLMChunkEvent( - chunk=LLMResultChunk( - model=self._task_state.llm_result.model, - prompt_messages=self._task_state.llm_result.prompt_messages, - delta=LLMResultChunkDelta( - index=0, - message=AssistantPromptMessage(content=self._task_state.llm_result.message.content), - ), - ) - ), - PublishFrom.TASK_PIPELINE, - ) - - self.queue_manager.publish( - QueueStopEvent(stopped_by=QueueStopEvent.StopBy.OUTPUT_MODERATION), PublishFrom.TASK_PIPELINE - ) - return True - else: - self.output_moderation_handler.append_new_token(text) - - return False diff --git a/api/core/app/task_pipeline/message_cycle_manager.py b/api/core/app/task_pipeline/message_cycle_manager.py deleted file mode 100644 index b41563a7de4b4e..00000000000000 --- a/api/core/app/task_pipeline/message_cycle_manager.py +++ /dev/null @@ -1,288 +0,0 @@ -import hashlib -import logging -from threading import Thread, Timer - -from flask import Flask, current_app -from sqlalchemy import select -from sqlalchemy.orm import Session - -from configs import dify_config -from core.app.entities.app_invoke_entities import ( - AdvancedChatAppGenerateEntity, - AgentChatAppGenerateEntity, - ChatAppGenerateEntity, - CompletionAppGenerateEntity, -) -from core.app.entities.queue_entities import ( - QueueAnnotationReplyEvent, - QueueMessageFileEvent, - QueueRetrieverResourcesEvent, -) -from core.app.entities.task_entities import ( - AnnotationReply, - AnnotationReplyAccount, - EasyUITaskState, - MessageFileStreamResponse, - MessageReplaceStreamResponse, - MessageStreamResponse, - StreamEvent, - WorkflowTaskState, -) -from core.db.session_factory import session_factory -from core.llm_generator.llm_generator import LLMGenerator -from core.tools.signature import sign_tool_file -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from models.enums import MessageFileBelongsTo -from models.model import App, AppMode, Conversation, MessageAnnotation, MessageFile -from services.annotation_service import AppAnnotationService - -logger = logging.getLogger(__name__) - - -class MessageCycleManager: - def __init__( - self, - *, - application_generate_entity: ChatAppGenerateEntity - | CompletionAppGenerateEntity - | AgentChatAppGenerateEntity - | AdvancedChatAppGenerateEntity, - task_state: EasyUITaskState | WorkflowTaskState, - ): - self._application_generate_entity = application_generate_entity - self._task_state = task_state - self._message_has_file: set[str] = set() - - def get_message_event_type(self, message_id: str) -> StreamEvent: - # Fast path: cached determination from prior QueueMessageFileEvent - if message_id in self._message_has_file: - return StreamEvent.MESSAGE_FILE - - # Use SQLAlchemy 2.x style session.scalar(select(...)) - with session_factory.create_session() as session: - message_file = session.scalar( - select(MessageFile).where(MessageFile.message_id == message_id, MessageFile.belongs_to == "assistant") - ) - - if message_file: - self._message_has_file.add(message_id) - return StreamEvent.MESSAGE_FILE - - return StreamEvent.MESSAGE - - def generate_conversation_name( - self, *, conversation_id: str, query: str, message_id: str | None = None - ) -> Thread | None: - """ - Generate conversation name. - :param conversation_id: conversation id - :param query: query - :return: thread - """ - if isinstance(self._application_generate_entity, CompletionAppGenerateEntity): - return None - - is_first_message = self._application_generate_entity.is_new_conversation - extras = self._application_generate_entity.extras - auto_generate_conversation_name = extras.get("auto_generate_conversation_name", True) - - thread: Thread | None = None - if auto_generate_conversation_name and is_first_message: - # start generate thread - # time.sleep not block other logic - thread = Timer( - 1, - self._generate_conversation_name_worker, - kwargs={ - "flask_app": current_app._get_current_object(), # type: ignore - "conversation_id": conversation_id, - "query": query, - "message_id": message_id, - }, - ) - thread.daemon = True - thread.start() - - if is_first_message: - self._application_generate_entity.is_new_conversation = False - - return thread - - def _generate_conversation_name_worker( - self, - flask_app: Flask, - conversation_id: str, - query: str, - message_id: str | None = None, - ): - with flask_app.app_context(): - with session_factory.create_session() as session: - # get conversation and message - stmt = select(Conversation).where(Conversation.id == conversation_id) - conversation = session.scalar(stmt) - - if not conversation: - return - - if conversation.mode != AppMode.COMPLETION: - app_model = session.get(App, conversation.app_id) - if not app_model: - return - - # generate conversation name - query_hash = hashlib.md5(query.encode()).hexdigest()[:16] - cache_key = f"conv_name:{conversation_id}:{query_hash}" - - cached_name = redis_client.get(cache_key) - if cached_name: - name = cached_name.decode("utf-8") - else: - try: - name = LLMGenerator.generate_conversation_name( - app_model.tenant_id, - query, - conversation_id, - conversation.app_id, - message_id=message_id, - ) - redis_client.setex(cache_key, 3600, name) - except Exception: - if dify_config.DEBUG: - logger.exception( - "generate conversation name failed, conversation_id: %s", conversation_id - ) - name = query[:47] + "..." if len(query) > 50 else query - conversation.name = name - session.commit() - - def handle_annotation_reply(self, event: QueueAnnotationReplyEvent, session: Session) -> MessageAnnotation | None: - """ - Handle annotation reply. - :param event: event - :return: - """ - reply = AppAnnotationService.get_annotation_reply_by_id(event.message_annotation_id, session) - if reply: - annotation, account_name = reply - self._task_state.metadata.annotation_reply = AnnotationReply( - id=annotation.id, - account=AnnotationReplyAccount( - id=annotation.account_id, - name=account_name if account_name is not None else "Dify user", - ), - ) - - return annotation - - return None - - def handle_retriever_resources(self, event: QueueRetrieverResourcesEvent): - """ - Handle retriever resources. - :param event: event - :return: - """ - if not self._application_generate_entity.app_config.additional_features: - raise ValueError("Additional features not found") - if self._application_generate_entity.app_config.additional_features.show_retrieve_source: - merged_resources = [r for r in self._task_state.metadata.retriever_resources or [] if r] - existing_ids = {(r.dataset_id, r.document_id) for r in merged_resources if r.dataset_id and r.document_id} - receipt_ids = {r.knowledge_fs_citation.id for r in merged_resources if r.knowledge_fs_citation} - - # Add new unique resources from the event - for resource in event.retriever_resources or []: - if not resource: - continue - - is_duplicate = ( - resource.knowledge_fs_citation.id in receipt_ids - if resource.knowledge_fs_citation - else ( - resource.knowledge_fs_citation is None - and resource.dataset_id - and resource.document_id - and (resource.dataset_id, resource.document_id) in existing_ids - ) - ) - - if not is_duplicate: - merged_resources.append(resource) - if resource.knowledge_fs_citation: - receipt_ids.add(resource.knowledge_fs_citation.id) - - for i, resource in enumerate(merged_resources, 1): - resource.position = i - - self._task_state.metadata.retriever_resources = merged_resources - - def message_file_to_stream_response(self, event: QueueMessageFileEvent) -> MessageFileStreamResponse | None: - """ - Message file to stream response. - :param event: event - :return: - """ - with Session(db.engine, expire_on_commit=False) as session: - message_file = session.scalar(select(MessageFile).where(MessageFile.id == event.message_file_id)) - - if message_file and message_file.url is not None: - self._message_has_file.add(message_file.message_id) - - # get tool file id - tool_file_id = message_file.url.split("/")[-1] - # trim extension - tool_file_id = tool_file_id.split(".")[0] - - # get extension - if "." in message_file.url: - extension = f".{message_file.url.split('.')[-1]}" - if len(extension) > 10: - extension = ".bin" - else: - extension = ".bin" - # add sign url to local file - if message_file.url.startswith("http"): - url = message_file.url - else: - url = sign_tool_file(tool_file_id=tool_file_id, extension=extension) - - return MessageFileStreamResponse( - task_id=self._application_generate_entity.task_id, - id=message_file.id, - type=message_file.type, - belongs_to=message_file.belongs_to or MessageFileBelongsTo.USER, - url=url, - ) - - return None - - def message_to_stream_response( - self, - answer: str, - message_id: str, - from_variable_selector: list[str] | None = None, - event_type: StreamEvent | None = None, - ) -> MessageStreamResponse: - """ - Message to stream response. - :param answer: answer - :param message_id: message id - :return: - """ - return MessageStreamResponse( - task_id=self._application_generate_entity.task_id, - id=message_id, - answer=answer, - from_variable_selector=from_variable_selector or [], - event=event_type or StreamEvent.MESSAGE, - ) - - def message_replace_to_stream_response(self, answer: str, reason: str = "") -> MessageReplaceStreamResponse: - """ - Message replace to stream response. - :param answer: answer - :return: - """ - return MessageReplaceStreamResponse( - task_id=self._application_generate_entity.task_id, answer=answer, reason=reason - ) diff --git a/api/core/plugin/backwards_invocation/app.py b/api/core/plugin/backwards_invocation/app.py deleted file mode 100644 index f955bc87afcd1b..00000000000000 --- a/api/core/plugin/backwards_invocation/app.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Compatibility exports for the service-owned plugin App invocation boundary. - -New callers must import from :mod:`services.plugin.app_invocation`. This module -only keeps the historical import path working until the remaining legacy core -plugin callers are retired. -""" - -from services.plugin.app_invocation import AppScopedEndUserProvisioner, PluginAppBackwardsInvocation - -__all__ = ["AppScopedEndUserProvisioner", "PluginAppBackwardsInvocation"] diff --git a/api/core/plugin/backwards_invocation/tool.py b/api/core/plugin/backwards_invocation/tool.py deleted file mode 100644 index b6f6e2c625795e..00000000000000 --- a/api/core/plugin/backwards_invocation/tool.py +++ /dev/null @@ -1,61 +0,0 @@ -from collections.abc import Generator -from typing import TYPE_CHECKING, Any - -from sqlalchemy.orm import Session - -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.plugin.backwards_invocation.base import BaseBackwardsInvocation -from core.tools.entities.tool_entities import ToolInvokeMessage, ToolProviderType -from core.tools.tool_engine import ToolEngine -from core.tools.tool_manager import ToolManager -from core.tools.utils.message_transformer import ToolFileMessageTransformer - -if TYPE_CHECKING: - from core.app.apps.workflow_app_runner import WorkflowRunDriver - - -class PluginToolBackwardsInvocation(BaseBackwardsInvocation): - """ - Backwards invocation for plugin tools. - """ - - @classmethod - def invoke_tool( - cls, - session: Session, - tenant_id: str, - user_id: str, - tool_type: ToolProviderType, - provider: str, - tool_name: str, - tool_parameters: dict[str, Any], - credential_id: str | None = None, - *, - execution_driver: "WorkflowRunDriver", - ) -> Generator[ToolInvokeMessage]: - """ - invoke tool - """ - # get tool runtime - try: - tool_runtime = ToolManager.get_tool_runtime_from_plugin( - tool_type, - tenant_id, - provider, - tool_name, - tool_parameters, - user_id=user_id, - credential_id=credential_id, - execution_driver=execution_driver, - ) - response = ToolEngine.generic_invoke( - session, tool_runtime, tool_parameters, user_id, DifyWorkflowCallbackHandler(), workflow_call_depth=1 - ) - - response = ToolFileMessageTransformer.transform_tool_invoke_messages( - response, user_id=user_id, tenant_id=tenant_id - ) - - return response - except Exception as e: - raise e diff --git a/api/core/tools/mcp_tool/tool.py b/api/core/tools/mcp_tool/tool.py index 6fbeff73750967..e645a93967b11d 100644 --- a/api/core/tools/mcp_tool/tool.py +++ b/api/core/tools/mcp_tool/tool.py @@ -326,7 +326,7 @@ def invoke_remote_mcp_tool( from sqlalchemy.orm import Session from extensions.ext_database import db - from services.tools.legacy_mcp_tools_manage_service import MCPToolManageService + from services.tools.mcp_tools_manage_service import MCPToolManageService # Step 1: Load provider entity and credentials in a short-lived session # This minimizes database connection hold time diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py deleted file mode 100644 index 47210354367e15..00000000000000 --- a/api/core/tools/tool_engine.py +++ /dev/null @@ -1,396 +0,0 @@ -import contextlib -import json -import logging -from collections.abc import Generator, Iterable -from copy import deepcopy -from datetime import UTC, datetime -from mimetypes import guess_type -from typing import Any, cast - -from sqlalchemy.orm import Session, sessionmaker -from yarl import URL - -from core.app.entities.app_invoke_entities import InvokeFrom -from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.ops.ops_trace_manager import TraceQueueManager -from core.tools.__base.tool import Tool -from core.tools.entities.tool_entities import ( - ToolInvokeMessage, - ToolInvokeMessageBinary, - ToolInvokeMeta, - ToolParameter, -) -from core.tools.errors import ( - ToolEngineInvokeError, - ToolInvokeError, - ToolNotFoundError, - ToolNotSupportedError, - ToolParameterValidationError, - ToolProviderCredentialValidationError, - ToolProviderNotFoundError, -) -from core.tools.utils.message_transformer import ToolFileMessageTransformer, safe_json_value -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.file import FileTransferMethod, FileType -from models.enums import CreatorUserRole, MessageFileBelongsTo -from models.model import Message, MessageFile - -logger = logging.getLogger(__name__) - - -class ToolEngine: - """ - Tool runtime engine take care of the tool executions. - """ - - @staticmethod - def agent_invoke( - session: Session, - tool: Tool, - tool_parameters: str | dict[str, Any], - user_id: str, - tenant_id: str, - message: Message, - invoke_from: InvokeFrom, - agent_tool_callback: DifyAgentCallbackHandler, - trace_manager: TraceQueueManager | None = None, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> tuple[str, list[str], ToolInvokeMeta]: - """ - Agent invokes the tool with the given arguments. - """ - # check if arguments is a string - if isinstance(tool_parameters, str): - # check if this tool has only one parameter - parameters = [ - parameter - for parameter in tool.get_runtime_parameters() - if parameter.form == ToolParameter.ToolParameterForm.LLM - ] - if parameters and len(parameters) == 1: - tool_parameters = {parameters[0].name: tool_parameters} - else: - with contextlib.suppress(Exception): - tool_parameters = json.loads(tool_parameters) - if not isinstance(tool_parameters, dict): - raise ValueError(f"tool_parameters should be a dict, but got a string: {tool_parameters}") - - try: - # hit the callback handler - agent_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) - - messages = ToolEngine._invoke(session, tool, tool_parameters, user_id, conversation_id, app_id, message_id) - invocation_meta_dict: dict[str, ToolInvokeMeta] = {} - - def message_callback( - invocation_meta_dict: dict[str, ToolInvokeMeta], - messages: Generator[ToolInvokeMessage | ToolInvokeMeta], - ): - for message in messages: - if isinstance(message, ToolInvokeMeta): - invocation_meta_dict["meta"] = message - else: - yield message - - messages = ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=message_callback(invocation_meta_dict, messages), - user_id=user_id, - tenant_id=tenant_id, - conversation_id=message.conversation_id, - ) - - message_list = list(messages) - - # extract binary data from tool invoke message - binary_files = ToolEngine._extract_tool_response_binary_and_text(message_list) - # create message file - message_files = ToolEngine._create_message_files( - tool_messages=binary_files, agent_message=message, invoke_from=invoke_from, user_id=user_id - ) - - plain_text = ToolEngine.tool_response_to_str(message_list) - - meta = invocation_meta_dict["meta"] - - # hit the callback handler - agent_tool_callback.on_tool_end( - tool_name=tool.entity.identity.name, - tool_inputs=tool_parameters, - tool_outputs=plain_text, - message_id=message.id, - trace_manager=trace_manager, - ) - - # transform tool invoke message to get LLM friendly message - return plain_text, message_files, meta - except ToolProviderCredentialValidationError as e: - logger.error(e, exc_info=True) - error_response = "Please check your tool provider credentials" - agent_tool_callback.on_tool_error(e) - except (ToolNotFoundError, ToolNotSupportedError, ToolProviderNotFoundError) as e: - error_response = f"there is not a tool named {tool.entity.identity.name}" - logger.error(e, exc_info=True) - agent_tool_callback.on_tool_error(e) - except ToolParameterValidationError as e: - error_response = f"tool parameters validation error: {e}, please check your tool parameters" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - except ToolInvokeError as e: - error_response = f"tool invoke error: {e}" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - except ToolEngineInvokeError as e: - meta = e.meta - error_response = f"tool invoke error: {meta.error}" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - return error_response, [], meta - except Exception as e: - error_response = f"unknown error: {e}" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - - return error_response, [], ToolInvokeMeta.error_instance(error_response) - - @staticmethod - def generic_invoke( - session: Session, - tool: Tool, - tool_parameters: dict[str, Any], - user_id: str, - workflow_tool_callback: DifyWorkflowCallbackHandler, - workflow_call_depth: int, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> Generator[ToolInvokeMessage]: - """ - Workflow invokes the tool with the given arguments. - """ - try: - # hit the callback handler - workflow_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) - - if isinstance(tool, WorkflowTool): - tool.workflow_call_depth = workflow_call_depth + 1 - - if tool.runtime and tool.runtime.runtime_parameters: - tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters} - - response = tool.invoke( - session=session, - user_id=user_id, - tool_parameters=tool_parameters, - conversation_id=conversation_id, - app_id=app_id, - message_id=message_id, - ) - - # hit the callback handler - response = workflow_tool_callback.on_tool_execution( - tool_name=tool.entity.identity.name, - tool_inputs=tool_parameters, - tool_outputs=response, - ) - - return response - except Exception as e: - workflow_tool_callback.on_tool_error(e) - raise e - - @staticmethod - def _invoke( - session: Session, - tool: Tool, - tool_parameters: dict[str, Any], - user_id: str, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> Generator[ToolInvokeMessage | ToolInvokeMeta]: - """ - Invoke the tool with the given arguments. - """ - started_at = datetime.now(UTC) - meta = ToolInvokeMeta( - time_cost=0.0, - error=None, - tool_config={ - "tool_name": tool.entity.identity.name, - "tool_provider": tool.entity.identity.provider, - "tool_provider_type": tool.tool_provider_type().value, - "tool_parameters": deepcopy(tool.runtime.runtime_parameters), - "tool_icon": tool.entity.identity.icon, - }, - ) - try: - yield from tool.invoke(session, user_id, tool_parameters, conversation_id, app_id, message_id) - except Exception as e: - meta.error = str(e) - raise ToolEngineInvokeError(meta) from e - finally: - ended_at = datetime.now(UTC) - meta.time_cost = (ended_at - started_at).total_seconds() - yield meta - - @staticmethod - def tool_response_to_str(tool_response: list[ToolInvokeMessage]) -> str: - """Convert tool invoke messages into the plain-text observation shown to the model/user.""" - parts: list[str] = [] - json_parts: list[str] = [] - - for response in tool_response: - if response.type == ToolInvokeMessage.MessageType.TEXT: - parts.append(cast(ToolInvokeMessage.TextMessage, response.message).text) - elif response.type == ToolInvokeMessage.MessageType.LINK: - parts.append( - f"result link: {cast(ToolInvokeMessage.TextMessage, response.message).text}." - + " please tell user to check it." - ) - elif response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: - parts.append( - "image has been created and sent to user already, " - + "you do not need to create it, just tell the user to check it now." - ) - elif response.type == ToolInvokeMessage.MessageType.JSON: - json_message = cast(ToolInvokeMessage.JsonMessage, response.message) - if json_message.suppress_output: - continue - json_parts.append( - json.dumps( - safe_json_value(cast(ToolInvokeMessage.JsonMessage, response.message).json_object), - ensure_ascii=False, - ) - ) - elif response.type == ToolInvokeMessage.MessageType.VARIABLE: - continue - else: - parts.append(str(response.message)) - - # Add JSON parts, avoiding duplicates from text parts. - if json_parts: - - def normalize_json(text: str) -> str: - try: - return json.dumps(json.loads(text), ensure_ascii=False, sort_keys=True, separators=(",", ":")) - except (ValueError, RecursionError): - return text - - existing_parts = {normalize_json(p) for p in parts} - parts.extend(p for p in json_parts if normalize_json(p) not in existing_parts) - - return "".join(parts) - - @staticmethod - def _extract_tool_response_binary_and_text( - tool_response: list[ToolInvokeMessage], - ) -> Generator[ToolInvokeMessageBinary]: - """ - Extract tool response binary - """ - for response in tool_response: - if response.type in { - ToolInvokeMessage.MessageType.IMAGE_LINK, - ToolInvokeMessage.MessageType.IMAGE, - ToolInvokeMessage.MessageType.BINARY_LINK, - }: - mimetype = None - if not response.meta: - raise ValueError("missing meta data") - if response.meta.get("mime_type"): - mimetype = response.meta.get("mime_type") - else: - with contextlib.suppress(Exception): - url = URL(cast(ToolInvokeMessage.TextMessage, response.message).text) - extension = url.suffix - guess_type_result, _ = guess_type(f"a{extension}") - if guess_type_result: - mimetype = guess_type_result - - if not mimetype: - mimetype = ( - "image/jpeg" - if response.type != ToolInvokeMessage.MessageType.BINARY_LINK - else "application/octet-stream" - ) - - yield ToolInvokeMessageBinary( - mimetype=response.meta.get("mime_type", mimetype), - url=cast(ToolInvokeMessage.TextMessage, response.message).text, - ) - elif response.type == ToolInvokeMessage.MessageType.BLOB: - if not response.meta: - raise ValueError("missing meta data") - - yield ToolInvokeMessageBinary( - mimetype=response.meta.get("mime_type", "application/octet-stream"), - url=cast(ToolInvokeMessage.TextMessage, response.message).text, - ) - elif response.type == ToolInvokeMessage.MessageType.LINK: - # check if there is a mime type in meta - if response.meta and "mime_type" in response.meta: - yield ToolInvokeMessageBinary( - mimetype=response.meta.get("mime_type", "application/octet-stream") - if response.meta - else "application/octet-stream", - url=cast(ToolInvokeMessage.TextMessage, response.message).text, - ) - - @staticmethod - def _create_message_files( - tool_messages: Iterable[ToolInvokeMessageBinary], - agent_message: Message, - invoke_from: InvokeFrom, - user_id: str, - ) -> list[str]: - """ - Create message files produced by a tool call. - - Tool file persistence is a side effect of agent execution. Use an - independent transaction so this helper never commits or closes the - caller's request-scoped session. - - :return: message file ids - """ - result = [] - - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - for message in tool_messages: - # extract tool file id from url - tool_file_id = message.url.split("/")[-1].split(".")[0] - message_file = MessageFile( - message_id=agent_message.id, - type=ToolEngine._resolve_tool_file_type(message), - transfer_method=FileTransferMethod.TOOL_FILE, - belongs_to=MessageFileBelongsTo.ASSISTANT, - url=message.url, - upload_file_id=tool_file_id, - created_by_role=( - CreatorUserRole.ACCOUNT - if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} - else CreatorUserRole.END_USER - ), - created_by=user_id, - ) - - session.add(message_file) - result.append(message_file.id) - - return result - - @staticmethod - def _resolve_tool_file_type(message: ToolInvokeMessageBinary) -> FileType: - if "image" in message.mimetype: - return FileType.IMAGE - elif "video" in message.mimetype: - return FileType.VIDEO - elif "audio" in message.mimetype: - return FileType.AUDIO - elif "text" in message.mimetype or "pdf" in message.mimetype: - return FileType.DOCUMENT - else: - return FileType.CUSTOM diff --git a/api/core/tools/tool_label_manager.py b/api/core/tools/tool_label_manager.py deleted file mode 100644 index 301ed3e1084d1c..00000000000000 --- a/api/core/tools/tool_label_manager.py +++ /dev/null @@ -1,135 +0,0 @@ -from sqlalchemy import delete, select -from sqlalchemy.orm import Session, sessionmaker - -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.entities.values import default_tool_label_name_list -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from extensions.ext_database import db -from models.tools import ToolLabelBinding - - -class ToolLabelManager: - @classmethod - def filter_tool_labels(cls, tool_labels: list[str]) -> list[str]: - """ - Filter tool labels - """ - tool_labels = [label for label in tool_labels if label in default_tool_label_name_list] - return list(set(tool_labels)) - - @classmethod - def update_tool_labels( - cls, controller: ToolProviderController, labels: list[str], session: Session | None = None - ) -> None: - """ - Update tool labels - - :param controller: tool provider controller - :param labels: list of tool labels - :param session: database session, if None, a new session will be created - :return: None - """ - - labels = cls.filter_tool_labels(labels) - - if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): - provider_id = controller.provider_id - else: - raise ValueError("Unsupported tool type") - - if session is not None: - cls._update_tool_labels_logics(session, provider_id, controller, labels) - else: - with sessionmaker(db.engine).begin() as _session: - cls._update_tool_labels_logics(_session, provider_id, controller, labels) - - @classmethod - def _update_tool_labels_logics( - cls, session: Session, provider_id: str, controller: ToolProviderController, labels: list[str] - ) -> None: - """ - Update tool labels logics - - :param session: database session - :param provider_id: tool provider ID - :param controller: tool provider controller - :param labels: list of tool labels - :return: None - """ - - # delete old labels - _ = session.execute( - delete(ToolLabelBinding).where( - ToolLabelBinding.tool_id == provider_id, ToolLabelBinding.tool_type == controller.provider_type - ) - ) - - # insert new labels - for label in labels: - session.add(ToolLabelBinding(tool_id=provider_id, tool_type=controller.provider_type, label_name=label)) - - @classmethod - def get_tool_labels(cls, controller: ToolProviderController) -> list[str]: - """ - Get tool labels - - :param controller: tool provider controller - :return: list of tool labels (str) - """ - - match controller: - case ApiToolProviderController() | WorkflowToolProviderController(): - provider_id = controller.provider_id - case BuiltinToolProviderController(): - return controller.tool_labels - case _: - raise ValueError("Unsupported tool type") - stmt = select(ToolLabelBinding.label_name).where( - ToolLabelBinding.tool_id == provider_id, - ToolLabelBinding.tool_type == controller.provider_type, - ) - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - labels: list[str] = list(_session.scalars(stmt).all()) - - return labels - - @classmethod - def get_tools_labels(cls, tool_providers: list[ToolProviderController]) -> dict[str, list[str]]: - """ - Get tools labels - - :param tool_providers: list of tool providers - - :return: dict of tool labels - :key: tool id - :value: list of tool labels - """ - if not tool_providers: - return {} - - provider_ids: list[str] = [] - provider_types: set[str] = set() - - for controller in tool_providers: - if not isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): - raise ValueError("Unsupported tool type") - provider_ids.append(controller.provider_id) - provider_types.add(controller.provider_type) - - labels: list[ToolLabelBinding] = [] - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - stmt = select(ToolLabelBinding).where( - ToolLabelBinding.tool_id.in_(provider_ids), ToolLabelBinding.tool_type.in_(list(provider_types)) - ) - labels = list(_session.scalars(stmt).all()) - - tool_labels: dict[str, list[str]] = {label.tool_id: [] for label in labels} - - for label in labels: - tool_labels[label.tool_id].append(label.label_name) - - return tool_labels diff --git a/api/core/tools/tool_manager.py b/api/core/tools/tool_manager.py deleted file mode 100644 index 87f70b18b74b18..00000000000000 --- a/api/core/tools/tool_manager.py +++ /dev/null @@ -1,1161 +0,0 @@ -import logging -import mimetypes -from collections.abc import Generator, Mapping -from os import listdir, path -from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol, cast - -import sqlalchemy as sa -from pydantic import TypeAdapter -from sqlalchemy import select -from sqlalchemy.orm import Session -from typing_extensions import TypedDict -from yarl import URL - -import contexts -from configs import dify_config -from core.agent.entities import AgentToolEntity -from core.app.entities.app_invoke_entities import InvokeFrom -from core.helper.module_import_helper import load_single_subclass_from_source -from core.helper.position_helper import is_filtered -from core.plugin.entities.plugin_daemon import CredentialType -from core.plugin.impl.exc import PluginDaemonNotFoundError, PluginNotFoundError -from core.plugin.impl.tool import PluginToolManager -from core.tools.__base.tool import Tool -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort -from core.tools.builtin_tool.tool import BuiltinTool -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.custom_tool.tool import ApiTool -from core.tools.entities.api_entities import ToolProviderApiEntity, ToolProviderTypeApiLiteral -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ApiProviderSchemaType, - EmojiIconDict, - ToolInvokeFrom, - ToolParameter, - ToolProviderType, - emoji_icon_adapter, -) -from core.tools.errors import ToolProviderNotFoundError -from core.tools.mcp_tool.provider import MCPToolProviderController -from core.tools.mcp_tool.tool import MCPTool -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.plugin_tool.tool import PluginTool -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.utils.configuration import ToolParameterConfigurationManager -from core.tools.utils.encryption import create_tool_provider_encrypter -from core.tools.utils.uuid_utils import is_valid_uuid -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.runtime import VariablePool -from graphon.variables.template_resolution import convert_template -from models.provider_ids import ToolProviderID -from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider -from services.tools.legacy_mcp_tools_manage_service import MCPToolManageService -from services.tools.legacy_tools_transform_service import ToolTransformService - -if TYPE_CHECKING: - from core.app.apps.workflow_app_runner import WorkflowRunDriver - -logger = logging.getLogger(__name__) - - -class ApiProviderControllerItem(TypedDict): - provider: ApiToolProvider - controller: ApiToolProviderController - - -_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any]) - - -class WorkflowToolRuntimeSpec(Protocol): - @property - def provider_type(self) -> ToolProviderType: ... - - @property - def provider_id(self) -> str: ... - - @property - def tool_name(self) -> str: ... - - @property - def tool_configurations(self) -> Mapping[str, Any]: ... - - @property - def credential_id(self) -> str | None: ... - - -class ToolCredentialResolver(Protocol): - def __call__( - self, - *, - controller: BuiltinToolProviderController, - provider_id: str, - workspace_id: str, - credential_id: str | None, - ) -> tuple[Mapping[str, object], CredentialType]: ... - - -class ToolManager: - _builtin_provider_lock = Lock() - _hardcoded_providers: dict[str, BuiltinToolProviderController] = {} - _external_hardcoded_provider_classes: list[type[BuiltinToolProviderController]] = [] - _builtin_providers_loaded = False - _builtin_tools_labels: dict[str, I18nObject | None] = {} - _credential_resolver: ClassVar[ToolCredentialResolver | None] = None - - @classmethod - def configure_credential_resolver(cls, resolver: ToolCredentialResolver) -> None: - """Register stateless composition; it must not capture an app, database or provider.""" - cls._credential_resolver = resolver - - @classmethod - def get_hardcoded_provider(cls, provider: str) -> BuiltinToolProviderController: - """ - - get the hardcoded provider - - """ - - if len(cls._hardcoded_providers) == 0: - # init the builtin providers - cls.load_hardcoded_providers_cache() - - return cls._hardcoded_providers[provider] - - @classmethod - def get_builtin_provider( - cls, provider: str, tenant_id: str - ) -> BuiltinToolProviderController | PluginToolProviderController: - """ - get the builtin provider - - :param provider: the name of the provider - :param tenant_id: the id of the tenant - :return: the provider - """ - # split provider to - - if len(cls._hardcoded_providers) == 0: - # init the builtin providers - cls.load_hardcoded_providers_cache() - - if provider not in cls._hardcoded_providers: - # get plugin provider - plugin_provider = cls.get_plugin_provider(provider, tenant_id) - if plugin_provider: - return plugin_provider - - return cls._hardcoded_providers[provider] - - @classmethod - def get_plugin_provider(cls, provider: str, tenant_id: str) -> PluginToolProviderController: - """ - get the plugin provider - """ - # check if context is set - - try: - contexts.plugin_tool_providers.get() - except LookupError: - contexts.plugin_tool_providers.set({}) - contexts.plugin_tool_providers_lock.set(Lock()) - - plugin_tool_providers = contexts.plugin_tool_providers.get() - if provider in plugin_tool_providers: - return plugin_tool_providers[provider] - - with contexts.plugin_tool_providers_lock.get(): - # double check - plugin_tool_providers = contexts.plugin_tool_providers.get() - if provider in plugin_tool_providers: - return plugin_tool_providers[provider] - - manager = PluginToolManager() - try: - provider_entity = manager.fetch_tool_provider(tenant_id, provider) - except (PluginNotFoundError, PluginDaemonNotFoundError) as exc: - # The plugin daemon rejected the provider name (probably an - # unknown or non-builtin identifier). Translate to the - # console's domain error so the API returns a 4xx instead - # of leaking the daemon's ``PluginNotFoundError`` as a 500. - raise ToolProviderNotFoundError(f"plugin provider {provider} not found") from exc - if not provider_entity: - raise ToolProviderNotFoundError(f"plugin provider {provider} not found") - - controller = PluginToolProviderController( - entity=provider_entity.declaration, - plugin_id=provider_entity.plugin_id, - plugin_unique_identifier=provider_entity.plugin_unique_identifier, - tenant_id=tenant_id, - ) - - plugin_tool_providers[provider] = controller - return controller - - @classmethod - def get_tool_runtime( - cls, - provider_type: ToolProviderType, - provider_id: str, - tool_name: str, - tenant_id: str, - user_id: str | None = None, - invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, - tool_invoke_from: ToolInvokeFrom = ToolInvokeFrom.AGENT, - credential_id: str | None = None, - execution_driver: "WorkflowRunDriver | None" = None, - ) -> BuiltinTool | PluginTool | ApiTool | WorkflowTool | MCPTool: - """ - get the tool runtime - - :param provider_type: the type of the provider - :param provider_id: the id of the provider - :param tool_name: the name of the tool - :param tenant_id: the tenant id - :param user_id: the caller id bound to runtime-scoped model/tool lookups - :param invoke_from: invoke from - :param tool_invoke_from: the tool invoke from - :param credential_id: the credential id - - :return: the tool - """ - match provider_type: - case ToolProviderType.BUILT_IN: - provider_controller = cls.get_builtin_provider(provider_id, tenant_id) - - builtin_tool = provider_controller.get_tool(tool_name) - if not builtin_tool: - raise ToolProviderNotFoundError(f"builtin tool {tool_name} not found") - - if not provider_controller.need_credentials: - return builtin_tool.fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials={}, - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - selected_id = ( - credential_id - if isinstance(provider_controller, PluginToolProviderController) and is_valid_uuid(credential_id) - else None - ) - if cls._credential_resolver is None: - raise RuntimeError("Tool credential resolver is not configured") - credentials, credential_type = cls._credential_resolver( - controller=provider_controller, - provider_id=provider_id, - workspace_id=tenant_id, - credential_id=selected_id, - ) - - return builtin_tool.fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials=dict(credentials), - credential_type=credential_type, - runtime_parameters={}, - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - - case ToolProviderType.API: - api_provider, credentials = cls.get_api_provider_controller(tenant_id, provider_id) - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=api_provider, - ) - return api_provider.get_tool(tool_name).fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials=dict(encrypter.decrypt(credentials)), - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - case ToolProviderType.WORKFLOW: - workflow_provider_stmt = select(WorkflowToolProvider).where( - WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id - ) - with Session(db.engine, expire_on_commit=False) as session, session.begin(): - workflow_provider = session.scalar(workflow_provider_stmt) - - if workflow_provider is None: - raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - - controller = ToolTransformService.workflow_provider_to_controller(db_provider=workflow_provider) - controller_tools: list[WorkflowTool] = controller.get_tools(tenant_id=workflow_provider.tenant_id) - if controller_tools is None or len(controller_tools) == 0: - raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - - workflow_tool = controller.get_tools(tenant_id=workflow_provider.tenant_id)[0].fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials={}, - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - workflow_tool.execution_driver = execution_driver - return workflow_tool - case ToolProviderType.APP: - raise NotImplementedError("app provider not implemented") - case ToolProviderType.PLUGIN: - plugin_tool = cls.get_plugin_provider(provider_id, tenant_id).get_tool(tool_name) - runtime = getattr(plugin_tool, "runtime", None) - if runtime is not None: - runtime.user_id = user_id - runtime.invoke_from = invoke_from - runtime.tool_invoke_from = tool_invoke_from - return plugin_tool - case ToolProviderType.MCP: - mcp_tool = cls.get_mcp_provider_controller(tenant_id, provider_id).get_tool(tool_name) - runtime = getattr(mcp_tool, "runtime", None) - if runtime is not None: - runtime.user_id = user_id - runtime.invoke_from = invoke_from - runtime.tool_invoke_from = tool_invoke_from - return mcp_tool - case ToolProviderType.DATASET_RETRIEVAL: - raise ToolProviderNotFoundError(f"provider type {provider_type.value} not found") - case _: - raise ToolProviderNotFoundError(f"provider type {provider_type} not found") - - @classmethod - def get_agent_tool_runtime( - cls, - tenant_id: str, - app_id: str, - agent_tool: AgentToolEntity, - user_id: str | None = None, - invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, - variable_pool: "VariablePool | None" = None, - allow_file_parameters: bool = False, - use_default_for_missing_form_parameters: bool = False, - execution_driver: "WorkflowRunDriver | None" = None, - ) -> Tool: - """ - get the agent tool runtime - """ - tool_entity = cls.get_tool_runtime( - provider_type=agent_tool.provider_type, - provider_id=agent_tool.provider_id, - tool_name=agent_tool.tool_name, - tenant_id=tenant_id, - user_id=user_id, - invoke_from=invoke_from, - tool_invoke_from=ToolInvokeFrom.AGENT, - credential_id=agent_tool.credential_id, - execution_driver=execution_driver, - ) - runtime_parameters: dict[str, Any] = {} - parameters = tool_entity.get_merged_runtime_parameters() - runtime_parameters = cls._convert_tool_parameters_type( - parameters, - variable_pool, - agent_tool.tool_parameters, - typ="agent", - allow_file_parameters=allow_file_parameters, - use_default_for_missing_form_parameters=use_default_for_missing_form_parameters, - ) - # decrypt runtime parameters - encryption_manager = ToolParameterConfigurationManager( - tenant_id=tenant_id, - tool_runtime=tool_entity, - provider_name=agent_tool.provider_id, - provider_type=agent_tool.provider_type, - identity_id=f"AGENT.{app_id}", - ) - runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) - if tool_entity.runtime is None or tool_entity.runtime.runtime_parameters is None: - raise ValueError("runtime not found or runtime parameters not found") - - tool_entity.runtime.runtime_parameters.update(runtime_parameters) - return tool_entity - - @classmethod - def get_workflow_tool_runtime( - cls, - tenant_id: str, - app_id: str, - node_id: str, - workflow_tool: WorkflowToolRuntimeSpec, - user_id: str | None = None, - invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, - variable_pool: "VariablePool | None" = None, - execution_driver: "WorkflowRunDriver | None" = None, - workflow_tool_definition: WorkflowTool | None = None, - ) -> Tool: - """ - get the workflow tool runtime - """ - - if workflow_tool_definition is None: - tool_runtime = cls.get_tool_runtime( - provider_type=workflow_tool.provider_type, - provider_id=workflow_tool.provider_id, - tool_name=workflow_tool.tool_name, - tenant_id=tenant_id, - user_id=user_id, - invoke_from=invoke_from, - tool_invoke_from=ToolInvokeFrom.WORKFLOW, - credential_id=workflow_tool.credential_id, - execution_driver=execution_driver, - ) - else: - tool_runtime = workflow_tool_definition.fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials={}, - invoke_from=invoke_from, - tool_invoke_from=ToolInvokeFrom.WORKFLOW, - ) - ) - tool_runtime.execution_driver = execution_driver - - parameters = tool_runtime.get_merged_runtime_parameters() - runtime_parameters = cls._convert_tool_parameters_type( - parameters, variable_pool, workflow_tool.tool_configurations, typ="workflow" - ) - # decrypt runtime parameters - encryption_manager = ToolParameterConfigurationManager( - tenant_id=tenant_id, - tool_runtime=tool_runtime, - provider_name=workflow_tool.provider_id, - provider_type=workflow_tool.provider_type, - identity_id=f"WORKFLOW.{app_id}.{node_id}", - ) - - if runtime_parameters: - runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) - - tool_runtime.runtime.runtime_parameters.update(runtime_parameters) - return tool_runtime - - @classmethod - def get_tool_runtime_from_plugin( - cls, - tool_type: ToolProviderType, - tenant_id: str, - provider: str, - tool_name: str, - tool_parameters: dict[str, Any], - user_id: str | None = None, - credential_id: str | None = None, - execution_driver: "WorkflowRunDriver | None" = None, - ) -> Tool: - """ - get tool runtime from plugin - """ - tool_entity = cls.get_tool_runtime( - provider_type=tool_type, - provider_id=provider, - tool_name=tool_name, - tenant_id=tenant_id, - user_id=user_id, - invoke_from=InvokeFrom.SERVICE_API, - tool_invoke_from=ToolInvokeFrom.PLUGIN, - credential_id=credential_id, - execution_driver=execution_driver, - ) - runtime_parameters: dict[str, Any] = {} - parameters = tool_entity.get_merged_runtime_parameters() - for parameter in parameters: - if parameter.form == ToolParameter.ToolParameterForm.FORM: - # save tool parameter to tool entity memory - value = parameter.init_frontend_parameter(tool_parameters.get(parameter.name)) - runtime_parameters[parameter.name] = value - - tool_entity.runtime.runtime_parameters.update(runtime_parameters) - return tool_entity - - @classmethod - def get_hardcoded_provider_icon(cls, provider: str) -> tuple[str, str]: - """ - get the absolute path of the icon of the hardcoded provider - - :param provider: the name of the provider - :return: the absolute path of the icon, the mime type of the icon - """ - # get provider - provider_controller = cls.get_hardcoded_provider(provider) - - absolute_path = path.join( - provider_controller.provider_directory, - "_assets", - provider_controller.entity.identity.icon, - ) - # check if the icon exists - if not path.exists(absolute_path): - raise ToolProviderNotFoundError(f"builtin provider {provider} icon not found") - - # get the mime type - mime_type, _ = mimetypes.guess_type(absolute_path) - mime_type = mime_type or "application/octet-stream" - - return absolute_path, mime_type - - @classmethod - def list_hardcoded_providers(cls): - # use cache first - if cls._builtin_providers_loaded: - yield from list(cls._hardcoded_providers.values()) - return - - with cls._builtin_provider_lock: - if cls._builtin_providers_loaded: - yield from list(cls._hardcoded_providers.values()) - return - - yield from cls._list_hardcoded_providers() - - @classmethod - def list_plugin_providers(cls, tenant_id: str) -> list[PluginToolProviderController]: - """ - list all the plugin providers - """ - - manager = PluginToolManager() - provider_entities = manager.fetch_tool_providers(tenant_id) - return [ - PluginToolProviderController( - entity=provider.declaration, - plugin_id=provider.plugin_id, - plugin_unique_identifier=provider.plugin_unique_identifier, - tenant_id=tenant_id, - ) - for provider in provider_entities - ] - - @classmethod - def list_builtin_providers( - cls, tenant_id: str - ) -> Generator[BuiltinToolProviderController | PluginToolProviderController]: - """ - list all the builtin providers - """ - yield from cls.list_hardcoded_providers() - # get plugin providers - yield from cls.list_plugin_providers(tenant_id) - - @classmethod - def _list_hardcoded_providers(cls) -> Generator[BuiltinToolProviderController]: - """ - list all the builtin providers - """ - for provider_path in listdir(path.join(path.dirname(path.realpath(__file__)), "builtin_tool", "providers")): - if provider_path.startswith("__"): - continue - - if path.isdir(path.join(path.dirname(path.realpath(__file__)), "builtin_tool", "providers", provider_path)): - if provider_path.startswith("__"): - continue - - # init provider - try: - provider_class = load_single_subclass_from_source( - module_name=f"core.tools.builtin_tool.providers.{provider_path}.{provider_path}", - script_path=path.join( - path.dirname(path.realpath(__file__)), - "builtin_tool", - "providers", - provider_path, - f"{provider_path}.py", - ), - parent_type=BuiltinToolProviderController, - ) - yield cls._cache_hardcoded_provider(provider_class()) - - except Exception: - logger.exception("load builtin provider %s", provider_path) - continue - for external_provider_class in cls._external_hardcoded_provider_classes: - try: - yield cls._cache_hardcoded_provider(external_provider_class()) - except Exception: - logger.exception("load builtin provider %s", external_provider_class.__module__) - # set builtin providers loaded - cls._builtin_providers_loaded = True - - @classmethod - def _cache_hardcoded_provider(cls, provider: BuiltinToolProviderController) -> BuiltinToolProviderController: - cls._hardcoded_providers[provider.entity.identity.name] = provider - for tool in provider.get_tools(): - cls._builtin_tools_labels[tool.entity.identity.name] = tool.entity.identity.label - return provider - - @classmethod - def register_hardcoded_provider_class(cls, provider_class: type[BuiltinToolProviderController]) -> None: - """Register a builtin provider owned outside ``core`` from the application composition root.""" - with cls._builtin_provider_lock: - if provider_class not in cls._external_hardcoded_provider_classes: - cls._external_hardcoded_provider_classes.append(provider_class) - cls._builtin_providers_loaded = False - - @classmethod - def load_hardcoded_providers_cache(cls): - for _ in cls.list_hardcoded_providers(): - pass - - @classmethod - def clear_hardcoded_providers_cache(cls): - cls._hardcoded_providers = {} - cls._builtin_providers_loaded = False - - @classmethod - def get_tool_label(cls, tool_name: str) -> I18nObject | None: - """ - get the tool label - - :param tool_name: the name of the tool - - :return: the label of the tool - """ - if len(cls._builtin_tools_labels) == 0: - # init the builtin providers - cls.load_hardcoded_providers_cache() - - if tool_name not in cls._builtin_tools_labels: - return None - - return cls._builtin_tools_labels[tool_name] - - @classmethod - def list_default_builtin_providers(cls, tenant_id: str) -> list[BuiltinToolProvider]: - """ - list all the builtin providers - """ - # according to multi credentials, select the one with is_default=True first, then created_at oldest - # for compatibility with old version - if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql": - # PostgreSQL: Use DISTINCT ON - sql = """ - SELECT DISTINCT ON (tenant_id, provider) id - FROM tool_builtin_providers - WHERE tenant_id = :tenant_id - ORDER BY tenant_id, provider, is_default DESC, created_at DESC - """ - else: - # MySQL: Use window function to achieve same result - sql = """ - SELECT id FROM ( - SELECT id, - ROW_NUMBER() OVER ( - PARTITION BY tenant_id, provider - ORDER BY is_default DESC, created_at DESC - ) as rn - FROM tool_builtin_providers - WHERE tenant_id = :tenant_id - ) ranked WHERE rn = 1 - """ - - with Session(db.engine, autoflush=False) as session: - ids = [row.id for row in session.execute(sa.text(sql), {"tenant_id": tenant_id}).all()] - return list(session.scalars(select(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)))) - - @classmethod - def list_providers_from_api( - cls, user_id: str, tenant_id: str, typ: ToolProviderTypeApiLiteral - ) -> list[ToolProviderApiEntity]: - result_providers: dict[str, ToolProviderApiEntity] = {} - - filters = [] - if not typ: - filters.extend(["builtin", "api", "workflow", "mcp"]) - else: - filters.append(typ) - - # Use a single session for all database operations to reduce connection overhead - with Session(db.engine) as session: - if "builtin" in filters: - builtin_providers = list(cls.list_builtin_providers(tenant_id)) - - # key: provider name, value: provider - db_builtin_providers = { - str(ToolProviderID(provider.provider)): provider - for provider in cls.list_default_builtin_providers(tenant_id) - } - - # append builtin providers - for provider in builtin_providers: - # handle include, exclude - if is_filtered( - include_set=dify_config.POSITION_TOOL_INCLUDES_SET, - exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET, - data=provider, - name_func=lambda x: x.entity.identity.name, - ): - continue - user_provider = ToolTransformService.builtin_provider_to_user_provider( - provider_controller=provider, - db_provider=db_builtin_providers.get(provider.entity.identity.name), - decrypt_credentials=False, - ) - - if isinstance(provider, PluginToolProviderController): - result_providers[f"plugin_provider.{user_provider.name}"] = user_provider - else: - result_providers[f"builtin_provider.{user_provider.name}"] = user_provider - - # get db api providers - if "api" in filters: - db_api_providers = session.scalars( - select(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id) - ).all() - - # Batch create controllers - api_provider_controllers: list[ApiProviderControllerItem] = [] - for api_provider in db_api_providers: - try: - controller = ToolTransformService.api_provider_to_controller(api_provider) - api_provider_controllers.append({"provider": api_provider, "controller": controller}) - except Exception: - # Skip invalid providers but continue processing others - logger.warning("Failed to create controller for API provider %s", api_provider.id) - - # Batch get labels for all API providers - if api_provider_controllers: - controllers = cast( - list[ToolProviderController], [item["controller"] for item in api_provider_controllers] - ) - labels = ToolLabelManager.get_tools_labels(controllers) - - for item in api_provider_controllers: - provider_controller = item["controller"] - db_provider = item["provider"] - provider_labels = labels.get(provider_controller.provider_id, []) - user_provider = ToolTransformService.api_provider_to_user_provider( - provider_controller=provider_controller, - db_provider=db_provider, - decrypt_credentials=False, - labels=provider_labels, - session=db.session(), - ) - result_providers[f"api_provider.{user_provider.name}"] = user_provider - - if "workflow" in filters: - # get workflow providers - workflow_providers = session.scalars( - select(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id) - ).all() - - workflow_provider_controllers: list[WorkflowToolProviderController] = [] - for workflow_provider in workflow_providers: - try: - workflow_controller: WorkflowToolProviderController = ( - ToolTransformService.workflow_provider_to_controller(db_provider=workflow_provider) - ) - workflow_provider_controllers.append(workflow_controller) - except Exception: - # app has been deleted - logger.exception("Failed to transform workflow provider %s to controller", workflow_provider.id) - continue - # Batch get labels for workflow providers - if workflow_provider_controllers: - workflow_controllers: list[ToolProviderController] = [ - cast(ToolProviderController, controller) for controller in workflow_provider_controllers - ] - labels = ToolLabelManager.get_tools_labels(workflow_controllers) - - for workflow_provider_controller in workflow_provider_controllers: - provider_labels = labels.get(workflow_provider_controller.provider_id, []) - user_provider = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=workflow_provider_controller, - labels=provider_labels, - ) - result_providers[f"workflow_provider.{user_provider.name}"] = user_provider - - if "mcp" in filters: - mcp_service = MCPToolManageService(session=session) - mcp_providers = mcp_service.list_providers(tenant_id=tenant_id) - for mcp_provider in mcp_providers: - result_providers[f"mcp_provider.{mcp_provider.name}"] = mcp_provider - - return BuiltinToolProviderSort.sort(list(result_providers.values())) - - @classmethod - def get_api_provider_controller( - cls, tenant_id: str, provider_id: str - ) -> tuple[ApiToolProviderController, dict[str, Any]]: - """ - get the api provider - - :param tenant_id: the id of the tenant - :param provider_id: the id of the provider - - :return: the provider controller, the credentials - """ - provider: ApiToolProvider | None = db.session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.id == provider_id, - ApiToolProvider.tenant_id == tenant_id, - ) - .limit(1) - ) - - if provider is None: - raise ToolProviderNotFoundError(f"api provider {provider_id} not found") - - auth_type = ApiProviderAuthType.NONE - provider_auth_type = provider.credentials.get("auth_type") - if provider_auth_type in ("api_key_header", "api_key"): # backward compatibility - auth_type = ApiProviderAuthType.API_KEY_HEADER - elif provider_auth_type == "api_key_query": - auth_type = ApiProviderAuthType.API_KEY_QUERY - - controller = ApiToolProviderController.from_db( - provider, - auth_type, - session=db.session(), - ) - controller.load_bundled_tools(provider.tools) - - return controller, provider.credentials - - @classmethod - def get_mcp_provider_controller(cls, tenant_id: str, provider_id: str) -> MCPToolProviderController: - """ - get the api provider - - :param tenant_id: the id of the tenant - :param provider_id: the persisted reference of the provider, normally its - server identifier, or the primary key for graphs written before that - convention - - :return: the provider controller, the credentials - """ - with Session(db.engine) as session: - mcp_service = MCPToolManageService(session=session) - try: - provider = mcp_service.get_provider_by_persisted_reference( - id_or_server_identifier=provider_id, tenant_id=tenant_id - ) - except ValueError as caught_error: - raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found") from caught_error - - controller = MCPToolProviderController.from_db(provider) - - return controller - - @classmethod - def user_get_api_provider(cls, provider: str, tenant_id: str, mask: bool = True): - """ - get api provider - """ - provider_name = provider - provider_obj: ApiToolProvider | None = db.session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider, - ) - .limit(1) - ) - - if provider_obj is None: - raise ValueError(f"you have not added provider {provider_name}") - - try: - credentials = _credentials_adapter.validate_json(provider_obj.credentials_str) or {} - except Exception: - credentials = {} - - # package tool provider controller - auth_type = ApiProviderAuthType.NONE - credentials_auth_type = credentials.get("auth_type") - if credentials_auth_type in ("api_key_header", "api_key"): # backward compatibility - auth_type = ApiProviderAuthType.API_KEY_HEADER - elif credentials_auth_type == "api_key_query": - auth_type = ApiProviderAuthType.API_KEY_QUERY - - controller = ApiToolProviderController.from_db( - provider_obj, - auth_type, - session=db.session(), - ) - # init tool configuration - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=controller, - ) - if mask: - masked_credentials = encrypter.mask_plugin_credentials(encrypter.decrypt(credentials)) - else: - masked_credentials = encrypter.decrypt(credentials) - - try: - icon = emoji_icon_adapter.validate_json(provider_obj.icon) - except Exception: - icon = {"background": "#252525", "content": "\ud83d\ude01"} - - # add tool labels - labels = ToolLabelManager.get_tool_labels(controller) - schema_type = provider_obj.schema_type - schema_type_value = schema_type.value if isinstance(schema_type, ApiProviderSchemaType) else schema_type - - return { - "schema_type": schema_type_value, - "schema": provider_obj.schema, - "tools": [tool.model_dump(mode="json") for tool in provider_obj.tools], - "icon": icon, - "description": provider_obj.description, - "credentials": masked_credentials, - "privacy_policy": provider_obj.privacy_policy, - "custom_disclaimer": provider_obj.custom_disclaimer, - "labels": labels, - } - - @classmethod - def generate_builtin_tool_icon_url(cls, provider_id: str) -> str: - return str( - URL(dify_config.CONSOLE_API_URL or "/") - / "console" - / "api" - / "workspaces" - / "current" - / "tool-provider" - / "builtin" - / provider_id - / "icon" - ) - - @classmethod - def generate_plugin_tool_icon_url(cls, tenant_id: str, filename: str) -> str: - return str( - URL(dify_config.CONSOLE_API_URL or "/") - / "console" - / "api" - / "workspaces" - / "current" - / "plugin" - / "icon" - % {"tenant_id": tenant_id, "filename": filename} - ) - - @classmethod - def generate_workflow_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict: - try: - # Use a short-lived session to avoid holding a database transaction - # during long-running nested workflow execution. - # Fixes: idle in transaction when Workflow Tool runs (#36902) - with Session(db.engine, expire_on_commit=False) as session: - workflow_provider: WorkflowToolProvider | None = session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) - .limit(1) - ) - - if workflow_provider is None: - raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - - icon = emoji_icon_adapter.validate_json(workflow_provider.icon) - return icon - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - - @classmethod - def generate_api_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict: - try: - # Use a short-lived session to avoid holding a database transaction - # during long-running tool execution. - with Session(db.engine, expire_on_commit=False) as session: - api_provider: ApiToolProvider | None = session.scalar( - select(ApiToolProvider) - .where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id) - .limit(1) - ) - - if api_provider is None: - raise ToolProviderNotFoundError(f"api provider {provider_id} not found") - - icon = emoji_icon_adapter.validate_json(api_provider.icon) - return icon - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - - @classmethod - def generate_mcp_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict | str: - try: - with Session(db.engine) as session: - mcp_service = MCPToolManageService(session=session) - try: - mcp_provider = mcp_service.get_provider_entity_by_persisted_reference( - id_or_server_identifier=provider_id, tenant_id=tenant_id - ) - return cast(EmojiIconDict | str, mcp_provider.provider_icon) - except ValueError as caught_error: - raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found") from caught_error - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - - @classmethod - def get_tool_icon( - cls, - tenant_id: str, - provider_type: ToolProviderType, - provider_id: str, - ) -> str | EmojiIconDict: - """ - get the tool icon - - :param tenant_id: the id of the tenant - :param provider_type: the type of the provider - :param provider_id: the id of the provider - :return: - """ - match provider_type: - case ToolProviderType.BUILT_IN: - provider = ToolManager.get_builtin_provider(provider_id, tenant_id) - if isinstance(provider, PluginToolProviderController): - try: - return cls.generate_plugin_tool_icon_url(tenant_id, provider.entity.identity.icon) - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - return cls.generate_builtin_tool_icon_url(provider_id) - case ToolProviderType.API: - return cls.generate_api_tool_icon_url(tenant_id, provider_id) - case ToolProviderType.WORKFLOW: - return cls.generate_workflow_tool_icon_url(tenant_id, provider_id) - case ToolProviderType.PLUGIN: - provider = ToolManager.get_plugin_provider(provider_id, tenant_id) - try: - return cls.generate_plugin_tool_icon_url(tenant_id, provider.entity.identity.icon) - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - case ToolProviderType.MCP: - return cls.generate_mcp_tool_icon_url(tenant_id, provider_id) - case ToolProviderType.APP | ToolProviderType.DATASET_RETRIEVAL: - raise ValueError(f"provider type {provider_type} not found") - case _: - raise ValueError(f"provider type {provider_type} not found") - - @classmethod - def _convert_tool_parameters_type( - cls, - parameters: list[ToolParameter], - variable_pool: "VariablePool | None", - tool_configurations: Mapping[str, Any], - typ: Literal["agent", "workflow", "tool"] = "workflow", - allow_file_parameters: bool = False, - use_default_for_missing_form_parameters: bool = False, - ) -> dict[str, Any]: - """ - Convert tool parameters type - """ - from graphon.nodes.tool.entities import ToolNodeData - from graphon.nodes.tool.exc import ToolParameterError - - runtime_parameters: dict[str, Any] = {} - for parameter in parameters: - if ( - parameter.type - in { - ToolParameter.ToolParameterType.SYSTEM_FILES, - ToolParameter.ToolParameterType.FILE, - ToolParameter.ToolParameterType.FILES, - } - and parameter.required - and typ == "agent" - and not allow_file_parameters - ): - raise ValueError(f"file type parameter {parameter.name} not supported in agent") - # save tool parameter to tool entity memory - if parameter.form == ToolParameter.ToolParameterForm.FORM: - if variable_pool: - config = tool_configurations.get(parameter.name, {}) - - selector_value = cls._extract_runtime_selector_value(parameter, config) - if selector_value is not None: - # Selector parameters carry structured dictionaries, not scalar ToolInput values. - runtime_parameters[parameter.name] = selector_value - continue - - if not (config and isinstance(config, dict) and config.get("value") is not None): - continue - tool_input = ToolNodeData.ToolInput.model_validate(tool_configurations.get(parameter.name, {})) - if tool_input.type == "variable": - variable_selector = tool_input.value - if not isinstance(variable_selector, list) or not all( - isinstance(selector_part, str) for selector_part in variable_selector - ): - raise ToolParameterError("Variable tool input must be a variable selector") - variable = variable_pool.get(variable_selector) - if variable is None: - raise ToolParameterError(f"Variable {tool_input.value} does not exist") - parameter_value = variable.value - elif tool_input.type == "constant": - parameter_value = tool_input.value - elif tool_input.type == "mixed": - segment_group = convert_template(variable_pool, str(tool_input.value)) - parameter_value = segment_group.text - else: - raise ToolParameterError(f"Unknown tool input type '{tool_input.type}'") - runtime_parameters[parameter.name] = parameter_value - - else: - parameter_value = tool_configurations.get(parameter.name) - if use_default_for_missing_form_parameters and parameter_value is None: - if parameter.default is not None: - parameter_value = parameter.default - elif ( - parameter.required - and parameter.type == ToolParameter.ToolParameterType.SELECT - and parameter.options - ): - parameter_value = parameter.options[0].value - else: - continue - value = parameter.init_frontend_parameter(parameter_value) - runtime_parameters[parameter.name] = value - return runtime_parameters - - @classmethod - def _extract_runtime_selector_value(cls, parameter: ToolParameter, config: Any) -> dict[str, Any] | None: - if parameter.type not in { - ToolParameter.ToolParameterType.MODEL_SELECTOR, - ToolParameter.ToolParameterType.APP_SELECTOR, - }: - return None - if not isinstance(config, dict): - return None - - input_value = config.get("value") - if isinstance(input_value, dict) and cls._is_selector_value(parameter, input_value): - return cast("dict[str, Any]", parameter.init_frontend_parameter(input_value)) - - if cls._is_selector_value(parameter, config): - selector_value = dict(config) - selector_value.pop("type", None) - selector_value.pop("value", None) - return cast("dict[str, Any]", parameter.init_frontend_parameter(selector_value)) - - return None - - @classmethod - def _is_selector_value(cls, parameter: ToolParameter, value: Mapping[str, Any]) -> bool: - if parameter.type == ToolParameter.ToolParameterType.MODEL_SELECTOR: - return ( - isinstance(value.get("provider"), str) - and isinstance(value.get("model"), str) - and isinstance(value.get("model_type"), str) - ) - if parameter.type == ToolParameter.ToolParameterType.APP_SELECTOR: - return isinstance(value.get("app_id"), str) - return False - - -ToolManager.load_hardcoded_providers_cache() diff --git a/api/core/tools/workflow_as_tool/provider.py b/api/core/tools/workflow_as_tool/provider.py deleted file mode 100644 index d613fc0ffc47cf..00000000000000 --- a/api/core/tools/workflow_as_tool/provider.py +++ /dev/null @@ -1,259 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import override - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager -from core.db.session_factory import session_factory -from core.plugin.entities.parameters import PluginParameterOption -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ( - ToolDescription, - ToolEntity, - ToolIdentity, - ToolParameter, - ToolProviderEntity, - ToolProviderIdentity, - ToolProviderType, -) -from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.variables.input_entities import VariableEntity, VariableEntityType -from models.account import Account -from models.model import App, AppMode -from models.tools import WorkflowToolProvider -from models.workflow import Workflow - -VARIABLE_TO_PARAMETER_TYPE_MAPPING = { - VariableEntityType.TEXT_INPUT: ToolParameter.ToolParameterType.STRING, - VariableEntityType.PARAGRAPH: ToolParameter.ToolParameterType.STRING, - VariableEntityType.SELECT: ToolParameter.ToolParameterType.SELECT, - VariableEntityType.NUMBER: ToolParameter.ToolParameterType.NUMBER, - VariableEntityType.CHECKBOX: ToolParameter.ToolParameterType.BOOLEAN, - VariableEntityType.FILE: ToolParameter.ToolParameterType.FILE, - VariableEntityType.FILE_LIST: ToolParameter.ToolParameterType.FILES, - VariableEntityType.JSON_OBJECT: ToolParameter.ToolParameterType.OBJECT, -} - - -class WorkflowToolProviderController(ToolProviderController[ToolProviderEntity, WorkflowTool | None]): - provider_id: str - tools: list[WorkflowTool] | None - - def __init__(self, entity: ToolProviderEntity, provider_id: str): - super().__init__(entity=entity) - self.provider_id = provider_id - self.tools = None - - @classmethod - def from_db(cls, db_provider: WorkflowToolProvider) -> WorkflowToolProviderController: - with session_factory.create_session() as session, session.begin(): - app = session.get(App, db_provider.app_id) - if not app: - raise ValueError("app not found") - - user = session.get(Account, db_provider.user_id) if db_provider.user_id else None - controller = WorkflowToolProviderController( - entity=ToolProviderEntity( - identity=ToolProviderIdentity( - author=user.name if user else "", - name=db_provider.label, - label=I18nObject(en_US=db_provider.label, zh_Hans=db_provider.label), - description=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), - icon=db_provider.icon, - ), - credentials_schema=[], - plugin_id=None, - ), - provider_id=db_provider.id, - ) - - controller.tools = [ - controller._get_db_provider_tool(db_provider, session=session, user=user), - ] - - return controller - - @property - @override - def provider_type(self) -> ToolProviderType: - return ToolProviderType.WORKFLOW - - def _get_db_provider_tool( - self, - db_provider: WorkflowToolProvider, - *, - session: Session, - user: Account | None = None, - ) -> WorkflowTool: - """ - get db provider tool - :param db_provider: the db provider - :return: the tool - """ - workflow: Workflow | None = session.scalar( - select(Workflow) - .where(Workflow.app_id == db_provider.app_id, Workflow.version == db_provider.version) - .limit(1) - ) - - if not workflow: - raise ValueError("workflow not found") - - return self.build_tool(db_provider, workflow, author=user.name if user else "") - - @staticmethod - def build_tool(db_provider: WorkflowToolProvider, workflow: Workflow, *, author: str) -> WorkflowTool: - """Build a tool from an already-loaded provider and its pinned definition.""" - # fetch start node - graph: Mapping = workflow.graph_dict - features_dict: Mapping = workflow.features_dict - features = WorkflowAppConfigManager.convert_features(config_dict=features_dict, app_mode=AppMode.WORKFLOW) - - parameters = db_provider.parameter_configurations - variables = WorkflowToolConfigurationUtils.get_workflow_graph_variables(graph) - - def fetch_workflow_variable(variable_name: str) -> VariableEntity | None: - return next(filter(lambda x: x.variable == variable_name, variables), None) - - workflow_tool_parameters = [] - for parameter in parameters: - variable = fetch_workflow_variable(parameter.name) - if variable: - parameter_type = None - options = [] - if variable.type not in VARIABLE_TO_PARAMETER_TYPE_MAPPING: - raise ValueError(f"unsupported variable type {variable.type}") - parameter_type = VARIABLE_TO_PARAMETER_TYPE_MAPPING[variable.type] - - if variable.type == VariableEntityType.SELECT and variable.options: - options = [ - PluginParameterOption(value=option, label=I18nObject(en_US=option, zh_Hans=option)) - for option in variable.options - ] - - workflow_tool_parameters.append( - ToolParameter( - name=parameter.name, - label=I18nObject(en_US=variable.label, zh_Hans=variable.label), - human_description=I18nObject(en_US=parameter.description, zh_Hans=parameter.description), - type=parameter_type, - form=parameter.form, - llm_description=parameter.description, - required=variable.required, - default=variable.default, - options=options, - placeholder=I18nObject(en_US="", zh_Hans=""), - ) - ) - elif features.file_upload: - workflow_tool_parameters.append( - ToolParameter( - name=parameter.name, - label=I18nObject(en_US=parameter.name, zh_Hans=parameter.name), - human_description=I18nObject(en_US=parameter.description, zh_Hans=parameter.description), - type=ToolParameter.ToolParameterType.SYSTEM_FILES, - llm_description=parameter.description, - required=False, - form=parameter.form, - placeholder=I18nObject(en_US="", zh_Hans=""), - ) - ) - else: - raise ValueError("variable not found") - - # get output schema from workflow - outputs = WorkflowToolConfigurationUtils.get_workflow_graph_output(graph) - - reserved_keys = {"json", "text", "files"} - - properties = {} - for output in outputs: - if output.variable not in reserved_keys: - properties[output.variable] = { - "type": output.value_type, - "description": "", - } - output_schema = {"type": "object", "properties": properties} - - return WorkflowTool( - workflow_as_tool_id=db_provider.id, - entity=ToolEntity( - identity=ToolIdentity( - author=author, - name=db_provider.name, - label=I18nObject(en_US=db_provider.label, zh_Hans=db_provider.label), - provider=db_provider.id, - icon=db_provider.icon, - ), - description=ToolDescription( - human=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), - llm=db_provider.description, - ), - parameters=workflow_tool_parameters, - output_schema=output_schema, - ), - runtime=ToolRuntime( - tenant_id=db_provider.tenant_id, - ), - workflow_app_id=workflow.app_id, - workflow_id=workflow.id, - version=db_provider.version, - workflow_call_depth=0, - label=db_provider.label, - ) - - def get_tools(self, tenant_id: str) -> list[WorkflowTool]: - """ - fetch tools from database - - :param tenant_id: the tenant id - :return: the tools - """ - if self.tools is not None: - return self.tools - - with Session(db.engine, expire_on_commit=False) as session, session.begin(): - db_provider: WorkflowToolProvider | None = session.scalar( - select(WorkflowToolProvider) - .where( - WorkflowToolProvider.tenant_id == tenant_id, - WorkflowToolProvider.id == self.provider_id, - ) - .limit(1) - ) - - if not db_provider: - return [] - - app = session.get(App, db_provider.app_id) - if not app: - raise ValueError("app not found") - - user = session.get(Account, db_provider.user_id) if db_provider.user_id else None - self.tools = [self._get_db_provider_tool(db_provider, session=session, user=user)] - - return self.tools - - @override - def get_tool(self, tool_name: str) -> WorkflowTool | None: - """ - get tool by name - - :param tool_name: the name of the tool - :return: the tool - """ - if self.tools is None: - return None - - for tool in self.tools: - if tool.entity.identity.name == tool_name: - return tool - - return None diff --git a/api/core/tools/workflow_as_tool/tool.py b/api/core/tools/workflow_as_tool/tool.py deleted file mode 100644 index bf4c990d06d9da..00000000000000 --- a/api/core/tools/workflow_as_tool/tool.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Compatibility wrapper for the service-owned nested workflow tool. - -New callers must import :class:`services.tools.workflow.tool.WorkflowTool` and -provide its execution dependencies directly. The legacy core tool stack still -constructs workflow tools while materializing provider metadata, so this -wrapper temporarily preserves that constructor and its ``execution_driver`` -attribute until the core stack is retired. -""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import TYPE_CHECKING, override - -from core.app.apps.draft_variable_saver import DraftVariableSaverFactory -from core.helper.trace_id_helper import ParentTraceContext -from core.tools.__base.tool import Tool -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.tool_entities import ToolEntity -from graphon.model_runtime.entities.llm_entities import LLMUsage -from models import Account -from models.tool_runtime_contracts import WorkflowToolQueries -from services.tools.workflow.tool import WorkflowTool as ServiceWorkflowTool -from services.workflow.execution.ports import WorkflowRuntime - -if TYPE_CHECKING: - from services.workflow.execution.contracts import WorkflowRunDriver - - -class WorkflowTool(ServiceWorkflowTool): - """Keep the legacy construction surface without restoring hidden dependencies. - - Metadata-only legacy callers may omit the new execution ports. Invocation - remains guarded by the canonical implementation, which raises when no - ``WorkflowRuntime`` was injected. - """ - - def __init__( - self, - workflow_app_id: str, - workflow_id: str, - workflow_as_tool_id: str, - version: str, - workflow_call_depth: int, - entity: ToolEntity, - runtime: ToolRuntime, - label: str = "Workflow", - execution_driver: WorkflowRunDriver | None = None, - *, - queries: WorkflowToolQueries | None = None, - draft_variable_saver: Callable[[str, Account], DraftVariableSaverFactory] | None = None, - workflow_runtime: WorkflowRuntime | None = None, - ) -> None: - if queries is None and workflow_runtime is not None: - queries = workflow_runtime.tools - - if queries is not None: - super().__init__( - workflow_app_id=workflow_app_id, - workflow_id=workflow_id, - workflow_as_tool_id=workflow_as_tool_id, - version=version, - workflow_call_depth=workflow_call_depth, - entity=entity, - runtime=runtime, - label=label, - queries=queries, - draft_variable_saver=draft_variable_saver, - workflow_runtime=workflow_runtime, - ) - else: - # The legacy provider also materializes tool metadata for list and - # configuration APIs. It has no execution composition root, so do - # not invent one; initialize only the state those APIs require. - self.workflow_app_id = workflow_app_id - self.workflow_id = workflow_id - self.workflow_as_tool_id = workflow_as_tool_id - self.version = version - self.workflow_call_depth = workflow_call_depth - self.label = label - self._draft_variable_saver = draft_variable_saver - self._workflow_runtime = None - self._latest_usage = LLMUsage.empty_usage() - self._parent_trace_context: ParentTraceContext | None = None - self._trace_session_id: str | None = None - Tool.__init__(self, entity=entity, runtime=runtime) - - self._legacy_queries = queries - self._legacy_execution_driver = execution_driver - - @property - def execution_driver(self) -> WorkflowRunDriver | None: - """Preserve the mutable attribute used by the retiring core manager.""" - if self._workflow_runtime is not None: - return self._workflow_runtime.execution_driver - return self._legacy_execution_driver - - @execution_driver.setter - def execution_driver(self, execution_driver: WorkflowRunDriver | None) -> None: - self._legacy_execution_driver = execution_driver - - @override - def fork_tool_runtime(self, runtime: ToolRuntime) -> WorkflowTool: - forked = self.__class__( - queries=self._legacy_queries, - entity=self.entity.model_copy(), - runtime=runtime, - workflow_app_id=self.workflow_app_id, - workflow_id=self.workflow_id, - workflow_as_tool_id=self.workflow_as_tool_id, - workflow_call_depth=self.workflow_call_depth, - version=self.version, - label=self.label, - execution_driver=self._legacy_execution_driver, - draft_variable_saver=self._draft_variable_saver, - workflow_runtime=self._workflow_runtime, - ) - forked._parent_trace_context = self._parent_trace_context.model_copy() if self._parent_trace_context else None - forked._trace_session_id = self._trace_session_id - return forked - - -__all__ = ["WorkflowTool"] diff --git a/api/extensions/application_services/workflow.py b/api/extensions/application_services/workflow.py index 285a45187ebfd7..da933383c68830 100644 --- a/api/extensions/application_services/workflow.py +++ b/api/extensions/application_services/workflow.py @@ -42,9 +42,6 @@ ) from services.workflow.variable_contracts import WorkflowExecutionVariables -if TYPE_CHECKING: - from core.app.apps.workflow_app_runner import WorkflowRunDriver - class _SessionBoundWorkflowToolInvoker(WorkflowToolInvoker): """Keep a transaction open until a workflow tool's lazy output is fully consumed.""" diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index dbfa99a892d003..f28ef214307cf1 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -20,7 +20,6 @@ from core.ops.base_trace_instance import bind_application_tracing_hooks from core.schemas.schema_manager import SchemaManager from core.tools.tool_file_manager import ToolFileManager -from core.tools.tool_manager import ToolManager as LegacyToolManager from enums import DeploymentEdition, WebAppAccessMode from extensions.application_services.account import AccountServices, build_account_services from extensions.application_services.agent import AgentAppServices, build_agent_app_services @@ -134,6 +133,7 @@ from services.api_based_extension_application_service import APIBasedExtensionApplicationService from services.app.advanced_prompt_template_service import AdvancedPromptTemplateService from services.app.api_key_service import AppApiKeyService +from services.app.generation.adapters.completion import CompletionAppGenerator from services.app.generation.runtime import AppGenerationRuntime from services.app.mcp_server_service import AppMCPServerService from services.app.workflow_comment_mention_gateway import CeleryWorkflowCommentMentionNotifier @@ -647,6 +647,7 @@ def build_application_services( repository=messages, ) + generation_runtime = build_workflow_execution_dependencies(database_client) return ApplicationServices( tools=tools, human_input_debug=build_human_input_debug_service(database_client=database_client), @@ -656,7 +657,7 @@ def build_application_services( console_workflow_variables=build_console_workflow_variables( database_client=database_client, variables=workflow_variables ), - workflow_runtime=build_workflow_execution_dependencies(database_client), + workflow_runtime=generation_runtime, workflow_suggestions=build_workflow_suggestions(database_client), workflow_agent_bindings=WorkflowAgentBindingResolver(database_client), snippet_generation=build_snippet_generation_service( @@ -829,7 +830,14 @@ def build_application_services( ), message_more_like_this=MessageMoreLikeThisService( repository=messages, - generator=MessageMoreLikeThisGenerator(session_factory=database_client), + generator=MessageMoreLikeThisGenerator( + session_factory=database_client, + completion_generator=CompletionAppGenerator( + retrieval=generation_runtime.retrieval, + records=generation_runtime.chat_records, + annotations=generation_runtime.annotation_replies, + ), + ), ), message_suggested_questions=MessageSuggestedQuestionsService( queries=suggested_questions, @@ -938,9 +946,6 @@ def init_app(app: Flask) -> None: app.extensions["knowledge_fs_node_dependencies"] = build_knowledge_fs_node_dependencies ToolManager.register_hardcoded_provider_class(KnowledgeFSProvider) - # Legacy workflow execution and saved trace extras still resolve providers here. - LegacyToolManager.register_hardcoded_provider_class(KnowledgeFSProvider) - LegacyToolManager.configure_credential_resolver(resolve_tool_credentials) ToolManager.configure_credential_resolver(resolve_tool_credentials) services = build_application_services( database_client=get_session_maker(), diff --git a/api/models/workflow.py b/api/models/workflow.py index f6ec17173f0e99..476cc725fe7121 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -24,7 +24,7 @@ ) from sqlalchemy.orm import Mapped, mapped_column -from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE, is_trigger_node_type +from core.trigger.constants import is_trigger_node_type from core.workflow.environment_variables import load_environment_variables from core.workflow.human_input_adapter import adapt_node_config_for_graph from core.workflow.llm_environment_variable import LLMEnvironmentVariable, dump_environment_variable @@ -50,7 +50,6 @@ BuiltinNodeTypes, NodeType, WorkflowExecutionStatus, - WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus, ) from graphon.file import File @@ -1174,41 +1173,6 @@ def execution_metadata_dict(self) -> dict[str, Any]: # cases where metadata is absent. return json.loads(self.execution_metadata) if self.execution_metadata else {} - @property - def extras(self) -> dict[str, Any]: - from core.tools.entities.tool_entities import ToolProviderType - from core.tools.tool_manager import ToolManager - from core.trigger.trigger_manager import TriggerManager - - extras: dict[str, Any] = {} - execution_metadata = self.execution_metadata_dict - if execution_metadata: - if self.node_type == BuiltinNodeTypes.TOOL and "tool_info" in execution_metadata: - tool_info: dict[str, Any] = execution_metadata["tool_info"] - provider_type = ToolProviderType(tool_info["provider_type"]) - if provider_type == ToolProviderType.WORKFLOW: - extras["workflow_tool"] = True - extras["icon"] = ToolManager.get_tool_icon( - tenant_id=self.tenant_id, - provider_type=provider_type, - provider_id=tool_info["provider_id"], - ) - elif self.node_type == BuiltinNodeTypes.DATASOURCE and "datasource_info" in execution_metadata: - datasource_info = execution_metadata["datasource_info"] - extras["icon"] = datasource_info.get("icon") - elif ( - self.node_type == TRIGGER_PLUGIN_NODE_TYPE - and WorkflowNodeExecutionMetadataKey.TRIGGER_INFO in execution_metadata - ): - trigger_info = execution_metadata[WorkflowNodeExecutionMetadataKey.TRIGGER_INFO] or {} - provider_id = trigger_info.get("provider_id") - if provider_id: - extras["icon"] = TriggerManager.get_trigger_plugin_icon( - tenant_id=self.tenant_id, - provider_id=provider_id, - ) - return extras - def _get_offload_by_type(self, type_: ExecutionOffLoadType) -> "WorkflowNodeExecutionOffload | None": return next(iter([i for i in self.offload_data if i.type_ == type_]), None) diff --git a/api/services/agent/workspace_service.py b/api/services/agent/workspace_service.py index 3e25cc12ca1fd5..70fa4cf3e90ad1 100644 --- a/api/services/agent/workspace_service.py +++ b/api/services/agent/workspace_service.py @@ -421,22 +421,6 @@ def collect_retired_workspace(cls, *, tenant_id: str, workspace_id: str) -> None ) session.commit() - @staticmethod - def validate_binding_generation( - binding: AgentWorkspaceBinding, - *, - base_home_snapshot_id: str | None, - agent_config_version_id: str, - agent_config_version_kind: AgentConfigVersionKind, - ) -> None: - """Validate a legacy execution caller against the repository-owned generation policy.""" - AgentWorkspaceRepository.validate_binding_generation( - binding, - base_home_snapshot_id=base_home_snapshot_id, - agent_config_version_id=agent_config_version_id, - agent_config_version_kind=agent_config_version_kind, - ) - @staticmethod def _client() -> Client: base_url = dify_config.AGENT_BACKEND_BASE_URL diff --git a/api/services/agent_tool_inner_service.py b/api/services/agent_tool_inner_service.py deleted file mode 100644 index 085f37d471c5f9..00000000000000 --- a/api/services/agent_tool_inner_service.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Service layer for Agent-backend core tool invocations. - -The `dify.core.tools` layer executes inside the API service boundary so Dify -Agent can expose API-owned tool providers (`plugin`, `builtin`, `api`, -`workflow`, and `mcp`) without learning their storage, credential, or runtime -internals. Production workflow and agent-app request builders still keep -existing plugin tool configs on the direct `dify.plugin.tools` route by -default; this service is the lower-level API execution path used when a caller -explicitly submits a `dify.core.tools` declaration. The service validates -tenant/app ownership, reuses `ToolManager.get_agent_tool_runtime(...)` to build -the runtime with `ToolInvokeFrom.AGENT`, invokes through -`ToolEngine.generic_invoke(...)`, and formats observations through the shared -public `ToolEngine.tool_response_to_str` helper so Agent and workflow paths -stay consistent. -""" - -from __future__ import annotations - -from sqlalchemy.orm import Session - -from core.agent.entities import AgentToolEntity -from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, UserFrom -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.tools.entities.tool_entities import ToolInvokeFrom, ToolProviderType -from core.tools.errors import ( - ToolInvokeError, - ToolNotFoundError, - ToolNotSupportedError, - ToolParameterValidationError, - ToolProviderCredentialValidationError, - ToolProviderNotFoundError, -) -from core.tools.tool_engine import ToolEngine -from core.tools.tool_manager import ToolManager -from core.tools.utils.message_transformer import ToolFileMessageTransformer -from models.model import App -from services.credentials.refresh import CredentialRefreshUnavailableError -from services.entities.agent_tool_inner import AgentToolInvokeRequest, AgentToolInvokeResponse -from services.errors.agent_tool_inner import AgentToolInnerServiceError -from services.workflow_run_agg import WorkflowRunAgg - - -class AgentToolInnerService: - """Invoke one API-owned Agent tool declaration, including explicit plugin-via-core calls.""" - - def invoke(self, request: AgentToolInvokeRequest, *, session: Session) -> AgentToolInvokeResponse: - app = session.get(App, request.caller.app_id) - if app is None: - raise AgentToolInnerServiceError( - error_code="app_not_found", - description="App not found.", - status_code=404, - ) - if app.tenant_id != request.caller.tenant_id: - raise AgentToolInnerServiceError( - error_code="app_tenant_mismatch", - description="App does not belong to the caller tenant.", - status_code=403, - ) - - agent_tool = AgentToolEntity( - provider_type=ToolProviderType.value_of(request.tool.provider_type), - provider_id=request.tool.provider_id, - tool_name=request.tool.tool_name, - tool_parameters=dict(request.tool.runtime_parameters), - credential_id=request.tool.credential_id, - ) - try: - dify_run_context = DifyRunContext( - tenant_id=request.caller.tenant_id, - app_id=request.caller.app_id, - user_id=request.caller.user_id, - user_from=UserFrom(request.caller.user_from), - invoke_from=InvokeFrom.value_of(request.caller.invoke_from), - trace_session_id=request.caller.trace_session_id, - ) - tool_runtime = ToolManager.get_agent_tool_runtime( - execution_driver=WorkflowRunAgg.run, - tenant_id=request.caller.tenant_id, - app_id=request.caller.app_id, - agent_tool=agent_tool, - user_id=request.caller.user_id, - invoke_from=dify_run_context.invoke_from, - variable_pool=None, - allow_file_parameters=True, - use_default_for_missing_form_parameters=True, - ) - tool_runtime.runtime.dify_run_context = dify_run_context - tool_runtime.runtime.tool_invoke_from = ( - ToolInvokeFrom.WORKFLOW if request.caller.workflow_id is not None else ToolInvokeFrom.AGENT - ) - messages = ToolEngine.generic_invoke( - session=session, - tool=tool_runtime, - tool_parameters=dict(request.tool.tool_parameters), - user_id=request.caller.user_id, - workflow_tool_callback=DifyWorkflowCallbackHandler(), - workflow_call_depth=0, - conversation_id=request.caller.conversation_id, - app_id=request.caller.app_id, - ) - transformed_messages = list( - ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=messages, - user_id=request.caller.user_id, - tenant_id=request.caller.tenant_id, - conversation_id=request.caller.conversation_id, - ) - ) - except CredentialRefreshUnavailableError: - raise - except ToolProviderNotFoundError as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_declaration_not_found", - description=str(exc), - status_code=404, - ) from exc - except ToolProviderCredentialValidationError as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_credential_invalid", - description=str(exc), - status_code=422, - ) from exc - except ToolParameterValidationError as exc: - raise AgentToolInnerServiceError( - error_code="tool_parameters_invalid", - description=str(exc), - status_code=422, - ) from exc - except (ToolInvokeError, ToolNotFoundError, ToolNotSupportedError) as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_invoke_failed", - description=str(exc), - status_code=422, - ) from exc - except ValueError as exc: - raise _map_value_error(exc) from exc - except Exception as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_invoke_unexpected_error", - description=str(exc), - status_code=500, - ) from exc - - return AgentToolInvokeResponse( - messages=[message.model_dump(mode="json") for message in transformed_messages], - observation=ToolEngine.tool_response_to_str(transformed_messages), - metadata={ - "provider_type": request.tool.provider_type, - "provider_id": request.tool.provider_id, - "tool_name": request.tool.tool_name, - }, - ) - - -def _map_value_error(error: ValueError) -> AgentToolInnerServiceError: - description = str(error) - if description == "app not found": - return AgentToolInnerServiceError( - error_code="app_not_found", - description="App not found.", - status_code=404, - ) - return AgentToolInnerServiceError( - error_code="agent_tool_invoke_failed", - description=description, - status_code=422, - ) diff --git a/api/services/app/generation/adapters/completion.py b/api/services/app/generation/adapters/completion.py index 05140779e8d518..9013d7d6711132 100644 --- a/api/services/app/generation/adapters/completion.py +++ b/api/services/app/generation/adapters/completion.py @@ -22,7 +22,7 @@ from core.ops.ops_trace_manager import TraceQueueManager from factories import file_factory from graphon.model_runtime.errors.invoke import InvokeAuthorizationError -from models import Account, App, EndUser +from models import Account, App, Conversation, EndUser, Message from models.annotation_reply import AnnotationReplies from services.app.generation.adapters.completion_runner import CompletionAppRunner from services.app.generation.message_records import MessageBasedAppGenerator @@ -183,50 +183,71 @@ def generate( session.expunge_all() session.close() - # init generate records - (conversation, message) = self._init_generate_records( - application_generate_entity, - ) - - # 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, - ) - - context = contextvars.copy_context() + return self.generate_from_entity(application_generate_entity, user=user) - # new thread with request context - @copy_current_request_context - def worker_with_context(): - return context.run( - self._generate_worker, - flask_app=current_app._get_current_object(), # type: ignore - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - message_id=message.id, - conversation_id=conversation.id, - ) + def generate_from_entity( + self, + application_generate_entity: CompletionAppGenerateEntity, + *, + user: Account | EndUser, + ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any]]: + """Persist one prepared completion before starting provider work. - worker_thread = threading.Thread(target=worker_with_context) + The caller owns preparation of the historical configuration, inputs, + files and admitted actor. The injected records port commits those + values in its own transaction; a failed record write therefore rolls + back before a worker can start. + """ + conversation, message = self._init_generate_records(application_generate_entity) + return self._start_generation( + application_generate_entity, + conversation=conversation, + message=message, + user=user, + ) - worker_thread.start() + def _start_generation( + self, + application_generate_entity: CompletionAppGenerateEntity, + *, + conversation: Conversation, + message: Message, + user: Account | EndUser, + ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any]]: + 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, + ) + context = contextvars.copy_context() - # return response or stream generator - response = self._handle_response( + @copy_current_request_context + def worker_with_context() -> None: + context.run( + self._generate_worker, + flask_app=current_app._get_current_object(), # type: ignore application_generate_entity=application_generate_entity, queue_manager=queue_manager, - conversation=conversation, - message=message, - user=user, - stream=streaming, + message_id=message.id, + conversation_id=conversation.id, ) - return CompletionAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) + threading.Thread(target=worker_with_context).start() + response = self._handle_response( + application_generate_entity=application_generate_entity, + queue_manager=queue_manager, + conversation=conversation, + message=message, + user=user, + stream=application_generate_entity.stream, + ) + return CompletionAppGenerateResponseConverter.convert( + response=response, + invoke_from=application_generate_entity.invoke_from, + ) def _generate_worker( self, @@ -376,48 +397,4 @@ def generate_more_like_this( session.expunge_all() session.close() - - # init generate records - (conversation, message) = self._init_generate_records( - application_generate_entity, - ) - - # 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, - ) - - context = contextvars.copy_context() - - # new thread with request context - @copy_current_request_context - def worker_with_context(): - return context.run( - self._generate_worker, - flask_app=current_app._get_current_object(), # type: ignore - application_generate_entity=application_generate_entity, - queue_manager=queue_manager, - message_id=message.id, - conversation_id=conversation.id, - ) - - worker_thread = threading.Thread(target=worker_with_context) - - 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=stream, - ) - - return CompletionAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) + return self.generate_from_entity(application_generate_entity, user=user) diff --git a/api/services/knowledge/dataset_service.py b/api/services/knowledge/dataset_service.py index 1dfd6e683e7891..f38130d976002a 100644 --- a/api/services/knowledge/dataset_service.py +++ b/api/services/knowledge/dataset_service.py @@ -1296,11 +1296,6 @@ class DocumentService: } DOCUMENT_BATCH_DOWNLOAD_ZIP_FILENAME_EXTENSION = ".zip" - @staticmethod - def get_documents_position(dataset_id: str, session: Session) -> int: - """Return the next document position for legacy pipeline execution callers.""" - return next_document_position(dataset_id, session) - @classmethod def normalize_display_status(cls, status: str | None) -> str | None: if not status: diff --git a/api/services/message_more_like_this_generator.py b/api/services/message_more_like_this_generator.py index b55092b50d9225..008d893bb4c4f5 100644 --- a/api/services/message_more_like_this_generator.py +++ b/api/services/message_more_like_this_generator.py @@ -17,7 +17,6 @@ from core.app.app_config.entities import EasyUIBasedAppModelConfigFrom from core.app.app_config.features.file_upload.manager import FileUploadConfigManager from core.app.apps.completion.app_config_manager import CompletionAppConfigManager -from core.app.apps.completion.app_generator import CompletionAppGenerator from core.app.entities.app_invoke_entities import CompletionAppGenerateEntity, InvokeFrom, UserFrom from core.app.file_access import DatabaseFileAccessController, FileAccessScope, bind_file_access_scope from extensions.ext_database import db @@ -26,8 +25,9 @@ from graphon.file.constants import FILE_MODEL_IDENTITY from libs.orjson import orjson_dumps from libs.stream import close_stream -from models.model import AppMode, AppModelConfigDict +from models.model import Account, AppMode, AppModelConfigDict, EndUser from models.utils.file_input_compat import build_file_from_stored_mapping +from services.app.generation.adapters.completion import CompletionAppGenerator from services.entities.message_entities import MessageAccount, MessageActor from services.message_more_like_this_service import ( MoreLikeThisFile, @@ -65,8 +65,14 @@ def close(self) -> None: class MessageMoreLikeThisGenerator: - def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + def __init__( + self, + *, + session_factory: sessionmaker[Session], + completion_generator: CompletionAppGenerator, + ) -> None: self._session_factory = session_factory + self._completion_generator = completion_generator self._file_access_controller = DatabaseFileAccessController() def generate( @@ -127,9 +133,22 @@ def generate( invoke_from=invoke_from, extras={}, ) - response = CompletionAppGenerator().generate_from_entity(entity, session_factory=self._session_factory) + prepared_user = self._load_prepared_user(actor) + response = self._completion_generator.generate_from_entity(entity, user=prepared_user) return response if isinstance(response, Mapping) else _MoreLikeThisEventStream(response) + def _load_prepared_user(self, actor: MessageActor) -> Account | EndUser: + """Load the admitted actor required by the shared response boundary.""" + with self._session_factory() as session: + if isinstance(actor, MessageAccount): + user = session.get(Account, actor.account_id) + else: + user = session.get(EndUser, actor.end_user_id) + if user is None: + raise ValueError("The admitted more-like-this actor no longer exists.") + session.expunge(user) + return user + def _file_mapping(file: MoreLikeThisFile) -> dict[str, object]: mapping: dict[str, object] = { diff --git a/api/services/snippet_generate_service.py b/api/services/snippet_generate_service.py index 0480638c64b572..8cca7ba58cba31 100644 --- a/api/services/snippet_generate_service.py +++ b/api/services/snippet_generate_service.py @@ -24,8 +24,6 @@ from dataclasses import replace from typing import Any, Protocol, cast -from sqlalchemy.orm import make_transient - from core.app.app_config.features.file_upload.manager import FileUploadConfigManager from core.app.entities.app_invoke_entities import InvokeFrom from core.app.file_access import DatabaseFileAccessController @@ -230,18 +228,6 @@ def _ensure_start_node(cls, workflow: WorkflowSnapshot, snippet: CustomizedSnipp return replace(workflow, graph=json.dumps(modified_graph)) - @classmethod - def ensure_start_node_for_worker(cls, workflow: Workflow, snippet: CustomizedSnippet) -> Workflow: - """Inject the virtual Start node after a worker reloads a snippet workflow.""" - snapshot = workflow_snapshot(workflow) - updated_snapshot = cls._ensure_start_node(snapshot, snippet) - if updated_snapshot.graph == snapshot.graph: - return workflow - - make_transient(workflow) - workflow.graph = updated_snapshot.graph - return workflow - @classmethod def _inject_virtual_start_node( cls, diff --git a/api/services/tools/legacy_api_tools_manage_service.py b/api/services/tools/legacy_api_tools_manage_service.py deleted file mode 100644 index 416045f80569ea..00000000000000 --- a/api/services/tools/legacy_api_tools_manage_service.py +++ /dev/null @@ -1,570 +0,0 @@ -import json -import logging -from typing import Any, Literal, TypedDict, cast - -from pydantic import TypeAdapter -from sqlalchemy import select -from sqlalchemy.orm import sessionmaker - -from core.entities.provider_entities import ProviderConfig, ProviderConfigType -from core.helper import ssrf_proxy -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ApiProviderSchemaType, -) -from core.tools.errors import ApiToolProviderNotFoundError -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.tool_manager import ToolManager -from core.tools.utils.encryption import create_tool_provider_encrypter -from core.tools.utils.parser import ApiBasedToolSchemaParser -from extensions.ext_database import db -from models.tools import ApiToolProvider -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class ApiSchemaParseResult(TypedDict): - schema_type: str - parameters_schema: list[dict[str, Any]] - credentials_schema: list[dict[str, Any]] - warning: dict[str, str] - - -class ApiToolPreviewResult(TypedDict, total=False): - result: str - error: str - - -class RemoteSchemaResult(TypedDict): - schema: str - - -class SimpleSuccessResult(TypedDict): - result: Literal["success"] - - -def _dump_api_tool_bundles(tool_bundles: list[ApiToolBundle]) -> list[dict[str, Any]]: - return cast(list[dict[str, Any]], TypeAdapter(list[ApiToolBundle]).dump_python(tool_bundles, mode="json")) - - -def _dump_provider_configs(configs: list[ProviderConfig]) -> list[dict[str, Any]]: - return cast(list[dict[str, Any]], TypeAdapter(list[ProviderConfig]).dump_python(configs, mode="json")) - - -class ApiToolManageService: - @staticmethod - def parser_api_schema(schema: str) -> ApiSchemaParseResult: - """ - parse api schema to tool bundle - """ - try: - warnings: dict[str, str] = {} - try: - tool_bundles, schema_type = ApiBasedToolSchemaParser.auto_parse_to_tool_bundle(schema, warning=warnings) - except Exception as e: - raise ValueError(f"invalid schema: {str(e)}") from e - - credentials_schema = [ - ProviderConfig( - name="auth_type", - type=ProviderConfigType.SELECT, - required=True, - default="none", - options=[ - ProviderConfig.Option(value="none", label=I18nObject(en_US="None", zh_Hans="无")), - ProviderConfig.Option(value="api_key", label=I18nObject(en_US="Api Key", zh_Hans="Api Key")), - ], - placeholder=I18nObject(en_US="Select auth type", zh_Hans="选择认证方式"), - ), - ProviderConfig( - name="api_key_header", - type=ProviderConfigType.TEXT_INPUT, - required=False, - placeholder=I18nObject(en_US="Enter api key header", zh_Hans="输入 api key header，如：X-API-KEY"), - default="api_key", - help=I18nObject(en_US="HTTP header name for api key", zh_Hans="HTTP 头部字段名，用于传递 api key"), - ), - ProviderConfig( - name="api_key_value", - type=ProviderConfigType.TEXT_INPUT, - required=False, - placeholder=I18nObject(en_US="Enter api key", zh_Hans="输入 api key"), - default="", - ), - ] - - return cast( - ApiSchemaParseResult, - { - "schema_type": schema_type.value, - "parameters_schema": _dump_api_tool_bundles(tool_bundles), - "credentials_schema": _dump_provider_configs(credentials_schema), - "warning": warnings, - }, - ) - except Exception as e: - raise ValueError(f"invalid schema: {str(e)}") from e - - @staticmethod - def convert_schema_to_tool_bundles( - schema: str, extra_info: dict[str, Any] | None = None - ) -> tuple[list[ApiToolBundle], ApiProviderSchemaType]: - """ - convert schema to tool bundles - - :return: the list of tool bundles, description - """ - try: - return ApiBasedToolSchemaParser.auto_parse_to_tool_bundle(schema, extra_info=extra_info) - except Exception as e: - raise ValueError(f"invalid schema: {str(e)}") from e - - @staticmethod - def create_api_tool_provider( - user_id: str, - tenant_id: str, - provider_name: str, - icon: dict[str, Any], - credentials: dict[str, Any], - schema_type: ApiProviderSchemaType, - schema: str, - privacy_policy: str, - custom_disclaimer: str, - labels: list[str], - ) -> SimpleSuccessResult: - """ - Create a new API tool provider. - - :param user_id: The ID of the user creating the provider. - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The name of the API tool provider. - :param icon: The icon configuration for the provider. - :param credentials: The credentials for the provider. - :param schema_type: The type of schema (e.g., OpenAPI). - :param schema: The raw schema string. - :param privacy_policy: The privacy policy URL or text. - :param custom_disclaimer: Custom disclaimer text. - :param labels: A list of labels for the provider. - :return: A dictionary indicating the result status. - """ - - provider_name = provider_name.strip() - - # check if the provider exists - # Create new session with automatic transaction management - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider: ApiToolProvider | None = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is not None: - raise ValueError(f"provider {provider_name} already exists") - - # parse openapi to tool bundle - extra_info: dict[str, str] = {} - # extra info like description will be set here - tool_bundles, schema_type = ApiToolManageService.convert_schema_to_tool_bundles(schema, extra_info) - - if len(tool_bundles) > 100: - raise ValueError("the number of apis should be less than 100") - - # create API tool provider - api_tool_provider = ApiToolProvider( - tenant_id=tenant_id, - user_id=user_id, - name=provider_name, - icon=json.dumps(icon), - schema=schema, - description=extra_info.get("description", ""), - schema_type_str=schema_type, - tools_str=json.dumps(_dump_api_tool_bundles(tool_bundles)), - credentials_str="{}", - privacy_policy=privacy_policy, - custom_disclaimer=custom_disclaimer, - ) - - if "auth_type" not in credentials: - raise ValueError("auth_type is required") - - # get auth type, none or api key - auth_type = ApiProviderAuthType.value_of(credentials["auth_type"]) - - # create provider entity - provider_controller = ApiToolProviderController.from_db(api_tool_provider, auth_type, session=_session) - # load tools into provider entity - provider_controller.load_bundled_tools(tool_bundles) - - # encrypt credentials - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=provider_controller, - ) - api_tool_provider.credentials_str = json.dumps(encrypter.encrypt(credentials)) - - _session.add(api_tool_provider) - - # update labels - ToolLabelManager.update_tool_labels(provider_controller, labels, _session) - - return {"result": "success"} - - @staticmethod - def get_api_tool_provider_remote_schema(user_id: str, tenant_id: str, url: str) -> RemoteSchemaResult: - """ - get api tool provider remote schema - """ - headers = { - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko)" - " Chrome/120.0.0.0 Safari/537.36 Edg/120.0.0.0", - "Accept": "*/*", - } - - try: - response = ssrf_proxy.get(url, headers=headers, timeout=10) - if response.status_code != 200: - raise ValueError(f"Got status code {response.status_code}") - schema = response.text - - # try to parse schema, avoid SSRF attack - ApiToolManageService.parser_api_schema(schema) - except Exception as e: - logger.exception("parse api schema error") - raise ValueError("invalid schema, please check the url you provided") from e - - return {"schema": schema} - - @staticmethod - def list_api_tool_provider_tools(user_id: str, tenant_id: str, provider_name: str) -> list[ToolApiEntity]: - """ - List tools provided by a specific API tool provider. - - :param user_id: The ID of the user requesting the list. - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The name of the API tool provider. - :return: A list of ToolApiEntity objects. - """ - - # create new session with automatic transaction management - provider: ApiToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is None: - raise ValueError(f"you have not added provider {provider_name}") - - controller = ToolTransformService.api_provider_to_controller(db_provider=provider) - labels = ToolLabelManager.get_tool_labels(controller) - - return [ - ToolTransformService.convert_tool_entity_to_api_entity( - tool_bundle, - tenant_id=tenant_id, - labels=labels, - ) - for tool_bundle in provider.tools - ] - - @staticmethod - def update_api_tool_provider( - user_id: str, - tenant_id: str, - provider_name: str, - original_provider: str, - icon: dict[str, Any], - credentials: dict[str, Any], - _schema_type: ApiProviderSchemaType, - schema: str, - privacy_policy: str | None, - custom_disclaimer: str, - labels: list[str], - ) -> SimpleSuccessResult: - """ - Update an existing API tool provider. - - :param user_id: The ID of the user updating the provider. - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The new name of the API tool provider. - :param original_provider: The original name of the API tool provider. - :param icon: The icon configuration for the provider. - :param credentials: The credentials for the provider. - :param _schema_type: The type of schema (e.g., OpenAPI). - :param schema: The raw schema string. - :param privacy_policy: The privacy policy URL or text. - :param custom_disclaimer: Custom disclaimer text. - :param labels: A list of labels for the provider. - :return: A dictionary indicating the result status. - """ - - provider_name = provider_name.strip() - - # check if the provider exists - # create new session with automatic transaction management - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider: ApiToolProvider | None = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == original_provider, - ) - .limit(1) - ) - - if provider is None: - raise ApiToolProviderNotFoundError(provider_name=original_provider, tenant_id=tenant_id) - - # parse openapi to tool bundle - extra_info: dict[str, str] = {} - # extra info like description will be set here - tool_bundles, schema_type = ApiToolManageService.convert_schema_to_tool_bundles(schema, extra_info) - - # update db provider - provider.name = provider_name - provider.icon = json.dumps(icon) - provider.schema = schema - provider.description = extra_info.get("description", "") - provider.schema_type_str = schema_type - provider.tools_str = json.dumps(_dump_api_tool_bundles(tool_bundles)) - provider.privacy_policy = privacy_policy - provider.custom_disclaimer = custom_disclaimer - - if "auth_type" not in credentials: - raise ValueError("auth_type is required") - - # get auth type, none or api key - auth_type = ApiProviderAuthType.value_of(credentials["auth_type"]) - - # create provider entity - provider_controller = ApiToolProviderController.from_db(provider, auth_type, session=_session) - # load tools into provider entity - provider_controller.load_bundled_tools(tool_bundles) - - # get original credentials if exists - encrypter, cache = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=provider_controller, - ) - - original_credentials = encrypter.decrypt(provider.credentials) - masked_credentials = encrypter.mask_plugin_credentials(original_credentials) - - # check if the credential has changed, save the original credential - for name, value in credentials.items(): - if name in masked_credentials and value == masked_credentials[name]: - credentials[name] = original_credentials[name] - - credentials = dict(encrypter.encrypt(credentials)) - provider.credentials_str = json.dumps(credentials) - - _session.add(provider) - - # update labels - ToolLabelManager.update_tool_labels(provider_controller, labels, _session) - - # delete cache - cache.delete() - - return {"result": "success"} - - @staticmethod - def delete_api_tool_provider(user_id: str, tenant_id: str, provider_name: str) -> SimpleSuccessResult: - """ - Delete an API tool provider. - - :param user_id: The ID of the user performing the deletion operation. - :param tenant_id: The ID of the workspace/tenant where the provider belongs. - :param provider_name: The unique name of the API tool provider to be deleted. - :raises ValueError: If the specified provider does not exist in the tenant. - :return: A dictionary indicating the result status. - """ - - # create new session with automatic transaction management - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider: ApiToolProvider | None = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is None: - raise ValueError(f"you have not added provider {provider_name}") - - _session.delete(provider) - - return {"result": "success"} - - @staticmethod - def get_api_tool_provider(user_id: str, tenant_id: str, provider: str) -> dict[str, Any]: - """ - Get API tool provider details. - - :param user_id: The ID of the user requesting the provider. - :param tenant_id: The ID of the workspace/tenant. - :param provider: The name of the API tool provider. - :return: A dictionary containing the provider details. - """ - return ToolManager.user_get_api_provider(provider=provider, tenant_id=tenant_id) - - @staticmethod - def test_api_tool_preview( - tenant_id: str, - provider_name: str, - tool_name: str, - credentials: dict[str, Any], - parameters: dict[str, Any], - schema_type: ApiProviderSchemaType | str, - schema: str, - ) -> ApiToolPreviewResult: - """ - Test an API tool before adding the API tool provider. - - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The name of the API tool provider. - :param tool_name: The name of the specific tool to test. - :param credentials: The credentials for the provider. - :param parameters: The parameters to pass to the tool. - :param schema_type: The type of schema (e.g., OpenAPI). - :param schema: The raw schema string. - :return: A dictionary containing the result or error message. - """ - - if schema_type not in [member.value for member in ApiProviderSchemaType]: - raise ValueError(f"invalid schema type {schema_type}") - - try: - tool_bundles, _ = ApiBasedToolSchemaParser.auto_parse_to_tool_bundle(schema) - except Exception as e: - raise ValueError("invalid schema") from e - - # get tool bundle - tool_bundle = next(filter(lambda tb: tb.operation_id == tool_name, tool_bundles), None) - if tool_bundle is None: - raise ValueError(f"invalid tool name {tool_name}") - - # create new session with automatic transaction management to get the provider - provider: ApiToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is None: - # create a fake db provider - provider = ApiToolProvider( - tenant_id="", - user_id="", - name="", - icon="", - schema=schema, - description="", - schema_type_str=ApiProviderSchemaType.OPENAPI, - tools_str=json.dumps(_dump_api_tool_bundles(tool_bundles)), - credentials_str=json.dumps(credentials), - ) - - if "auth_type" not in credentials: - raise ValueError("auth_type is required") - - # get auth type, none or api key - auth_type = ApiProviderAuthType.value_of(credentials["auth_type"]) - - # create provider entity - provider_controller = ApiToolProviderController.from_db(provider, auth_type, session=db.session()) - # load tools into provider entity - provider_controller.load_bundled_tools(tool_bundles) - - # decrypt credentials - if provider.id: - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=provider_controller, - ) - decrypted_credentials = encrypter.decrypt(credentials) - # check if the credential has changed, save the original credential - masked_credentials = encrypter.mask_plugin_credentials(decrypted_credentials) - for name, value in credentials.items(): - if name in masked_credentials and value == masked_credentials[name]: - credentials[name] = decrypted_credentials[name] - - try: - provider_controller.validate_credentials_format(credentials) - # get tool - tool = provider_controller.get_tool(tool_name) - tool = tool.fork_tool_runtime( - runtime=ToolRuntime( - credentials=credentials, - tenant_id=tenant_id, - ) - ) - result = tool.validate_credentials(credentials, parameters) - except Exception as e: - return {"error": str(e)} - - return {"result": result or "empty response"} - - @staticmethod - def list_api_tools(tenant_id: str) -> list[ToolProviderApiEntity]: - """ - List all API tools for a specific tenant. - - :param tenant_id: The ID of the workspace/tenant. - :return: A list of ToolProviderApiEntity objects. - """ - # get all api providers - # create new session with automatic transaction management - providers: list[ApiToolProvider] = [] - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - providers = list( - _session.scalars(select(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id)).all() - ) - - result: list[ToolProviderApiEntity] = [] - for provider in providers: - # convert provider controller to user provider - provider_controller = ToolTransformService.api_provider_to_controller(db_provider=provider) - labels = ToolLabelManager.get_tool_labels(provider_controller) - user_provider = ToolTransformService.api_provider_to_user_provider( - provider_controller, db_provider=provider, decrypt_credentials=True, session=db.session() - ) - user_provider.labels = labels - - # add icon - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_provider) - - tools = provider_controller.get_tools(tenant_id=tenant_id) - - for tool in tools or []: - user_provider.tools.append( - ToolTransformService.convert_tool_entity_to_api_entity( - tenant_id=tenant_id, tool=tool, labels=labels - ) - ) - - result.append(user_provider) - - return result diff --git a/api/services/tools/legacy_builtin_tools_manage_service.py b/api/services/tools/legacy_builtin_tools_manage_service.py deleted file mode 100644 index 443b16820e69fb..00000000000000 --- a/api/services/tools/legacy_builtin_tools_manage_service.py +++ /dev/null @@ -1,806 +0,0 @@ -import json -import logging -from collections.abc import Mapping -from pathlib import Path -from typing import Any - -from sqlalchemy import delete, exists, func, select, update -from sqlalchemy.orm import Session, sessionmaker - -from configs import dify_config -from constants import HIDDEN_VALUE, UNKNOWN_VALUE -from core.helper.name_generator import generate_incremental_name -from core.helper.position_helper import is_filtered -from core.helper.provider_cache import NoOpProviderCredentialCache, ToolProviderCredentialsCache -from core.plugin.entities.plugin_daemon import CredentialType -from core.plugin.plugin_service import PluginService -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort -from core.tools.entities.api_entities import ( - ToolApiEntity, - ToolProviderApiEntity, - ToolProviderCredentialApiEntity, - ToolProviderCredentialInfoApiEntity, -) -from core.tools.errors import ToolProviderNotFoundError -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.tool_manager import ToolManager -from core.tools.utils.encryption import ProviderConfigCache, create_provider_encrypter -from core.tools.utils.system_encryption import decrypt_system_params -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from models.account import Account -from models.enums import PermissionEnum -from models.provider_ids import ToolProviderID -from models.tools import BuiltinToolProvider, ToolOAuthSystemClient, ToolOAuthTenantClient -from services.credentials.query import CredentialQuery, ToolCredentialRecord -from services.plugin.dispatch_identity import PluginOAuthClientConfig -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class BuiltinToolManageService: - __MAX_BUILTIN_TOOL_PROVIDER_COUNT__ = 100 - - @staticmethod - def delete_custom_oauth_client_params(tenant_id: str, provider: str): - """ - delete custom oauth client params - """ - tool_provider = ToolProviderID(provider) - with sessionmaker(bind=db.engine).begin() as session: - session.execute( - delete(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ) - .execution_options(synchronize_session=False) - ) - return {"result": "success"} - - @staticmethod - def get_builtin_tool_provider_oauth_client_schema(tenant_id: str, provider_name: str): - """ - get builtin tool provider oauth client schema - """ - provider = ToolManager.get_builtin_provider(provider_name, tenant_id) - verified = not isinstance(provider, PluginToolProviderController) or PluginService.is_plugin_verified( - tenant_id, provider.plugin_unique_identifier - ) - - is_oauth_custom_client_enabled = BuiltinToolManageService.is_oauth_custom_client_enabled( - tenant_id, provider_name - ) - is_system_oauth_params_exists = verified and BuiltinToolManageService.is_oauth_system_client_exists( - provider_name - ) - result = { - "schema": provider.get_oauth_client_schema(), - "is_oauth_custom_client_enabled": is_oauth_custom_client_enabled, - "is_system_oauth_params_exists": is_system_oauth_params_exists, - "client_params": BuiltinToolManageService.get_custom_oauth_client_params(tenant_id, provider_name), - "redirect_uri": f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_name}/tool/callback", - } - return result - - @staticmethod - def list_builtin_tool_provider_tools(tenant_id: str, provider: str) -> list[ToolApiEntity]: - """ - list builtin tool provider tools - - :param tenant_id: the id of the tenant - :param provider: the name of the provider - - :return: the list of tools - """ - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - tools = provider_controller.get_tools() - - result: list[ToolApiEntity] = [] - for tool in tools or []: - result.append( - ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool, - tenant_id=tenant_id, - labels=ToolLabelManager.get_tool_labels(provider_controller), - ) - ) - - return result - - @staticmethod - def get_builtin_tool_provider_info(tenant_id: str, provider: str): - """ - get builtin tool provider info - """ - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - # check if user has added the provider - builtin_provider = BuiltinToolManageService.get_builtin_provider(provider, tenant_id) - if builtin_provider is None: - raise ValueError(f"you have not added provider {provider}") - - entity = ToolTransformService.builtin_provider_to_user_provider( - provider_controller=provider_controller, - db_provider=builtin_provider, - decrypt_credentials=True, - ) - - entity.original_credentials = {} - return entity - - @staticmethod - def list_builtin_provider_credentials_schema(provider_name: str, credential_type: CredentialType, tenant_id: str): - """ - list builtin provider credentials schema - - :param credential_type: credential type - :param provider_name: the name of the provider - :param tenant_id: the id of the tenant - :return: the list of tool providers - """ - provider = ToolManager.get_builtin_provider(provider_name, tenant_id) - return provider.get_credentials_schema_by_type(credential_type) - - @staticmethod - def update_builtin_tool_provider( - user_id: str, - tenant_id: str, - provider: str, - credential_id: str, - credentials: dict[str, Any] | None = None, - name: str | None = None, - ): - """ - update builtin tool provider - """ - cache: ProviderConfigCache | None = None - with sessionmaker(bind=db.engine).begin() as session: - # get if the provider exists - db_provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.id == credential_id, - ) - .limit(1) - ) - if db_provider is None: - raise ValueError(f"you have not added provider {provider}") - - try: - if CredentialType.of(db_provider.credential_type).is_editable() and credentials: - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller.need_credentials: - raise ValueError(f"provider {provider} does not need credentials") - - encrypter, cache = BuiltinToolManageService.create_tool_encrypter( - tenant_id, db_provider, provider, provider_controller - ) - - original_credentials = encrypter.decrypt(db_provider.credentials) - new_credentials: dict[str, Any] = { - key: value if value != HIDDEN_VALUE else original_credentials.get(key, UNKNOWN_VALUE) - for key, value in credentials.items() - } - - if CredentialType.of(db_provider.credential_type).is_validate_allowed(): - provider_controller.validate_credentials(user_id, new_credentials) - - # encrypt credentials - db_provider.encrypted_credentials = json.dumps(encrypter.encrypt(new_credentials)) - - # update name if provided - if name and name != db_provider.name: - # check if the name is already used - if session.scalar( - select( - exists().where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.name == name, - ) - ) - ): - raise ValueError(f"the credential name '{name}' is already used") - - db_provider.name = name - - # Visibility is immutable after creation — no update path. To change scope, - # create a new credential. partial-member access is handled by RBAC. - - except Exception as e: - raise ValueError(str(e)) from e - if cache is not None: - cache.delete() - return {"result": "success"} - - @staticmethod - def add_builtin_tool_provider( - user_id: str, - api_type: CredentialType, - tenant_id: str, - provider: str, - credentials: dict[str, Any], - expires_at: int = -1, - name: str | None = None, - visibility: str | None = None, - ): - """ - add builtin tool provider - """ - try: - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller.need_credentials: - raise ValueError(f"provider {provider} does not need credentials") - - # Validation may call external services, so it runs before the lock is taken - # to keep its duration independent of the lock timeout. - if CredentialType.of(api_type).is_validate_allowed(): - provider_controller.validate_credentials(user_id, credentials) - - lock = f"builtin_tool_provider_create_lock:{tenant_id}_{provider}" - # The transaction commits before the lock is released, so concurrent requests - # see this credential when they check the count limit and the name. - with redis_client.lock(lock, timeout=20): - with sessionmaker(bind=db.engine).begin() as session: - provider_count = ( - session.scalar( - select(func.count(BuiltinToolProvider.id)).where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - ) - ) - or 0 - ) - - # check if the provider count is reached the limit - if provider_count >= BuiltinToolManageService.__MAX_BUILTIN_TOOL_PROVIDER_COUNT__: - raise ValueError(f"you have reached the maximum number of providers for {provider}") - - # generate name if not provided - if name is None or name == "": - name = BuiltinToolManageService.generate_builtin_tool_provider_name( - session=session, tenant_id=tenant_id, provider=provider, credential_type=api_type - ) - else: - # check if the name is already used - if session.scalar( - select( - exists().where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.name == name, - ) - ) - ): - raise ValueError(f"the credential name '{name}' is already used") - - # create encrypter - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type(api_type) - ], - cache=NoOpProviderCredentialCache(), - ) - - visibility_enum = PermissionEnum(visibility) if visibility else PermissionEnum.ALL_TEAM - # Plugin credentials only expose only_me / all_team_members at creation; - # partial-member access is handled by workspace RBAC, not per-credential. - if visibility_enum == PermissionEnum.PARTIAL_TEAM: - raise ValueError("partial_members visibility is no longer supported for plugin credentials") - - db_provider = BuiltinToolProvider( - tenant_id=tenant_id, - user_id=user_id, - provider=provider, - encrypted_credentials=json.dumps(encrypter.encrypt(credentials)), - credential_type=api_type, - name=name, - expires_at=expires_at if expires_at is not None else -1, - visibility=visibility_enum, - ) - - session.add(db_provider) - session.flush() - credential_id = db_provider.id - except Exception as e: - raise ValueError(str(e)) from e - - return {"result": "success", "id": credential_id} - - @staticmethod - def create_tool_encrypter( - tenant_id: str, - db_provider: BuiltinToolProvider | ToolCredentialRecord, - provider: str, - provider_controller: BuiltinToolProviderController, - ): - encrypter, cache = create_provider_encrypter( - tenant_id=tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type(db_provider.credential_type) - ], - cache=ToolProviderCredentialsCache(tenant_id=tenant_id, provider=provider, credential_id=db_provider.id), - ) - return encrypter, cache - - @staticmethod - def generate_builtin_tool_provider_name( - tenant_id: str, provider: str, credential_type: CredentialType, *, session: Session - ) -> str: - db_providers = session.scalars( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.credential_type == credential_type, - ) - .order_by(BuiltinToolProvider.created_at.desc()) - ).all() - return generate_incremental_name( - [provider.name for provider in db_providers], - f"{credential_type.get_name()}", - ) - - @staticmethod - def get_builtin_tool_provider_credentials( - tenant_id: str, - provider_name: str, - user: Account | None = None, - include_credential_ids: list[str] | None = None, - *, - credential_query: CredentialQuery, - ) -> list[ToolProviderCredentialApiEntity]: - """ - get builtin tool provider credentials, filtered by visibility. - - ``user`` is used to filter the result list by per-credential visibility - (only_me / all_team_members / legacy partial_members). When ``None`` the - query returns every credential for the tenant — meant for internal / - background callers that don't act on behalf of a specific user. - - ``include_credential_ids`` lets callers request specific credential IDs that should be - returned even if the visibility filter would normally hide them (e.g. an only_me credential - owned by another member which the current workflow/agent node still references). Those - rows are marked with ``from_other_member=True`` so the UI can render them as - borrowed-from-teammate (selectable but not editable). - """ - providers = credential_query.list_tools( - workspace_id=tenant_id, - provider=provider_name, - actor_id=user.id if user else None, - include_credential_ids=include_credential_ids or (), - ) - if not providers: - return [] - provider_controller = ToolManager.get_builtin_provider(providers[0].provider, tenant_id) - credentials: list[ToolProviderCredentialApiEntity] = [] - for provider in providers: - encrypter, _ = BuiltinToolManageService.create_tool_encrypter( - tenant_id, provider, provider.provider, provider_controller - ) - decrypted = encrypter.mask_plugin_credentials(encrypter.decrypt(provider.credentials)) - credentials.append( - ToolProviderCredentialApiEntity( - id=provider.id, - name=provider.name, - provider=provider.provider, - credential_type=provider.credential_type, - is_default=provider.is_default, - credentials=dict(decrypted), - visibility=provider.visibility, - created_by=provider.created_by, - partial_member_list=list(provider.partial_member_ids), - from_other_member=provider.from_other_member, - ) - ) - return credentials - - @staticmethod - def get_builtin_tool_provider_credential_info( - tenant_id: str, - provider: str, - user: Account | None = None, - include_credential_ids: list[str] | None = None, - *, - credential_query: CredentialQuery, - ) -> ToolProviderCredentialInfoApiEntity: - """ - get builtin tool provider credential info - """ - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - supported_credential_types = provider_controller.get_supported_credential_types() - credentials = BuiltinToolManageService.get_builtin_tool_provider_credentials( - tenant_id, - provider, - user=user, - include_credential_ids=include_credential_ids, - credential_query=credential_query, - ) - credential_info = ToolProviderCredentialInfoApiEntity( - supported_credential_types=supported_credential_types, - is_oauth_custom_client_enabled=BuiltinToolManageService.is_oauth_custom_client_enabled(tenant_id, provider), - credentials=credentials, - ) - - return credential_info - - @staticmethod - def delete_builtin_tool_provider(tenant_id: str, provider: str, credential_id: str): - """ - delete tool provider - """ - with sessionmaker(bind=db.engine).begin() as session: - db_provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.id == credential_id, - ) - .limit(1) - ) - - if db_provider is None: - raise ValueError(f"you have not added provider {provider}") - - session.delete(db_provider) - - # delete cache - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - _, cache = BuiltinToolManageService.create_tool_encrypter( - tenant_id, db_provider, provider, provider_controller - ) - cache.delete() - - return {"result": "success"} - - @staticmethod - def set_default_provider(tenant_id: str, provider: str, id: str): - """ - set default provider - """ - with sessionmaker(bind=db.engine).begin() as session: - # get provider - target_provider = session.scalar( - select(BuiltinToolProvider) - .where(BuiltinToolProvider.id == id, BuiltinToolProvider.tenant_id == tenant_id) - .limit(1) - ) - if target_provider is None: - raise ValueError("provider not found") - - # clear default provider (tenant-scoped: only one default per provider per workspace) - session.execute( - update(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.is_default.is_(True), - ) - .values(is_default=False) - .execution_options(synchronize_session=False) - ) - - # set new default provider - target_provider.is_default = True - - return {"result": "success"} - - @staticmethod - def is_oauth_system_client_exists(provider_name: str) -> bool: - """ - check if oauth system client exists - """ - tool_provider = ToolProviderID(provider_name) - with Session(db.engine, autoflush=False) as session: - system_client = session.scalar( - select(ToolOAuthSystemClient) - .where( - ToolOAuthSystemClient.plugin_id == tool_provider.plugin_id, - ToolOAuthSystemClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - return system_client is not None - - @staticmethod - def is_oauth_custom_client_enabled(tenant_id: str, provider: str) -> bool: - """ - check if oauth custom client is enabled - """ - tool_provider = ToolProviderID(provider) - with Session(db.engine, autoflush=False) as session: - user_client = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.enabled.is_(True), - ) - .limit(1) - ) - return user_client is not None and user_client.enabled - - @staticmethod - def get_oauth_client( - tenant_id: str, provider: str, *, provider_controller: PluginToolProviderController | None = None - ) -> PluginOAuthClientConfig | None: - """ - get builtin tool provider - """ - tool_provider = ToolProviderID(provider) - if provider_controller is None: - provider_controller = ToolManager.get_plugin_provider(provider, tenant_id) - plugin_unique_identifier = provider_controller.plugin_unique_identifier - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()], - cache=NoOpProviderCredentialCache(), - ) - with Session(db.engine, autoflush=False) as session: - user_client = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.enabled.is_(True), - ) - .limit(1) - ) - oauth_params: Mapping[str, Any] | None = None - if user_client: - oauth_params = encrypter.decrypt(user_client.oauth_params) - return PluginOAuthClientConfig(plugin_unique_identifier, dict(oauth_params)) - - # only verified provider can use official oauth client - is_verified = PluginService.is_plugin_verified(tenant_id, plugin_unique_identifier) - if not is_verified: - return None - - system_client = session.scalar( - select(ToolOAuthSystemClient) - .where( - ToolOAuthSystemClient.plugin_id == tool_provider.plugin_id, - ToolOAuthSystemClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - if system_client: - try: - oauth_params = decrypt_system_params(system_client.encrypted_oauth_params) - except Exception as e: - raise ValueError(f"Error decrypting system oauth params: {e}") from e - - return ( - PluginOAuthClientConfig(plugin_unique_identifier, dict(oauth_params)) - if oauth_params is not None - else None - ) - - @staticmethod - def get_builtin_tool_provider_icon(provider: str): - """ - get tool provider icon and it's mimetype - """ - icon_path, mime_type = ToolManager.get_hardcoded_provider_icon(provider) - icon_bytes = Path(icon_path).read_bytes() - - return icon_bytes, mime_type - - @staticmethod - def list_builtin_tools(user_id: str, tenant_id: str) -> list[ToolProviderApiEntity]: - """ - list builtin tools - """ - # get all builtin providers - provider_controllers = ToolManager.list_builtin_providers(tenant_id) - - # get all user added providers - db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id) - - # rewrite db_providers - for db_provider in db_providers: - db_provider.provider = str(ToolProviderID(db_provider.provider)) - - # find provider - def find_provider(provider): - return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None) - - result: list[ToolProviderApiEntity] = [] - - for provider_controller in provider_controllers: - try: - # handle include, exclude - if is_filtered( - include_set=dify_config.POSITION_TOOL_INCLUDES_SET, - exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET, - data=provider_controller, - name_func=lambda x: x.entity.identity.name, - ): - continue - - # convert provider controller to user provider - user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider( - provider_controller=provider_controller, - db_provider=find_provider(provider_controller.entity.identity.name), - decrypt_credentials=True, - ) - - # add icon - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider) - - tools = provider_controller.get_tools() - for tool in tools or []: - user_builtin_provider.tools.append( - ToolTransformService.convert_tool_entity_to_api_entity( - tenant_id=tenant_id, - tool=tool, - labels=ToolLabelManager.get_tool_labels(provider_controller), - ) - ) - - result.append(user_builtin_provider) - except Exception as e: - raise e - - return BuiltinToolProviderSort.sort(result) - - @staticmethod - def get_builtin_provider(provider_name: str, tenant_id: str) -> BuiltinToolProvider | None: - """ - This method is used to fetch the builtin provider from the database - 1.if the default provider exists, return the default provider - 2.if the default provider does not exist, return the oldest provider - """ - with Session(db.engine, autoflush=False) as session: - try: - full_provider_name = provider_name - provider_id_entity = ToolProviderID(provider_name) - provider_name = provider_id_entity.provider_name - - if provider_id_entity.organization != "langgenius": - provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == full_provider_name, - ) - .order_by( - BuiltinToolProvider.is_default.desc(), # default=True first - BuiltinToolProvider.created_at.asc(), # oldest first - ) - .limit(1) - ) - else: - provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - (BuiltinToolProvider.provider == provider_name) - | (BuiltinToolProvider.provider == full_provider_name), - ) - .order_by( - BuiltinToolProvider.is_default.desc(), # default=True first - BuiltinToolProvider.created_at.asc(), # oldest first - ) - .limit(1) - ) - - if provider is None: - return None - - provider.provider = ToolProviderID(provider.provider).to_string() - return provider - except Exception: - # it's an old provider without organization - return session.scalar( - select(BuiltinToolProvider) - .where(BuiltinToolProvider.tenant_id == tenant_id, BuiltinToolProvider.provider == provider_name) - .order_by( - BuiltinToolProvider.is_default.desc(), # default=True first - BuiltinToolProvider.created_at.asc(), # oldest first - ) - .limit(1) - ) - - @staticmethod - def save_custom_oauth_client_params( - tenant_id: str, - provider: str, - client_params: dict[str, Any] | None = None, - enable_oauth_custom_client: bool | None = None, - ): - """ - setup oauth custom client - """ - if client_params is None and enable_oauth_custom_client is None: - return {"result": "success"} - - tool_provider = ToolProviderID(provider) - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller: - raise ToolProviderNotFoundError(f"Provider {provider} not found") - - if not isinstance(provider_controller, (BuiltinToolProviderController, PluginToolProviderController)): - raise ValueError(f"Provider {provider} is not a builtin or plugin provider") - - with sessionmaker(bind=db.engine).begin() as session: - custom_client_params = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - - # if the record does not exist, create a basic record - if custom_client_params is None: - custom_client_params = ToolOAuthTenantClient( - tenant_id=tenant_id, - plugin_id=tool_provider.plugin_id, - provider=tool_provider.provider_name, - ) - session.add(custom_client_params) - - if client_params is not None: - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()], - cache=NoOpProviderCredentialCache(), - ) - original_params = encrypter.decrypt(custom_client_params.oauth_params) - new_params = { - key: value if value != HIDDEN_VALUE else original_params.get(key, UNKNOWN_VALUE) - for key, value in client_params.items() - } - custom_client_params.encrypted_oauth_params = json.dumps(encrypter.encrypt(new_params)) - - if enable_oauth_custom_client is not None: - custom_client_params.enabled = enable_oauth_custom_client - - return {"result": "success"} - - @staticmethod - def get_custom_oauth_client_params(tenant_id: str, provider: str): - """ - get custom oauth client params - """ - with Session(db.engine) as session: - tool_provider = ToolProviderID(provider) - custom_oauth_client_params = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - if custom_oauth_client_params is None: - return {} - - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller: - raise ToolProviderNotFoundError(f"Provider {provider} not found") - - if not isinstance(provider_controller, BuiltinToolProviderController): - raise ValueError(f"Provider {provider} is not a builtin or plugin provider") - - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()], - cache=NoOpProviderCredentialCache(), - ) - - return encrypter.mask_plugin_credentials(encrypter.decrypt(custom_oauth_client_params.oauth_params)) diff --git a/api/services/tools/legacy_mcp_tools_manage_service.py b/api/services/tools/legacy_mcp_tools_manage_service.py deleted file mode 100644 index 0d90406b0eebf2..00000000000000 --- a/api/services/tools/legacy_mcp_tools_manage_service.py +++ /dev/null @@ -1,884 +0,0 @@ -import hashlib -import json -import logging -import uuid -from collections.abc import Mapping -from datetime import datetime -from enum import StrEnum -from typing import Any -from urllib.parse import urlparse - -from pydantic import BaseModel, Field -from sqlalchemy import Select, or_, select -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session - -from core.entities.mcp_provider import IdentityMode, MCPAuthentication, MCPConfiguration, MCPProviderEntity -from core.helper import encrypter -from core.helper.provider_cache import NoOpProviderCredentialCache -from core.mcp.auth.auth_flow import auth -from core.mcp.auth_client import MCPClientWithAuthRetry -from core.mcp.entities import AuthActionType, AuthResult -from core.mcp.error import MCPAuthError, MCPError -from core.mcp.types import Tool as MCPTool -from core.tools.entities.api_entities import ToolProviderApiEntity -from core.tools.utils.encryption import ProviderConfigEncrypter -from models.tools import MCPToolProvider -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - -# Constants -UNCHANGED_SERVER_URL_PLACEHOLDER = "[__HIDDEN__]" -CLIENT_NAME = "Dify" -EMPTY_TOOLS_JSON = "[]" -EMPTY_CREDENTIALS_JSON = "{}" - - -class OAuthDataType(StrEnum): - """Types of OAuth data that can be saved.""" - - TOKENS = "tokens" - CLIENT_INFO = "client_info" - CODE_VERIFIER = "code_verifier" - MIXED = "mixed" - - -class ReconnectResult(BaseModel): - """Result of reconnecting to an MCP provider""" - - authed: bool = Field(description="Whether the provider is authenticated") - tools: str = Field(description="JSON string of tool list") - encrypted_credentials: str = Field(description="JSON string of encrypted credentials") - - -class ServerUrlValidationResult(BaseModel): - """Result of server URL validation check""" - - needs_validation: bool - validation_passed: bool = False - reconnect_result: ReconnectResult | None = None - encrypted_server_url: str | None = None - server_url_hash: str | None = None - - @property - def should_update_server_url(self) -> bool: - """Check if server URL should be updated based on validation result""" - return self.needs_validation and self.validation_passed and self.reconnect_result is not None - - -class ProviderUrlValidationData(BaseModel): - """Data required for URL validation, extracted from database to perform network operations outside of session""" - - current_server_url_hash: str - headers: dict[str, str] - timeout: float | None - sse_read_timeout: float | None - - -class MCPToolManageService: - """Service class for managing MCP tools and providers.""" - - def __init__(self, session: Session): - self._session = session - - # ========== Provider CRUD Operations ========== - - def get_provider_by_id(self, *, provider_id: str, tenant_id: str) -> MCPToolProvider: - """ - Get MCP provider by its primary key. - - Args: - provider_id: Provider primary key (UUID). Server identifiers are not - accepted here — use get_provider_by_server_identifier instead. - tenant_id: Tenant ID - - Returns: - MCPToolProvider instance - - Raises: - ValueError: If provider_id is not a UUID, or the provider does not exist - """ - try: - uuid.UUID(provider_id) - except (AttributeError, TypeError, ValueError) as exc: - raise ValueError("expected a valid UUID, but got: " + provider_id) from exc - - stmt = select(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.id == provider_id) - return self._scalar_one_provider(stmt) - - def get_provider_by_server_identifier(self, *, server_identifier: str, tenant_id: str) -> MCPToolProvider: - """ - Get MCP provider by its tenant-scoped server identifier. - - Args: - server_identifier: Server identifier (the runtime-facing reference) - tenant_id: Tenant ID - - Returns: - MCPToolProvider instance - - Raises: - ValueError: If provider not found - """ - stmt = select(MCPToolProvider).where( - MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == server_identifier - ) - return self._scalar_one_provider(stmt) - - def _scalar_one_provider(self, stmt: Select[tuple[MCPToolProvider]]) -> MCPToolProvider: - provider = self._session.scalar(stmt) - if not provider: - raise ValueError("MCP tool not found") - return provider - - def get_provider_entity_by_id(self, *, provider_id: str, tenant_id: str) -> MCPProviderEntity: - """Get provider entity by primary key.""" - return self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id).to_entity() - - def get_provider_entity_by_server_identifier(self, *, server_identifier: str, tenant_id: str) -> MCPProviderEntity: - """Get provider entity by server identifier.""" - return self.get_provider_by_server_identifier( - server_identifier=server_identifier, tenant_id=tenant_id - ).to_entity() - - def get_provider_by_persisted_reference(self, *, id_or_server_identifier: str, tenant_id: str) -> MCPToolProvider: - """ - Resolve an MCP provider from a reference persisted in a graph or agent config. - - For backward compatibility. Previously, the id returned by the - server is non-deterministic, depending on the query parameter, - so we accept both primary key (id) and server identifier for any existing DSLs - - Raises: - ValueError: If no provider matches the reference - """ - try: - uuid.UUID(id_or_server_identifier) - except (AttributeError, TypeError, ValueError): - # Not uuid-shaped, so it can only ever be a server identifier. Skipping - # the primary-key lookup also keeps its "expected a valid UUID" error, - # which would misdescribe a legitimate reference, out of this path. - return self.get_provider_by_server_identifier( - server_identifier=id_or_server_identifier, tenant_id=tenant_id - ) - - try: - return self.get_provider_by_id(provider_id=id_or_server_identifier, tenant_id=tenant_id) - except ValueError: - # A server identifier is free text, so it may itself look like a uuid. - return self.get_provider_by_server_identifier( - server_identifier=id_or_server_identifier, tenant_id=tenant_id - ) - - def get_provider_entity_by_persisted_reference( - self, *, id_or_server_identifier: str, tenant_id: str - ) -> MCPProviderEntity: - """Resolve a provider entity from a persisted graph or agent config reference. - - Accepts both meanings — see get_provider_by_persisted_reference. - """ - return self.get_provider_by_persisted_reference( - id_or_server_identifier=id_or_server_identifier, tenant_id=tenant_id - ).to_entity() - - def create_provider( - self, - *, - tenant_id: str, - name: str, - server_url: str, - user_id: str, - icon: str, - icon_type: str, - icon_background: str, - server_identifier: str, - configuration: MCPConfiguration, - authentication: MCPAuthentication | None = None, - headers: dict[str, str] | None = None, - identity_mode: IdentityMode = IdentityMode.OFF, - ) -> ToolProviderApiEntity: - """Create a new MCP provider.""" - # Validate URL format - if not self._is_valid_url(server_url): - raise ValueError("Server URL is not valid.") - - server_url_hash = hashlib.sha256(server_url.encode()).hexdigest() - - # Check for existing provider - self._check_provider_exists(tenant_id, name, server_url_hash, server_identifier) - - # Encrypt sensitive data - encrypted_server_url = encrypter.encrypt_token(tenant_id, server_url) - encrypted_headers = self._prepare_encrypted_dict(headers, tenant_id) if headers else None - encrypted_credentials = None - if authentication is not None and authentication.client_id: - encrypted_credentials = self._build_and_encrypt_credentials( - authentication.client_id, authentication.client_secret, tenant_id - ) - - # Create provider - mcp_tool = MCPToolProvider( - tenant_id=tenant_id, - name=name, - server_url=encrypted_server_url, - server_url_hash=server_url_hash, - user_id=user_id, - authed=False, - tools=EMPTY_TOOLS_JSON, - icon=self._prepare_icon(icon, icon_type, icon_background), - server_identifier=server_identifier, - timeout=configuration.timeout, - sse_read_timeout=configuration.sse_read_timeout, - encrypted_headers=encrypted_headers, - encrypted_credentials=encrypted_credentials, - identity_mode=identity_mode, - ) - - self._session.add(mcp_tool) - self._session.flush() - - mcp_providers = ToolTransformService.mcp_provider_to_user_provider(mcp_tool) - return mcp_providers - - def update_provider( - self, - *, - tenant_id: str, - provider_id: str, - name: str, - server_url: str, - icon: str, - icon_type: str, - icon_background: str, - server_identifier: str, - headers: dict[str, str] | None = None, - configuration: MCPConfiguration, - authentication: MCPAuthentication | None = None, - validation_result: ServerUrlValidationResult | None = None, - identity_mode: IdentityMode = IdentityMode.OFF, - ) -> None: - """ - Update an MCP provider. - - Args: - validation_result: Pre-validation result from validate_server_url_standalone. - If provided and contains reconnect_result, it will be used - instead of performing network operations. - """ - mcp_provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - - # Check for duplicate name (excluding current provider) - if name != mcp_provider.name: - stmt = select(MCPToolProvider).where( - MCPToolProvider.tenant_id == tenant_id, - MCPToolProvider.name == name, - MCPToolProvider.id != provider_id, - ) - existing_provider = self._session.scalar(stmt) - if existing_provider: - raise ValueError(f"MCP tool {name} already exists") - - # Get URL update data from validation result - encrypted_server_url = None - server_url_hash = None - reconnect_result = None - - if validation_result and validation_result.encrypted_server_url: - # Use all data from validation result - encrypted_server_url = validation_result.encrypted_server_url - server_url_hash = validation_result.server_url_hash - reconnect_result = validation_result.reconnect_result - - try: - # Update basic fields - mcp_provider.updated_at = datetime.now() - mcp_provider.name = name - mcp_provider.icon = self._prepare_icon(icon, icon_type, icon_background) - mcp_provider.server_identifier = server_identifier - - # Update server URL if changed - if encrypted_server_url and server_url_hash: - mcp_provider.server_url = encrypted_server_url - mcp_provider.server_url_hash = server_url_hash - - if reconnect_result: - mcp_provider.authed = reconnect_result.authed - mcp_provider.tools = reconnect_result.tools - mcp_provider.encrypted_credentials = reconnect_result.encrypted_credentials - - # Update optional configuration fields - self._update_optional_fields(mcp_provider, configuration) - - # Update headers if provided - if headers is not None: - mcp_provider.encrypted_headers = self._process_headers(headers, mcp_provider, tenant_id) - - # Update credentials if provided - if authentication and authentication.client_id: - mcp_provider.encrypted_credentials = self._process_credentials(authentication, mcp_provider, tenant_id) - - # Update user-identity forwarding mode. The controller has already - # resolved "leave unchanged" and applied the Enterprise-edition gate, - # so this is always a concrete, vetted value. - mcp_provider.identity_mode = identity_mode - - # Flush changes to database - self._session.flush() - - except IntegrityError as e: - self._handle_integrity_error(e, name, server_url, server_identifier) - - def delete_provider(self, *, tenant_id: str, provider_id: str) -> None: - """Delete an MCP provider.""" - mcp_tool = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - self._session.delete(mcp_tool) - - def list_providers(self, *, tenant_id: str, include_sensitive: bool = True) -> list[ToolProviderApiEntity]: - """List all MCP providers for a tenant. - - Args: - tenant_id: Tenant ID - include_sensitive: If False, skip expensive decryption operations (default: True for backward compatibility) - """ - from models.account import Account - - stmt = select(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant_id).order_by(MCPToolProvider.name) - mcp_providers = self._session.scalars(stmt).all() - - if not mcp_providers: - return [] - - # Batch query all users to avoid N+1 problem - user_ids = {provider.user_id for provider in mcp_providers} - users = self._session.scalars(select(Account).where(Account.id.in_(user_ids))).all() - user_name_map = {user.id: user.name for user in users} - - return [ - ToolTransformService.mcp_provider_to_user_provider( - provider, - user_name=user_name_map.get(provider.user_id), - include_sensitive=include_sensitive, - ) - for provider in mcp_providers - ] - - # ========== Tool Operations ========== - - def list_provider_tools(self, *, tenant_id: str, provider_id: str) -> ToolProviderApiEntity: - """List tools from remote MCP server.""" - # Load provider and convert to entity - db_provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - provider_entity = db_provider.to_entity() - - # Verify authentication - if not provider_entity.authed: - raise ValueError("Please auth the tool first") - - # Prepare headers with auth token - headers = self._prepare_auth_headers(provider_entity) - - # Retrieve tools from remote server - server_url = provider_entity.decrypt_server_url() - try: - tools = self._retrieve_remote_mcp_tools(server_url, headers, provider_entity) - except MCPError as e: - raise ValueError(f"Failed to connect to MCP server: {e}") from e - - # Update database with retrieved tools (ensure description is a non-null string) - tools_payload = [] - for tool in tools: - data = tool.model_dump() - if data.get("description") is None: - data["description"] = "" - tools_payload.append(data) - db_provider.tools = json.dumps(tools_payload) - db_provider.authed = True - db_provider.updated_at = datetime.now() - self._session.flush() - - # Build API response - return self._build_tool_provider_response(db_provider, provider_entity, tools) - - # ========== OAuth and Credentials Operations ========== - - def update_provider_credentials( - self, *, provider_id: str, tenant_id: str, credentials: dict[str, Any], authed: bool | None = None - ) -> None: - """ - Update provider credentials with encryption. - - Args: - provider_id: Provider ID - tenant_id: Tenant ID - credentials: Credentials to save - authed: Whether provider is authenticated (None means keep current state) - """ - from core.tools.mcp_tool.provider import MCPToolProviderController - - # Get provider from current session - provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - - # Encrypt new credentials - provider_controller = MCPToolProviderController.from_db(provider) - tool_configuration = ProviderConfigEncrypter( - tenant_id=provider.tenant_id, - config=list(provider_controller.get_credentials_schema()), - provider_config_cache=NoOpProviderCredentialCache(), - ) - encrypted_credentials = tool_configuration.encrypt(credentials) - - # Update provider - provider.updated_at = datetime.now() - provider.encrypted_credentials = json.dumps({**provider.credentials, **encrypted_credentials}) - - if authed is not None: - provider.authed = authed - if not authed: - provider.tools = EMPTY_TOOLS_JSON - - # Flush changes to database - self._session.flush() - - def save_oauth_data( - self, provider_id: str, tenant_id: str, data: dict[str, Any], data_type: OAuthDataType = OAuthDataType.MIXED - ) -> None: - """ - Save OAuth-related data (tokens, client info, code verifier). - - Args: - provider_id: Provider ID - tenant_id: Tenant ID - data: Data to save (tokens, client info, or code verifier) - data_type: Type of OAuth data to save - """ - # Determine if this makes the provider authenticated - authed = ( - data_type == OAuthDataType.TOKENS or (data_type == OAuthDataType.MIXED and "access_token" in data) or None - ) - - # update_provider_credentials will validate provider existence - self.update_provider_credentials(provider_id=provider_id, tenant_id=tenant_id, credentials=data, authed=authed) - - def clear_provider_credentials(self, *, provider_id: str, tenant_id: str) -> None: - """ - Clear all credentials for a provider. - - Args: - provider_id: Provider ID - tenant_id: Tenant ID - """ - # Get provider from current session - provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - - provider.tools = EMPTY_TOOLS_JSON - provider.encrypted_credentials = EMPTY_CREDENTIALS_JSON - provider.updated_at = datetime.now() - provider.authed = False - - # ========== Private Helper Methods ========== - - def _check_provider_exists(self, tenant_id: str, name: str, server_url_hash: str, server_identifier: str) -> None: - """Check if provider with same attributes already exists.""" - stmt = select(MCPToolProvider).where( - MCPToolProvider.tenant_id == tenant_id, - or_( - MCPToolProvider.name == name, - MCPToolProvider.server_url_hash == server_url_hash, - MCPToolProvider.server_identifier == server_identifier, - ), - ) - existing_provider = self._session.scalar(stmt) - - if existing_provider: - if existing_provider.name == name: - raise ValueError(f"MCP tool {name} already exists") - if existing_provider.server_url_hash == server_url_hash: - raise ValueError("MCP tool with this server URL already exists") - if existing_provider.server_identifier == server_identifier: - raise ValueError(f"MCP tool {server_identifier} already exists") - - def _prepare_icon(self, icon: str, icon_type: str, icon_background: str) -> str: - """Prepare icon data for storage.""" - if icon_type == "emoji": - return json.dumps({"content": icon, "background": icon_background}) - return icon - - def _encrypt_dict_fields(self, data: dict[str, Any], secret_fields: list[str], tenant_id: str) -> Mapping[str, str]: - """Encrypt specified fields in a dictionary. - - Args: - data: Dictionary containing data to encrypt - secret_fields: List of field names to encrypt - tenant_id: Tenant ID for encryption - - Returns: - JSON string of encrypted data - """ - from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType - from core.tools.utils.encryption import create_provider_encrypter - - # Create config for secret fields - config = [BasicProviderConfig(type=ProviderConfigType.SECRET_INPUT, name=field) for field in secret_fields] - - encrypter_instance, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=config, - cache=NoOpProviderCredentialCache(), - ) - - encrypted_data = encrypter_instance.encrypt(data) - return encrypted_data - - def _prepare_encrypted_dict(self, headers: dict[str, str], tenant_id: str) -> str: - """Encrypt headers and prepare for storage.""" - # All headers are treated as secret - return json.dumps(self._encrypt_dict_fields(headers, list(headers.keys()), tenant_id)) - - def _prepare_auth_headers(self, provider_entity: MCPProviderEntity) -> dict[str, str]: - """Prepare headers with OAuth token if available.""" - headers = provider_entity.decrypt_headers() - tokens = provider_entity.retrieve_tokens() - if tokens: - headers["Authorization"] = f"{tokens.token_type.capitalize()} {tokens.access_token}" - return headers - - def _retrieve_remote_mcp_tools( - self, - server_url: str, - headers: dict[str, str], - provider_entity: MCPProviderEntity, - ): - """Retrieve tools from remote MCP server.""" - with MCPClientWithAuthRetry( - server_url=server_url, - headers=headers, - timeout=provider_entity.timeout, - sse_read_timeout=provider_entity.sse_read_timeout, - provider_entity=provider_entity, - ) as mcp_client: - return mcp_client.list_tools() - - _ACTION_TO_OAUTH: dict[AuthActionType, OAuthDataType] = { - AuthActionType.SAVE_CLIENT_INFO: OAuthDataType.CLIENT_INFO, - AuthActionType.SAVE_TOKENS: OAuthDataType.TOKENS, - AuthActionType.SAVE_CODE_VERIFIER: OAuthDataType.CODE_VERIFIER, - } - - def execute_auth_actions(self, auth_result: AuthResult) -> dict[str, str]: - """ - Execute the actions returned by the auth function. - - This method processes the AuthResult and performs the necessary database operations. - - Args: - auth_result: The result from the auth function - - Returns: - The response from the auth result - """ - for action in auth_result.actions: - if action.provider_id is None or action.tenant_id is None: - continue - - oauth_type = self._ACTION_TO_OAUTH.get(action.action_type) - if oauth_type is not None: - self.save_oauth_data(action.provider_id, action.tenant_id, action.data, oauth_type) - - return auth_result.response - - def auth_with_actions( - self, - provider_entity: MCPProviderEntity, - authorization_code: str | None = None, - resource_metadata_url: str | None = None, - scope_hint: str | None = None, - ) -> dict[str, str]: - """ - Perform authentication and execute all resulting actions. - - This method is used by MCPClientWithAuthRetry for automatic re-authentication. - - Args: - provider_entity: The MCP provider entity - authorization_code: Optional authorization code - resource_metadata_url: Optional Protected Resource Metadata URL from WWW-Authenticate - scope_hint: Optional scope hint from WWW-Authenticate header - - Returns: - Response dictionary from auth result - """ - auth_result = auth( - provider_entity, - authorization_code, - resource_metadata_url=resource_metadata_url, - scope_hint=scope_hint, - ) - return self.execute_auth_actions(auth_result) - - def get_provider_for_url_validation(self, *, tenant_id: str, provider_id: str) -> ProviderUrlValidationData: - """ - Get provider data required for URL validation. - This method performs database read and should be called within a session. - - Returns: - ProviderUrlValidationData: Data needed for standalone URL validation - """ - provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - provider_entity = provider.to_entity() - return ProviderUrlValidationData( - current_server_url_hash=provider.server_url_hash, - headers=provider_entity.headers, - timeout=provider_entity.timeout, - sse_read_timeout=provider_entity.sse_read_timeout, - ) - - @staticmethod - def validate_server_url_standalone( - *, - tenant_id: str, - new_server_url: str, - validation_data: ProviderUrlValidationData, - ) -> ServerUrlValidationResult: - """ - Validate server URL change by attempting to connect to the new server. - This method performs network operations and MUST be called OUTSIDE of any database session - to avoid holding locks during network I/O. - - Args: - tenant_id: Tenant ID for encryption - new_server_url: The new server URL to validate - validation_data: Provider data obtained from get_provider_for_url_validation - - Returns: - ServerUrlValidationResult: Validation result with connection status and tools if successful - """ - # Handle hidden/unchanged URL - if UNCHANGED_SERVER_URL_PLACEHOLDER in new_server_url: - return ServerUrlValidationResult(needs_validation=False) - - # Validate URL format - parsed = urlparse(new_server_url) - if not all([parsed.scheme, parsed.netloc]) or parsed.scheme not in ["http", "https"]: - raise ValueError("Server URL is not valid.") - - # Always encrypt and hash the URL - encrypted_server_url = encrypter.encrypt_token(tenant_id, new_server_url) - new_server_url_hash = hashlib.sha256(new_server_url.encode()).hexdigest() - - # Check if URL is actually different - if new_server_url_hash == validation_data.current_server_url_hash: - # URL hasn't changed, but still return the encrypted data - return ServerUrlValidationResult( - needs_validation=False, - encrypted_server_url=encrypted_server_url, - server_url_hash=new_server_url_hash, - ) - - # Perform network validation - this is the expensive operation that should be outside session - reconnect_result = MCPToolManageService._reconnect_with_url( - server_url=new_server_url, - headers=validation_data.headers, - timeout=validation_data.timeout, - sse_read_timeout=validation_data.sse_read_timeout, - ) - return ServerUrlValidationResult( - needs_validation=True, - validation_passed=True, - reconnect_result=reconnect_result, - encrypted_server_url=encrypted_server_url, - server_url_hash=new_server_url_hash, - ) - - @staticmethod - def reconnect_with_url( - *, - server_url: str, - headers: dict[str, str], - timeout: float | None, - sse_read_timeout: float | None, - ) -> ReconnectResult: - return MCPToolManageService._reconnect_with_url( - server_url=server_url, - headers=headers, - timeout=timeout, - sse_read_timeout=sse_read_timeout, - ) - - @staticmethod - def _reconnect_with_url( - *, - server_url: str, - headers: dict[str, str], - timeout: float | None, - sse_read_timeout: float | None, - ) -> ReconnectResult: - """ - Attempt to connect to MCP server with given URL. - This is a static method that performs network I/O without database access. - """ - from core.mcp.mcp_client import MCPClient - - try: - with MCPClient( - server_url=server_url, - headers=headers, - timeout=timeout, - sse_read_timeout=sse_read_timeout, - ) as mcp_client: - tools = mcp_client.list_tools() - # Ensure tool descriptions are non-null in payload - tools_payload = [] - for t in tools: - d = t.model_dump() - if d.get("description") is None: - d["description"] = "" - tools_payload.append(d) - return ReconnectResult( - authed=True, - tools=json.dumps(tools_payload), - encrypted_credentials=EMPTY_CREDENTIALS_JSON, - ) - except MCPAuthError: - return ReconnectResult(authed=False, tools=EMPTY_TOOLS_JSON, encrypted_credentials=EMPTY_CREDENTIALS_JSON) - except MCPError as e: - raise ValueError(f"Failed to re-connect MCP server: {e}") from e - - def _build_tool_provider_response( - self, db_provider: MCPToolProvider, provider_entity: MCPProviderEntity, tools: list[MCPTool] - ) -> ToolProviderApiEntity: - """Build API response for tool provider.""" - user = db_provider.load_user(self._session) - user_name = user.name if user else None - response = provider_entity.to_api_response( - user_name=user_name, - ) - response["tools"] = ToolTransformService.mcp_tool_to_user_tool(db_provider, tools, user_name=user_name) - response["plugin_unique_identifier"] = provider_entity.server_identifier - return ToolProviderApiEntity(**response) - - def _handle_integrity_error( - self, error: IntegrityError, name: str, server_url: str, server_identifier: str - ) -> None: - """Handle database integrity errors with user-friendly messages.""" - error_msg = str(error.orig) - if "unique_mcp_provider_name" in error_msg: - raise ValueError(f"MCP tool {name} already exists") - if "unique_mcp_provider_server_url" in error_msg: - raise ValueError(f"MCP tool {server_url} already exists") - if "unique_mcp_provider_server_identifier" in error_msg: - raise ValueError(f"MCP tool {server_identifier} already exists") - raise error - - def _is_valid_url(self, url: str) -> bool: - """Validate URL format.""" - if not url: - return False - try: - parsed = urlparse(url) - return all([parsed.scheme, parsed.netloc]) and parsed.scheme in ["http", "https"] - except (ValueError, TypeError): - return False - - def _update_optional_fields(self, mcp_provider: MCPToolProvider, configuration: MCPConfiguration) -> None: - """Update optional configuration fields using setattr for cleaner code.""" - field_mapping = {"timeout": configuration.timeout, "sse_read_timeout": configuration.sse_read_timeout} - - for field, value in field_mapping.items(): - if value is not None: - setattr(mcp_provider, field, value) - - def _process_headers(self, headers: dict[str, str], mcp_provider: MCPToolProvider, tenant_id: str) -> str | None: - """Process headers update, handling empty dict to clear headers.""" - if not headers: - return None - - # Merge with existing headers to preserve masked values - final_headers = self._merge_headers_with_masked(incoming_headers=headers, mcp_provider=mcp_provider) - return self._prepare_encrypted_dict(final_headers, tenant_id) - - def _process_credentials( - self, authentication: MCPAuthentication, mcp_provider: MCPToolProvider, tenant_id: str - ) -> str: - """Process credentials update, handling masked values.""" - # Merge with existing credentials - final_client_id, final_client_secret = self._merge_credentials_with_masked( - authentication.client_id, authentication.client_secret, mcp_provider - ) - - # Build and encrypt - return self._build_and_encrypt_credentials(final_client_id, final_client_secret, tenant_id) - - def _merge_headers_with_masked( - self, incoming_headers: dict[str, str], mcp_provider: MCPToolProvider - ) -> dict[str, str]: - """Merge incoming headers with existing ones, preserving unchanged masked values. - - Args: - incoming_headers: Headers from frontend (may contain masked values) - mcp_provider: The MCP provider instance - - Returns: - Final headers dict with proper values (original for unchanged masked, new for changed) - """ - mcp_provider_entity = mcp_provider.to_entity() - existing_decrypted = mcp_provider_entity.decrypt_headers() - existing_masked = mcp_provider_entity.masked_headers() - - return { - key: (str(existing_decrypted[key]) if key in existing_masked and value == existing_masked[key] else value) - for key, value in incoming_headers.items() - if key in existing_decrypted or value != existing_masked.get(key) - } - - def _merge_credentials_with_masked( - self, - client_id: str, - client_secret: str | None, - mcp_provider: MCPToolProvider, - ) -> tuple[ - str, - str | None, - ]: - """Merge incoming credentials with existing ones, preserving unchanged masked values. - - Args: - client_id: Client ID from frontend (may be masked) - client_secret: Client secret from frontend (may be masked) - mcp_provider: The MCP provider instance - - Returns: - Tuple of (final_client_id, final_client_secret) - """ - mcp_provider_entity = mcp_provider.to_entity() - existing_decrypted = mcp_provider_entity.decrypt_credentials() - existing_masked = mcp_provider_entity.masked_credentials() - - # Check if client_id is masked and unchanged - final_client_id = client_id - if existing_masked.get("client_id") and client_id == existing_masked["client_id"]: - # Use existing decrypted value - final_client_id = existing_decrypted.get("client_id", client_id) - - # Check if client_secret is masked and unchanged - final_client_secret = client_secret - if existing_masked.get("client_secret") and client_secret == existing_masked["client_secret"]: - # Use existing decrypted value - final_client_secret = existing_decrypted.get("client_secret", client_secret) - - return final_client_id, final_client_secret - - def _build_and_encrypt_credentials(self, client_id: str, client_secret: str | None, tenant_id: str) -> str: - """Build credentials and encrypt sensitive fields.""" - # Create a flat structure with all credential data - credentials_data = { - "client_id": client_id, - "client_name": CLIENT_NAME, - "is_dynamic_registration": False, - } - secret_fields = [] - if client_secret is not None: - credentials_data["encrypted_client_secret"] = client_secret - secret_fields = ["encrypted_client_secret"] - client_info = self._encrypt_dict_fields(credentials_data, secret_fields, tenant_id) - return json.dumps({"client_information": client_info}) diff --git a/api/services/tools/legacy_tools_manage_service.py b/api/services/tools/legacy_tools_manage_service.py deleted file mode 100644 index 44db6a97e36a59..00000000000000 --- a/api/services/tools/legacy_tools_manage_service.py +++ /dev/null @@ -1,26 +0,0 @@ -import logging - -from core.tools.entities.api_entities import ToolProviderTypeApiLiteral -from core.tools.tool_manager import ToolManager -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class ToolCommonService: - @staticmethod - def list_tool_providers(user_id: str, tenant_id: str, typ: ToolProviderTypeApiLiteral | None = None): - """ - list tool providers - - :return: the list of tool providers - """ - providers = ToolManager.list_providers_from_api(user_id, tenant_id, typ) - - # add icon - for provider in providers: - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=provider) - - result = [provider.to_dict() for provider in providers] - - return result diff --git a/api/services/tools/legacy_tools_transform_service.py b/api/services/tools/legacy_tools_transform_service.py deleted file mode 100644 index f9310e1f72df44..00000000000000 --- a/api/services/tools/legacy_tools_transform_service.py +++ /dev/null @@ -1,539 +0,0 @@ -import logging -from collections.abc import Mapping -from typing import Any - -from pydantic import TypeAdapter, ValidationError -from sqlalchemy.orm import Session -from yarl import URL - -from configs import dify_config -from core.helper.provider_cache import ToolProviderCredentialsCache -from core.mcp.types import Tool as MCPTool -from core.plugin.entities.plugin_daemon import CredentialType, PluginDatasourceProviderEntity -from core.plugin.plugin_service import PluginService -from core.tools.__base.tool import Tool -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity, ToolProviderCredentialApiEntity -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ToolParameter, - ToolProviderType, - emoji_icon_adapter, -) -from core.tools.mcp_tool.tool import schema_accepts_multiple_types -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider - -logger = logging.getLogger(__name__) - -_mcp_tools_adapter: TypeAdapter[list[MCPTool]] = TypeAdapter(list[MCPTool]) - - -class ToolTransformService: - _MCP_SCHEMA_TYPE_RESOLUTION_MAX_DEPTH = 10 - - @classmethod - def get_tool_provider_icon_url( - cls, provider_type: str, provider_name: str, icon: str | Mapping[str, str] - ) -> str | Mapping[str, str]: - """ - get tool provider icon url - """ - url_prefix = ( - URL(dify_config.CONSOLE_API_URL or "/") / "console" / "api" / "workspaces" / "current" / "tool-provider" - ) - - match provider_type: - case ToolProviderType.BUILT_IN: - return str(url_prefix / "builtin" / provider_name / "icon") - case ToolProviderType.API | ToolProviderType.WORKFLOW: - try: - if isinstance(icon, str): - parsed = emoji_icon_adapter.validate_json(icon) - return {"background": parsed["background"], "content": parsed["content"]} - return {"background": icon["background"], "content": icon["content"]} - except (ValueError, KeyError): - return {"background": "#252525", "content": "\ud83d\ude01"} - case ToolProviderType.MCP: - if isinstance(icon, Mapping): - return {"background": icon.get("background", ""), "content": icon.get("content", "")} - return icon - case ToolProviderType.PLUGIN | ToolProviderType.APP | ToolProviderType.DATASET_RETRIEVAL: - return "" - case _: - return "" - - @staticmethod - def repack_provider( - tenant_id: str, provider: dict[str, Any] | ToolProviderApiEntity | PluginDatasourceProviderEntity - ): - """ - repack provider - - :param tenant_id: the tenant id - :param provider: the provider dict - """ - match provider: - case dict() if "icon" in provider: - provider["icon"] = ToolTransformService.get_tool_provider_icon_url( - provider_type=provider["type"], provider_name=provider["name"], icon=provider["icon"] - ) - case ToolProviderApiEntity(): - if provider.plugin_id: - if isinstance(provider.icon, str): - provider.icon = PluginService.get_plugin_icon_url(tenant_id=tenant_id, filename=provider.icon) - if isinstance(provider.icon_dark, str) and provider.icon_dark: - provider.icon_dark = PluginService.get_plugin_icon_url( - tenant_id=tenant_id, filename=provider.icon_dark - ) - else: - provider.icon = ToolTransformService.get_tool_provider_icon_url( - provider_type=provider.type.value, provider_name=provider.name, icon=provider.icon - ) - if provider.icon_dark: - provider.icon_dark = ToolTransformService.get_tool_provider_icon_url( - provider_type=provider.type.value, provider_name=provider.name, icon=provider.icon_dark - ) - case PluginDatasourceProviderEntity(): - if provider.plugin_id: - if isinstance(provider.declaration.identity.icon, str): - provider.declaration.identity.icon = PluginService.get_plugin_icon_url( - tenant_id=tenant_id, filename=provider.declaration.identity.icon - ) - - @classmethod - def builtin_provider_to_user_provider( - cls, - provider_controller: BuiltinToolProviderController | PluginToolProviderController, - db_provider: BuiltinToolProvider | None, - decrypt_credentials: bool = True, - ) -> ToolProviderApiEntity: - """ - convert provider controller to user provider - """ - result = ToolProviderApiEntity( - id=provider_controller.entity.identity.name, - author=provider_controller.entity.identity.author, - name=provider_controller.entity.identity.name, - description=provider_controller.entity.identity.description, - icon=provider_controller.entity.identity.icon, - icon_dark=provider_controller.entity.identity.icon_dark or "", - label=provider_controller.entity.identity.label, - type=ToolProviderType.BUILT_IN, - masked_credentials={}, - is_team_authorization=False, - plugin_id=None, - tools=[], - labels=provider_controller.tool_labels, - ) - - if isinstance(provider_controller, PluginToolProviderController): - result.plugin_id = provider_controller.plugin_id - result.plugin_unique_identifier = provider_controller.plugin_unique_identifier - - # get credentials schema - schema = { - x.to_basic_provider_config().name: x - for x in provider_controller.get_credentials_schema_by_type( - CredentialType.of(db_provider.credential_type) if db_provider else CredentialType.API_KEY - ) - } - - masked_creds = {} - for name in schema: - masked_creds[name] = "" - result.masked_credentials = masked_creds - - # check if the provider need credentials - if not provider_controller.need_credentials: - result.is_team_authorization = True - result.allow_delete = False - elif db_provider: - result.is_team_authorization = True - - if decrypt_credentials: - credentials = db_provider.credentials - if not db_provider.tenant_id: - raise ValueError(f"Required tenant_id is missing for BuiltinToolProvider with id {db_provider.id}") - # init tool configuration - encrypter, _ = create_provider_encrypter( - tenant_id=db_provider.tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type( - CredentialType.of(db_provider.credential_type) - ) - ], - cache=ToolProviderCredentialsCache( - tenant_id=db_provider.tenant_id, - provider=db_provider.provider, - credential_id=db_provider.id, - ), - ) - # decrypt the credentials and mask the credentials - decrypted_credentials = encrypter.decrypt(data=credentials) - masked_credentials = encrypter.mask_plugin_credentials(data=decrypted_credentials) - - result.masked_credentials = masked_credentials - result.original_credentials = decrypted_credentials - - return result - - @staticmethod - def api_provider_to_controller( - db_provider: ApiToolProvider, - ) -> ApiToolProviderController: - """ - convert provider controller to user provider - """ - # package tool provider controller - auth_type = ApiProviderAuthType.NONE - credentials_auth_type = db_provider.credentials.get("auth_type") - if credentials_auth_type in ("api_key_header", "api_key"): # backward compatibility - auth_type = ApiProviderAuthType.API_KEY_HEADER - elif credentials_auth_type == "api_key_query": - auth_type = ApiProviderAuthType.API_KEY_QUERY - - controller = ApiToolProviderController.from_db( - db_provider=db_provider, - auth_type=auth_type, - session=db.session(), - ) - - return controller - - @staticmethod - def workflow_provider_to_controller(db_provider: WorkflowToolProvider) -> WorkflowToolProviderController: - """ - convert provider controller to provider - """ - return WorkflowToolProviderController.from_db(db_provider) - - @staticmethod - def workflow_provider_to_user_provider( - provider_controller: WorkflowToolProviderController, - labels: list[str] | None = None, - workflow_app_id: str | None = None, - ): - """ - convert provider controller to user provider - """ - return ToolProviderApiEntity( - id=provider_controller.provider_id, - author=provider_controller.entity.identity.author, - name=provider_controller.entity.identity.name, - description=provider_controller.entity.identity.description, - icon=provider_controller.entity.identity.icon, - icon_dark=provider_controller.entity.identity.icon_dark or "", - label=provider_controller.entity.identity.label, - type=ToolProviderType.WORKFLOW, - masked_credentials={}, - is_team_authorization=True, - plugin_id=None, - plugin_unique_identifier=None, - tools=[], - labels=labels or [], - workflow_app_id=workflow_app_id, - ) - - @staticmethod - def mcp_provider_to_user_provider( - db_provider: MCPToolProvider, - user_name: str | None = None, - include_sensitive: bool = True, - ) -> ToolProviderApiEntity: - from core.entities.mcp_provider import MCPConfiguration - - # Use provided user_name to avoid N+1 query, fallback to load_user() if not provided - if user_name is None: - user = db_provider.load_user(db.session()) - user_name = user.name if user else None - - # Convert to entity and use its API response method - provider_entity = db_provider.to_entity() - - response = provider_entity.to_api_response(user_name=user_name, include_sensitive=include_sensitive) - try: - mcp_tools = _mcp_tools_adapter.validate_json(db_provider.tools) - except (ValidationError, ValueError): - mcp_tools = [] - # Add additional fields specific to the transform - response["id"] = db_provider.id - response["tools"] = ToolTransformService.mcp_tool_to_user_tool(db_provider, mcp_tools, user_name=user_name) - response["server_identifier"] = db_provider.server_identifier - - # Convert configuration dict to MCPConfiguration object - if "configuration" in response and isinstance(response["configuration"], dict): - response["configuration"] = MCPConfiguration( - timeout=float(response["configuration"]["timeout"]), - sse_read_timeout=float(response["configuration"]["sse_read_timeout"]), - ) - - return ToolProviderApiEntity(**response) - - @staticmethod - def mcp_tool_to_user_tool( - mcp_provider: MCPToolProvider, tools: list[MCPTool], user_name: str | None = None - ) -> list[ToolApiEntity]: - # Use provided user_name to avoid N+1 query, fallback to load_user() if not provided - if user_name is None: - user = mcp_provider.load_user(db.session()) - user_name = user.name if user else "Anonymous" - - return [ - ToolApiEntity( - author=user_name or "Anonymous", - name=tool.name, - # Prefer the server-provided human-readable title and fall back - # to the programmatic name when the server omitted the title - # or returned ``title: null`` (e.g. Exa's MCP server). - label=I18nObject(en_US=tool.title or tool.name, zh_Hans=tool.title or tool.name), - description=I18nObject(en_US=tool.description or "", zh_Hans=tool.description or ""), - parameters=ToolTransformService.convert_mcp_schema_to_parameter(tool.inputSchema), - labels=[], - output_schema=tool.outputSchema or {}, - ) - for tool in tools - ] - - @classmethod - def api_provider_to_user_provider( - cls, - provider_controller: ApiToolProviderController, - db_provider: ApiToolProvider, - decrypt_credentials: bool = True, - labels: list[str] | None = None, - *, - session: Session, - ) -> ToolProviderApiEntity: - """ - convert provider controller to user provider - """ - username = "Anonymous" - user = db_provider.user(session=session) - if user is None: - raise ValueError(f"user is None for api provider {db_provider.id}") - try: - if not user: - raise ValueError("user not found") - - username = user.name - except Exception: - logger.exception("failed to get user name for api provider %s", db_provider.id) - # add provider into providers - credentials = db_provider.credentials - result = ToolProviderApiEntity( - id=db_provider.id, - author=username, - name=db_provider.name, - description=I18nObject( - en_US=db_provider.description, - zh_Hans=db_provider.description, - ), - icon=db_provider.icon, - label=I18nObject( - en_US=db_provider.name, - zh_Hans=db_provider.name, - ), - type=ToolProviderType.API, - plugin_id=None, - plugin_unique_identifier=None, - masked_credentials={}, - is_team_authorization=True, - tools=[], - labels=labels or [], - ) - - if decrypt_credentials: - # init tool configuration - encrypter, _ = create_tool_provider_encrypter( - tenant_id=db_provider.tenant_id, - controller=provider_controller, - ) - - # decrypt the credentials and mask the credentials - decrypted_credentials = encrypter.decrypt(data=credentials) - masked_credentials = encrypter.mask_plugin_credentials(data=decrypted_credentials) - - result.masked_credentials = masked_credentials - - return result - - @staticmethod - def convert_tool_entity_to_api_entity( - tool: ApiToolBundle | WorkflowTool | Tool, - tenant_id: str, - labels: list[str] | None = None, - ) -> ToolApiEntity: - """ - convert tool to user tool - """ - if isinstance(tool, Tool): - # fork tool runtime - tool = tool.fork_tool_runtime( - runtime=ToolRuntime( - credentials={}, - tenant_id=tenant_id, - ) - ) - - # get tool parameters - base_parameters = tool.entity.parameters or [] - # get tool runtime parameters - runtime_parameters = tool.get_runtime_parameters() - - # merge parameters using a functional approach to avoid type issues - merged_parameters: list[ToolParameter] = [] - - # create a mapping of runtime parameters for quick lookup - runtime_param_map = {(rp.name, rp.form): rp for rp in runtime_parameters} - - # process base parameters, replacing with runtime versions if they exist - for base_param in base_parameters: - key = (base_param.name, base_param.form) - if key in runtime_param_map: - merged_parameters.append(runtime_param_map[key]) - else: - merged_parameters.append(base_param) - - # add any runtime parameters that weren't in base parameters - for runtime_parameter in runtime_parameters: - if runtime_parameter.form == ToolParameter.ToolParameterForm.FORM: - # check if this parameter is already in merged_parameters - already_exists = any( - p.name == runtime_parameter.name and p.form == runtime_parameter.form for p in merged_parameters - ) - if not already_exists: - merged_parameters.append(runtime_parameter) - - return ToolApiEntity( - author=tool.entity.identity.author, - name=tool.entity.identity.name, - label=tool.entity.identity.label, - description=tool.entity.description.human if tool.entity.description else I18nObject(en_US=""), - output_schema=tool.entity.output_schema, - parameters=merged_parameters, - labels=labels or [], - ) - else: - assert tool.operation_id - return ToolApiEntity( - author=tool.author, - name=tool.operation_id or "", - label=I18nObject(en_US=tool.operation_id, zh_Hans=tool.operation_id), - description=I18nObject(en_US=tool.summary or "", zh_Hans=tool.summary or ""), - output_schema=tool.output_schema, - parameters=tool.parameters, - labels=labels or [], - ) - - @staticmethod - def convert_builtin_provider_to_credential_entity( - provider: BuiltinToolProvider, credentials: dict[str, Any] - ) -> ToolProviderCredentialApiEntity: - return ToolProviderCredentialApiEntity( - id=provider.id, - name=provider.name, - provider=provider.provider, - credential_type=provider.credential_type, - is_default=provider.is_default, - credentials=credentials, - ) - - @staticmethod - def convert_mcp_schema_to_parameter(schema: dict[str, Any]) -> list["ToolParameter"]: - """ - Convert MCP JSON schema to tool parameters - - :param schema: JSON schema dictionary - :return: list of ToolParameter instances - """ - - def resolve_property_type(prop: dict[str, Any], depth: int = 0) -> str: - """ - Resolve a JSON schema property type while guarding against cyclic or deeply nested unions. - """ - if depth >= ToolTransformService._MCP_SCHEMA_TYPE_RESOLUTION_MAX_DEPTH: - return "string" - prop_type = prop.get("type") - match prop_type: - case list(): - non_null_types = [type_name for type_name in prop_type if type_name != "null"] - if non_null_types: - return non_null_types[0] - if prop_type: - return "string" - case str(): - if prop_type == "null": - return "string" - return prop_type - - for union_key in ("anyOf", "oneOf"): - union_schemas = prop.get(union_key) - if not isinstance(union_schemas, list): - continue - - for union_schema in union_schemas: - if not isinstance(union_schema, dict): - continue - union_type = resolve_property_type(union_schema, depth + 1) - if union_type != "null": - return union_type - - all_of_schemas = prop.get("allOf") - if isinstance(all_of_schemas, list): - for all_of_schema in all_of_schemas: - if not isinstance(all_of_schema, dict): - continue - all_of_type = resolve_property_type(all_of_schema, depth + 1) - if all_of_type != "null": - return all_of_type - return "string" - - def create_parameter( - name: str, description: str, param_type: str, required: bool, input_schema: dict[str, Any] | None = None - ) -> ToolParameter: - """Create a ToolParameter instance with given attributes""" - input_schema_dict: dict[str, Any] = {"input_schema": input_schema} if input_schema else {} - return ToolParameter( - name=name, - llm_description=description, - label=I18nObject(en_US=name), - form=ToolParameter.ToolParameterForm.LLM, - required=required, - type=ToolParameter.ToolParameterType(param_type), - human_description=I18nObject(en_US=description), - **input_schema_dict, - ) - - def process_properties( - props: dict[str, dict[str, Any]], required: list[str], prefix: str = "" - ) -> list[ToolParameter]: - """Process properties recursively""" - type_mapping = {"integer": "number", "float": "number"} - complex_types = ["array", "object"] - - parameters = [] - for name, prop in props.items(): - current_description = prop.get("description", "") - prop_type = resolve_property_type(prop) - if prop_type in type_mapping: - prop_type = type_mapping[prop_type] - input_schema = prop if prop_type in complex_types or schema_accepts_multiple_types(prop) else None - parameters.append( - create_parameter(name, current_description, prop_type, name in required, input_schema) - ) - - return parameters - - if schema.get("type") == "object" and "properties" in schema: - return process_properties(schema["properties"], schema.get("required", [])) - return [] diff --git a/api/services/tools/legacy_workflow_tools_manage_service.py b/api/services/tools/legacy_workflow_tools_manage_service.py deleted file mode 100644 index c48de89944fd35..00000000000000 --- a/api/services/tools/legacy_workflow_tools_manage_service.py +++ /dev/null @@ -1,412 +0,0 @@ -import json -import logging -from datetime import datetime -from typing import Any - -from sqlalchemy import delete, or_, select -from sqlalchemy.orm import sessionmaker - -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity -from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration, emoji_icon_adapter -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.model_runtime.utils.encoders import jsonable_encoder -from models.model import App -from models.tools import WorkflowToolProvider -from models.workflow import Workflow -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class WorkflowToolManageService: - """ - Service class for managing workflow tools. - """ - - @staticmethod - def create_workflow_tool( - *, - user_id: str, - tenant_id: str, - workflow_app_id: str, - name: str, - label: str, - icon: dict[str, Any], - description: str, - parameters: list[WorkflowToolParameterConfiguration], - privacy_policy: str = "", - labels: list[str] | None = None, - import_id: str = "", - ): - # check if the name is unique - existing_workflow_tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - # query if the name or app_id exists - existing_workflow_tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where( - WorkflowToolProvider.tenant_id == tenant_id, - # name or app_id - or_(WorkflowToolProvider.name == name, WorkflowToolProvider.app_id == workflow_app_id), - ) - .limit(1) - ) - - # if the name or app_id exists raise error - if existing_workflow_tool_provider is not None: - raise ValueError(f"Tool with name {name} or app_id {workflow_app_id} already exists") - - # query the app and its published workflow in the same session - app: App | None = None - workflow: Workflow | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - app = _session.scalar(select(App).where(App.id == workflow_app_id, App.tenant_id == tenant_id).limit(1)) - if app is not None: - workflow = app.workflow_with_session(session=_session) - - # if not found raise error - if app is None: - raise ValueError(f"App {workflow_app_id} not found") - - # if not found raise error - if workflow is None: - raise ValueError(f"Workflow not found for app {workflow_app_id}") - - # create workflow tool provider - workflow_tool_provider = WorkflowToolProvider( - tenant_id=tenant_id, - user_id=user_id, - app_id=workflow_app_id, - name=name, - label=label, - icon=json.dumps(icon), - description=description, - parameter_configuration=json.dumps([p.model_dump() for p in parameters]), - privacy_policy=privacy_policy, - version=workflow.version, - ) - if import_id: - workflow_tool_provider.id = import_id - try: - WorkflowToolProviderController.from_db(workflow_tool_provider) - except Exception as e: - logger.warning(e, exc_info=True) - raise ValueError(str(e)) from e - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - _session.add(workflow_tool_provider) - - # keep the session open to make orm instances in the same session - if labels is not None: - ToolLabelManager.update_tool_labels( - ToolTransformService.workflow_provider_to_controller(workflow_tool_provider), labels - ) - - return {"result": "success"} - - @classmethod - def update_workflow_tool( - cls, - user_id: str, - tenant_id: str, - workflow_tool_id: str, - name: str, - label: str, - icon: dict[str, Any], - description: str, - parameters: list[WorkflowToolParameterConfiguration], - privacy_policy: str = "", - labels: list[str] | None = None, - ): - """ - Update a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: workflow tool id - :param name: name - :param label: label - :param icon: icon - :param description: description - :param parameters: parameters - :param privacy_policy: privacy policy - :param labels: labels - :return: the updated tool - """ - - existing_workflow_tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - # query if the name exists for other tools - existing_workflow_tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where( - WorkflowToolProvider.tenant_id == tenant_id, - WorkflowToolProvider.name == name, - WorkflowToolProvider.id != workflow_tool_id, - ) - .limit(1) - ) - - # if the name exists raise error - if existing_workflow_tool_provider is not None: - raise ValueError(f"Tool with name {name} already exists") - - # query the workflow tool provider - workflow_tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - workflow_tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id) - .limit(1) - ) - - # if not found raise error - if workflow_tool_provider is None: - raise ValueError(f"Tool {workflow_tool_id} not found") - - # query the app and its published workflow in the same session - app: App | None = None - workflow: Workflow | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - app = _session.scalar( - select(App).where(App.id == workflow_tool_provider.app_id, App.tenant_id == tenant_id).limit(1) - ) - if app is not None: - workflow = app.workflow_with_session(session=_session) - - # if not found raise error - if app is None: - raise ValueError(f"App {workflow_tool_provider.app_id} not found") - - # if not found raise error - if workflow is None: - raise ValueError(f"Workflow not found for app {workflow_tool_provider.app_id}") - - with sessionmaker(db.engine).begin() as _session: - _session.add(workflow_tool_provider) - - # update workflow tool provider - workflow_tool_provider.name = name - workflow_tool_provider.label = label - workflow_tool_provider.icon = json.dumps(icon) - workflow_tool_provider.description = description - workflow_tool_provider.parameter_configuration = json.dumps([p.model_dump() for p in parameters]) - workflow_tool_provider.privacy_policy = privacy_policy - workflow_tool_provider.version = workflow.version - workflow_tool_provider.updated_at = datetime.now() - - try: - WorkflowToolProviderController.from_db(workflow_tool_provider) - except Exception as e: - raise ValueError(str(e)) from e - - if labels is not None: - ToolLabelManager.update_tool_labels( - ToolTransformService.workflow_provider_to_controller(workflow_tool_provider), - labels, - session=_session, - ) - - return {"result": "success"} - - @classmethod - def list_tenant_workflow_tools(cls, user_id: str, tenant_id: str) -> list[ToolProviderApiEntity]: - """ - List workflow tools. - - :param user_id: the user id - :param tenant_id: the tenant id - :return: the list of tools - """ - - providers: list[WorkflowToolProvider] = [] - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - providers = list( - _session.scalars(select(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id)).all() - ) - - # Create a mapping from provider_id to app_id - provider_id_to_app_id = {provider.id: provider.app_id for provider in providers} - - tools: list[WorkflowToolProviderController] = [] - for provider in providers: - try: - tools.append(ToolTransformService.workflow_provider_to_controller(provider)) - except Exception: - # skip deleted tools - logger.exception("Failed to load workflow tool provider %s", provider.id) - - labels = ToolLabelManager.get_tools_labels([tool for tool in tools if isinstance(tool, ToolProviderController)]) - - result: list[ToolProviderApiEntity] = [] - - for tool in tools: - workflow_app_id = provider_id_to_app_id.get(tool.provider_id) - user_tool_provider = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=tool, - labels=labels.get(tool.provider_id, []), - workflow_app_id=workflow_app_id, - ) - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_tool_provider) - user_tool_provider.tools = [ - ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool.get_tools(tenant_id)[0], - labels=labels.get(tool.provider_id, []), - tenant_id=tenant_id, - ) - ] - result.append(user_tool_provider) - - return result - - @classmethod - def delete_workflow_tool(cls, user_id: str, tenant_id: str, workflow_tool_id: str): - """ - Delete a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: the workflow tool id - """ - - with sessionmaker(db.engine).begin() as _session: - _ = _session.execute( - delete(WorkflowToolProvider).where( - WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id - ) - ) - - return {"result": "success"} - - @classmethod - def get_workflow_tool_by_tool_id(cls, user_id: str, tenant_id: str, workflow_tool_id: str): - """ - Get a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: the workflow tool id - :return: the tool - """ - - tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id) - .limit(1) - ) - - return cls._get_workflow_tool(tenant_id, tool_provider) - - @classmethod - def get_workflow_tool_by_app_id(cls, user_id: str, tenant_id: str, workflow_app_id: str): - """ - Get a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_app_id: the workflow app id - :return: the tool - """ - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - tool_provider: WorkflowToolProvider | None = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.app_id == workflow_app_id) - .limit(1) - ) - - return cls._get_workflow_tool(tenant_id, tool_provider) - - @classmethod - def _get_workflow_tool(cls, tenant_id: str, db_tool: WorkflowToolProvider | None): - """ - Get a workflow tool. - - :db_tool: the database tool - :return: the tool - """ - if db_tool is None: - raise ValueError("Tool not found") - - workflow_app: App | None = None - workflow: Workflow | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - workflow_app = _session.scalar( - select(App).where(App.id == db_tool.app_id, App.tenant_id == db_tool.tenant_id).limit(1) - ) - if workflow_app is not None: - workflow = workflow_app.workflow_with_session(session=_session) - - if workflow_app is None: - raise ValueError(f"App {db_tool.app_id} not found") - - if not workflow: - raise ValueError("Workflow not found") - - tool = ToolTransformService.workflow_provider_to_controller(db_tool) - workflow_tools: list[WorkflowTool] = tool.get_tools(tenant_id) - if len(workflow_tools) == 0: - raise ValueError(f"Tool {db_tool.id} not found") - - tool_entity = workflow_tools[0].entity - # get output schema from workflow tool entity - output_schema = tool_entity.output_schema - - return { - "name": db_tool.name, - "label": db_tool.label, - "workflow_tool_id": db_tool.id, - "workflow_app_id": db_tool.app_id, - "icon": emoji_icon_adapter.validate_json(db_tool.icon), - "description": db_tool.description, - "parameters": jsonable_encoder(db_tool.parameter_configurations), - "output_schema": output_schema, - "tool": ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool.get_tools(db_tool.tenant_id)[0], - labels=ToolLabelManager.get_tool_labels(tool), - tenant_id=tenant_id, - ), - "synced": workflow.version == db_tool.version, - "privacy_policy": db_tool.privacy_policy, - } - - @classmethod - def list_single_workflow_tools(cls, user_id: str, tenant_id: str, workflow_tool_id: str) -> list[ToolApiEntity]: - """ - List workflow tool provider tools. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: the workflow tool id - :return: the list of tools - """ - - provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id) - .limit(1) - ) - - if provider is None: - raise ValueError(f"Tool {workflow_tool_id} not found") - - tool = ToolTransformService.workflow_provider_to_controller(provider) - workflow_tools: list[WorkflowTool] = tool.get_tools(tenant_id) - if len(workflow_tools) == 0: - raise ValueError(f"Tool {workflow_tool_id} not found") - - return [ - ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool.get_tools(provider.tenant_id)[0], - labels=ToolLabelManager.get_tool_labels(tool), - tenant_id=tenant_id, - ) - ] diff --git a/api/services/workflow/execution/adapters/persistence.py b/api/services/workflow/execution/adapters/persistence.py index ccc5575d8ee8f2..b8cec78f8d7d72 100644 --- a/api/services/workflow/execution/adapters/persistence.py +++ b/api/services/workflow/execution/adapters/persistence.py @@ -28,6 +28,7 @@ GraphRunAbortedEvent, GraphRunFailedEvent, GraphRunPartialSucceededEvent, + GraphRunPausedEvent, GraphRunStartedEvent, GraphRunSucceededEvent, NodeEvent, @@ -168,6 +169,8 @@ def on_event(self, event: EngineEvent) -> None: self._handle_graph_run_failed(event) case GraphRunAbortedEvent(): self._handle_graph_run_aborted(event) + case GraphRunPausedEvent(): + self._handle_graph_run_paused(event) case NodeRunRetryEvent(): self._handle_node_retry(event) case NodeRunStartedEvent(): @@ -245,6 +248,13 @@ def _finish_workflow_execution(self, execution: WorkflowExecution) -> None: self._enqueue_trace_task(execution) _inspector_publish_workflow_completed(workflow_run_id=execution.id_, status=str(execution.status.value)) + def _handle_graph_run_paused(self, event: GraphRunPausedEvent) -> None: + execution = self._get_workflow_execution() + execution.status = WorkflowExecutionStatus.PAUSED + execution.outputs = event.outputs + self._populate_completion_statistics(execution, update_finished=False) + self._workflow_execution_repository.save(execution) + # ------------------------------------------------------------------ # Node-level handlers # ------------------------------------------------------------------ @@ -442,9 +452,10 @@ def _merge_retry_history( merged_process_data[RETRY_HISTORY_PROCESS_DATA_KEY] = raw_history return merged_process_data - def _populate_completion_statistics(self, execution: WorkflowExecution) -> None: + def _populate_completion_statistics(self, execution: WorkflowExecution, *, update_finished: bool = True) -> None: execution.accumulated_elapsed_time = self._application_generate_entity.workflow_execution_timing.finish() - execution.finished_at = naive_utc_now() + if update_finished: + execution.finished_at = naive_utc_now() runtime_state = self.runtime_state execution.total_tokens = runtime_state.total_tokens execution.total_steps = runtime_state.node_run_steps diff --git a/api/services/workflow_event_snapshot_service.py b/api/services/workflow_event_snapshot_service.py index a15310d0226347..fe165364d6aa46 100644 --- a/api/services/workflow_event_snapshot_service.py +++ b/api/services/workflow_event_snapshot_service.py @@ -53,8 +53,10 @@ from repositories.entities.workflow_pause import WorkflowPauseEntity from repositories.factory import DifyAPIRepositoryFactory from services.app_task_service import AppTaskService +from services.tools.provider_queries import ToolProviderIcons from services.workflow.execution.adapters.response_converter import WorkflowResponseConverter from services.workflow.execution.adapters.response_stream import WorkflowEventStream +from services.workflow.node_execution_extras import node_execution_extras from services.workflow_pause_service import WorkflowPauseSnapshot, load_workflow_pause_snapshot logger = logging.getLogger(__name__) @@ -91,6 +93,7 @@ def build_workflow_event_stream( tenant_id: str, app_id: str, session_maker: sessionmaker[Session], + tool_providers: ToolProviderIcons, human_input_surface: HumanInputSurface | None = None, idle_timeout: float = 300, ping_interval: float = 10.0, @@ -207,6 +210,7 @@ def _generate() -> Generator[Mapping[str, Any] | str]: resumption_context=resumption_context, session_maker=session_maker, human_input_surface=human_input_surface, + tool_providers=tool_providers, ) else: snapshot_events = _build_snapshot_events( @@ -215,6 +219,7 @@ def _generate() -> Generator[Mapping[str, Any] | str]: task_id=task_id, message_context=message_context, node_executions=node_executions, + tool_providers=tool_providers, ) for event in snapshot_events: @@ -355,6 +360,7 @@ def _build_snapshot_events( task_id: str, message_context: MessageContext | None, node_executions: Mapping[str, WorkflowNodeExecutionModel], + tool_providers: ToolProviderIcons, ) -> list[Mapping[str, Any]]: events: list[Mapping[str, Any]] = [] @@ -377,6 +383,7 @@ def _build_snapshot_events( task_id=task_id, snapshot=snapshot, execution=execution, + tool_providers=tool_providers, ) _apply_message_context(node_started, message_context) events.append(node_started) @@ -404,6 +411,7 @@ def _build_paused_snapshot_events( node_executions: Mapping[str, WorkflowNodeExecutionModel], resumption_context: WorkflowResumptionContext | None, session_maker: sessionmaker[Session], + tool_providers: ToolProviderIcons, human_input_surface: HumanInputSurface | None = None, ) -> list[Mapping[str, Any]]: events = _build_snapshot_events( @@ -412,6 +420,7 @@ def _build_paused_snapshot_events( task_id=task_id, message_context=message_context, node_executions=node_executions, + tool_providers=tool_providers, ) pause_snapshot = load_workflow_pause_snapshot( pause_entity, session_factory=session_maker, context=resumption_context @@ -505,6 +514,7 @@ def _build_node_started_event( task_id: str, snapshot: WorkflowNodeExecutionSnapshot, execution: WorkflowNodeExecutionModel | None, + tool_providers: ToolProviderIcons, ) -> dict[str, Any]: created_at = int(snapshot.created_at.timestamp()) if snapshot.created_at else 0 response = NodeStartStreamResponse( @@ -520,7 +530,7 @@ def _build_node_started_event( inputs=execution.inputs_dict if execution else None, inputs_truncated=execution.inputs_truncated if execution else False, created_at=created_at, - extras=execution.extras if execution else {}, + extras=node_execution_extras(execution, tool_providers=tool_providers) if execution else {}, iteration_id=snapshot.iteration_id, loop_id=snapshot.loop_id, ), diff --git a/api/services/workflow_persistence.py b/api/services/workflow_persistence.py deleted file mode 100644 index ccc5575d8ee8f2..00000000000000 --- a/api/services/workflow_persistence.py +++ /dev/null @@ -1,549 +0,0 @@ -"""Persist workflow execution events on behalf of the upper run orchestrator. - -The aggregate constructs this observer and supplies its repositories, indexes, -and history before the engine starts. Engine events remain storage-agnostic. -""" - -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime -from typing import Any, override - -from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity -from core.app.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt -from core.helper.trace_id_helper import ParentTraceContext -from core.ops.base_trace_instance import application_tracing_hooks -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository -from core.tools.workflow_as_tool.repository import WorkflowToolSource -from core.workflow.node_execution import WorkflowNodeExecution as DifyWorkflowNodeExecution -from core.workflow.node_execution_process_data import WORKFLOW_TOOL_ROOT_APP_ID_KEY, keep_agent_and_tool_ids -from core.workflow.system_variables import SystemVariableKey -from core.workflow.variable_prefixes import SYSTEM_VARIABLE_NODE_ID -from core.workflow.workflow_run_outputs import project_node_outputs_for_workflow_run -from graphon.engine.layer import Layer -from graphon.engine_events import ( - EngineEvent, - GraphRunAbortedEvent, - GraphRunFailedEvent, - GraphRunPartialSucceededEvent, - GraphRunStartedEvent, - GraphRunSucceededEvent, - NodeEvent, - NodeRunExceptionEvent, - NodeRunFailedEvent, - NodeRunPauseRequestedEvent, - NodeRunRetryEvent, - NodeRunStartedEvent, - NodeRunSucceededEvent, -) -from graphon.entities import WorkflowNodeExecution -from graphon.enums import ( - BuiltinNodeTypes, - WorkflowExecutionStatus, - WorkflowNodeExecutionMetadataKey, - WorkflowNodeExecutionStatus, - WorkflowType, -) -from graphon.node_events import NodeRunResult -from libs.datetime_utils import naive_utc_now -from repositories.entities.workflow_execution import WorkflowExecution -from services.workflow.inspector_events import ( - publish_node_changed as _inspector_publish_node_changed, -) -from services.workflow.inspector_events import ( - publish_workflow_completed as _inspector_publish_workflow_completed, -) - - -@dataclass(slots=True) -class PersistenceWorkflowInfo: - """Static workflow metadata required for persistence.""" - - workflow_id: str - workflow_type: WorkflowType - version: str - graph_data: Mapping[str, Any] - - -class WorkflowPersistenceLayer(Layer): - """Engine layer that persists workflow and node execution state.""" - - def __init__( - self, - *, - application_generate_entity: AdvancedChatAppGenerateEntity | WorkflowAppGenerateEntity, - workflow_info: PersistenceWorkflowInfo, - workflow_execution_repository: WorkflowExecutionRepository, - workflow_node_execution_repository: WorkflowNodeExecutionRepository, - trace_manager: TraceQueueManager | None = None, - ) -> None: - super().__init__() - self._application_generate_entity = application_generate_entity - self._workflow_info = workflow_info - self._workflow_execution_repository = workflow_execution_repository - self._workflow_node_execution_repository = workflow_node_execution_repository - self._trace_manager = trace_manager - - self._workflow_execution: WorkflowExecution | None = None - self._node_execution_cache: dict[str, WorkflowNodeExecution] = {} - self._node_run_indices: Mapping[str, int] = {} - self._workflow_tool_layers: dict[tuple[str, str], WorkflowPersistenceLayer] = {} - - @property - def workflow_execution(self) -> WorkflowExecution: - return self._get_workflow_execution() - - def set_node_run_indices(self, indices: Mapping[str, int]) -> None: - """Consume the upper layer's indices, allocated before engine event delivery.""" - self._node_run_indices = indices - - def set_node_execution_history(self, executions: Sequence[WorkflowNodeExecution]) -> None: - """Restore materialized history before execution starts or a source frame is scheduled.""" - self._node_execution_cache = { - execution.node_execution_id or execution.id: execution for execution in executions - } - - def create_workflow_tool_event_listener( - self, - source: WorkflowToolSource, - caller_workflow_id: str, - caller_node_execution_id: str, - *, - node_executions: Sequence[WorkflowNodeExecution] = (), - ) -> Callable[[NodeEvent], None]: - """Persist hidden tool nodes under their source app, without creating another run.""" - key = (source.app_id, source.workflow_id) - tool_layer = self._workflow_tool_layers.get(key) - if tool_layer is None: - tool_layer = WorkflowPersistenceLayer( - application_generate_entity=self._application_generate_entity, - workflow_info=PersistenceWorkflowInfo( - workflow_id=source.workflow_id, - workflow_type=WorkflowType.WORKFLOW, - version="", - graph_data=source.graph_config, - ), - workflow_execution_repository=self._workflow_execution_repository, - workflow_node_execution_repository=self._workflow_node_execution_repository.for_workflow_tool( - source.app_id - ), - ) - tool_layer.set_node_run_indices(self._node_run_indices) - tool_layer.set_node_execution_history(node_executions) - self._workflow_tool_layers[key] = tool_layer - - def on_node_event(event: NodeEvent) -> None: - tool_layer._workflow_execution = self._get_workflow_execution() - # Lifecycle ownership must survive independently of the configured run-log store. - event.node_run_result.process_data = { - **event.node_run_result.process_data, - WORKFLOW_TOOL_ROOT_APP_ID_KEY: self._application_generate_entity.app_config.app_id, - } - if isinstance(event, NodeRunStartedEvent) and not isinstance(event, NodeRunRetryEvent): - tool_layer._handle_node_started( - event, - caller_workflow_id=caller_workflow_id, - caller_node_execution_id=caller_node_execution_id, - ) - else: - tool_layer.on_event(event) - - return on_node_event - - # ------------------------------------------------------------------ - # Layer lifecycle - # ------------------------------------------------------------------ - @override - def on_event(self, event: EngineEvent) -> None: - match event: - case GraphRunStartedEvent(): - self._handle_graph_run_started() - case GraphRunSucceededEvent(): - self._handle_graph_run_succeeded(event) - case GraphRunPartialSucceededEvent(): - self._handle_graph_run_partial_succeeded(event) - case GraphRunFailedEvent(): - self._handle_graph_run_failed(event) - case GraphRunAbortedEvent(): - self._handle_graph_run_aborted(event) - case NodeRunRetryEvent(): - self._handle_node_retry(event) - case NodeRunStartedEvent(): - self._handle_node_started(event) - case NodeRunSucceededEvent(): - self._handle_node_succeeded(event) - case NodeRunFailedEvent(): - self._handle_node_failed(event) - case NodeRunExceptionEvent(): - self._handle_node_exception(event) - case NodeRunPauseRequestedEvent(): - self._handle_node_pause_requested(event) - - # ------------------------------------------------------------------ - # Graph-level handlers - # ------------------------------------------------------------------ - def _handle_graph_run_started(self) -> None: - execution_id = self._get_execution_id() - timing = self._application_generate_entity.workflow_execution_timing - timing.start() - workflow_execution = WorkflowExecution( - id_=execution_id, - workflow_id=self._workflow_info.workflow_id, - workflow_type=self._workflow_info.workflow_type, - workflow_version=self._workflow_info.version, - graph=self._workflow_info.graph_data, - inputs=self._prepare_workflow_inputs(), - started_at=naive_utc_now(), - accumulated_elapsed_time=timing.elapsed_time, - ) - - self._workflow_execution_repository.save_synchronously(workflow_execution) - self._workflow_execution = workflow_execution - tracing_hooks = application_tracing_hooks() - if tracing_hooks is not None: - tracing_hooks.freeze_workflow_run( - self._application_generate_entity, - workflow_id=self._workflow_info.workflow_id, - workflow_run_id=execution_id, - conversation_id=self._system_variables().get(SystemVariableKey.CONVERSATION_ID.value), - graph_data=self._workflow_info.graph_data, - ) - - def _handle_graph_run_succeeded(self, event: GraphRunSucceededEvent) -> None: - execution = self._get_workflow_execution() - execution.outputs = event.outputs - execution.status = WorkflowExecutionStatus.SUCCEEDED - self._finish_workflow_execution(execution) - - def _handle_graph_run_partial_succeeded(self, event: GraphRunPartialSucceededEvent) -> None: - execution = self._get_workflow_execution() - execution.outputs = event.outputs - execution.status = WorkflowExecutionStatus.PARTIAL_SUCCEEDED - execution.exceptions_count = event.exceptions_count - self._finish_workflow_execution(execution) - - def _handle_graph_run_failed(self, event: GraphRunFailedEvent) -> None: - execution = self._get_workflow_execution() - execution.status = WorkflowExecutionStatus.FAILED - execution.error_message = event.error - execution.exceptions_count = event.exceptions_count - self._finish_workflow_execution(execution) - - def _handle_graph_run_aborted(self, event: GraphRunAbortedEvent) -> None: - execution = self._get_workflow_execution() - execution.status = WorkflowExecutionStatus.STOPPED - execution.error_message = event.reason or "Workflow execution aborted" - self._finish_workflow_execution(execution) - - def _finish_workflow_execution(self, execution: WorkflowExecution) -> None: - self._populate_completion_statistics(execution) - if execution.status in (WorkflowExecutionStatus.FAILED, WorkflowExecutionStatus.STOPPED): - self._fail_running_node_executions(error_message=execution.error_message or "") - self._workflow_execution_repository.save_synchronously(execution) - self._enqueue_trace_task(execution) - _inspector_publish_workflow_completed(workflow_run_id=execution.id_, status=str(execution.status.value)) - - # ------------------------------------------------------------------ - # Node-level handlers - # ------------------------------------------------------------------ - def _handle_node_started( - self, - event: NodeRunStartedEvent, - *, - caller_workflow_id: str | None = None, - caller_node_execution_id: str | None = None, - ) -> None: - execution = self._get_workflow_execution() - iteration_id = event.node_run_result.metadata.get(WorkflowNodeExecutionMetadataKey.ITERATION_ID) - loop_id = event.node_run_result.metadata.get(WorkflowNodeExecutionMetadataKey.LOOP_ID) - - metadata: dict[WorkflowNodeExecutionMetadataKey, Any] = { - **event.node_run_result.metadata, - WorkflowNodeExecutionMetadataKey.ITERATION_ID: iteration_id, - WorkflowNodeExecutionMetadataKey.LOOP_ID: loop_id, - } - if event.node_type == BuiltinNodeTypes.TOOL and event.provider_type: - metadata[WorkflowNodeExecutionMetadataKey.TOOL_INFO] = { - "provider_type": event.provider_type, - "provider_id": event.provider_id, - } - - previous = self._node_execution_cache.get(event.id) - domain_execution = DifyWorkflowNodeExecution( - id=previous.id if previous is not None else event.id, - node_execution_id=event.id, - workflow_id=self._workflow_info.workflow_id, - triggered_from_workflow_id=caller_workflow_id, - triggered_from_node_execution_id=caller_node_execution_id, - workflow_execution_id=execution.id_, - predecessor_node_id=event.predecessor_node_id, - index=self._node_run_indices[event.id], - node_id=event.node_id, - node_type=event.node_type, - title=event.node_title, - status=WorkflowNodeExecutionStatus.RUNNING, - process_data=event.node_run_result.process_data or None, - metadata=metadata, - created_at=previous.created_at if previous is not None else event.start_at, - ) - - self._node_execution_cache[event.id] = domain_execution - if event.node_type == BuiltinNodeTypes.AGENT and event.node_version == "2": - self._workflow_node_execution_repository.save_synchronously(domain_execution) - else: - self._workflow_node_execution_repository.save(domain_execution) - - _inspector_publish_node_changed(workflow_run_id=execution.id_, node_id=event.node_id, status="running") - - def _handle_node_retry(self, event: NodeRunRetryEvent) -> None: - domain_execution = self._get_node_execution(event.id) - domain_execution.status = WorkflowNodeExecutionStatus.RETRY - domain_execution.error = event.error - self._append_retry_history(domain_execution, event) - self._workflow_node_execution_repository.save(domain_execution) - self._workflow_node_execution_repository.save_execution_data(domain_execution) - _inspector_publish_node_changed( - workflow_run_id=self._get_workflow_execution().id_, - node_id=domain_execution.node_id, - status="retry", - ) - - def _handle_node_succeeded(self, event: NodeRunSucceededEvent) -> None: - domain_execution = self._get_node_execution(event.id) - self._update_node_execution( - domain_execution, - event.node_run_result, - WorkflowNodeExecutionStatus.SUCCEEDED, - finished_at=event.finished_at, - ) - _inspector_publish_node_changed( - workflow_run_id=self._get_workflow_execution().id_, - node_id=domain_execution.node_id, - status="succeeded", - ) - - def _handle_node_failed(self, event: NodeRunFailedEvent) -> None: - domain_execution = self._get_node_execution(event.id) - self._update_node_execution( - domain_execution, - event.node_run_result, - WorkflowNodeExecutionStatus.FAILED, - error=event.error, - finished_at=event.finished_at, - ) - _inspector_publish_node_changed( - workflow_run_id=self._get_workflow_execution().id_, - node_id=domain_execution.node_id, - status="failed", - ) - - def _handle_node_exception(self, event: NodeRunExceptionEvent) -> None: - domain_execution = self._get_node_execution(event.id) - self._update_node_execution( - domain_execution, - event.node_run_result, - WorkflowNodeExecutionStatus.EXCEPTION, - error=event.error, - finished_at=event.finished_at, - ) - _inspector_publish_node_changed( - workflow_run_id=self._get_workflow_execution().id_, - node_id=domain_execution.node_id, - status="exception", - ) - - def _handle_node_pause_requested(self, event: NodeRunPauseRequestedEvent) -> None: - domain_execution = self._get_node_execution(event.id) - self._update_node_execution( - domain_execution, - event.node_run_result, - WorkflowNodeExecutionStatus.PAUSED, - error="", - update_outputs=False, - ) - - # ------------------------------------------------------------------ - # Helpers - # ------------------------------------------------------------------ - def _get_execution_id(self) -> str: - workflow_execution_id = self._system_variables().get(SystemVariableKey.WORKFLOW_EXECUTION_ID) - if not workflow_execution_id: - raise ValueError("workflow_execution_id must be provided in system variables for pause/resume flows") - return str(workflow_execution_id) - - def _prepare_workflow_inputs(self) -> Mapping[str, Any]: - inputs = {**self._application_generate_entity.inputs} - for field_name, value in self._system_variables().items(): - if field_name == SystemVariableKey.CONVERSATION_ID.value: - # Conversation IDs are tied to the current session; omit them so persisted - # workflow inputs stay reusable without binding future runs to this conversation. - continue - inputs[f"sys.{field_name}"] = value - # Local import to avoid circular dependency during app bootstrapping. - from services.workflow.execution.adapters.workflow_entry import WorkflowEntry - - handled = WorkflowEntry.handle_special_values(inputs) - return handled or {} - - def _get_workflow_execution(self) -> WorkflowExecution: - if self._workflow_execution is None: - raise ValueError("workflow execution not initialized") - return self._workflow_execution - - def _get_node_execution(self, node_execution_id: str) -> WorkflowNodeExecution: - if node_execution_id not in self._node_execution_cache: - raise ValueError(f"Node execution not found for id={node_execution_id}") - return self._node_execution_cache[node_execution_id] - - def _append_retry_history(self, execution: WorkflowNodeExecution, event: NodeRunRetryEvent) -> None: - """Append a validated full attempt before repository truncation or offload.""" - finished_at = naive_utc_now() - process_data = keep_agent_and_tool_ids( - event.node_run_result.process_data, - execution.process_data, - ) - process_data = dict(process_data or {}) - raw_history = process_data.get(RETRY_HISTORY_PROCESS_DATA_KEY) - history = list(raw_history) if isinstance(raw_history, list) else [] - projected_outputs = project_node_outputs_for_workflow_run( - node_type=execution.node_type, - inputs=event.node_run_result.inputs, - outputs=event.node_run_result.outputs, - ) - attempt = WorkflowNodeRetryAttempt( - retry_index=event.retry_index, - inputs=event.node_run_result.inputs, - process_data=event.node_run_result.process_data, - outputs=projected_outputs, - error=event.node_run_result.error or event.error, - elapsed_time=max((finished_at - event.start_at).total_seconds(), 0.0), - execution_metadata={key.value: value for key, value in event.node_run_result.metadata.items()}, - created_at=int(event.start_at.timestamp()), - finished_at=int(finished_at.timestamp()), - ) - history.append(attempt.model_dump(mode="json")) - process_data[RETRY_HISTORY_PROCESS_DATA_KEY] = history - execution.process_data = process_data - - @staticmethod - def _merge_retry_history( - existing_process_data: Mapping[str, Any] | None, - next_process_data: Mapping[str, Any] | None, - ) -> Mapping[str, Any] | None: - """Keep internal retry history while replacing node-specific Process Data.""" - merged_process_data = keep_agent_and_tool_ids(existing_process_data, next_process_data) - raw_history = (existing_process_data or {}).get(RETRY_HISTORY_PROCESS_DATA_KEY) - if not isinstance(raw_history, list) or not raw_history: - return merged_process_data - - merged_process_data = dict(merged_process_data or {}) - merged_process_data[RETRY_HISTORY_PROCESS_DATA_KEY] = raw_history - return merged_process_data - - def _populate_completion_statistics(self, execution: WorkflowExecution) -> None: - execution.accumulated_elapsed_time = self._application_generate_entity.workflow_execution_timing.finish() - execution.finished_at = naive_utc_now() - runtime_state = self.runtime_state - execution.total_tokens = runtime_state.total_tokens - execution.total_steps = runtime_state.node_run_steps - execution.outputs = execution.outputs or runtime_state.outputs - execution.exceptions_count = max(execution.exceptions_count, runtime_state.exceptions_count) - - def _update_node_execution( - self, - domain_execution: WorkflowNodeExecution, - node_result: NodeRunResult, - status: WorkflowNodeExecutionStatus, - *, - error: str | None = None, - update_outputs: bool = True, - finished_at: datetime | None = None, - ) -> None: - actual_finished_at = finished_at or naive_utc_now() - domain_execution.status = status - domain_execution.finished_at = actual_finished_at - domain_execution.elapsed_time = max((actual_finished_at - domain_execution.created_at).total_seconds(), 0.0) - - if error: - domain_execution.error = error - - if update_outputs: - projected_outputs = project_node_outputs_for_workflow_run( - node_type=domain_execution.node_type, - inputs=node_result.inputs, - outputs=node_result.outputs, - ) - process_data = self._merge_retry_history(domain_execution.process_data, node_result.process_data) - domain_execution.update_from_mapping( - inputs=node_result.inputs, - process_data=process_data, - outputs=projected_outputs, - metadata={**(domain_execution.metadata or {}), **node_result.metadata}, - ) - else: - domain_execution.process_data = keep_agent_and_tool_ids( - node_result.process_data, - domain_execution.process_data, - ) - - self._workflow_node_execution_repository.save(domain_execution) - self._workflow_node_execution_repository.save_execution_data(domain_execution) - - def persist_unfinished_nodes_synchronously(self) -> None: - """Make paused participants durable before exposing a resumable pause.""" - for tool_layer in self._workflow_tool_layers.values(): - tool_layer.persist_unfinished_nodes_synchronously() - for execution in self._node_execution_cache.values(): - if execution.status in ( - WorkflowNodeExecutionStatus.RUNNING, - WorkflowNodeExecutionStatus.PAUSED, - WorkflowNodeExecutionStatus.RETRY, - ): - self._workflow_node_execution_repository.save_synchronously(execution) - - def _fail_running_node_executions(self, *, error_message: str) -> None: - for tool_layer in self._workflow_tool_layers.values(): - tool_layer._fail_running_node_executions(error_message=error_message) - now = naive_utc_now() - unfinished = [] - for execution in self._node_execution_cache.values(): - if execution.status in (WorkflowNodeExecutionStatus.RUNNING, WorkflowNodeExecutionStatus.RETRY): - execution.status = WorkflowNodeExecutionStatus.FAILED - execution.error = error_message - execution.finished_at = now - execution.elapsed_time = max((now - execution.created_at).total_seconds(), 0.0) - unfinished.append(execution) - self._workflow_node_execution_repository.save_many(unfinished) - - def _enqueue_trace_task(self, execution: WorkflowExecution) -> None: - if not self._trace_manager: - return - - conversation_id = self._system_variables().get(SystemVariableKey.CONVERSATION_ID.value) - external_trace_id = None - trace_session_id = None - parent_trace_context = None - if isinstance(self._application_generate_entity, (WorkflowAppGenerateEntity, AdvancedChatAppGenerateEntity)): - extras = self._application_generate_entity.extras - external_trace_id = extras.get("external_trace_id") - trace_session_id = extras.get("trace_session_id") - parent_trace_context = extras.get("parent_trace_context") - if isinstance(parent_trace_context, ParentTraceContext): - parent_trace_context = parent_trace_context.model_dump(exclude_none=True) - - trace_task = TraceTask( - TraceTaskName.WORKFLOW_TRACE, - workflow_execution=execution, - conversation_id=conversation_id, - user_id=self._trace_manager.user_id, - external_trace_id=external_trace_id, - trace_session_id=trace_session_id, - parent_trace_context=parent_trace_context, - ) - self._trace_manager.add_trace_task(trace_task) - - def _system_variables(self) -> Mapping[str, Any]: - runtime_state = self.runtime_state - return runtime_state.variable_pool.get_by_prefix(SYSTEM_VARIABLE_NODE_ID) diff --git a/api/services/workflow_run_agg.py b/api/services/workflow_run_agg.py index c2177f73f3d680..a95cee78b68e02 100644 --- a/api/services/workflow_run_agg.py +++ b/api/services/workflow_run_agg.py @@ -56,8 +56,8 @@ from services.workflow.execution.adapters.human_input_events import build_human_input_pause_reason from services.workflow.execution.adapters.node_factory import DifyNodeFactory from services.workflow.execution.adapters.node_runtime import resolve_dify_run_context +from services.workflow.execution.adapters.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from services.workflow.execution.contracts import PreparedWorkflowRun, WorkflowRunRunner -from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from services.workflow_run_index import WorkflowRunIndex from services.workflow_tool_source_service import WorkflowToolSnapshot, WorkflowToolSourceService diff --git a/api/tests/e2e_tests/containers/services/workflow/test_workflow_event_snapshot_service.py b/api/tests/e2e_tests/containers/services/workflow/test_workflow_event_snapshot_service.py index 4a2ada49a14b6b..f7d6a59f5a4fd4 100644 --- a/api/tests/e2e_tests/containers/services/workflow/test_workflow_event_snapshot_service.py +++ b/api/tests/e2e_tests/containers/services/workflow/test_workflow_event_snapshot_service.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import override +from unittest.mock import Mock from uuid import uuid4 from sqlalchemy import Engine, delete @@ -173,6 +174,7 @@ def test_build_snapshot_events_resolves_variable_select_options(db_session_with_ resumption_context=_build_resumption_context(workflow_run_id, form.id), session_maker=session_maker, node_executions={}, + tool_providers=Mock(), ) human_input_events = [event for event in events if event.get("event") == "human_input_required"] diff --git a/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py b/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py index 6ffc67e61bce0d..478069ca51f211 100644 --- a/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py +++ b/api/tests/e2e_tests/contracts/core/app/apps/completion/test_prepared_completion.py @@ -19,9 +19,13 @@ PromptTemplateEntity, ) from core.app.apps.completion.app_config_manager import CompletionAppConfig -from core.app.apps.completion.app_generator import CompletionAppGenerator from core.app.entities.app_invoke_entities import CompletionAppGenerateEntity, InvokeFrom -from models import AppMode, Conversation, Message +from extensions.application_services.annotation import build_annotation_replies +from extensions.application_services.retrieval import build_dataset_retrieval +from models import App, AppMode, Conversation, EndUser, Message +from models.enums import EndUserType +from repositories.app.generation_repository import AppGenerationRepository +from services.app.generation.adapters.completion import CompletionAppGenerator from tests.helpers.unit.core.model_fixtures import make_model_config @@ -50,13 +54,35 @@ def _entity(invoke_from: InvokeFrom) -> CompletionAppGenerateEntity: ) +def _generator(session_factory: sessionmaker[Session]) -> CompletionAppGenerator: + return CompletionAppGenerator( + retrieval=build_dataset_retrieval(session_factory), + records=AppGenerationRepository(session_factory), + annotations=build_annotation_replies(session_factory), + ) + + +def _save_app(entity: CompletionAppGenerateEntity, session_factory: sessionmaker[Session]) -> None: + with session_factory.begin() as session: + session.add( + App( + id=entity.app_config.app_id, + tenant_id=entity.app_config.tenant_id, + name="Prepared completion", + mode=AppMode.COMPLETION, + enable_site=True, + enable_api=True, + ) + ) + + @pytest.mark.parametrize("invoke_from", [InvokeFrom.WEB_APP, InvokeFrom.EXPLORE]) def test_prepared_records_preserve_input_and_actor_ownership( sqlite_session_factory: sessionmaker[Session], invoke_from: InvokeFrom ) -> None: entity = _entity(invoke_from) - with sqlite_session_factory(expire_on_commit=False) as session: - conversation, message = CompletionAppGenerator()._init_generate_records(entity, session=session) + _save_app(entity, sqlite_session_factory) + conversation, message = _generator(sqlite_session_factory)._init_generate_records(entity) with sqlite_session_factory() as verification: saved = verification.get(Message, message.id) saved_conversation = verification.get(Conversation, conversation.id) @@ -83,9 +109,17 @@ def test_invalid_prepared_input_rolls_back_without_starting_a_worker( sqlite_session_factory: sessionmaker[Session], ) -> None: entity = _entity(InvokeFrom.WEB_APP) + _save_app(entity, sqlite_session_factory) entity.inputs = {**entity.inputs, "invalid_json": object()} + user = EndUser( + id=entity.user_id, + tenant_id=entity.app_config.tenant_id, + app_id=entity.app_config.app_id, + type=EndUserType.BROWSER, + session_id=str(uuid4()), + ) with pytest.raises(StatementError) as error: - CompletionAppGenerator().generate_from_entity(entity, session_factory=sqlite_session_factory) + _generator(sqlite_session_factory).generate_from_entity(entity, user=user) assert isinstance(error.value.orig, TypeError) with sqlite_session_factory() as verification: assert verification.scalar(select(func.count()).select_from(Message)) == 0 diff --git a/api/tests/unit_tests/controllers/console/test_human_input_form.py b/api/tests/unit_tests/controllers/console/test_human_input_form.py index 7a210601398029..3ae56a835da8ed 100644 --- a/api/tests/unit_tests/controllers/console/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/console/test_human_input_form.py @@ -430,6 +430,12 @@ def get_workflow_run_by_id_and_tenant_id(self, **_kwargs): return workflow_run snapshot_calls: list[dict[str, object]] = [] + tool_providers = Mock() + monkeypatch.setattr( + human_input_form_module, + "application_services", + lambda: SimpleNamespace(tools=SimpleNamespace(tool_providers=tool_providers)), + ) def snapshot_events() -> Iterator[dict[str, str]]: assert has_request_context() @@ -470,6 +476,7 @@ def snapshot_builder(**kwargs: object) -> Iterator[dict[str, str]]: "tenant_id": "t1", "app_id": "app-1", "session_maker": ANY, + "tool_providers": tool_providers, "human_input_surface": HumanInputSurface.CONSOLE, "include_node_details": include_node_details, "close_on_pause": False, diff --git a/api/tests/unit_tests/controllers/openapi/test_workflow_events.py b/api/tests/unit_tests/controllers/openapi/test_workflow_events.py index a226fc2611f2c9..2c853ded7c96b9 100644 --- a/api/tests/unit_tests/controllers/openapi/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/openapi/test_workflow_events.py @@ -11,6 +11,7 @@ from __future__ import annotations import sys +from collections.abc import Iterator from types import SimpleNamespace from unittest.mock import Mock @@ -245,7 +246,19 @@ def test_stream_reads_everything_off_the_context_before_streaming( msg_gen_mock = Mock() msg_gen_mock.retrieve_events.return_value = iter(["a", "b"]) monkeypatch.setattr(module, "WorkflowEventStream", msg_gen_mock) - monkeypatch.setattr(module, "build_workflow_event_stream", lambda **_kwargs: iter(["a", "b"])) + snapshot_calls: list[dict[str, object]] = [] + + def snapshot_builder(**kwargs: object) -> Iterator[str]: + snapshot_calls.append(kwargs) + return iter(["a", "b"]) + + monkeypatch.setattr(module, "build_workflow_event_stream", snapshot_builder) + tool_providers = Mock() + monkeypatch.setattr( + module, + "application_services", + lambda: SimpleNamespace(tools=SimpleNamespace(tool_providers=tool_providers)), + ) ctx = _context(_make_account(), CreatorUserRole.ACCOUNT) api = OpenApiWorkflowEventsApi() @@ -257,6 +270,10 @@ def test_stream_reads_everything_off_the_context_before_streaming( assert resp.mimetype == "text/event-stream" assert body == "event: a\n\nevent: b\n\n" + if include_state_snapshot: + assert snapshot_calls[0]["tool_providers"] is tool_providers + else: + assert snapshot_calls == [] @pytest.mark.parametrize("continue_on_pause", [False, True], ids=["close", "keep_open"]) def test_continue_on_pause_comes_off_the_declared_query_model( diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py index 8a139a73689442..486e55136e8f54 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow_events.py @@ -400,6 +400,12 @@ def test_workflow_events_snapshot_continue_on_pause_keeps_pause_open( ) msg_generator = Mock() snapshot_calls: list[dict[str, object]] = [] + tool_providers = Mock() + monkeypatch.setattr( + workflow_events_module, + "application_services", + lambda: SimpleNamespace(tools=SimpleNamespace(tool_providers=tool_providers)), + ) def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: snapshot_calls.append(kwargs) @@ -428,6 +434,7 @@ def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: "tenant_id": "tenant-1", "app_id": "app-1", "session_maker": ANY, + "tool_providers": tool_providers, "human_input_surface": HumanInputSurface.SERVICE_API, "close_on_pause": False, } @@ -827,6 +834,7 @@ def test_snapshot_events_include_pause_payload_contract( resumption_context=resumption_context, session_maker=sqlite_session_maker, node_executions={}, + tool_providers=Mock(), ) assert [event["event"] for event in events] == [ @@ -1017,6 +1025,12 @@ def test_running_run_with_snapshot( ) msg_generator = Mock() snapshot_calls: list[dict[str, object]] = [] + tool_providers = Mock() + monkeypatch.setattr( + workflow_events_module, + "application_services", + lambda: SimpleNamespace(tools=SimpleNamespace(tool_providers=tool_providers)), + ) def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: snapshot_calls.append(kwargs) diff --git a/api/tests/unit_tests/controllers/web/test_workflow_events.py b/api/tests/unit_tests/controllers/web/test_workflow_events.py index ea7d4ba9fb6602..aca28a0f223597 100644 --- a/api/tests/unit_tests/controllers/web/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/web/test_workflow_events.py @@ -3,6 +3,7 @@ from __future__ import annotations from datetime import datetime +from types import SimpleNamespace from unittest.mock import ANY, MagicMock, patch import pytest @@ -147,6 +148,11 @@ def test_snapshot_stream_can_continue_across_pauses( mock_factory.create_api_workflow_run_repository.return_value = mock_repo snapshot_calls: list[dict[str, object]] = [] + tool_providers = MagicMock() + monkeypatch.setattr( + "controllers.web.workflow_events.application_services", + lambda: SimpleNamespace(tools=SimpleNamespace(tool_providers=tool_providers)), + ) def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: snapshot_calls.append(kwargs) @@ -165,6 +171,7 @@ def snapshot_builder(**kwargs: object) -> list[dict[str, str]]: "tenant_id": "tenant-1", "app_id": "app-1", "session_maker": ANY, + "tool_providers": tool_providers, "close_on_pause": False, } ] diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index 476fe7e09a54a8..e1d15d414ac190 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py @@ -23,7 +23,6 @@ from constants import UUID_NIL from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig from core.app.apps.exc import GenerateTaskStoppedError -from core.app.apps.workflow import app_generator as wf_app_gen_module from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom from core.app.file_access import DatabaseFileAccessController from core.app.workflow.file_runtime import DifyWorkflowFileRuntime @@ -69,6 +68,7 @@ from services.workflow.execution.adapters.chatflow.app_generator import AdvancedChatAppGenerator from services.workflow.execution.adapters.chatflow.record_preparation import prepare_chat_records from services.workflow.execution.adapters.node_factory import DifyNodeFactory +from services.workflow.execution.adapters.workflow import app_generator as wf_app_gen_module from services.workflow.execution.ports import WorkflowRuntime from tests.helpers.unit.config_override import apply_config_overrides from tests.helpers.unit.model_factories import make_upload_file diff --git a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py index 39d3e42582dcf3..67b1da962b61e9 100644 --- a/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py +++ b/api/tests/unit_tests/core/app/apps/common/test_workflow_response_converter.py @@ -46,9 +46,11 @@ from models.human_input_entities import ParagraphInputConfig, SelectInputConfig, StringListSource, UserActionConfig from models.model import AppMode from models.workflow import WorkflowNodeExecutionModel +from services.tools.provider_queries import ToolProviderIcons from services.workflow.execution.adapters import response_converter as workflow_response_converter from services.workflow.execution.adapters.events import WorkflowEventPublisher from services.workflow.execution.adapters.response_converter import WorkflowResponseConverter +from services.workflow.execution.ports import ExecutionContexts class TestWorkflowResponseConverterCases: @@ -441,7 +443,7 @@ def test_human_input_form_timeout_stream_response_contains_timeout_metadata( def _workflow_response_converter_resumption_build_converter( - *, workflow_contexts, tool_providers + *, workflow_contexts: ExecutionContexts, tool_providers: ToolProviderIcons ) -> WorkflowResponseConverter: """Construct a minimal WorkflowResponseConverter for testing.""" system_variables = build_system_variables( @@ -478,7 +480,9 @@ def _workflow_response_converter_resumption_build_converter( class TestWorkflowResponseConverterResumptionCases: @staticmethod - def test_workflow_start_stream_response_carries_resumption_reason(*, workflow_contexts, tool_providers) -> None: + def test_workflow_start_stream_response_carries_resumption_reason( + *, workflow_contexts: ExecutionContexts, tool_providers: ToolProviderIcons + ) -> None: converter = _workflow_response_converter_resumption_build_converter( workflow_contexts=workflow_contexts, tool_providers=tool_providers ) @@ -491,7 +495,9 @@ def test_workflow_start_stream_response_carries_resumption_reason(*, workflow_co assert resp.data.reason is WorkflowStartReason.RESUMPTION @staticmethod - def test_workflow_start_stream_response_carries_initial_reason(*, workflow_contexts, tool_providers) -> None: + def test_workflow_start_stream_response_carries_initial_reason( + *, workflow_contexts: ExecutionContexts, tool_providers: ToolProviderIcons + ) -> None: converter = _workflow_response_converter_resumption_build_converter( workflow_contexts=workflow_contexts, tool_providers=tool_providers ) @@ -505,8 +511,8 @@ def test_workflow_start_stream_response_carries_initial_reason(*, workflow_conte @staticmethod def test_node_start_preserves_the_allocated_index_when_events_arrive_out_of_order( - *, workflow_contexts, tool_providers - ): + *, workflow_contexts: ExecutionContexts, tool_providers: ToolProviderIcons + ) -> None: converter = _workflow_response_converter_resumption_build_converter( workflow_contexts=workflow_contexts, tool_providers=tool_providers ) @@ -536,7 +542,6 @@ def test_workflow_tool_detail_hint_in_live_and_saved_traces( monkeypatch, provider_type, *, workflow_contexts, tool_providers ): monkeypatch.setattr(workflow_response_converter.ToolManager, "get_tool_icon", lambda **_kwargs: "icon") - monkeypatch.setattr("core.tools.tool_manager.ToolManager.get_tool_icon", lambda **_kwargs: "icon") converter = _workflow_response_converter_resumption_build_converter( workflow_contexts=workflow_contexts, tool_providers=tool_providers ) @@ -562,13 +567,17 @@ def test_workflow_tool_detail_hint_in_live_and_saved_traces( assert response is not None from services.workflow.node_execution_extras import node_execution_extras - assert response.data.extras == saved.extras assert response.data.extras == node_execution_extras(saved, tool_providers=tool_providers) assert response.data.extras.get("workflow_tool", False) == (provider_type == "workflow") @staticmethod @pytest.mark.parametrize("next_node_already_persisted", [False, True]) - def test_resumed_node_metadata_and_sequence(next_node_already_persisted, *, workflow_contexts, tool_providers): + def test_resumed_node_metadata_and_sequence( + next_node_already_persisted: bool, + *, + workflow_contexts: ExecutionContexts, + tool_providers: ToolProviderIcons, + ) -> None: started_at = datetime(2026, 9, 7) converter = _workflow_response_converter_resumption_build_converter( workflow_contexts=workflow_contexts, tool_providers=tool_providers @@ -577,7 +586,7 @@ def test_resumed_node_metadata_and_sequence(next_node_already_persisted, *, work task_id="task-1", workflow_run_id="run-1", workflow_id="wf-1", reason=WorkflowStartReason.INITIAL ) - def node_start(execution_id, title, node_type, index): + def node_start(execution_id: str, title: str, node_type: str, index: int) -> QueueNodeStartedEvent: return QueueNodeStartedEvent( node_execution_id=execution_id, node_id=execution_id, diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py index 9c3186d79aa235..8d9418aaeed749 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py @@ -10,12 +10,12 @@ from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.execution_coordinator import AppExecutionState from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager -from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueErrorEvent, QueueTextChunkEvent from graphon.engine.command import AbortCommand from models import Tenant from models.model import AppMode +from services.workflow.execution.adapters.workflow.app_queue_manager import WorkflowAppQueueManager class DummyQueueManager(AppQueueManager): diff --git a/api/tests/unit_tests/core/app/layers/test_conversation_variable_persist_layer.py b/api/tests/unit_tests/core/app/layers/test_conversation_variable_persist_layer.py index b4d5808f4c58d0..022aa5f439a6f7 100644 --- a/api/tests/unit_tests/core/app/layers/test_conversation_variable_persist_layer.py +++ b/api/tests/unit_tests/core/app/layers/test_conversation_variable_persist_layer.py @@ -1,6 +1,5 @@ from unittest.mock import Mock -from core.app.layers.conversation_variable_persist_layer import ConversationVariablePersistenceLayer from core.workflow.system_variables import build_system_variables from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID from graphon.engine.command import InMemoryChannel @@ -10,6 +9,7 @@ from graphon.runtime import ReadOnlyRuntimeStateWrapper, RuntimeState, VariablePool from graphon.variables import StringVariable from libs.datetime_utils import naive_utc_now +from services.workflow.execution.adapters.chatflow.conversation_variables import ConversationVariablePersistenceLayer def _build_graph_runtime_state(conversation_id: str) -> ReadOnlyRuntimeStateWrapper: diff --git a/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py b/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py index a603f7b501b2c6..643d8fa2363a90 100644 --- a/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py +++ b/api/tests/unit_tests/core/app/task_pipeline/test_based_generate_task_pipeline.py @@ -1,53 +1,15 @@ -from decimal import Decimal from types import SimpleNamespace from unittest.mock import Mock import pytest from dify_agent.protocol import RunFailureType -from sqlalchemy.orm import Session from clients.agent_backend.errors import AgentBackendRunFailedError from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter -from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueErrorEvent from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline from core.errors.error import QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError, InvokeRateLimitError -from models.enums import ConversationFromSource, MessageStatus -from models.model import AppMode, Message - - -def _persist_message(session: Session, *, message_id: str) -> Message: - message = Message( - id=message_id, - app_id="app-1", - model_provider=None, - model_id=None, - override_model_configs=None, - conversation_id="conversation-1", - _inputs={}, - query="query", - message={}, - message_unit_price=Decimal(0), - answer="", - answer_unit_price=Decimal(0), - parent_message_id=None, - total_price=None, - currency="USD", - status=MessageStatus.NORMAL, - error=None, - message_metadata=None, - invoke_from=InvokeFrom.WEB_APP, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id=None, - from_account_id="account-1", - workflow_run_id=None, - app_mode=AppMode.COMPLETION, - ) - session.add(message) - session.commit() - session.expunge_all() - return message class TestBasedGenerateTaskPipeline: @@ -96,34 +58,6 @@ def test_handle_error_preserves_agent_backend_run_failed_error(self, pipeline): assert "Knowledge retrieval failed" in str(err) assert "agent_run_id=run-1" in str(err) - @pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) - def test_handle_error_updates_message_when_found(self, pipeline, sqlite_session: Session): - event = QueueErrorEvent(error=ValueError("oops")) - _persist_message(sqlite_session, message_id="msg-1") - - err = pipeline.handle_error(event=event, session=sqlite_session, message_id="msg-1") - - assert err is event.error - sqlite_session.flush() - sqlite_session.expire_all() - updated_message = sqlite_session.get(Message, "msg-1") - assert updated_message is not None - assert updated_message.status == MessageStatus.ERROR - assert updated_message.error == "oops" - - @pytest.mark.parametrize("sqlite_session", [(Message,)], indirect=True) - def test_handle_error_leaves_other_messages_unchanged(self, pipeline, sqlite_session: Session): - event = QueueErrorEvent(error=ValueError("oops")) - _persist_message(sqlite_session, message_id="other-message") - - err = pipeline.handle_error(event=event, session=sqlite_session, message_id="missing-message") - - assert err is event.error - untouched_message = sqlite_session.get(Message, "other-message") - assert untouched_message is not None - assert untouched_message.status == MessageStatus.NORMAL - assert untouched_message.error is None - def test_error_to_stream_response_and_ping(self, pipeline): error_response = pipeline.error_to_stream_response(ValueError("boom")) ping_response = pipeline.ping_stream_response() diff --git a/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py b/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py index 64d68758e433d1..09e6be7adb3656 100644 --- a/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py +++ b/api/tests/unit_tests/core/tools/mcp_tool/test_tool.py @@ -2,16 +2,16 @@ import base64 import contextlib -from collections.abc import Iterator +from collections.abc import Callable from decimal import Decimal from typing import Any from unittest.mock import patch import pytest -from sqlalchemy.engine import Engine from sqlalchemy.orm import Session from core.app.entities.app_invoke_entities import InvokeFrom +from core.entities.mcp_provider import IdentityMode from core.mcp.types import ( AudioContent, BlobResourceContents, @@ -63,7 +63,7 @@ def _build_mcp_tool(*, with_output_schema: bool = True) -> MCPTool: # ----- M2/M3 user-identity forwarding --------------------------------------- -def _build_forwarding_tool(*, mode: str = "idp_token") -> MCPTool: +def _build_forwarding_tool(*, mode: IdentityMode = IdentityMode.IDP_TOKEN) -> MCPTool: """Helper that builds an MCPTool with the identity_mode set.""" entity = ToolEntity( identity=ToolIdentity( @@ -88,7 +88,7 @@ def _build_forwarding_tool(*, mode: str = "idp_token") -> MCPTool: class TestMcpToolCases: @staticmethod - def test_mcp_tool_provider_type_and_fork_runtime(): + def test_mcp_tool_provider_type_and_fork_runtime() -> None: tool = _build_mcp_tool() assert tool.tool_provider_type() == ToolProviderType.MCP @@ -98,7 +98,7 @@ def test_mcp_tool_provider_type_and_fork_runtime(): assert forked.server_identifier == "mcp_server" @staticmethod - def test_mcp_tool_text_and_json_processing_helpers(): + def test_mcp_tool_text_and_json_processing_helpers() -> None: tool = _build_mcp_tool() json_messages = list(tool._process_text_content(TextContent(type="text", text='{"a": 1}'))) @@ -106,6 +106,7 @@ def test_mcp_tool_text_and_json_processing_helpers(): plain_messages = list(tool._process_text_content(TextContent(type="text", text="not-json"))) assert plain_messages[0].type == ToolInvokeMessage.MessageType.TEXT + assert isinstance(plain_messages[0].message, ToolInvokeMessage.TextMessage) assert plain_messages[0].message.text == "not-json" list_messages = list(tool._process_json_content([{"k": 1}, {"k": 2}])) @@ -119,10 +120,11 @@ def test_mcp_tool_text_and_json_processing_helpers(): assert mixed_list_messages[0].type == ToolInvokeMessage.MessageType.TEXT primitive_messages = list(tool._process_json_content(123)) + assert isinstance(primitive_messages[0].message, ToolInvokeMessage.TextMessage) assert primitive_messages[0].message.text == "123" @staticmethod - def test_mcp_tool_usage_extraction_helpers(): + def test_mcp_tool_usage_extraction_helpers() -> None: usage = MCPTool._extract_usage_dict({"usage": {"total_tokens": 9}}) assert usage == {"total_tokens": 9} @@ -145,7 +147,7 @@ def test_mcp_tool_usage_extraction_helpers(): assert derived.total_tokens == 0 @staticmethod - def test_mcp_tool_invoke_handles_content_types_and_structured_output(unbound_session: Session): + def test_mcp_tool_invoke_handles_content_types_and_structured_output(unbound_session: Session) -> None: tool = _build_mcp_tool() img_data = base64.b64encode(b"img").decode() blob_data = base64.b64encode(b"blob").decode() @@ -181,7 +183,7 @@ def test_mcp_tool_invoke_handles_content_types_and_structured_output(unbound_ses assert tool.latest_usage.total_tokens == 5 @staticmethod - def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource(unbound_session: Session): + def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource(unbound_session: Session) -> None: tool = _build_mcp_tool() # Use model_construct to bypass pydantic validation and force unsupported resource path. bad_resource = EmbeddedResource.model_construct(type="resource", resource=object()) @@ -192,7 +194,7 @@ def test_mcp_tool_invoke_raises_for_unsupported_embedded_resource(unbound_sessio list(tool.invoke(session=unbound_session, user_id="user-1", tool_parameters={})) @staticmethod - def test_mcp_tool_handle_none_parameter_filters_empty_values(): + def test_mcp_tool_handle_none_parameter_filters_empty_values() -> None: tool = _build_mcp_tool() cleaned = tool._handle_none_parameter({"a": 1, "b": None, "c": "", "d": " ", "e": "ok"}) assert cleaned == {"a": 1, "e": "ok"} @@ -212,10 +214,14 @@ def test_mcp_tool_handle_none_parameter_filters_empty_values(): ) @pytest.mark.parametrize("domain", [["light", "cover"], "light"]) def test_mcp_tool_invoke_preserves_values_of_mixed_type_union_parameters( - sqlite_session: Session, domain_schema, domain - ): + sqlite_session: Session, domain_schema: dict[str, Any], domain: str | list[str] + ) -> None: tool = _build_mcp_tool() - schema = {"type": "object", "properties": {"domain": domain_schema, "name": _STRING}, "required": []} + schema: dict[str, Any] = { + "type": "object", + "properties": {"domain": domain_schema, "name": _STRING}, + "required": [], + } tool.entity.parameters = ToolTransformService.convert_mcp_schema_to_parameter(schema) result = CallToolResult(content=[TextContent(type="text", text="ok")]) @@ -228,7 +234,7 @@ def test_mcp_tool_invoke_preserves_values_of_mixed_type_union_parameters( assert sent["name"] == "1" @staticmethod - def test_inject_forwarded_identity_stamps_custom_header(): + def test_inject_forwarded_identity_stamps_custom_header() -> None: """The minted SSO token must be placed in X-Dify-SSO-Token; the workspace-scoped Authorization header and any other custom headers must pass through untouched so provider credentials keep working.""" @@ -250,7 +256,7 @@ def test_inject_forwarded_identity_stamps_custom_header(): assert headers["X-Other"] == "keep" @staticmethod - def test_inject_forwarded_identity_translates_token_error_to_invoke_error(): + def test_inject_forwarded_identity_translates_token_error_to_invoke_error() -> None: """EnterpriseService failures must surface as ToolInvokeError so the workflow halts loudly instead of proceeding without identity.""" from core.tools.mcp_tool.tool import FORWARDED_IDENTITY_HEADER @@ -273,7 +279,7 @@ def test_inject_forwarded_identity_translates_token_error_to_invoke_error(): assert "Authorization" not in headers @staticmethod - def test_inject_forwarded_identity_sends_end_user_type_for_webapp(): + def test_inject_forwarded_identity_sends_end_user_type_for_webapp() -> None: """A WEB_APP run forwards user_type=end_user so enterprise routes to the published-webapp token store.""" tool = _build_forwarding_tool() @@ -291,7 +297,7 @@ def test_inject_forwarded_identity_sends_end_user_type_for_webapp(): assert issue.call_args.kwargs["user_type"] == "end_user" @staticmethod - def test_inject_forwarded_identity_sends_account_type_for_debugger(): + def test_inject_forwarded_identity_sends_account_type_for_debugger() -> None: """A DEBUGGER/console run forwards user_type=account (the existing behaviour).""" tool = _build_forwarding_tool() # built with InvokeFrom.DEBUGGER headers: dict[str, str] = {} @@ -307,7 +313,9 @@ def test_inject_forwarded_identity_sends_account_type_for_debugger(): assert issue.call_args.kwargs["user_type"] == "account" @staticmethod - def test_invoke_remote_mcp_tool_fails_closed_when_user_id_missing(config_overrides): + def test_invoke_remote_mcp_tool_fails_closed_when_user_id_missing( + config_overrides: Callable[..., None], + ) -> None: """When forwarding is enabled AND the deployment is enterprise, missing user_id must raise — never silently invoke as the static identity.""" tool = _build_forwarding_tool() @@ -317,7 +325,9 @@ def test_invoke_remote_mcp_tool_fails_closed_when_user_id_missing(config_overrid tool.invoke_remote_mcp_tool({}, user_id=None, app_id=None) @staticmethod - def test_invoke_skips_forwarding_outside_enterprise_edition(config_overrides): + def test_invoke_skips_forwarding_outside_enterprise_edition( + config_overrides: Callable[..., None], + ) -> None: """Non-enterprise deployments treat the DB selector as a no-op: a stale `identity_mode="idp_token"` row must NOT raise (fail-closed) AND must NOT call the enterprise inner API. The runtime falls through to the @@ -364,13 +374,6 @@ def _make_mcp_tool(output_schema: dict[str, Any] | None = None) -> MCPTool: class TestMcpToolCasesTools: - @staticmethod - @pytest.fixture - def orm_session(sqlite_engine: Engine) -> Iterator[Session]: - """Use a real ORM session while MCP transport remains mocked.""" - with Session(sqlite_engine) as session: - yield session - class TestMCPToolInvoke: @pytest.mark.parametrize( ("content_factory", "mime_type"), @@ -385,7 +388,12 @@ class TestMCPToolInvoke: ), ], ) - def test_invoke_image_or_audio_yields_blob(self, content_factory, mime_type, orm_session: Session) -> None: + def test_invoke_image_or_audio_yields_blob( + self, + content_factory: Callable[[str, str], ImageContent | AudioContent], + mime_type: str, + unbound_session: Session, + ) -> None: tool = _make_mcp_tool() raw = b"\x00\x01test-bytes\x02" b64 = base64.b64encode(raw).decode() @@ -393,7 +401,7 @@ def test_invoke_image_or_audio_yields_blob(self, content_factory, mime_type, orm result = CallToolResult(content=[content]) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + messages = list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) assert len(messages) == 1 msg = messages[0] @@ -402,14 +410,14 @@ def test_invoke_image_or_audio_yields_blob(self, content_factory, mime_type, orm assert msg.message.blob == raw assert msg.meta == {"mime_type": mime_type} - def test_invoke_embedded_text_resource_yields_text(self, orm_session: Session) -> None: + def test_invoke_embedded_text_resource_yields_text(self, unbound_session: Session) -> None: tool = _make_mcp_tool() text_resource = TextResourceContents(uri="file://test.txt", mimeType="text/plain", text="hello world") content = EmbeddedResource(type="resource", resource=text_resource) result = CallToolResult(content=[content]) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + messages = list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) assert len(messages) == 1 msg = messages[0] @@ -422,7 +430,7 @@ def test_invoke_embedded_text_resource_yields_text(self, orm_session: Session) - [("application/pdf", "application/pdf"), (None, "application/octet-stream")], ) def test_invoke_embedded_blob_resource_yields_blob( - self, mime_type, expected_mime, orm_session: Session + self, mime_type: str | None, expected_mime: str, unbound_session: Session ) -> None: tool = _make_mcp_tool() raw = b"binary-data" @@ -432,7 +440,7 @@ def test_invoke_embedded_blob_resource_yields_blob( result = CallToolResult(content=[content]) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + messages = list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) assert len(messages) == 1 msg = messages[0] @@ -441,27 +449,33 @@ def test_invoke_embedded_blob_resource_yields_blob( assert msg.message.blob == raw assert msg.meta == {"mime_type": expected_mime} - def test_invoke_yields_variables_when_structured_content_and_schema(self, orm_session: Session) -> None: + def test_invoke_yields_variables_when_structured_content_and_schema(self, unbound_session: Session) -> None: tool = _make_mcp_tool(output_schema={"type": "object"}) result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"}) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + messages = list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) # Expect two variable messages corresponding to keys a and b assert len(messages) == 2 - var_msgs = [m for m in messages if isinstance(m.message, ToolInvokeMessage.VariableMessage)] - assert {m.message.variable_name for m in var_msgs} == {"a", "b"} + variable_messages = [ + message.message + for message in messages + if isinstance(message.message, ToolInvokeMessage.VariableMessage) + ] + assert {message.variable_name for message in variable_messages} == {"a", "b"} # Validate values - values = {m.message.variable_name: m.message.variable_value for m in var_msgs} + values = {message.variable_name: message.variable_value for message in variable_messages} assert values == {"a": 1, "b": "x"} - def test_invoke_yields_json_when_structured_content_has_no_output_schema(self, orm_session: Session) -> None: + def test_invoke_yields_json_when_structured_content_has_no_output_schema( + self, unbound_session: Session + ) -> None: tool = _make_mcp_tool() result = CallToolResult(content=[], structuredContent={"a": 1, "b": "x"}) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - messages = list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + messages = list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) assert len(messages) == 1 msg = messages[0] @@ -609,7 +623,7 @@ def test_derive_usage_from_result_calculates_total_tokens(self) -> None: assert usage.prompt_tokens == 100 assert usage.completion_tokens == 50 - def test_invoke_sets_latest_usage_from_meta(self, orm_session: Session) -> None: + def test_invoke_sets_latest_usage_from_meta(self, unbound_session: Session) -> None: """Test that _invoke sets _latest_usage from result meta.""" tool = _make_mcp_tool() meta = { @@ -624,7 +638,7 @@ def test_invoke_sets_latest_usage_from_meta(self, orm_session: Session) -> None: result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=meta) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) # Verify latest_usage was set correctly assert tool.latest_usage.prompt_tokens == 200 @@ -632,13 +646,13 @@ def test_invoke_sets_latest_usage_from_meta(self, orm_session: Session) -> None: assert tool.latest_usage.total_tokens == 300 assert tool.latest_usage.total_price == Decimal("0.003") - def test_invoke_with_no_meta_returns_empty_usage(self, orm_session: Session) -> None: + def test_invoke_with_no_meta_returns_empty_usage(self, unbound_session: Session) -> None: """Test that _invoke returns empty usage when no meta is present.""" tool = _make_mcp_tool() result = CallToolResult(content=[TextContent(type="text", text="test")], _meta=None) with patch.object(tool, "invoke_remote_mcp_tool", return_value=result): - list(tool._invoke(session=orm_session, user_id="test_user", tool_parameters={})) + list(tool._invoke(session=unbound_session, user_id="test_user", tool_parameters={})) # Verify latest_usage is empty assert tool.latest_usage.total_tokens == 0 @@ -679,7 +693,7 @@ def test_initial_usage_is_empty(self) -> None: {"level1": {"level2": {"usage": {"total_tokens": 200}}}}, ], ) - def test_various_meta_formats(self, meta_data) -> None: + def test_various_meta_formats(self, meta_data: dict[str, Any]) -> None: """Test that various meta formats are correctly parsed.""" result = CallToolResult(content=[], _meta=meta_data) usage = MCPTool._derive_usage_from_result(result) diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py index e415d6cde3449f..3fef6d98ddbbf4 100644 --- a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py +++ b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py @@ -24,7 +24,6 @@ from clients.agent_backend import FakeAgentBackendRunClient from context import capture_current_context from context.flask_app_context import capture_flask_context -from core.app.apps.workflow.app_config_manager import WorkflowAppConfig from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, WorkflowAppGenerateEntity from core.app.file_access import FileAccessScope, bind_file_access_scope from core.app.workflow.file_runtime import init_app @@ -138,12 +137,13 @@ from services.workflow.execution.adapters.human_input_events import resolve_human_input_node_id from services.workflow.execution.adapters.node_factory import DifyNodeFactory from services.workflow.execution.adapters.node_runtime import DifyToolFileManager, DifyToolNodeRuntime +from services.workflow.execution.adapters.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer +from services.workflow.execution.adapters.workflow.app_config_manager import WorkflowAppConfig from services.workflow.execution.adapters.workflow_entry import WorkflowEntry from services.workflow.execution.adapters.workflow_tool_container_handler import ( WorkflowToolContainerHandler, WorkflowToolNestedContainerHandler, ) -from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from services.workflow_run_index import WorkflowRunIndex from tests.helpers.unit.config_override import apply_config_overrides from tests.workflow_test_utils import build_test_graph_init_params, build_test_run_context diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index e4228085076562..c8cef41b4c1509 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -258,7 +258,6 @@ def test_init_app_registers_services_for_the_current_app( from extensions.application_services.tools import resolve_tool_credentials - assert ext_application_services.LegacyToolManager._credential_resolver is resolve_tool_credentials assert ext_application_services.ToolManager._credential_resolver is resolve_tool_credentials with app.app_context(): diff --git a/api/tests/unit_tests/services/app/test_message_persistence.py b/api/tests/unit_tests/services/app/test_message_persistence.py index 04a27ca03424f3..3de1dd6981dfef 100644 --- a/api/tests/unit_tests/services/app/test_message_persistence.py +++ b/api/tests/unit_tests/services/app/test_message_persistence.py @@ -273,18 +273,40 @@ def completed(saved: Message, **_kwargs: object) -> None: assert len(session.scalars(select(Message)).all()) == 1 -def test_error_event_updates_the_same_injected_message( +def test_error_event_updates_only_the_same_injected_message( message_store: MessageStore, annotation_replies: AnnotationReplies ) -> None: sessions, records = message_store identity, pipeline = initialize(records, AppMode.CHAT, ChatAppGenerateEntity, annotation_replies) + with sessions.begin() as session: + session.add( + make_message( + message_id="other-message", + app_id=identity.app_id, + conversation_id=identity.conversation_id, + inputs={}, + query="other question", + message={}, + answer="", + status=MessageStatus.NORMAL, + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ) + ) + responses = list(pipeline._process_stream_response(publisher=None)) assert len(responses) == 1 with sessions() as session: message = session.get(Message, identity.message_id) + other_message = session.get(Message, "other-message") assert message is not None + assert other_message is not None assert message.status == MessageStatus.ERROR assert message.error == "generation failed" + assert other_message.status == MessageStatus.NORMAL + assert other_message.error is None engine: Engine = sessions.kw["bind"] assert checked_out_connections(engine) == 0 diff --git a/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py b/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py index 90bee38bd4960f..83f3dcda9ecf9b 100644 --- a/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py +++ b/api/tests/unit_tests/services/human_input_v2/test_form_delivery_runtime.py @@ -1589,6 +1589,7 @@ def test_pause_persistence_and_reconnect_preserve_v2_form_and_initiator_only_tok pause_entity=pause, resumption_context=None, session_maker=context.sessions, + tool_providers=tool_providers, ) required = next(event["data"] for event in events if event["event"] == "human_input_required") paused = next(event["data"] for event in events if event["event"] == "workflow_paused") diff --git a/api/tests/unit_tests/services/test_message_more_like_this_generator.py b/api/tests/unit_tests/services/test_message_more_like_this_generator.py index 77327078a7e1d0..06daf9315c66ec 100644 --- a/api/tests/unit_tests/services/test_message_more_like_this_generator.py +++ b/api/tests/unit_tests/services/test_message_more_like_this_generator.py @@ -23,6 +23,8 @@ from core.hosting_configuration import HostingProvider, TrialHostingQuota from core.workflow.file_reference import resolve_file_record_id from extensions import ext_hosting_provider +from extensions.application_services.annotation import build_annotation_replies +from extensions.application_services.retrieval import build_dataset_retrieval from extensions.ext_database import db from extensions.storage.storage_type import StorageType from graphon.file import File, FileTransferMethod, FileType @@ -31,7 +33,9 @@ from models import App, AppMode, AppModelConfig, Conversation, EndUser, Message, MessageFile, UploadFile from models.enums import ConversationFromSource, CreatorUserRole, EndUserType, MessageFileBelongsTo from models.provider import Provider, ProviderModel, ProviderType +from repositories.app.generation_repository import AppGenerationRepository from repositories.message_repository import MessageRepository +from services.app.generation.adapters.completion import CompletionAppGenerator from services.entities.message_entities import MessageEndUser from services.message_more_like_this_generator import ( MessageMoreLikeThisGenerator, @@ -113,9 +117,17 @@ def runtime(sqlite_engine: Engine, sqlite_session_factory: sessionmaker[Session] session.flush() factory = sessionmaker(bind=sqlite_engine) + completion_generator = CompletionAppGenerator( + retrieval=build_dataset_retrieval(factory), + records=AppGenerationRepository(factory), + annotations=build_annotation_replies(factory), + ) service = MessageMoreLikeThisService( repository=MessageRepository(session_factory=factory), - generator=MessageMoreLikeThisGenerator(session_factory=factory), + generator=MessageMoreLikeThisGenerator( + session_factory=factory, + completion_generator=completion_generator, + ), ) app = Flask(__name__) app.config["SQLALCHEMY_DATABASE_URI"] = str(sqlite_engine.url) diff --git a/api/tests/unit_tests/services/test_message_more_like_this_service.py b/api/tests/unit_tests/services/test_message_more_like_this_service.py index b8698fa1da2bad..8e9fab173e92ae 100644 --- a/api/tests/unit_tests/services/test_message_more_like_this_service.py +++ b/api/tests/unit_tests/services/test_message_more_like_this_service.py @@ -12,12 +12,16 @@ from sqlalchemy import Connection, Engine, event, func, select from sqlalchemy.orm import Session, SessionTransaction, sessionmaker +from extensions.application_services.annotation import build_annotation_replies +from extensions.application_services.retrieval import build_dataset_retrieval from extensions.ext_database import db from graphon.file import FileTransferMethod, FileType from models import Account, App, AppMode, AppModelConfig, Conversation, EndUser, Message, MessageFile from models.enums import ConversationFromSource, CreatorUserRole, EndUserType, MessageFileBelongsTo from models.provider import Provider +from repositories.app.generation_repository import AppGenerationRepository from repositories.message_repository import MessageRepository +from services.app.generation.adapters.completion import CompletionAppGenerator from services.app_definition_query_service import AppDefinitionUnavailableError from services.entities.message_entities import MessageAccount, MessageActor, MessageEndUser from services.errors.app import MoreLikeThisDisabledError @@ -157,10 +161,19 @@ def record_scoped(session: Session, _transaction: SessionTransaction, _connectio event.listen(db.session.session_factory, "after_begin", record_scoped) repository = MessageRepository(session_factory=factory) + completion_generator = CompletionAppGenerator( + retrieval=build_dataset_retrieval(factory), + records=AppGenerationRepository(factory), + annotations=build_annotation_replies(factory), + ) with flask_app.test_request_context("/more-like-this"): yield _Harness( service=MessageMoreLikeThisService( - repository=repository, generator=MessageMoreLikeThisGenerator(session_factory=factory) + repository=repository, + generator=MessageMoreLikeThisGenerator( + session_factory=factory, + completion_generator=completion_generator, + ), ), repository=repository, app_id=app.id, diff --git a/api/tests/unit_tests/services/test_workflow_event_snapshot_service.py b/api/tests/unit_tests/services/test_workflow_event_snapshot_service.py index 8a3e13b8c44e63..3087e01351d47c 100644 --- a/api/tests/unit_tests/services/test_workflow_event_snapshot_service.py +++ b/api/tests/unit_tests/services/test_workflow_event_snapshot_service.py @@ -64,6 +64,8 @@ _build_paused_snapshot_events as _build_checkpoint_snapshot_events, ) +_TEST_TOOL_PROVIDERS = MagicMock() + @dataclass(frozen=True) class _FakePauseEntity(WorkflowPauseEntity): @@ -198,6 +200,7 @@ def test_build_snapshot_events_applies_message_context() -> None: task_id="task-1", message_context=message_context, node_executions={}, + tool_providers=_TEST_TOOL_PROVIDERS, ) assert [event["event"] for event in events] == [ @@ -821,6 +824,8 @@ def get_executions(tenant_id: str, app_id: str, workflow_run_id: str) -> Sequenc buffer_state = BufferState(queue.Queue(), Event(), Event(), Event(), task_id_hint="task-1") buffer_state.done_event.set() monkeypatch.setattr(service_module, "_start_buffering", lambda _sub: buffer_state) + extras_builder = MagicMock(return_value={"icon": "node-icon"}) + monkeypatch.setattr(service_module, "node_execution_extras", extras_builder) events = list( build_workflow_event_stream( @@ -829,7 +834,8 @@ def get_executions(tenant_id: str, app_id: str, workflow_run_id: str) -> Sequenc tenant_id="tenant-1", app_id="app-1", session_maker=sqlite_session_factory, - **({"include_node_details": True} if include_node_details else {}), + tool_providers=_TEST_TOOL_PROVIDERS, + include_node_details=include_node_details, ) ) started, finished = [ @@ -845,10 +851,17 @@ def get_executions(tenant_id: str, app_id: str, workflow_run_id: str) -> Sequenc assert finished["process_data"] == {"action": "approve"} assert finished["outputs"] == {"approved_text": "approved"} assert finished["execution_metadata"] == {"total_tokens": 7} + assert started["extras"] == {"icon": "node-icon"} + extras_builder.assert_called_once() + replayed_execution = extras_builder.call_args.args[0] + assert replayed_execution.id == execution.id + assert extras_builder.call_args.kwargs == {"tool_providers": _TEST_TOOL_PROVIDERS} else: assert detail_queries == [] assert started["inputs"] is None + assert started["extras"] == {} assert all(finished[key] is None for key in ("inputs", "process_data", "outputs", "execution_metadata")) + extras_builder.assert_not_called() def test_build_workflow_event_stream_should_emit_ping_and_terminal_snapshot_event( @@ -908,6 +921,7 @@ def test_build_workflow_event_stream_should_emit_ping_and_terminal_snapshot_even tenant_id="tenant-1", app_id="app-1", session_maker=session_maker, + tool_providers=_TEST_TOOL_PROVIDERS, ) ) @@ -984,6 +998,7 @@ def test_build_advanced_chat_snapshot_requires_conversation_context( tenant_id="tenant-1", app_id="app-1", session_maker=unbound_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, ) conversation_lookup.assert_not_called() app_lookup.assert_not_called() @@ -1022,6 +1037,7 @@ def test_build_non_suspended_advanced_chat_snapshot_uses_app_scoped_fallback( tenant_id="tenant-1", app_id="app-1", session_maker=session_maker, + tool_providers=_TEST_TOOL_PROVIDERS, ) # Assert @@ -1081,6 +1097,7 @@ def get(self, timeout: int = 1) -> None: tenant_id="tenant-1", app_id="app-1", session_maker=sqlite_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, idle_timeout=20.0, ping_interval=5.0, ) @@ -1127,6 +1144,7 @@ def test_build_workflow_event_stream_should_exit_when_buffer_done_and_empty( tenant_id="tenant-1", app_id="app-1", session_maker=sqlite_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, ) ) @@ -1171,6 +1189,7 @@ def test_build_workflow_event_stream_should_continue_when_pause_loading_fails( tenant_id="tenant-1", app_id="app-1", session_maker=sqlite_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, ) ) @@ -1219,6 +1238,7 @@ def test_build_snapshot_events_preserves_public_form_token( resumption_context=resumption_context, session_maker=sqlite_session_factory, node_executions={}, + tool_providers=_TEST_TOOL_PROVIDERS, ) assert events[-2]["event"] == StreamEvent.HUMAN_INPUT_REQUIRED @@ -1260,6 +1280,7 @@ def _build_recipient_snapshot_events( session_maker=session_maker, human_input_surface=HumanInputSurface.OPENAPI, node_executions={}, + tool_providers=_TEST_TOOL_PROVIDERS, ) @@ -1371,6 +1392,7 @@ def test_build_snapshot_events_resolves_pause_reason_select_options( resumption_context=resumption_context, session_maker=sqlite_session_factory, node_executions={}, + tool_providers=_TEST_TOOL_PROVIDERS, ) human_input_event = events[-2] @@ -1419,6 +1441,7 @@ def test_build_workflow_event_stream_loads_pause_tokens_without_flask_app_contex tenant_id="tenant-1", app_id="app-1", session_maker=sqlite_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, ) ) @@ -1461,6 +1484,7 @@ def _build_paused_snapshot_events( resumption_context=_build_resumption_context("task-ctx", pause_reasons=pause_entity.pause_reasons), session_maker=session_maker, node_executions={}, + tool_providers=_TEST_TOOL_PROVIDERS, ) @@ -1563,6 +1587,7 @@ def test_stream_continues_with_live_events_when_all_forms_submitted( tenant_id="tenant-1", app_id="app-1", session_maker=sqlite_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, close_on_pause=True, ) ) @@ -1619,6 +1644,7 @@ def test_authorized_snapshot_indexes_legacy_pause_for_stop( tenant_id=run.tenant_id, app_id=run.app_id, session_maker=sqlite_session_factory, + tool_providers=_TEST_TOOL_PROVIDERS, ) stream.close() sqlite_session.expire_all() diff --git a/api/tests/unit_tests/services/test_workflow_run_index.py b/api/tests/unit_tests/services/test_workflow_run_index.py index 121ee501e860ca..adceb43ee3cc7c 100644 --- a/api/tests/unit_tests/services/test_workflow_run_index.py +++ b/api/tests/unit_tests/services/test_workflow_run_index.py @@ -20,7 +20,7 @@ from models.model import AppMode from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom from services.workflow.execution.adapters.node_factory import DifyNodeFactory -from services.workflow_persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer +from services.workflow.execution.adapters.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from services.workflow_run_index import WorkflowRunIndex from tests.workflow_test_utils import build_test_graph_init_params diff --git a/api/tests/unit_tests/services/workflow/test_node_execution_extras.py b/api/tests/unit_tests/services/workflow/test_node_execution_extras.py new file mode 100644 index 00000000000000..eb0a0fd143c556 --- /dev/null +++ b/api/tests/unit_tests/services/workflow/test_node_execution_extras.py @@ -0,0 +1,87 @@ +import json +from typing import cast + +import pytest + +from core.tools.entities.tool_entities import ToolProviderType +from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE +from core.trigger.trigger_manager import TriggerManager +from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey +from models.workflow import WorkflowNodeExecutionModel +from services.tools.provider_queries import ToolProviderIcons +from services.tools.tool_manager import ToolManager +from services.workflow.node_execution_extras import node_execution_extras + + +def _execution(*, node_type: str, metadata: dict[str, object]) -> WorkflowNodeExecutionModel: + return WorkflowNodeExecutionModel( + tenant_id="tenant-1", + node_type=node_type, + execution_metadata=json.dumps(metadata), + ) + + +@pytest.mark.parametrize( + ("provider_type", "expected"), + [ + (ToolProviderType.BUILT_IN, {"icon": "tool-icon"}), + (ToolProviderType.WORKFLOW, {"workflow_tool": True, "icon": "tool-icon"}), + ], +) +def test_tool_extras_use_injected_provider_icons( + monkeypatch: pytest.MonkeyPatch, + provider_type: ToolProviderType, + expected: dict[str, object], +) -> None: + tool_providers = cast(ToolProviderIcons, object()) + execution = _execution( + node_type=BuiltinNodeTypes.TOOL, + metadata={"tool_info": {"provider_type": provider_type, "provider_id": "provider-1"}}, + ) + + def get_tool_icon(**kwargs: object) -> str: + assert kwargs == { + "tenant_id": "tenant-1", + "provider_type": provider_type, + "provider_id": "provider-1", + "tool_providers": tool_providers, + } + return "tool-icon" + + monkeypatch.setattr(ToolManager, "get_tool_icon", get_tool_icon) + + assert node_execution_extras(execution, tool_providers=tool_providers) == expected + + +def test_datasource_extras_preserve_persisted_icon() -> None: + execution = _execution( + node_type=BuiltinNodeTypes.DATASOURCE, + metadata={"datasource_info": {"icon": "datasource-icon"}}, + ) + + assert node_execution_extras(execution, tool_providers=cast(ToolProviderIcons, object())) == { + "icon": "datasource-icon" + } + + +def test_trigger_extras_resolve_plugin_icon(monkeypatch: pytest.MonkeyPatch) -> None: + execution = _execution( + node_type=TRIGGER_PLUGIN_NODE_TYPE, + metadata={WorkflowNodeExecutionMetadataKey.TRIGGER_INFO: {"provider_id": "provider-1"}}, + ) + + def get_trigger_plugin_icon(*, tenant_id: str, provider_id: str) -> str: + assert (tenant_id, provider_id) == ("tenant-1", "provider-1") + return "trigger-icon" + + monkeypatch.setattr(TriggerManager, "get_trigger_plugin_icon", get_trigger_plugin_icon) + + assert node_execution_extras(execution, tool_providers=cast(ToolProviderIcons, object())) == { + "icon": "trigger-icon" + } + + +def test_extras_are_empty_without_supported_metadata() -> None: + execution = _execution(node_type=BuiltinNodeTypes.TOOL, metadata={}) + + assert node_execution_extras(execution, tool_providers=cast(ToolProviderIcons, object())) == {}