from __future__ import annotations import json import logging from collections.abc import Generator, Mapping, Sequence from typing import Any, cast from sqlalchemy import select from core.app.file_access import DatabaseFileAccessController from core.db.session_factory import session_factory from core.helper.trace_id_helper import ParentTraceContext, extract_parent_trace_context_from_args from core.tools.__base.tool import Tool from core.tools.__base.tool_runtime import ToolRuntime from core.tools.entities.tool_entities import ( ToolEntity, ToolInvokeMessage, ToolParameter, ToolProviderType, ) from core.tools.errors import ToolInvokeError from core.workflow.file_reference import resolve_file_record_id from factories.file_factory import build_from_mapping from graphon.file import FILE_MODEL_IDENTITY, File, FileTransferMethod from graphon.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata from models import Account, Tenant from models.model import App, EndUser from models.utils.file_input_compat import build_file_from_stored_mapping from models.workflow import Workflow logger = logging.getLogger(__name__) _file_access_controller = DatabaseFileAccessController() class WorkflowTool(Tool): """ Workflow tool. """ _parent_trace_context: ParentTraceContext | None def __init__( self, workflow_app_id: str, workflow_as_tool_id: str, version: str, workflow_entities: dict[str, Any], workflow_call_depth: int, entity: ToolEntity, runtime: ToolRuntime, label: str = "Workflow", ): self.workflow_app_id = workflow_app_id self.workflow_as_tool_id = workflow_as_tool_id self.version = version self.workflow_entities = workflow_entities self.workflow_call_depth = workflow_call_depth self.label = label self._latest_usage = LLMUsage.empty_usage() self._parent_trace_context = None super().__init__(entity=entity, runtime=runtime) def tool_provider_type(self) -> ToolProviderType: """ get the tool provider type :return: the tool provider type """ return ToolProviderType.WORKFLOW def _invoke( self, user_id: str, tool_parameters: dict[str, Any], conversation_id: str | None = None, app_id: str | None = None, message_id: str | None = None, ) -> Generator[ToolInvokeMessage, None, None]: """ invoke the tool """ app = self._get_app(app_id=self.workflow_app_id) workflow = self._get_workflow(app_id=self.workflow_app_id, version=self.version) # transform the tool parameters tool_parameters, files = self._transform_args(tool_parameters=tool_parameters) from core.app.apps.workflow.app_generator import WorkflowAppGenerator generator = WorkflowAppGenerator() assert self.runtime is not None assert self.runtime.invoke_from is not None user = self._resolve_user(user_id=user_id) if user is None: raise ToolInvokeError("User not found") self._latest_usage = LLMUsage.empty_usage() generator_args: dict[str, Any] = {"inputs": tool_parameters, "files": files} if self._parent_trace_context: generator_args.update( extract_parent_trace_context_from_args({"parent_trace_context": self._parent_trace_context}) ) result = generator.generate( app_model=app, workflow=workflow, user=user, args=generator_args, invoke_from=self.runtime.invoke_from, streaming=False, call_depth=self.workflow_call_depth + 1, # NOTE(QuantumGhost): We explicitly set `pause_state_config` to `None` # because workflow pausing mechanisms (such as HumanInput) are not # supported within WorkflowTool execution context. pause_state_config=None, ) assert isinstance(result, dict) data = result.get("data", {}) if err := data.get("error"): raise ToolInvokeError(err) outputs = data.get("outputs") if outputs is None: outputs = {} else: outputs, files = self._extract_files(outputs) # type: ignore for file in files: yield self.create_file_message(file) # type: ignore # traverse `outputs` field and create variable messages for key, value in outputs.items(): if key not in {"text", "json", "files"}: yield self.create_variable_message(variable_name=key, variable_value=value) self._latest_usage = self._derive_usage_from_result(data) yield self.create_text_message(json.dumps(outputs, ensure_ascii=False)) yield self.create_json_message(outputs, suppress_output=True) @property def latest_usage(self) -> LLMUsage: return self._latest_usage @classmethod def _derive_usage_from_result(cls, data: Mapping[str, Any]) -> LLMUsage: usage_dict = cls._extract_usage_dict(data) if usage_dict is not None: return LLMUsage.from_metadata(cast(LLMUsageMetadata, dict(usage_dict))) total_tokens = data.get("total_tokens") total_price = data.get("total_price") if total_tokens is None and total_price is None: return LLMUsage.empty_usage() usage_metadata: dict[str, Any] = {} if total_tokens is not None: try: usage_metadata["total_tokens"] = int(str(total_tokens)) except (TypeError, ValueError): pass if total_price is not None: usage_metadata["total_price"] = str(total_price) currency = data.get("currency") if currency is not None: usage_metadata["currency"] = currency if not usage_metadata: return LLMUsage.empty_usage() return LLMUsage.from_metadata(cast(LLMUsageMetadata, usage_metadata)) @classmethod def _extract_usage_dict(cls, payload: Mapping[str, Any]) -> Mapping[str, Any] | None: usage_candidate = payload.get("usage") if isinstance(usage_candidate, Mapping): return usage_candidate metadata_candidate = payload.get("metadata") if isinstance(metadata_candidate, Mapping): usage_candidate = metadata_candidate.get("usage") if isinstance(usage_candidate, Mapping): return usage_candidate for value in payload.values(): if isinstance(value, Mapping): found = cls._extract_usage_dict(value) if found is not None: return found elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): for item in value: if isinstance(item, Mapping): found = cls._extract_usage_dict(item) if found is not None: return found return None def fork_tool_runtime(self, runtime: ToolRuntime) -> WorkflowTool: """ fork a new tool with metadata :return: the new tool """ forked = self.__class__( entity=self.entity.model_copy(), runtime=runtime, workflow_app_id=self.workflow_app_id, workflow_as_tool_id=self.workflow_as_tool_id, workflow_entities=self.workflow_entities, workflow_call_depth=self.workflow_call_depth, version=self.version, label=self.label, ) forked._parent_trace_context = self._parent_trace_context.model_copy() if self._parent_trace_context else None return forked def set_parent_trace_context( self, *, parent_workflow_run_id: str, parent_node_execution_id: str, ) -> None: """Attach outer workflow trace context without exposing it as tool input.""" self._parent_trace_context = ParentTraceContext( parent_workflow_run_id=parent_workflow_run_id, parent_node_execution_id=parent_node_execution_id, ) def clear_parent_trace_context(self) -> None: """Remove parent trace context before invoking this tool outside a nested workflow.""" self._parent_trace_context = None def _resolve_user(self, user_id: str) -> Account | EndUser | None: """ Resolve user object in both HTTP and worker contexts. In HTTP context: dereference the current_user LocalProxy (can return Account or EndUser). In worker context: load Account(knowledge pipeline) or EndUser(trigger) from database by user_id. Returns: Account | EndUser | None: The resolved user object, or None if resolution fails. """ return self._resolve_user_from_database(user_id=user_id) def _resolve_user_from_database(self, user_id: str) -> Account | EndUser | None: """ Resolve user from database (worker/Celery context). """ with session_factory.create_session() as session: tenant_stmt = select(Tenant).where(Tenant.id == self.runtime.tenant_id) tenant = session.scalar(tenant_stmt) if not tenant: return None user_stmt = select(Account).where(Account.id == user_id) user = session.scalar(user_stmt) if user: user.current_tenant = tenant session.expunge(user) return user end_user_stmt = select(EndUser).where(EndUser.id == user_id, EndUser.tenant_id == tenant.id) end_user = session.scalar(end_user_stmt) if end_user: session.expunge(end_user) return end_user return None def _get_workflow(self, app_id: str, version: str) -> Workflow: """ get the workflow by app id and version """ with session_factory.create_session() as session, session.begin(): if not version: stmt = ( select(Workflow) .where(Workflow.app_id == app_id, Workflow.version != Workflow.VERSION_DRAFT) .order_by(Workflow.created_at.desc()) ) workflow = session.scalars(stmt).first() else: stmt = select(Workflow).where(Workflow.app_id == app_id, Workflow.version == version) workflow = session.scalar(stmt) if not workflow: raise ValueError("workflow not found or not published") session.expunge(workflow) return workflow def _get_app(self, app_id: str) -> App: """ get the app by app id """ stmt = select(App).where(App.id == app_id) with session_factory.create_session() as session, session.begin(): app = session.scalar(stmt) if not app: raise ValueError("app not found") session.expunge(app) return app def _transform_args(self, tool_parameters: dict[str, Any]) -> tuple[dict[str, Any], list[dict[str, str | None]]]: """ transform the tool parameters :param tool_parameters: the tool parameters :return: tool_parameters, files """ parameter_rules = self.get_merged_runtime_parameters() parameters_result = {} files = [] for parameter in parameter_rules: if parameter.type == ToolParameter.ToolParameterType.SYSTEM_FILES: file = tool_parameters.get(parameter.name) if file: try: file_var_list = [ build_file_from_stored_mapping( file_mapping=cast(Mapping[str, Any], f), tenant_id=str(self.runtime.tenant_id), ) for f in file if isinstance(f, Mapping) ] for file in file_var_list: file_dict: dict[str, str | None] = { "transfer_method": file.transfer_method.value, "type": file.type.value, } match file.transfer_method: case FileTransferMethod.TOOL_FILE: file_dict["tool_file_id"] = resolve_file_record_id(file.reference) case FileTransferMethod.LOCAL_FILE: file_dict["upload_file_id"] = resolve_file_record_id(file.reference) case FileTransferMethod.DATASOURCE_FILE: file_dict["datasource_file_id"] = resolve_file_record_id(file.reference) case FileTransferMethod.REMOTE_URL: file_dict["url"] = file.generate_url() files.append(file_dict) except Exception: logger.exception("Failed to transform file %s", file) else: parameters_result[parameter.name] = tool_parameters.get(parameter.name) return parameters_result, files def _extract_files(self, outputs: dict[str, Any]) -> tuple[dict[str, Any], list[File]]: """ extract files from the result :return: the result, files """ files: list[File] = [] result = {} for key, value in outputs.items(): if isinstance(value, list): for item in value: if isinstance(item, dict) and item.get("dify_model_identity") == FILE_MODEL_IDENTITY: item = self._update_file_mapping(item) file = build_from_mapping( mapping=item, tenant_id=str(self.runtime.tenant_id), access_controller=_file_access_controller, ) files.append(file) elif isinstance(value, dict) and value.get("dify_model_identity") == FILE_MODEL_IDENTITY: value = self._update_file_mapping(value) file = build_from_mapping( mapping=value, tenant_id=str(self.runtime.tenant_id), access_controller=_file_access_controller, ) files.append(file) result[key] = value return result, files def _update_file_mapping(self, file_dict: dict[str, Any]) -> dict[str, Any]: file_id = resolve_file_record_id(file_dict.get("reference") or file_dict.get("related_id")) transfer_method_value = file_dict.get("transfer_method") if not isinstance(transfer_method_value, str): raise ValueError("Workflow file mapping is missing a valid transfer_method") transfer_method = FileTransferMethod.value_of(transfer_method_value) match transfer_method: case FileTransferMethod.TOOL_FILE: file_dict["tool_file_id"] = file_id case FileTransferMethod.LOCAL_FILE: file_dict["upload_file_id"] = file_id case FileTransferMethod.REMOTE_URL | FileTransferMethod.DATASOURCE_FILE: pass return file_dict