You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

workflow_draft_variable_service.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746
  1. import dataclasses
  2. import datetime
  3. import logging
  4. from collections.abc import Mapping, Sequence
  5. from enum import StrEnum
  6. from typing import Any, ClassVar
  7. from sqlalchemy import Engine, orm
  8. from sqlalchemy.dialects.postgresql import insert
  9. from sqlalchemy.orm import Session, sessionmaker
  10. from sqlalchemy.sql.expression import and_, or_
  11. from core.app.entities.app_invoke_entities import InvokeFrom
  12. from core.file.models import File
  13. from core.variables import Segment, StringSegment, Variable
  14. from core.variables.consts import MIN_SELECTORS_LENGTH
  15. from core.variables.segments import ArrayFileSegment, FileSegment
  16. from core.variables.types import SegmentType
  17. from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, ENVIRONMENT_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
  18. from core.workflow.enums import SystemVariableKey
  19. from core.workflow.nodes import NodeType
  20. from core.workflow.nodes.variable_assigner.common.helpers import get_updated_variables
  21. from core.workflow.variable_loader import VariableLoader
  22. from factories.file_factory import StorageKeyLoader
  23. from factories.variable_factory import build_segment, segment_to_variable
  24. from models import App, Conversation
  25. from models.enums import DraftVariableType
  26. from models.workflow import Workflow, WorkflowDraftVariable, is_system_variable_editable
  27. from repositories.factory import DifyAPIRepositoryFactory
  28. _logger = logging.getLogger(__name__)
  29. @dataclasses.dataclass(frozen=True)
  30. class WorkflowDraftVariableList:
  31. variables: list[WorkflowDraftVariable]
  32. total: int | None = None
  33. class WorkflowDraftVariableError(Exception):
  34. pass
  35. class VariableResetError(WorkflowDraftVariableError):
  36. pass
  37. class UpdateNotSupportedError(WorkflowDraftVariableError):
  38. pass
  39. class DraftVarLoader(VariableLoader):
  40. # This implements the VariableLoader interface for loading draft variables.
  41. #
  42. # ref: core.workflow.variable_loader.VariableLoader
  43. # Database engine used for loading variables.
  44. _engine: Engine
  45. # Application ID for which variables are being loaded.
  46. _app_id: str
  47. _tenant_id: str
  48. _fallback_variables: Sequence[Variable]
  49. def __init__(
  50. self,
  51. engine: Engine,
  52. app_id: str,
  53. tenant_id: str,
  54. fallback_variables: Sequence[Variable] | None = None,
  55. ) -> None:
  56. self._engine = engine
  57. self._app_id = app_id
  58. self._tenant_id = tenant_id
  59. self._fallback_variables = fallback_variables or []
  60. def _selector_to_tuple(self, selector: Sequence[str]) -> tuple[str, str]:
  61. return (selector[0], selector[1])
  62. def load_variables(self, selectors: list[list[str]]) -> list[Variable]:
  63. if not selectors:
  64. return []
  65. # Map each selector (as a tuple via `_selector_to_tuple`) to its corresponding Variable instance.
  66. variable_by_selector: dict[tuple[str, str], Variable] = {}
  67. with Session(bind=self._engine, expire_on_commit=False) as session:
  68. srv = WorkflowDraftVariableService(session)
  69. draft_vars = srv.get_draft_variables_by_selectors(self._app_id, selectors)
  70. for draft_var in draft_vars:
  71. segment = draft_var.get_value()
  72. variable = segment_to_variable(
  73. segment=segment,
  74. selector=draft_var.get_selector(),
  75. id=draft_var.id,
  76. name=draft_var.name,
  77. description=draft_var.description,
  78. )
  79. selector_tuple = self._selector_to_tuple(variable.selector)
  80. variable_by_selector[selector_tuple] = variable
  81. # Important:
  82. files: list[File] = []
  83. for draft_var in draft_vars:
  84. value = draft_var.get_value()
  85. if isinstance(value, FileSegment):
  86. files.append(value.value)
  87. elif isinstance(value, ArrayFileSegment):
  88. files.extend(value.value)
  89. with Session(bind=self._engine) as session:
  90. storage_key_loader = StorageKeyLoader(session, tenant_id=self._tenant_id)
  91. storage_key_loader.load_storage_keys(files)
  92. return list(variable_by_selector.values())
  93. class WorkflowDraftVariableService:
  94. _session: Session
  95. def __init__(self, session: Session) -> None:
  96. """
  97. Initialize the WorkflowDraftVariableService with a SQLAlchemy session.
  98. Args:
  99. session (Session): The SQLAlchemy session used to execute database queries.
  100. The provided session must be bound to an `Engine` object, not a specific `Connection`.
  101. Raises:
  102. AssertionError: If the provided session is not bound to an `Engine` object.
  103. """
  104. self._session = session
  105. engine = session.get_bind()
  106. # Ensure the session is bound to a engine.
  107. assert isinstance(engine, Engine)
  108. session_maker = sessionmaker(bind=engine, expire_on_commit=False)
  109. self._api_node_execution_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository(
  110. session_maker
  111. )
  112. def get_variable(self, variable_id: str) -> WorkflowDraftVariable | None:
  113. return self._session.query(WorkflowDraftVariable).filter(WorkflowDraftVariable.id == variable_id).first()
  114. def get_draft_variables_by_selectors(
  115. self,
  116. app_id: str,
  117. selectors: Sequence[list[str]],
  118. ) -> list[WorkflowDraftVariable]:
  119. ors = []
  120. for selector in selectors:
  121. assert len(selector) >= MIN_SELECTORS_LENGTH, f"Invalid selector to get: {selector}"
  122. node_id, name = selector[:2]
  123. ors.append(and_(WorkflowDraftVariable.node_id == node_id, WorkflowDraftVariable.name == name))
  124. # NOTE(QuantumGhost): Although the number of `or` expressions may be large, as long as
  125. # each expression includes conditions on both `node_id` and `name` (which are covered by the unique index),
  126. # PostgreSQL can efficiently retrieve the results using a bitmap index scan.
  127. #
  128. # Alternatively, a `SELECT` statement could be constructed for each selector and
  129. # combined using `UNION` to fetch all rows.
  130. # Benchmarking indicates that both approaches yield comparable performance.
  131. variables = (
  132. self._session.query(WorkflowDraftVariable).where(WorkflowDraftVariable.app_id == app_id, or_(*ors)).all()
  133. )
  134. return variables
  135. def list_variables_without_values(self, app_id: str, page: int, limit: int) -> WorkflowDraftVariableList:
  136. criteria = WorkflowDraftVariable.app_id == app_id
  137. total = None
  138. query = self._session.query(WorkflowDraftVariable).filter(criteria)
  139. if page == 1:
  140. total = query.count()
  141. variables = (
  142. # Do not load the `value` field.
  143. query.options(orm.defer(WorkflowDraftVariable.value))
  144. .order_by(WorkflowDraftVariable.created_at.desc())
  145. .limit(limit)
  146. .offset((page - 1) * limit)
  147. .all()
  148. )
  149. return WorkflowDraftVariableList(variables=variables, total=total)
  150. def _list_node_variables(self, app_id: str, node_id: str) -> WorkflowDraftVariableList:
  151. criteria = (
  152. WorkflowDraftVariable.app_id == app_id,
  153. WorkflowDraftVariable.node_id == node_id,
  154. )
  155. query = self._session.query(WorkflowDraftVariable).filter(*criteria)
  156. variables = query.order_by(WorkflowDraftVariable.created_at.desc()).all()
  157. return WorkflowDraftVariableList(variables=variables)
  158. def list_node_variables(self, app_id: str, node_id: str) -> WorkflowDraftVariableList:
  159. return self._list_node_variables(app_id, node_id)
  160. def list_conversation_variables(self, app_id: str) -> WorkflowDraftVariableList:
  161. return self._list_node_variables(app_id, CONVERSATION_VARIABLE_NODE_ID)
  162. def list_system_variables(self, app_id: str) -> WorkflowDraftVariableList:
  163. return self._list_node_variables(app_id, SYSTEM_VARIABLE_NODE_ID)
  164. def get_conversation_variable(self, app_id: str, name: str) -> WorkflowDraftVariable | None:
  165. return self._get_variable(app_id=app_id, node_id=CONVERSATION_VARIABLE_NODE_ID, name=name)
  166. def get_system_variable(self, app_id: str, name: str) -> WorkflowDraftVariable | None:
  167. return self._get_variable(app_id=app_id, node_id=SYSTEM_VARIABLE_NODE_ID, name=name)
  168. def get_node_variable(self, app_id: str, node_id: str, name: str) -> WorkflowDraftVariable | None:
  169. return self._get_variable(app_id, node_id, name)
  170. def _get_variable(self, app_id: str, node_id: str, name: str) -> WorkflowDraftVariable | None:
  171. variable = (
  172. self._session.query(WorkflowDraftVariable)
  173. .where(
  174. WorkflowDraftVariable.app_id == app_id,
  175. WorkflowDraftVariable.node_id == node_id,
  176. WorkflowDraftVariable.name == name,
  177. )
  178. .first()
  179. )
  180. return variable
  181. def update_variable(
  182. self,
  183. variable: WorkflowDraftVariable,
  184. name: str | None = None,
  185. value: Segment | None = None,
  186. ) -> WorkflowDraftVariable:
  187. if not variable.editable:
  188. raise UpdateNotSupportedError(f"variable not support updating, id={variable.id}")
  189. if name is not None:
  190. variable.set_name(name)
  191. if value is not None:
  192. variable.set_value(value)
  193. variable.last_edited_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  194. self._session.flush()
  195. return variable
  196. def _reset_conv_var(self, workflow: Workflow, variable: WorkflowDraftVariable) -> WorkflowDraftVariable | None:
  197. conv_var_by_name = {i.name: i for i in workflow.conversation_variables}
  198. conv_var = conv_var_by_name.get(variable.name)
  199. if conv_var is None:
  200. self._session.delete(instance=variable)
  201. self._session.flush()
  202. _logger.warning(
  203. "Conversation variable not found for draft variable, id=%s, name=%s", variable.id, variable.name
  204. )
  205. return None
  206. variable.set_value(conv_var)
  207. variable.last_edited_at = None
  208. self._session.add(variable)
  209. self._session.flush()
  210. return variable
  211. def _reset_node_var_or_sys_var(
  212. self, workflow: Workflow, variable: WorkflowDraftVariable
  213. ) -> WorkflowDraftVariable | None:
  214. # If a variable does not allow updating, it makes no sence to resetting it.
  215. if not variable.editable:
  216. return variable
  217. # No execution record for this variable, delete the variable instead.
  218. if variable.node_execution_id is None:
  219. self._session.delete(instance=variable)
  220. self._session.flush()
  221. _logger.warning("draft variable has no node_execution_id, id=%s, name=%s", variable.id, variable.name)
  222. return None
  223. node_exec = self._api_node_execution_repo.get_execution_by_id(variable.node_execution_id)
  224. if node_exec is None:
  225. _logger.warning(
  226. "Node exectution not found for draft variable, id=%s, name=%s, node_execution_id=%s",
  227. variable.id,
  228. variable.name,
  229. variable.node_execution_id,
  230. )
  231. self._session.delete(instance=variable)
  232. self._session.flush()
  233. return None
  234. outputs_dict = node_exec.outputs_dict or {}
  235. # a sentinel value used to check the absent of the output variable key.
  236. absent = object()
  237. if variable.get_variable_type() == DraftVariableType.NODE:
  238. # Get node type for proper value extraction
  239. node_config = workflow.get_node_config_by_id(variable.node_id)
  240. node_type = workflow.get_node_type_from_node_config(node_config)
  241. # Note: Based on the implementation in `_build_from_variable_assigner_mapping`,
  242. # VariableAssignerNode (both v1 and v2) can only create conversation draft variables.
  243. # For consistency, we should simply return when processing VARIABLE_ASSIGNER nodes.
  244. #
  245. # This implementation must remain synchronized with the `_build_from_variable_assigner_mapping`
  246. # and `save` methods.
  247. if node_type == NodeType.VARIABLE_ASSIGNER:
  248. return variable
  249. output_value = outputs_dict.get(variable.name, absent)
  250. else:
  251. output_value = outputs_dict.get(f"sys.{variable.name}", absent)
  252. # We cannot use `is None` to check the existence of an output variable here as
  253. # the value of the output may be `None`.
  254. if output_value is absent:
  255. # If variable not found in execution data, delete the variable
  256. self._session.delete(instance=variable)
  257. self._session.flush()
  258. return None
  259. value_seg = WorkflowDraftVariable.build_segment_with_type(variable.value_type, output_value)
  260. # Extract variable value using unified logic
  261. variable.set_value(value_seg)
  262. variable.last_edited_at = None # Reset to indicate this is a reset operation
  263. self._session.flush()
  264. return variable
  265. def reset_variable(self, workflow: Workflow, variable: WorkflowDraftVariable) -> WorkflowDraftVariable | None:
  266. variable_type = variable.get_variable_type()
  267. if variable_type == DraftVariableType.SYS and not is_system_variable_editable(variable.name):
  268. raise VariableResetError(f"cannot reset system variable, variable_id={variable.id}")
  269. if variable_type == DraftVariableType.CONVERSATION:
  270. return self._reset_conv_var(workflow, variable)
  271. else:
  272. return self._reset_node_var_or_sys_var(workflow, variable)
  273. def delete_variable(self, variable: WorkflowDraftVariable):
  274. self._session.delete(variable)
  275. def delete_workflow_variables(self, app_id: str):
  276. (
  277. self._session.query(WorkflowDraftVariable)
  278. .filter(WorkflowDraftVariable.app_id == app_id)
  279. .delete(synchronize_session=False)
  280. )
  281. def delete_node_variables(self, app_id: str, node_id: str):
  282. return self._delete_node_variables(app_id, node_id)
  283. def _delete_node_variables(self, app_id: str, node_id: str):
  284. self._session.query(WorkflowDraftVariable).where(
  285. WorkflowDraftVariable.app_id == app_id,
  286. WorkflowDraftVariable.node_id == node_id,
  287. ).delete()
  288. def _get_conversation_id_from_draft_variable(self, app_id: str) -> str | None:
  289. draft_var = self._get_variable(
  290. app_id=app_id,
  291. node_id=SYSTEM_VARIABLE_NODE_ID,
  292. name=str(SystemVariableKey.CONVERSATION_ID),
  293. )
  294. if draft_var is None:
  295. return None
  296. segment = draft_var.get_value()
  297. if not isinstance(segment, StringSegment):
  298. _logger.warning(
  299. "sys.conversation_id variable is not a string: app_id=%s, id=%s",
  300. app_id,
  301. draft_var.id,
  302. )
  303. return None
  304. return segment.value
  305. def get_or_create_conversation(
  306. self,
  307. account_id: str,
  308. app: App,
  309. workflow: Workflow,
  310. ) -> str:
  311. """
  312. get_or_create_conversation creates and returns the ID of a conversation for debugging.
  313. If a conversation already exists, as determined by the following criteria, its ID is returned:
  314. - The system variable `sys.conversation_id` exists in the draft variable table, and
  315. - A corresponding conversation record is found in the database.
  316. If no such conversation exists, a new conversation is created and its ID is returned.
  317. """
  318. conv_id = self._get_conversation_id_from_draft_variable(workflow.app_id)
  319. if conv_id is not None:
  320. conversation = (
  321. self._session.query(Conversation)
  322. .filter(
  323. Conversation.id == conv_id,
  324. Conversation.app_id == workflow.app_id,
  325. )
  326. .first()
  327. )
  328. # Only return the conversation ID if it exists and is valid (has a correspond conversation record in DB).
  329. if conversation is not None:
  330. return conv_id
  331. conversation = Conversation(
  332. app_id=workflow.app_id,
  333. app_model_config_id=app.app_model_config_id,
  334. model_provider=None,
  335. model_id="",
  336. override_model_configs=None,
  337. mode=app.mode,
  338. name="Draft Debugging Conversation",
  339. inputs={},
  340. introduction="",
  341. system_instruction="",
  342. system_instruction_tokens=0,
  343. status="normal",
  344. invoke_from=InvokeFrom.DEBUGGER.value,
  345. from_source="console",
  346. from_end_user_id=None,
  347. from_account_id=account_id,
  348. )
  349. self._session.add(conversation)
  350. self._session.flush()
  351. return conversation.id
  352. def prefill_conversation_variable_default_values(self, workflow: Workflow):
  353. """"""
  354. draft_conv_vars: list[WorkflowDraftVariable] = []
  355. for conv_var in workflow.conversation_variables:
  356. draft_var = WorkflowDraftVariable.new_conversation_variable(
  357. app_id=workflow.app_id,
  358. name=conv_var.name,
  359. value=conv_var,
  360. description=conv_var.description,
  361. )
  362. draft_conv_vars.append(draft_var)
  363. _batch_upsert_draft_varaible(
  364. self._session,
  365. draft_conv_vars,
  366. policy=_UpsertPolicy.IGNORE,
  367. )
  368. class _UpsertPolicy(StrEnum):
  369. IGNORE = "ignore"
  370. OVERWRITE = "overwrite"
  371. def _batch_upsert_draft_varaible(
  372. session: Session,
  373. draft_vars: Sequence[WorkflowDraftVariable],
  374. policy: _UpsertPolicy = _UpsertPolicy.OVERWRITE,
  375. ) -> None:
  376. if not draft_vars:
  377. return None
  378. # Although we could use SQLAlchemy ORM operations here, we choose not to for several reasons:
  379. #
  380. # 1. The variable saving process involves writing multiple rows to the
  381. # `workflow_draft_variables` table. Batch insertion significantly improves performance.
  382. # 2. Using the ORM would require either:
  383. #
  384. # a. Checking for the existence of each variable before insertion,
  385. # resulting in 2n SQL statements for n variables and potential concurrency issues.
  386. # b. Attempting insertion first, then updating if a unique index violation occurs,
  387. # which still results in n to 2n SQL statements.
  388. #
  389. # Both approaches are inefficient and suboptimal.
  390. # 3. We do not need to retrieve the results of the SQL execution or populate ORM
  391. # model instances with the returned values.
  392. # 4. Batch insertion with `ON CONFLICT DO UPDATE` allows us to insert or update all
  393. # variables in a single SQL statement, avoiding the issues above.
  394. #
  395. # For these reasons, we use the SQLAlchemy query builder and rely on dialect-specific
  396. # insert operations instead of the ORM layer.
  397. stmt = insert(WorkflowDraftVariable).values([_model_to_insertion_dict(v) for v in draft_vars])
  398. if policy == _UpsertPolicy.OVERWRITE:
  399. stmt = stmt.on_conflict_do_update(
  400. index_elements=WorkflowDraftVariable.unique_app_id_node_id_name(),
  401. set_={
  402. # Refresh creation timestamp to ensure updated variables
  403. # appear first in chronologically sorted result sets.
  404. "created_at": stmt.excluded.created_at,
  405. "updated_at": stmt.excluded.updated_at,
  406. "last_edited_at": stmt.excluded.last_edited_at,
  407. "description": stmt.excluded.description,
  408. "value_type": stmt.excluded.value_type,
  409. "value": stmt.excluded.value,
  410. "visible": stmt.excluded.visible,
  411. "editable": stmt.excluded.editable,
  412. "node_execution_id": stmt.excluded.node_execution_id,
  413. },
  414. )
  415. elif _UpsertPolicy.IGNORE:
  416. stmt = stmt.on_conflict_do_nothing(index_elements=WorkflowDraftVariable.unique_app_id_node_id_name())
  417. else:
  418. raise Exception("Invalid value for update policy.")
  419. session.execute(stmt)
  420. def _model_to_insertion_dict(model: WorkflowDraftVariable) -> dict[str, Any]:
  421. d: dict[str, Any] = {
  422. "app_id": model.app_id,
  423. "last_edited_at": None,
  424. "node_id": model.node_id,
  425. "name": model.name,
  426. "selector": model.selector,
  427. "value_type": model.value_type,
  428. "value": model.value,
  429. "node_execution_id": model.node_execution_id,
  430. }
  431. if model.visible is not None:
  432. d["visible"] = model.visible
  433. if model.editable is not None:
  434. d["editable"] = model.editable
  435. if model.created_at is not None:
  436. d["created_at"] = model.created_at
  437. if model.updated_at is not None:
  438. d["updated_at"] = model.updated_at
  439. if model.description is not None:
  440. d["description"] = model.description
  441. return d
  442. def _build_segment_for_serialized_values(v: Any) -> Segment:
  443. """
  444. Reconstructs Segment objects from serialized values, with special handling
  445. for FileSegment and ArrayFileSegment types.
  446. This function should only be used when:
  447. 1. No explicit type information is available
  448. 2. The input value is in serialized form (dict or list)
  449. It detects potential file objects in the serialized data and properly rebuilds the
  450. appropriate segment type.
  451. """
  452. return build_segment(WorkflowDraftVariable.rebuild_file_types(v))
  453. class DraftVariableSaver:
  454. # _DUMMY_OUTPUT_IDENTITY is a placeholder output for workflow nodes.
  455. # Its sole possible value is `None`.
  456. #
  457. # This is used to signal the execution of a workflow node when it has no other outputs.
  458. _DUMMY_OUTPUT_IDENTITY: ClassVar[str] = "__dummy__"
  459. _DUMMY_OUTPUT_VALUE: ClassVar[None] = None
  460. # _EXCLUDE_VARIABLE_NAMES_MAPPING maps node types and versions to variable names that
  461. # should be excluded when saving draft variables. This prevents certain internal or
  462. # technical variables from being exposed in the draft environment, particularly those
  463. # that aren't meant to be directly edited or viewed by users.
  464. _EXCLUDE_VARIABLE_NAMES_MAPPING: dict[NodeType, frozenset[str]] = {
  465. NodeType.LLM: frozenset(["finish_reason"]),
  466. NodeType.LOOP: frozenset(["loop_round"]),
  467. }
  468. # Database session used for persisting draft variables.
  469. _session: Session
  470. # The application ID associated with the draft variables.
  471. # This should match the `Workflow.app_id` of the workflow to which the current node belongs.
  472. _app_id: str
  473. # The ID of the node for which DraftVariableSaver is saving output variables.
  474. _node_id: str
  475. # The type of the current node (see NodeType).
  476. _node_type: NodeType
  477. #
  478. _node_execution_id: str
  479. # _enclosing_node_id identifies the container node that the current node belongs to.
  480. # For example, if the current node is an LLM node inside an Iteration node
  481. # or Loop node, then `_enclosing_node_id` refers to the ID of
  482. # the containing Iteration or Loop node.
  483. #
  484. # If the current node is not nested within another node, `_enclosing_node_id` is
  485. # `None`.
  486. _enclosing_node_id: str | None
  487. def __init__(
  488. self,
  489. session: Session,
  490. app_id: str,
  491. node_id: str,
  492. node_type: NodeType,
  493. node_execution_id: str,
  494. enclosing_node_id: str | None = None,
  495. ):
  496. # Important: `node_execution_id` parameter refers to the primary key (`id`) of the
  497. # WorkflowNodeExecutionModel/WorkflowNodeExecution, not their `node_execution_id`
  498. # field. These are distinct database fields with different purposes.
  499. self._session = session
  500. self._app_id = app_id
  501. self._node_id = node_id
  502. self._node_type = node_type
  503. self._node_execution_id = node_execution_id
  504. self._enclosing_node_id = enclosing_node_id
  505. def _create_dummy_output_variable(self):
  506. return WorkflowDraftVariable.new_node_variable(
  507. app_id=self._app_id,
  508. node_id=self._node_id,
  509. name=self._DUMMY_OUTPUT_IDENTITY,
  510. node_execution_id=self._node_execution_id,
  511. value=build_segment(self._DUMMY_OUTPUT_VALUE),
  512. visible=False,
  513. editable=False,
  514. )
  515. def _should_save_output_variables_for_draft(self) -> bool:
  516. if self._enclosing_node_id is not None and self._node_type != NodeType.VARIABLE_ASSIGNER:
  517. # Currently we do not save output variables for nodes inside loop or iteration.
  518. return False
  519. return True
  520. def _build_from_variable_assigner_mapping(self, process_data: Mapping[str, Any]) -> list[WorkflowDraftVariable]:
  521. draft_vars: list[WorkflowDraftVariable] = []
  522. updated_variables = get_updated_variables(process_data) or []
  523. for item in updated_variables:
  524. selector = item.selector
  525. if len(selector) < MIN_SELECTORS_LENGTH:
  526. raise Exception("selector too short")
  527. # NOTE(QuantumGhost): only the following two kinds of variable could be updated by
  528. # VariableAssigner: ConversationVariable and iteration variable.
  529. # We only save conversation variable here.
  530. if selector[0] != CONVERSATION_VARIABLE_NODE_ID:
  531. continue
  532. segment = WorkflowDraftVariable.build_segment_with_type(segment_type=item.value_type, value=item.new_value)
  533. draft_vars.append(
  534. WorkflowDraftVariable.new_conversation_variable(
  535. app_id=self._app_id,
  536. name=item.name,
  537. value=segment,
  538. )
  539. )
  540. # Add a dummy output variable to indicate that this node is executed.
  541. draft_vars.append(self._create_dummy_output_variable())
  542. return draft_vars
  543. def _build_variables_from_start_mapping(self, output: Mapping[str, Any]) -> list[WorkflowDraftVariable]:
  544. draft_vars = []
  545. has_non_sys_variables = False
  546. for name, value in output.items():
  547. value_seg = _build_segment_for_serialized_values(value)
  548. node_id, name = self._normalize_variable_for_start_node(name)
  549. # If node_id is not `sys`, it means that the variable is a user-defined input field
  550. # in `Start` node.
  551. if node_id != SYSTEM_VARIABLE_NODE_ID:
  552. draft_vars.append(
  553. WorkflowDraftVariable.new_node_variable(
  554. app_id=self._app_id,
  555. node_id=self._node_id,
  556. name=name,
  557. node_execution_id=self._node_execution_id,
  558. value=value_seg,
  559. visible=True,
  560. editable=True,
  561. )
  562. )
  563. has_non_sys_variables = True
  564. else:
  565. if name == SystemVariableKey.FILES:
  566. # Here we know the type of variable must be `array[file]`, we
  567. # just build files from the value.
  568. files = [File.model_validate(v) for v in value]
  569. if files:
  570. value_seg = WorkflowDraftVariable.build_segment_with_type(SegmentType.ARRAY_FILE, files)
  571. else:
  572. value_seg = ArrayFileSegment(value=[])
  573. draft_vars.append(
  574. WorkflowDraftVariable.new_sys_variable(
  575. app_id=self._app_id,
  576. name=name,
  577. node_execution_id=self._node_execution_id,
  578. value=value_seg,
  579. editable=self._should_variable_be_editable(node_id, name),
  580. )
  581. )
  582. if not has_non_sys_variables:
  583. draft_vars.append(self._create_dummy_output_variable())
  584. return draft_vars
  585. def _normalize_variable_for_start_node(self, name: str) -> tuple[str, str]:
  586. if not name.startswith(f"{SYSTEM_VARIABLE_NODE_ID}."):
  587. return self._node_id, name
  588. _, name_ = name.split(".", maxsplit=1)
  589. return SYSTEM_VARIABLE_NODE_ID, name_
  590. def _build_variables_from_mapping(self, output: Mapping[str, Any]) -> list[WorkflowDraftVariable]:
  591. draft_vars = []
  592. for name, value in output.items():
  593. if not self._should_variable_be_saved(name):
  594. _logger.debug(
  595. "Skip saving variable as it has been excluded by its node_type, name=%s, node_type=%s",
  596. name,
  597. self._node_type,
  598. )
  599. continue
  600. if isinstance(value, Segment):
  601. value_seg = value
  602. else:
  603. value_seg = _build_segment_for_serialized_values(value)
  604. draft_vars.append(
  605. WorkflowDraftVariable.new_node_variable(
  606. app_id=self._app_id,
  607. node_id=self._node_id,
  608. name=name,
  609. node_execution_id=self._node_execution_id,
  610. value=value_seg,
  611. visible=self._should_variable_be_visible(self._node_id, self._node_type, name),
  612. )
  613. )
  614. return draft_vars
  615. def save(
  616. self,
  617. process_data: Mapping[str, Any] | None = None,
  618. outputs: Mapping[str, Any] | None = None,
  619. ):
  620. draft_vars: list[WorkflowDraftVariable] = []
  621. if outputs is None:
  622. outputs = {}
  623. if process_data is None:
  624. process_data = {}
  625. if not self._should_save_output_variables_for_draft():
  626. return
  627. if self._node_type == NodeType.VARIABLE_ASSIGNER:
  628. draft_vars = self._build_from_variable_assigner_mapping(process_data=process_data)
  629. elif self._node_type == NodeType.START:
  630. draft_vars = self._build_variables_from_start_mapping(outputs)
  631. else:
  632. draft_vars = self._build_variables_from_mapping(outputs)
  633. _batch_upsert_draft_varaible(self._session, draft_vars)
  634. @staticmethod
  635. def _should_variable_be_editable(node_id: str, name: str) -> bool:
  636. if node_id in (CONVERSATION_VARIABLE_NODE_ID, ENVIRONMENT_VARIABLE_NODE_ID):
  637. return False
  638. if node_id == SYSTEM_VARIABLE_NODE_ID and not is_system_variable_editable(name):
  639. return False
  640. return True
  641. @staticmethod
  642. def _should_variable_be_visible(node_id: str, node_type: NodeType, name: str) -> bool:
  643. if node_type in NodeType.IF_ELSE:
  644. return False
  645. if node_id == SYSTEM_VARIABLE_NODE_ID and not is_system_variable_editable(name):
  646. return False
  647. return True
  648. def _should_variable_be_saved(self, name: str) -> bool:
  649. exclude_var_names = self._EXCLUDE_VARIABLE_NAMES_MAPPING.get(self._node_type)
  650. if exclude_var_names is None:
  651. return True
  652. return name not in exclude_var_names