From ea89d2a17c6450c1c9751068322a52938db96f02 Mon Sep 17 00:00:00 2001 From: QuantumGhost Date: Tue, 17 Jun 2025 10:49:19 +0800 Subject: [PATCH] fix(api): fix the issue that draft var loader does not decode value properly. Refactor the segment loading logic. --- api/core/workflow/variable_loader.py | 2 +- api/models/workflow.py | 57 +++++++++++++++---- .../workflow_draft_variable_service.py | 55 +++++++----------- 3 files changed, 68 insertions(+), 46 deletions(-) diff --git a/api/core/workflow/variable_loader.py b/api/core/workflow/variable_loader.py index 4648a9b752..4842ee00a5 100644 --- a/api/core/workflow/variable_loader.py +++ b/api/core/workflow/variable_loader.py @@ -76,4 +76,4 @@ def load_into_variable_pool( variables_to_load.append(list(selector)) loaded = variable_loader.load_variables(variables_to_load) for var in loaded: - variable_pool.add(var.selector, var.value) + variable_pool.add(var.selector, var) diff --git a/api/models/workflow.py b/api/models/workflow.py index fac6ee712a..654d550d34 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -9,12 +9,12 @@ from uuid import uuid4 from flask_login import current_user from sqlalchemy import orm +from core.file.constants import maybe_file_object from core.file.models import File from core.variables import utils as variable_utils -from core.variables.segments import ArrayFileSegment, FileSegment from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID from core.workflow.nodes.enums import NodeType -from factories.variable_factory import build_segment +from factories.variable_factory import TypeMismatchError, build_segment_with_type from ._workflow_exc import NodeNotFoundError, WorkflowDataError @@ -1065,15 +1065,52 @@ class WorkflowDraftVariable(Base): def _loads_value(self) -> Segment: value = json.loads(self.value) - value_type = self.value_type - if value_type == SegmentType.FILE: - file = File.model_validate(value) - return FileSegment(value=file) - elif value_type == SegmentType.ARRAY_FILE: - files = [File.model_validate(i) for i in value] - return ArrayFileSegment(value=files) + return self.build_segment_with_type(self.value_type, value) + + @staticmethod + def rebuild_file_types(value: Any) -> Any: + # NOTE(QuantumGhost): Temporary workaround for structured data handling. + # By this point, `output` has been converted to dict by + # `WorkflowEntry.handle_special_values`, so we need to + # reconstruct File objects from their serialized form + # to maintain proper variable saving behavior. + # + # Ideally, we should work with structured data objects directly + # rather than their serialized forms. + # However, multiple components in the codebase depend on + # `WorkflowEntry.handle_special_values`, making a comprehensive migration challenging. + if isinstance(value, dict): + if not maybe_file_object(value): + return value + return File.model_validate(value) + elif isinstance(value, list) and value: + first = value[0] + if not maybe_file_object(first): + return value + return [File.model_validate(i) for i in value] else: - return build_segment(value) + return value + + @classmethod + def build_segment_with_type(cls, segment_type: SegmentType, value: Any) -> Segment: + # Extends `variable_factory.build_segment_with_type` functionality by + # reconstructing `FileSegment`` or `ArrayFileSegment`` objects from + # their serialized dictionary or list representations, respectively. + if segment_type == SegmentType.FILE: + if isinstance(value, File): + return build_segment_with_type(segment_type, value) + elif isinstance(value, dict): + file = cls.rebuild_file_types(value) + return build_segment_with_type(segment_type, file) + else: + raise TypeMismatchError(f"expected dict or File for FileSegment, got {type(value)}") + if segment_type == SegmentType.ARRAY_FILE: + if not isinstance(value, list): + raise TypeMismatchError(f"expected list for ArrayFileSegment, got {type(value)}") + file_list = cls.rebuild_file_types(value) + return build_segment_with_type(segment_type=segment_type, value=file_list) + + return build_segment_with_type(segment_type=segment_type, value=value) def get_value(self) -> Segment: """Decode the serialized value into its corresponding `Segment` object. diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py index e92e7d773d..4f4e749d41 100644 --- a/api/services/workflow_draft_variable_service.py +++ b/api/services/workflow_draft_variable_service.py @@ -11,17 +11,17 @@ from sqlalchemy.orm import Session from sqlalchemy.sql.expression import and_, or_ from core.app.entities.app_invoke_entities import InvokeFrom -from core.file.constants import maybe_file_object from core.file.models import File from core.variables import Segment, StringSegment, Variable from core.variables.consts import MIN_SELECTORS_LENGTH from core.variables.segments import ArrayFileSegment +from core.variables.types import SegmentType from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, ENVIRONMENT_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID from core.workflow.enums import SystemVariableKey from core.workflow.nodes import NodeType from core.workflow.nodes.variable_assigner.common.helpers import get_updated_variables from core.workflow.variable_loader import VariableLoader -from factories.variable_factory import build_segment, build_segment_with_type, segment_to_variable +from factories.variable_factory import build_segment, segment_to_variable from models import App, Conversation from models.enums import DraftVariableType from models.workflow import Workflow, WorkflowDraftVariable, WorkflowNodeExecutionModel, is_system_variable_editable @@ -83,9 +83,7 @@ class DraftVarLoader(VariableLoader): draft_vars = srv.get_draft_variables_by_selectors(self._app_id, selectors) for draft_var in draft_vars: - segment = _build_segment_for_value( - draft_var.value, - ) + segment = draft_var.get_value() variable = segment_to_variable( segment=segment, selector=draft_var.get_selector(), @@ -280,7 +278,7 @@ class WorkflowDraftVariableService: self._session.flush() return None value = outputs_dict[variable.name] - value_seg = build_segment_with_type(variable.value_type, _rebuild_file_types_from_dict(value)) + value_seg = WorkflowDraftVariable.build_segment_with_type(variable.value_type, value) # Extract variable value using unified logic variable.set_value(value_seg) variable.last_edited_at = None # Reset to indicate this is a reset operation @@ -480,32 +478,19 @@ def _model_to_insertion_dict(model: WorkflowDraftVariable) -> dict[str, Any]: return d -def _rebuild_file_types_from_dict(value: Any) -> Any: - # NOTE(QuantumGhost): Temporary workaround for structured data handling. - # By this point, `output` has been converted to dict by - # `WorkflowEntry.handle_special_values`, so we need to - # reconstruct File objects from their serialized form - # to maintain proper variable saving behavior. - # - # Ideally, we should work with structured data objects directly - # rather than their serialized forms. - # However, multiple components in the codebase depend on - # `WorkflowEntry.handle_special_values`, making a comprehensive migration challenging. - if isinstance(value, dict): - if not maybe_file_object(value): - return value - return File.model_validate(value) - elif isinstance(value, list) and value: - first = value[0] - if not maybe_file_object(first): - return value - return [File.model_validate(i) for i in value] - else: - return value +def _build_segment_for_serialized_values(v: Any) -> Segment: + """ + Reconstructs Segment objects from serialized values, with special handling + for FileSegment and ArrayFileSegment types. + This function should only be used when: + 1. No explicit type information is available + 2. The input value is in serialized form (dict or list) -def _build_segment_for_value(v: Any) -> Segment: - return build_segment(_rebuild_file_types_from_dict(v)) + It detects potential file objects in the serialized data and properly rebuilds the + appropriate segment type. + """ + return build_segment(WorkflowDraftVariable.rebuild_file_types(v)) class DraftVariableSaver: @@ -577,7 +562,7 @@ class DraftVariableSaver: node_id=self._node_id, name=self._DUMMY_OUTPUT_IDENTITY, node_execution_id=self._node_execution_id, - value=_build_segment_for_value(self._DUMMY_OUTPUT_VALUE), + value=build_segment(self._DUMMY_OUTPUT_VALUE), visible=False, editable=False, ) @@ -604,7 +589,7 @@ class DraftVariableSaver: # We only save conversation variable here. if selector[0] != CONVERSATION_VARIABLE_NODE_ID: continue - segment = _build_segment_for_value(item.new_value) + segment = WorkflowDraftVariable.build_segment_with_type(segment_type=item.value_type, value=item.new_value) draft_vars.append( WorkflowDraftVariable.new_conversation_variable( app_id=self._app_id, @@ -620,7 +605,7 @@ class DraftVariableSaver: draft_vars = [] has_non_sys_variables = False for name, value in output.items(): - value_seg = _build_segment_for_value(value) + value_seg = _build_segment_for_serialized_values(value) node_id, name = self._normalize_variable_for_start_node(name) # If node_id is not `sys`, it means that the variable is a user-defined input field # in `Start` node. @@ -643,7 +628,7 @@ class DraftVariableSaver: # just build files from the value. files = [File.model_validate(v) for v in value] if files: - value_seg = _build_segment_for_value(files) + value_seg = WorkflowDraftVariable.build_segment_with_type(SegmentType.ARRAY_FILE, files) else: value_seg = ArrayFileSegment(value=[]) @@ -677,7 +662,7 @@ class DraftVariableSaver: ) continue - value_seg = _build_segment_for_value(value) + value_seg = _build_segment_for_serialized_values(value) draft_vars.append( WorkflowDraftVariable.new_node_variable( app_id=self._app_id,