chore: remove pm2 from web docker

Co-authored-by: hyoban <38493346+hyoban@users.noreply.github.com>

api/.importlinter

@@ -110,7 +110,6 @@ ignore_imports =
     core.workflow.nodes.iteration.iteration_node -> core.app.workflow.node_factory
     core.workflow.nodes.knowledge_index.knowledge_index_node -> core.rag.index_processor.index_processor_factory
     core.workflow.nodes.llm.llm_utils -> configs
-    core.workflow.nodes.llm.llm_utils -> core.app.entities.app_invoke_entities
     core.workflow.nodes.llm.llm_utils -> core.model_manager
     core.workflow.nodes.llm.protocols -> core.model_manager
     core.workflow.nodes.llm.llm_utils -> core.model_runtime.model_providers.__base.large_language_model
@@ -129,13 +128,9 @@ ignore_imports =
     core.workflow.nodes.human_input.human_input_node -> core.app.entities.app_invoke_entities
     core.workflow.nodes.knowledge_index.knowledge_index_node -> core.app.entities.app_invoke_entities
     core.workflow.nodes.knowledge_retrieval.knowledge_retrieval_node -> core.app.app_config.entities
-    core.workflow.nodes.llm.node -> core.app.entities.app_invoke_entities
-    core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.app.entities.app_invoke_entities
     core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.prompt.advanced_prompt_transform
     core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.prompt.simple_prompt_transform
     core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.model_runtime.model_providers.__base.large_language_model
-    core.workflow.nodes.question_classifier.question_classifier_node -> core.app.entities.app_invoke_entities
-    core.workflow.nodes.question_classifier.question_classifier_node -> core.prompt.advanced_prompt_transform
     core.workflow.nodes.question_classifier.question_classifier_node -> core.prompt.simple_prompt_transform
     core.workflow.nodes.start.entities -> core.app.app_config.entities
     core.workflow.nodes.start.start_node -> core.app.app_config.entities
@@ -145,8 +140,6 @@ ignore_imports =
     core.workflow.nodes.llm.llm_utils -> core.entities.provider_entities
     core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.model_manager
     core.workflow.nodes.question_classifier.question_classifier_node -> core.model_manager
-    core.workflow.nodes.llm.llm_utils -> core.variables.segments
-    core.workflow.nodes.loop.entities -> core.variables.types
     core.workflow.nodes.tool.tool_node -> core.tools.utils.message_transformer
     core.workflow.nodes.tool.tool_node -> models
     core.workflow.nodes.agent.agent_node -> models.model
@@ -183,54 +176,6 @@ ignore_imports =
     core.workflow.nodes.llm.file_saver -> core.tools.signature
     core.workflow.nodes.llm.file_saver -> core.tools.tool_file_manager
     core.workflow.nodes.tool.tool_node -> core.tools.errors
-    core.workflow.conversation_variable_updater -> core.variables
-    core.workflow.graph_engine.entities.commands -> core.variables.variables
-    core.workflow.nodes.agent.agent_node -> core.variables.segments
-    core.workflow.nodes.answer.answer_node -> core.variables
-    core.workflow.nodes.code.code_node -> core.variables.segments
-    core.workflow.nodes.code.code_node -> core.variables.types
-    core.workflow.nodes.code.entities -> core.variables.types
-    core.workflow.nodes.document_extractor.node -> core.variables
-    core.workflow.nodes.document_extractor.node -> core.variables.segments
-    core.workflow.nodes.http_request.executor -> core.variables.segments
-    core.workflow.nodes.http_request.node -> core.variables.segments
-    core.workflow.nodes.human_input.entities -> core.variables.consts
-    core.workflow.nodes.iteration.iteration_node -> core.variables
-    core.workflow.nodes.iteration.iteration_node -> core.variables.segments
-    core.workflow.nodes.iteration.iteration_node -> core.variables.variables
-    core.workflow.nodes.knowledge_retrieval.knowledge_retrieval_node -> core.variables
-    core.workflow.nodes.knowledge_retrieval.knowledge_retrieval_node -> core.variables.segments
-    core.workflow.nodes.list_operator.node -> core.variables
-    core.workflow.nodes.list_operator.node -> core.variables.segments
-    core.workflow.nodes.llm.node -> core.variables
-    core.workflow.nodes.loop.loop_node -> core.variables
-    core.workflow.nodes.parameter_extractor.entities -> core.variables.types
-    core.workflow.nodes.parameter_extractor.exc -> core.variables.types
-    core.workflow.nodes.parameter_extractor.parameter_extractor_node -> core.variables.types
-    core.workflow.nodes.tool.tool_node -> core.variables.segments
-    core.workflow.nodes.tool.tool_node -> core.variables.variables
-    core.workflow.nodes.trigger_webhook.node -> core.variables.types
-    core.workflow.nodes.trigger_webhook.node -> core.variables.variables
-    core.workflow.nodes.variable_aggregator.entities -> core.variables.types
-    core.workflow.nodes.variable_aggregator.variable_aggregator_node -> core.variables.segments
-    core.workflow.nodes.variable_assigner.common.helpers -> core.variables
-    core.workflow.nodes.variable_assigner.common.helpers -> core.variables.consts
-    core.workflow.nodes.variable_assigner.common.helpers -> core.variables.types
-    core.workflow.nodes.variable_assigner.v1.node -> core.variables
-    core.workflow.nodes.variable_assigner.v2.helpers -> core.variables
-    core.workflow.nodes.variable_assigner.v2.node -> core.variables
-    core.workflow.nodes.variable_assigner.v2.node -> core.variables.consts
-    core.workflow.runtime.graph_runtime_state_protocol -> core.variables.segments
-    core.workflow.runtime.read_only_wrappers -> core.variables.segments
-    core.workflow.runtime.variable_pool -> core.variables
-    core.workflow.runtime.variable_pool -> core.variables.consts
-    core.workflow.runtime.variable_pool -> core.variables.segments
-    core.workflow.runtime.variable_pool -> core.variables.variables
-    core.workflow.utils.condition.processor -> core.variables
-    core.workflow.utils.condition.processor -> core.variables.segments
-    core.workflow.variable_loader -> core.variables
-    core.workflow.variable_loader -> core.variables.consts
-    core.workflow.workflow_type_encoder -> core.variables
     core.workflow.nodes.agent.agent_node -> extensions.ext_database
     core.workflow.nodes.knowledge_index.knowledge_index_node -> extensions.ext_database
     core.workflow.nodes.llm.file_saver -> extensions.ext_database

api/controllers/console/app/workflow_draft_variable.py

@@ -15,11 +15,11 @@ from controllers.console.app.error import (
 from controllers.console.app.wraps import get_app_model
 from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required
 from controllers.web.error import InvalidArgumentError, NotFoundError
-from core.variables.segment_group import SegmentGroup
-from core.variables.segments import ArrayFileSegment, FileSegment, Segment
-from core.variables.types import SegmentType
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
 from core.workflow.file import helpers as file_helpers
+from core.workflow.variables.segment_group import SegmentGroup
+from core.workflow.variables.segments import ArrayFileSegment, FileSegment, Segment
+from core.workflow.variables.types import SegmentType
 from extensions.ext_database import db
 from factories.file_factory import build_from_mapping, build_from_mappings
 from factories.variable_factory import build_segment_with_type

api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py

@@ -21,8 +21,8 @@ from controllers.console.app.workflow_draft_variable import (
 from controllers.console.datasets.wraps import get_rag_pipeline
 from controllers.console.wraps import account_initialization_required, setup_required
 from controllers.web.error import InvalidArgumentError, NotFoundError
-from core.variables.types import SegmentType
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
+from core.workflow.variables.types import SegmentType
 from extensions.ext_database import db
 from factories.file_factory import build_from_mapping, build_from_mappings
 from factories.variable_factory import build_segment_with_type

api/core/app/apps/advanced_chat/app_runner.py

@@ -25,7 +25,6 @@ from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, Workfl
 from core.db.session_factory import session_factory
 from core.moderation.base import ModerationError
 from core.moderation.input_moderation import InputModeration
-from core.variables.variables import Variable
 from core.workflow.enums import WorkflowType
 from core.workflow.graph_engine.command_channels.redis_channel import RedisChannel
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
@@ -34,6 +33,7 @@ from core.workflow.repositories.workflow_node_execution_repository import Workfl
 from core.workflow.runtime import GraphRuntimeState, VariablePool
 from core.workflow.system_variable import SystemVariable
 from core.workflow.variable_loader import VariableLoader
+from core.workflow.variables.variables import Variable
 from core.workflow.workflow_entry import WorkflowEntry
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client

api/core/app/apps/common/workflow_response_converter.py

@@ -49,7 +49,6 @@ 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.trigger_manager import TriggerManager
-from core.variables.segments import ArrayFileSegment, FileSegment, Segment
 from core.workflow.entities.pause_reason import HumanInputRequired
 from core.workflow.entities.workflow_start_reason import WorkflowStartReason
 from core.workflow.enums import (
@@ -62,6 +61,7 @@ from core.workflow.enums import (
 from core.workflow.file import FILE_MODEL_IDENTITY, File
 from core.workflow.runtime import GraphRuntimeState
 from core.workflow.system_variable import SystemVariable
+from core.workflow.variables.segments import ArrayFileSegment, FileSegment, Segment
 from core.workflow.workflow_entry import WorkflowEntry
 from core.workflow.workflow_type_encoder import WorkflowRuntimeTypeConverter
 from extensions.ext_database import db

api/core/app/apps/pipeline/pipeline_runner.py

@@ -11,7 +11,6 @@ from core.app.entities.app_invoke_entities import (
 )
 from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer
 from core.app.workflow.node_factory import DifyNodeFactory
-from core.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput
 from core.workflow.entities.graph_init_params import GraphInitParams
 from core.workflow.enums import WorkflowType
 from core.workflow.graph import Graph
@@ -21,6 +20,7 @@ from core.workflow.repositories.workflow_node_execution_repository import Workfl
 from core.workflow.runtime import GraphRuntimeState, VariablePool
 from core.workflow.system_variable import SystemVariable
 from core.workflow.variable_loader import VariableLoader
+from core.workflow.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput
 from core.workflow.workflow_entry import WorkflowEntry
 from extensions.ext_database import db
 from models.dataset import Document, Pipeline

api/core/app/layers/conversation_variable_persist_layer.py

@@ -1,12 +1,12 @@
 import logging
 
-from core.variables import VariableBase
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.conversation_variable_updater import ConversationVariableUpdater
 from core.workflow.enums import NodeType
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
 from core.workflow.graph_events import GraphEngineEvent, NodeRunSucceededEvent
 from core.workflow.nodes.variable_assigner.common import helpers as common_helpers
+from core.workflow.variables import VariableBase
 
 logger = logging.getLogger(__name__)
 

api/core/app/llm/model_access.py

@@ -83,14 +83,21 @@ def fetch_model_config(
         raise ModelNotExistError(f"Model {node_data_model.name} not exist.")
     provider_model.raise_for_status()
 
-    stop: list[str] = []
-    if "stop" in node_data_model.completion_params:
-        stop = node_data_model.completion_params.pop("stop")
+    completion_params = dict(node_data_model.completion_params)
+    stop = completion_params.pop("stop", [])
+    if not isinstance(stop, list):
+        stop = []
 
     model_schema = model_instance.model_type_instance.get_model_schema(node_data_model.name, credentials)
     if not model_schema:
         raise ModelNotExistError(f"Model {node_data_model.name} not exist.")
 
+    model_instance.provider = node_data_model.provider
+    model_instance.model_name = node_data_model.name
+    model_instance.credentials = credentials
+    model_instance.parameters = completion_params
+    model_instance.stop = tuple(stop)
+
     return model_instance, ModelConfigWithCredentialsEntity(
         provider=node_data_model.provider,
         model=node_data_model.name,
@@ -98,6 +105,6 @@ def fetch_model_config(
         mode=node_data_model.mode,
         provider_model_bundle=provider_model_bundle,
         credentials=credentials,
-        parameters=node_data_model.completion_params,
+        parameters=completion_params,
         stop=stop,
     )

api/core/app/workflow/node_factory.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, final
+from typing import TYPE_CHECKING, Any, cast, final
 
 from typing_extensions import override
 
@@ -9,6 +9,10 @@ from core.datasource.datasource_manager import DatasourceManager
 from core.helper.code_executor.code_executor import CodeExecutionError, CodeExecutor
 from core.helper.code_executor.code_node_provider import CodeNodeProvider
 from core.helper.ssrf_proxy import ssrf_proxy
+from core.model_manager import ModelInstance
+from core.model_runtime.entities.model_entities import ModelType
+from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
+from core.prompt.entities.advanced_prompt_entities import MemoryConfig
 from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
 from core.tools.tool_file_manager import ToolFileManager
 from core.workflow.entities.graph_config import NodeConfigDict
@@ -23,7 +27,11 @@ from core.workflow.nodes.datasource import DatasourceNode
 from core.workflow.nodes.document_extractor import DocumentExtractorNode, UnstructuredApiConfig
 from core.workflow.nodes.http_request import HttpRequestNode, build_http_request_config
 from core.workflow.nodes.knowledge_retrieval.knowledge_retrieval_node import KnowledgeRetrievalNode
+from core.workflow.nodes.llm import llm_utils
+from core.workflow.nodes.llm.entities import ModelConfig
+from core.workflow.nodes.llm.exc import LLMModeRequiredError, ModelNotExistError
 from core.workflow.nodes.llm.node import LLMNode
+from core.workflow.nodes.llm.protocols import PromptMessageMemory
 from core.workflow.nodes.node_mapping import LATEST_VERSION, NODE_TYPE_CLASSES_MAPPING
 from core.workflow.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode
 from core.workflow.nodes.question_classifier.question_classifier_node import QuestionClassifierNode
@@ -171,6 +179,8 @@ class DifyNodeFactory(NodeFactory):
             )
 
         if node_type == NodeType.LLM:
+            model_instance = self._build_model_instance_for_llm_node(node_data)
+            memory = self._build_memory_for_llm_node(node_data=node_data, model_instance=model_instance)
             return LLMNode(
                 id=node_id,
                 config=node_config,
@@ -178,6 +188,8 @@ class DifyNodeFactory(NodeFactory):
                 graph_runtime_state=self.graph_runtime_state,
                 credentials_provider=self._llm_credentials_provider,
                 model_factory=self._llm_model_factory,
+                model_instance=model_instance,
+                memory=memory,
             )
 
         if node_type == NodeType.DATASOURCE:
@@ -208,6 +220,7 @@ class DifyNodeFactory(NodeFactory):
             )
 
         if node_type == NodeType.QUESTION_CLASSIFIER:
+            model_instance = self._build_model_instance_for_llm_node(node_data)
             return QuestionClassifierNode(
                 id=node_id,
                 config=node_config,
@@ -215,9 +228,11 @@ class DifyNodeFactory(NodeFactory):
                 graph_runtime_state=self.graph_runtime_state,
                 credentials_provider=self._llm_credentials_provider,
                 model_factory=self._llm_model_factory,
+                model_instance=model_instance,
             )
 
         if node_type == NodeType.PARAMETER_EXTRACTOR:
+            model_instance = self._build_model_instance_for_llm_node(node_data)
             return ParameterExtractorNode(
                 id=node_id,
                 config=node_config,
@@ -225,6 +240,7 @@ class DifyNodeFactory(NodeFactory):
                 graph_runtime_state=self.graph_runtime_state,
                 credentials_provider=self._llm_credentials_provider,
                 model_factory=self._llm_model_factory,
+                model_instance=model_instance,
             )
 
         return node_class(
@@ -233,3 +249,55 @@ class DifyNodeFactory(NodeFactory):
             graph_init_params=self.graph_init_params,
             graph_runtime_state=self.graph_runtime_state,
         )
+
+    def _build_model_instance_for_llm_node(self, node_data: Mapping[str, Any]) -> ModelInstance:
+        node_data_model = ModelConfig.model_validate(node_data["model"])
+        if not node_data_model.mode:
+            raise LLMModeRequiredError("LLM mode is required.")
+
+        credentials = self._llm_credentials_provider.fetch(node_data_model.provider, node_data_model.name)
+        model_instance = self._llm_model_factory.init_model_instance(node_data_model.provider, node_data_model.name)
+        provider_model_bundle = model_instance.provider_model_bundle
+
+        provider_model = provider_model_bundle.configuration.get_provider_model(
+            model=node_data_model.name,
+            model_type=ModelType.LLM,
+        )
+        if provider_model is None:
+            raise ModelNotExistError(f"Model {node_data_model.name} not exist.")
+        provider_model.raise_for_status()
+
+        completion_params = dict(node_data_model.completion_params)
+        stop = completion_params.pop("stop", [])
+        if not isinstance(stop, list):
+            stop = []
+
+        model_schema = model_instance.model_type_instance.get_model_schema(node_data_model.name, credentials)
+        if not model_schema:
+            raise ModelNotExistError(f"Model {node_data_model.name} not exist.")
+
+        model_instance.provider = node_data_model.provider
+        model_instance.model_name = node_data_model.name
+        model_instance.credentials = credentials
+        model_instance.parameters = completion_params
+        model_instance.stop = tuple(stop)
+        model_instance.model_type_instance = cast(LargeLanguageModel, model_instance.model_type_instance)
+        return model_instance
+
+    def _build_memory_for_llm_node(
+        self,
+        *,
+        node_data: Mapping[str, Any],
+        model_instance: ModelInstance,
+    ) -> PromptMessageMemory | None:
+        raw_memory_config = node_data.get("memory")
+        if raw_memory_config is None:
+            return None
+
+        node_memory = MemoryConfig.model_validate(raw_memory_config)
+        return llm_utils.fetch_memory(
+            variable_pool=self.graph_runtime_state.variable_pool,
+            app_id=self.graph_init_params.app_id,
+            node_data_memory=node_memory,
+            model_instance=model_instance,
+        )

api/core/helper/code_executor/template_transformer.py

@@ -5,7 +5,7 @@ from base64 import b64encode
 from collections.abc import Mapping
 from typing import Any
 
-from core.variables.utils import dumps_with_segments
+from core.workflow.variables.utils import dumps_with_segments
 
 
 class TemplateTransformer(ABC):

api/core/model_manager.py

@@ -1,5 +1,5 @@
 import logging
-from collections.abc import Callable, Generator, Iterable, Sequence
+from collections.abc import Callable, Generator, Iterable, Mapping, Sequence
 from typing import IO, Any, Literal, Optional, Union, cast, overload
 
 from configs import dify_config
@@ -38,6 +38,9 @@ class ModelInstance:
         self.model_name = model
         self.provider = provider_model_bundle.configuration.provider.provider
         self.credentials = self._fetch_credentials_from_bundle(provider_model_bundle, model)
+        # Runtime LLM invocation fields.
+        self.parameters: Mapping[str, Any] = {}
+        self.stop: Sequence[str] = ()
         self.model_type_instance = self.provider_model_bundle.model_type_instance
         self.load_balancing_manager = self._get_load_balancing_manager(
             configuration=provider_model_bundle.configuration,

api/core/prompt/advanced_prompt_transform.py

@@ -4,6 +4,7 @@ from typing import cast
 from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.helper.code_executor.jinja2.jinja2_formatter import Jinja2Formatter
 from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_manager import ModelInstance
 from core.model_runtime.entities import (
     AssistantPromptMessage,
     PromptMessage,
@@ -44,7 +45,8 @@ class AdvancedPromptTransform(PromptTransform):
         context: str | None,
         memory_config: MemoryConfig | None,
         memory: TokenBufferMemory | None,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
         image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
     ) -> list[PromptMessage]:
         prompt_messages = []
@@ -59,6 +61,7 @@ class AdvancedPromptTransform(PromptTransform):
                 memory_config=memory_config,
                 memory=memory,
                 model_config=model_config,
+                model_instance=model_instance,
                 image_detail_config=image_detail_config,
             )
         elif isinstance(prompt_template, list) and all(isinstance(item, ChatModelMessage) for item in prompt_template):
@@ -71,6 +74,7 @@ class AdvancedPromptTransform(PromptTransform):
                 memory_config=memory_config,
                 memory=memory,
                 model_config=model_config,
+                model_instance=model_instance,
                 image_detail_config=image_detail_config,
             )
 
@@ -85,7 +89,8 @@ class AdvancedPromptTransform(PromptTransform):
         context: str | None,
         memory_config: MemoryConfig | None,
         memory: TokenBufferMemory | None,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
         image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
     ) -> list[PromptMessage]:
         """
@@ -111,6 +116,7 @@ class AdvancedPromptTransform(PromptTransform):
                     parser=parser,
                     prompt_inputs=prompt_inputs,
                     model_config=model_config,
+                    model_instance=model_instance,
                 )
 
             if query:
@@ -146,7 +152,8 @@ class AdvancedPromptTransform(PromptTransform):
         context: str | None,
         memory_config: MemoryConfig | None,
         memory: TokenBufferMemory | None,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
         image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
     ) -> list[PromptMessage]:
         """
@@ -198,8 +205,13 @@ class AdvancedPromptTransform(PromptTransform):
 
         prompt_message_contents: list[PromptMessageContentUnionTypes] = []
         if memory and memory_config:
-            prompt_messages = self._append_chat_histories(memory, memory_config, prompt_messages, model_config)
-
+            prompt_messages = self._append_chat_histories(
+                memory,
+                memory_config,
+                prompt_messages,
+                model_config=model_config,
+                model_instance=model_instance,
+            )
             if files and query is not None:
                 for file in files:
                     prompt_message_contents.append(
@@ -276,7 +288,8 @@ class AdvancedPromptTransform(PromptTransform):
         role_prefix: MemoryConfig.RolePrefix,
         parser: PromptTemplateParser,
         prompt_inputs: Mapping[str, str],
-        model_config: ModelConfigWithCredentialsEntity,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
     ) -> Mapping[str, str]:
         prompt_inputs = dict(prompt_inputs)
         if "#histories#" in parser.variable_keys:
@@ -286,7 +299,11 @@ class AdvancedPromptTransform(PromptTransform):
                 prompt_inputs = {k: inputs[k] for k in parser.variable_keys if k in inputs}
                 tmp_human_message = UserPromptMessage(content=parser.format(prompt_inputs))
 
-                rest_tokens = self._calculate_rest_token([tmp_human_message], model_config)
+                rest_tokens = self._calculate_rest_token(
+                    [tmp_human_message],
+                    model_config=model_config,
+                    model_instance=model_instance,
+                )
 
                 histories = self._get_history_messages_from_memory(
                     memory=memory,

api/core/prompt/agent_history_prompt_transform.py

@@ -41,7 +41,7 @@ class AgentHistoryPromptTransform(PromptTransform):
         if not self.memory:
             return prompt_messages
 
-        max_token_limit = self._calculate_rest_token(self.prompt_messages, self.model_config)
+        max_token_limit = self._calculate_rest_token(self.prompt_messages, model_config=self.model_config)
 
         model_type_instance = self.model_config.provider_model_bundle.model_type_instance
         model_type_instance = cast(LargeLanguageModel, model_type_instance)

api/core/prompt/prompt_transform.py

@@ -4,45 +4,83 @@ from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEnti
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities.message_entities import PromptMessage
-from core.model_runtime.entities.model_entities import ModelPropertyKey
+from core.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey
 from core.prompt.entities.advanced_prompt_entities import MemoryConfig
 
 
 class PromptTransform:
+    def _resolve_model_runtime(
+        self,
+        *,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
+    ) -> tuple[ModelInstance, AIModelEntity]:
+        if model_instance is None:
+            if model_config is None:
+                raise ValueError("Either model_config or model_instance must be provided.")
+            model_instance = ModelInstance(
+                provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
+            )
+            model_instance.credentials = model_config.credentials
+            model_instance.parameters = model_config.parameters
+            model_instance.stop = model_config.stop
+
+        model_schema = model_instance.model_type_instance.get_model_schema(
+            model=model_instance.model_name,
+            credentials=model_instance.credentials,
+        )
+        if model_schema is None:
+            if model_config is None:
+                raise ValueError("Model schema not found for the provided model instance.")
+            model_schema = model_config.model_schema
+
+        return model_instance, model_schema
+
     def _append_chat_histories(
         self,
         memory: TokenBufferMemory,
         memory_config: MemoryConfig,
         prompt_messages: list[PromptMessage],
-        model_config: ModelConfigWithCredentialsEntity,
+        *,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
     ) -> list[PromptMessage]:
-        rest_tokens = self._calculate_rest_token(prompt_messages, model_config)
+        rest_tokens = self._calculate_rest_token(
+            prompt_messages,
+            model_config=model_config,
+            model_instance=model_instance,
+        )
         histories = self._get_history_messages_list_from_memory(memory, memory_config, rest_tokens)
         prompt_messages.extend(histories)
 
         return prompt_messages
 
     def _calculate_rest_token(
-        self, prompt_messages: list[PromptMessage], model_config: ModelConfigWithCredentialsEntity
+        self,
+        prompt_messages: list[PromptMessage],
+        *,
+        model_config: ModelConfigWithCredentialsEntity | None = None,
+        model_instance: ModelInstance | None = None,
     ) -> int:
+        model_instance, model_schema = self._resolve_model_runtime(
+            model_config=model_config,
+            model_instance=model_instance,
+        )
+        model_parameters = model_instance.parameters
         rest_tokens = 2000
 
-        model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
+        model_context_tokens = model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
         if model_context_tokens:
-            model_instance = ModelInstance(
-                provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
-            )
-
             curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
 
             max_tokens = 0
-            for parameter_rule in model_config.model_schema.parameter_rules:
+            for parameter_rule in 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 "")
+                        model_parameters.get(parameter_rule.name)
+                        or model_parameters.get(parameter_rule.use_template or "")
                     ) or 0
 
             rest_tokens = model_context_tokens - max_tokens - curr_message_tokens

api/core/prompt/simple_prompt_transform.py

@@ -252,7 +252,7 @@ class SimplePromptTransform(PromptTransform):
         if memory:
             tmp_human_message = UserPromptMessage(content=prompt)
 
-            rest_tokens = self._calculate_rest_token([tmp_human_message], model_config)
+            rest_tokens = self._calculate_rest_token([tmp_human_message], model_config=model_config)
             histories = self._get_history_messages_from_memory(
                 memory=memory,
                 memory_config=MemoryConfig(

api/core/rag/index_processor/processor/paragraph_index_processor.py

@@ -469,7 +469,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
         if not isinstance(result, LLMResult):
             raise ValueError("Expected LLMResult when stream=False")
 
-        summary_content = getattr(result.message, "content", "")
+        summary_content = result.message.get_text_content()
         usage = result.usage
 
         # Deduct quota for summary generation (same as workflow nodes)

api/core/workflow/conversation_variable_updater.py

@@ -1,7 +1,7 @@
 import abc
 from typing import Protocol
 
-from core.variables import VariableBase
+from core.workflow.variables import VariableBase
 
 
 class ConversationVariableUpdater(Protocol):

api/core/workflow/graph_engine/entities/commands.py

@@ -11,7 +11,7 @@ from typing import Any
 
 from pydantic import BaseModel, Field
 
-from core.variables.variables import Variable
+from core.workflow.variables.variables import Variable
 
 
 class CommandType(StrEnum):

api/core/workflow/nodes/agent/agent_node.py

@@ -25,7 +25,6 @@ from core.tools.entities.tool_entities import (
 )
 from core.tools.tool_manager import ToolManager
 from core.tools.utils.message_transformer import ToolFileMessageTransformer
-from core.variables.segments import ArrayFileSegment, StringSegment
 from core.workflow.enums import (
     NodeType,
     SystemVariableKey,
@@ -44,6 +43,7 @@ from core.workflow.nodes.agent.entities import AgentNodeData, AgentOldVersionMod
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.base.variable_template_parser import VariableTemplateParser
 from core.workflow.runtime import VariablePool
+from core.workflow.variables.segments import ArrayFileSegment, StringSegment
 from extensions.ext_database import db
 from factories import file_factory
 from factories.agent_factory import get_plugin_agent_strategy

api/core/workflow/nodes/answer/answer_node.py

@@ -1,13 +1,13 @@
 from collections.abc import Mapping, Sequence
 from typing import Any
 
-from core.variables import ArrayFileSegment, FileSegment, Segment
 from core.workflow.enums import NodeExecutionType, NodeType, WorkflowNodeExecutionStatus
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.answer.entities import AnswerNodeData
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.base.template import Template
 from core.workflow.nodes.base.variable_template_parser import VariableTemplateParser
+from core.workflow.variables import ArrayFileSegment, FileSegment, Segment
 
 
 class AnswerNode(Node[AnswerNodeData]):

api/core/workflow/nodes/code/code_node.py

@@ -5,13 +5,13 @@ from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast
 from core.helper.code_executor.code_node_provider import CodeNodeProvider
 from core.helper.code_executor.javascript.javascript_code_provider import JavascriptCodeProvider
 from core.helper.code_executor.python3.python3_code_provider import Python3CodeProvider
-from core.variables.segments import ArrayFileSegment
-from core.variables.types import SegmentType
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.code.entities import CodeLanguage, CodeNodeData
 from core.workflow.nodes.code.limits import CodeNodeLimits
+from core.workflow.variables.segments import ArrayFileSegment
+from core.workflow.variables.types import SegmentType
 
 from .exc import (
     CodeNodeError,

api/core/workflow/nodes/code/entities.py

@@ -3,9 +3,9 @@ from typing import Annotated, Literal
 from pydantic import AfterValidator, BaseModel
 
 from core.helper.code_executor.code_executor import CodeLanguage
-from core.variables.types import SegmentType
 from core.workflow.nodes.base import BaseNodeData
 from core.workflow.nodes.base.entities import VariableSelector
+from core.workflow.variables.types import SegmentType
 
 _ALLOWED_OUTPUT_FROM_CODE = frozenset(
     [

api/core/workflow/nodes/document_extractor/node.py

@@ -21,12 +21,12 @@ from docx.table import Table
 from docx.text.paragraph import Paragraph
 
 from core.helper import ssrf_proxy
-from core.variables import ArrayFileSegment
-from core.variables.segments import ArrayStringSegment, FileSegment
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
 from core.workflow.file import File, FileTransferMethod, file_manager
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
+from core.workflow.variables import ArrayFileSegment
+from core.workflow.variables.segments import ArrayStringSegment, FileSegment
 
 from .entities import DocumentExtractorNodeData, UnstructuredApiConfig
 from .exc import DocumentExtractorError, FileDownloadError, TextExtractionError, UnsupportedFileTypeError

api/core/workflow/nodes/http_request/executor.py

@@ -11,10 +11,10 @@ import httpx
 from json_repair import repair_json
 
 from core.helper.ssrf_proxy import ssrf_proxy
-from core.variables.segments import ArrayFileSegment, FileSegment
 from core.workflow.file.enums import FileTransferMethod
 from core.workflow.file.file_manager import file_manager as default_file_manager
 from core.workflow.runtime import VariablePool
+from core.workflow.variables.segments import ArrayFileSegment, FileSegment
 
 from ..protocols import FileManagerProtocol, HttpClientProtocol
 from .entities import (

api/core/workflow/nodes/http_request/node.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any
 
 from core.helper.ssrf_proxy import ssrf_proxy
 from core.tools.tool_file_manager import ToolFileManager
-from core.variables.segments import ArrayFileSegment
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
 from core.workflow.file import File, FileTransferMethod
 from core.workflow.file.file_manager import file_manager as default_file_manager
@@ -15,6 +14,7 @@ from core.workflow.nodes.base.entities import VariableSelector
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.http_request.executor import Executor
 from core.workflow.nodes.protocols import FileManagerProtocol, HttpClientProtocol
+from core.workflow.variables.segments import ArrayFileSegment
 from factories import file_factory
 
 from .config import build_http_request_config, resolve_http_request_config

api/core/workflow/nodes/human_input/entities.py

@@ -10,10 +10,10 @@ from typing import Annotated, Any, ClassVar, Literal, Self
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
-from core.variables.consts import SELECTORS_LENGTH
 from core.workflow.nodes.base import BaseNodeData
 from core.workflow.nodes.base.variable_template_parser import VariableTemplateParser
 from core.workflow.runtime import VariablePool
+from core.workflow.variables.consts import SELECTORS_LENGTH
 
 from .enums import ButtonStyle, DeliveryMethodType, EmailRecipientType, FormInputType, PlaceholderType, TimeoutUnit
 

api/core/workflow/nodes/iteration/iteration_node.py

@@ -7,9 +7,6 @@ from typing import TYPE_CHECKING, Any, NewType, cast
 from typing_extensions import TypeIs
 
 from core.model_runtime.entities.llm_entities import LLMUsage
-from core.variables import IntegerVariable, NoneSegment
-from core.variables.segments import ArrayAnySegment, ArraySegment
-from core.variables.variables import Variable
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.enums import (
     NodeExecutionType,
@@ -36,6 +33,9 @@ from core.workflow.nodes.base import LLMUsageTrackingMixin
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData
 from core.workflow.runtime import VariablePool
+from core.workflow.variables import IntegerVariable, NoneSegment
+from core.workflow.variables.segments import ArrayAnySegment, ArraySegment
+from core.workflow.variables.variables import Variable
 from libs.datetime_utils import naive_utc_now
 
 from .exc import (

api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py

@@ -5,12 +5,6 @@ from typing import TYPE_CHECKING, Any, Literal
 from core.app.app_config.entities import DatasetRetrieveConfigEntity
 from core.model_runtime.entities.llm_entities import LLMUsage
 from core.model_runtime.utils.encoders import jsonable_encoder
-from core.variables import (
-    ArrayFileSegment,
-    FileSegment,
-    StringSegment,
-)
-from core.variables.segments import ArrayObjectSegment
 from core.workflow.entities import GraphInitParams
 from core.workflow.enums import (
     NodeType,
@@ -22,6 +16,12 @@ from core.workflow.nodes.base import LLMUsageTrackingMixin
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.llm.file_saver import FileSaverImpl, LLMFileSaver
 from core.workflow.repositories.rag_retrieval_protocol import KnowledgeRetrievalRequest, RAGRetrievalProtocol, Source
+from core.workflow.variables import (
+    ArrayFileSegment,
+    FileSegment,
+    StringSegment,
+)
+from core.workflow.variables.segments import ArrayObjectSegment
 
 from .entities import KnowledgeRetrievalNodeData
 from .exc import (

api/core/workflow/nodes/list_operator/node.py

@@ -1,12 +1,12 @@
 from collections.abc import Callable, Sequence
 from typing import Any, TypeAlias, TypeVar
 
-from core.variables import ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment
-from core.variables.segments import ArrayAnySegment, ArrayBooleanSegment, ArraySegment
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
 from core.workflow.file import File
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
+from core.workflow.variables import ArrayFileSegment, ArrayNumberSegment, ArrayStringSegment
+from core.workflow.variables.segments import ArrayAnySegment, ArrayBooleanSegment, ArraySegment
 
 from .entities import FilterOperator, ListOperatorNodeData, Order
 from .exc import InvalidConditionError, InvalidFilterValueError, InvalidKeyError, ListOperatorError

api/core/workflow/nodes/llm/llm_utils.py

@@ -5,21 +5,17 @@ from sqlalchemy import select, update
 from sqlalchemy.orm import Session
 
 from configs import dify_config
-from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.entities.provider_entities import ProviderQuotaType, QuotaUnit
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities.llm_entities import LLMUsage
-from core.model_runtime.entities.model_entities import ModelType
+from core.model_runtime.entities.model_entities import AIModelEntity
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
 from core.prompt.entities.advanced_prompt_entities import MemoryConfig
-from core.variables.segments import ArrayAnySegment, ArrayFileSegment, FileSegment, NoneSegment, StringSegment
 from core.workflow.enums import SystemVariableKey
 from core.workflow.file.models import File
-from core.workflow.nodes.llm.entities import ModelConfig
-from core.workflow.nodes.llm.exc import LLMModeRequiredError, ModelNotExistError
-from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory
 from core.workflow.runtime import VariablePool
+from core.workflow.variables.segments import ArrayAnySegment, ArrayFileSegment, FileSegment, NoneSegment, StringSegment
 from extensions.ext_database import db
 from libs.datetime_utils import naive_utc_now
 from models.model import Conversation
@@ -29,46 +25,14 @@ from models.provider_ids import ModelProviderID
 from .exc import InvalidVariableTypeError
 
 
-def fetch_model_config(
-    *,
-    node_data_model: ModelConfig,
-    credentials_provider: CredentialsProvider,
-    model_factory: ModelFactory,
-) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]:
-    if not node_data_model.mode:
-        raise LLMModeRequiredError("LLM mode is required.")
-
-    credentials = credentials_provider.fetch(node_data_model.provider, node_data_model.name)
-    model_instance = model_factory.init_model_instance(node_data_model.provider, node_data_model.name)
-    provider_model_bundle = model_instance.provider_model_bundle
-
-    provider_model = provider_model_bundle.configuration.get_provider_model(
-        model=node_data_model.name,
-        model_type=ModelType.LLM,
+def fetch_model_schema(*, model_instance: ModelInstance) -> AIModelEntity:
+    model_schema = cast(LargeLanguageModel, model_instance.model_type_instance).get_model_schema(
+        model_instance.model_name,
+        model_instance.credentials,
     )
-    if provider_model is None:
-        raise ModelNotExistError(f"Model {node_data_model.name} not exist.")
-    provider_model.raise_for_status()
-
-    stop: list[str] = []
-    if "stop" in node_data_model.completion_params:
-        stop = node_data_model.completion_params.pop("stop")
-
-    model_schema = model_instance.model_type_instance.get_model_schema(node_data_model.name, credentials)
     if not model_schema:
-        raise ModelNotExistError(f"Model {node_data_model.name} not exist.")
-
-    model_instance.model_type_instance = cast(LargeLanguageModel, model_instance.model_type_instance)
-    return model_instance, ModelConfigWithCredentialsEntity(
-        provider=node_data_model.provider,
-        model=node_data_model.name,
-        model_schema=model_schema,
-        mode=node_data_model.mode,
-        provider_model_bundle=provider_model_bundle,
-        credentials=credentials,
-        parameters=node_data_model.completion_params,
-        stop=stop,
-    )
+        raise ValueError(f"Model schema not found for {model_instance.model_name}")
+    return model_schema
 
 
 def fetch_files(variable_pool: VariablePool, selector: Sequence[str]) -> Sequence["File"]:

api/core/workflow/nodes/llm/node.py

@@ -11,11 +11,9 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from sqlalchemy import select
 
-from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.helper.code_executor import CodeExecutor, CodeLanguage
 from core.llm_generator.output_parser.errors import OutputParserError
 from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output
-from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities import (
     ImagePromptMessageContent,
@@ -38,20 +36,12 @@ from core.model_runtime.entities.message_entities import (
     SystemPromptMessage,
     UserPromptMessage,
 )
-from core.model_runtime.entities.model_entities import AIModelEntity, ModelFeature, ModelPropertyKey
+from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptTemplate, MemoryConfig
 from core.prompt.utils.prompt_message_util import PromptMessageUtil
 from core.rag.entities.citation_metadata import RetrievalSourceMetadata
 from core.tools.signature import sign_upload_file
-from core.variables import (
-    ArrayFileSegment,
-    ArraySegment,
-    FileSegment,
-    NoneSegment,
-    ObjectSegment,
-    StringSegment,
-)
 from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID
 from core.workflow.entities import GraphInitParams
 from core.workflow.enums import (
@@ -72,8 +62,16 @@ from core.workflow.node_events import (
 from core.workflow.nodes.base.entities import VariableSelector
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.base.variable_template_parser import VariableTemplateParser
-from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory
+from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory, PromptMessageMemory
 from core.workflow.runtime import VariablePool
+from core.workflow.variables import (
+    ArrayFileSegment,
+    ArraySegment,
+    FileSegment,
+    NoneSegment,
+    ObjectSegment,
+    StringSegment,
+)
 from extensions.ext_database import db
 from models.dataset import SegmentAttachmentBinding
 from models.model import UploadFile
@@ -83,7 +81,6 @@ from .entities import (
     LLMNodeChatModelMessage,
     LLMNodeCompletionModelPromptTemplate,
     LLMNodeData,
-    ModelConfig,
 )
 from .exc import (
     InvalidContextStructureError,
@@ -116,6 +113,8 @@ class LLMNode(Node[LLMNodeData]):
     _llm_file_saver: LLMFileSaver
     _credentials_provider: CredentialsProvider
     _model_factory: ModelFactory
+    _model_instance: ModelInstance
+    _memory: PromptMessageMemory | None
 
     def __init__(
         self,
@@ -126,6 +125,8 @@ class LLMNode(Node[LLMNodeData]):
         *,
         credentials_provider: CredentialsProvider,
         model_factory: ModelFactory,
+        model_instance: ModelInstance,
+        memory: PromptMessageMemory | None = None,
         llm_file_saver: LLMFileSaver | None = None,
     ):
         super().__init__(
@@ -139,6 +140,8 @@ class LLMNode(Node[LLMNodeData]):
 
         self._credentials_provider = credentials_provider
         self._model_factory = model_factory
+        self._model_instance = model_instance
+        self._memory = memory
 
         if llm_file_saver is None:
             llm_file_saver = FileSaverImpl(
@@ -202,29 +205,12 @@ class LLMNode(Node[LLMNodeData]):
                 node_inputs["#context_files#"] = [file.model_dump() for file in context_files]
 
             # fetch model config
-            model_instance, model_config = self._fetch_model_config(
-                node_data_model=self.node_data.model,
-            )
-            model_name = getattr(model_instance, "model_name", None)
-            if not isinstance(model_name, str):
-                model_name = model_config.model
-            model_provider = getattr(model_instance, "provider", None)
-            if not isinstance(model_provider, str):
-                model_provider = model_config.provider
-            model_schema = model_instance.model_type_instance.get_model_schema(
-                model_name,
-                model_instance.credentials,
-            )
-            if not model_schema:
-                raise ValueError(f"Model schema not found for {model_name}")
+            model_instance = self._model_instance
+            model_name = model_instance.model_name
+            model_provider = model_instance.provider
+            model_stop = model_instance.stop
 
-            # fetch memory
-            memory = llm_utils.fetch_memory(
-                variable_pool=variable_pool,
-                app_id=self.app_id,
-                node_data_memory=self.node_data.memory,
-                model_instance=model_instance,
-            )
+            memory = self._memory
 
             query: str | None = None
             if self.node_data.memory:
@@ -240,9 +226,7 @@ class LLMNode(Node[LLMNodeData]):
                 context=context,
                 memory=memory,
                 model_instance=model_instance,
-                model_schema=model_schema,
-                model_parameters=self.node_data.model.completion_params,
-                stop=model_config.stop,
+                stop=model_stop,
                 prompt_template=self.node_data.prompt_template,
                 memory_config=self.node_data.memory,
                 vision_enabled=self.node_data.vision.enabled,
@@ -254,7 +238,6 @@ class LLMNode(Node[LLMNodeData]):
 
             # handle invoke result
             generator = LLMNode.invoke_llm(
-                node_data_model=self.node_data.model,
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 stop=stop,
@@ -371,7 +354,6 @@ class LLMNode(Node[LLMNodeData]):
     @staticmethod
     def invoke_llm(
         *,
-        node_data_model: ModelConfig,
         model_instance: ModelInstance,
         prompt_messages: Sequence[PromptMessage],
         stop: Sequence[str] | None = None,
@@ -384,11 +366,10 @@ class LLMNode(Node[LLMNodeData]):
         node_type: NodeType,
         reasoning_format: Literal["separated", "tagged"] = "tagged",
     ) -> Generator[NodeEventBase | LLMStructuredOutput, None, None]:
-        model_schema = model_instance.model_type_instance.get_model_schema(
-            node_data_model.name, model_instance.credentials
-        )
-        if not model_schema:
-            raise ValueError(f"Model schema not found for {node_data_model.name}")
+        model_parameters = model_instance.parameters
+        invoke_model_parameters = dict(model_parameters)
+
+        model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
 
         if structured_output_enabled:
             output_schema = LLMNode.fetch_structured_output_schema(
@@ -402,7 +383,7 @@ class LLMNode(Node[LLMNodeData]):
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 json_schema=output_schema,
-                model_parameters=node_data_model.completion_params,
+                model_parameters=invoke_model_parameters,
                 stop=list(stop or []),
                 stream=True,
                 user=user_id,
@@ -412,7 +393,7 @@ class LLMNode(Node[LLMNodeData]):
 
             invoke_result = model_instance.invoke_llm(
                 prompt_messages=list(prompt_messages),
-                model_parameters=node_data_model.completion_params,
+                model_parameters=invoke_model_parameters,
                 stop=list(stop or []),
                 stream=True,
                 user=user_id,
@@ -771,33 +752,14 @@ class LLMNode(Node[LLMNodeData]):
 
         return None
 
-    def _fetch_model_config(
-        self,
-        *,
-        node_data_model: ModelConfig,
-    ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]:
-        model, model_config_with_cred = llm_utils.fetch_model_config(
-            node_data_model=node_data_model,
-            credentials_provider=self._credentials_provider,
-            model_factory=self._model_factory,
-        )
-        completion_params = model_config_with_cred.parameters
-
-        model_config_with_cred.parameters = completion_params
-        # NOTE(-LAN-): This line modify the `self.node_data.model`, which is used in `_invoke_llm()`.
-        node_data_model.completion_params = completion_params
-        return model, model_config_with_cred
-
     @staticmethod
     def fetch_prompt_messages(
         *,
         sys_query: str | None = None,
         sys_files: Sequence[File],
         context: str | None = None,
-        memory: TokenBufferMemory | None = None,
+        memory: PromptMessageMemory | None = None,
         model_instance: ModelInstance,
-        model_schema: AIModelEntity,
-        model_parameters: Mapping[str, Any],
         prompt_template: Sequence[LLMNodeChatModelMessage] | LLMNodeCompletionModelPromptTemplate,
         stop: Sequence[str] | None = None,
         memory_config: MemoryConfig | None = None,
@@ -808,6 +770,7 @@ class LLMNode(Node[LLMNodeData]):
         context_files: list[File] | None = None,
     ) -> tuple[Sequence[PromptMessage], Sequence[str] | None]:
         prompt_messages: list[PromptMessage] = []
+        model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
 
         if isinstance(prompt_template, list):
             # For chat model
@@ -826,8 +789,6 @@ class LLMNode(Node[LLMNodeData]):
                 memory=memory,
                 memory_config=memory_config,
                 model_instance=model_instance,
-                model_schema=model_schema,
-                model_parameters=model_parameters,
             )
             # Extend prompt_messages with memory messages
             prompt_messages.extend(memory_messages)
@@ -865,8 +826,6 @@ class LLMNode(Node[LLMNodeData]):
                 memory=memory,
                 memory_config=memory_config,
                 model_instance=model_instance,
-                model_schema=model_schema,
-                model_parameters=model_parameters,
             )
             # Insert histories into the prompt
             prompt_content = prompt_messages[0].content
@@ -1316,23 +1275,23 @@ def _calculate_rest_token(
     *,
     prompt_messages: list[PromptMessage],
     model_instance: ModelInstance,
-    model_schema: AIModelEntity,
-    model_parameters: Mapping[str, Any],
 ) -> int:
     rest_tokens = 2000
+    runtime_model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
+    runtime_model_parameters = model_instance.parameters
 
-    model_context_tokens = model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
+    model_context_tokens = runtime_model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
     if model_context_tokens:
         curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
 
         max_tokens = 0
-        for parameter_rule in model_schema.parameter_rules:
+        for parameter_rule in runtime_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_parameters.get(parameter_rule.name)
-                    or model_parameters.get(str(parameter_rule.use_template))
+                    runtime_model_parameters.get(parameter_rule.name)
+                    or runtime_model_parameters.get(str(parameter_rule.use_template))
                     or 0
                 )
 
@@ -1344,11 +1303,9 @@ def _calculate_rest_token(
 
 def _handle_memory_chat_mode(
     *,
-    memory: TokenBufferMemory | None,
+    memory: PromptMessageMemory | None,
     memory_config: MemoryConfig | None,
     model_instance: ModelInstance,
-    model_schema: AIModelEntity,
-    model_parameters: Mapping[str, Any],
 ) -> Sequence[PromptMessage]:
     memory_messages: Sequence[PromptMessage] = []
     # Get messages from memory for chat model
@@ -1356,8 +1313,6 @@ def _handle_memory_chat_mode(
         rest_tokens = _calculate_rest_token(
             prompt_messages=[],
             model_instance=model_instance,
-            model_schema=model_schema,
-            model_parameters=model_parameters,
         )
         memory_messages = memory.get_history_prompt_messages(
             max_token_limit=rest_tokens,
@@ -1368,11 +1323,9 @@ def _handle_memory_chat_mode(
 
 def _handle_memory_completion_mode(
     *,
-    memory: TokenBufferMemory | None,
+    memory: PromptMessageMemory | None,
     memory_config: MemoryConfig | None,
     model_instance: ModelInstance,
-    model_schema: AIModelEntity,
-    model_parameters: Mapping[str, Any],
 ) -> str:
     memory_text = ""
     # Get history text from memory for completion model
@@ -1380,20 +1333,51 @@ def _handle_memory_completion_mode(
         rest_tokens = _calculate_rest_token(
             prompt_messages=[],
             model_instance=model_instance,
-            model_schema=model_schema,
-            model_parameters=model_parameters,
         )
         if not memory_config.role_prefix:
             raise MemoryRolePrefixRequiredError("Memory role prefix is required for completion model.")
-        memory_text = memory.get_history_prompt_text(
+        memory_messages = memory.get_history_prompt_messages(
             max_token_limit=rest_tokens,
             message_limit=memory_config.window.size if memory_config.window.enabled else None,
+        )
+        memory_text = _convert_history_messages_to_text(
+            history_messages=memory_messages,
             human_prefix=memory_config.role_prefix.user,
             ai_prefix=memory_config.role_prefix.assistant,
         )
     return memory_text
 
 
+def _convert_history_messages_to_text(
+    *,
+    history_messages: Sequence[PromptMessage],
+    human_prefix: str,
+    ai_prefix: str,
+) -> str:
+    string_messages: list[str] = []
+    for message in history_messages:
+        if message.role == PromptMessageRole.USER:
+            role = human_prefix
+        elif message.role == PromptMessageRole.ASSISTANT:
+            role = ai_prefix
+        else:
+            continue
+
+        if isinstance(message.content, list):
+            content_parts = []
+            for content in message.content:
+                if isinstance(content, TextPromptMessageContent):
+                    content_parts.append(content.data)
+                elif isinstance(content, ImagePromptMessageContent):
+                    content_parts.append("[image]")
+
+            inner_msg = "\n".join(content_parts)
+            string_messages.append(f"{role}: {inner_msg}")
+        else:
+            string_messages.append(f"{role}: {message.content}")
+    return "\n".join(string_messages)
+
+
 def _handle_completion_template(
     *,
     template: LLMNodeCompletionModelPromptTemplate,

api/core/workflow/nodes/llm/protocols.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import Any, Protocol
 
 from core.model_manager import ModelInstance
+from core.model_runtime.entities import PromptMessage
 
 
 class CredentialsProvider(Protocol):
@@ -19,3 +21,13 @@ class ModelFactory(Protocol):
     def init_model_instance(self, provider_name: str, model_name: str) -> ModelInstance:
         """Create a model instance that is ready for schema lookup and invocation."""
         ...
+
+
+class PromptMessageMemory(Protocol):
+    """Port for loading memory as prompt messages for LLM nodes."""
+
+    def get_history_prompt_messages(
+        self, max_token_limit: int = 2000, message_limit: int | None = None
+    ) -> Sequence[PromptMessage]:
+        """Return historical prompt messages constrained by token/message limits."""
+        ...

api/core/workflow/nodes/loop/entities.py

@@ -3,9 +3,9 @@ from typing import Annotated, Any, Literal
 
 from pydantic import AfterValidator, BaseModel, Field, field_validator
 
-from core.variables.types import SegmentType
 from core.workflow.nodes.base import BaseLoopNodeData, BaseLoopState, BaseNodeData
 from core.workflow.utils.condition.entities import Condition
+from core.workflow.variables.types import SegmentType
 
 _VALID_VAR_TYPE = frozenset(
     [

api/core/workflow/nodes/loop/loop_node.py

@@ -6,7 +6,6 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 from core.model_runtime.entities.llm_entities import LLMUsage
-from core.variables import Segment, SegmentType
 from core.workflow.enums import (
     NodeExecutionType,
     NodeType,
@@ -31,6 +30,7 @@ from core.workflow.nodes.base import LLMUsageTrackingMixin
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.loop.entities import LoopCompletedReason, LoopNodeData, LoopVariableData
 from core.workflow.utils.condition.processor import ConditionProcessor
+from core.workflow.variables import Segment, SegmentType
 from factories.variable_factory import TypeMismatchError, build_segment_with_type, segment_to_variable
 from libs.datetime_utils import naive_utc_now
 

api/core/workflow/nodes/parameter_extractor/entities.py

@@ -8,9 +8,9 @@ from pydantic import (
 )
 
 from core.prompt.entities.advanced_prompt_entities import MemoryConfig
-from core.variables.types import SegmentType
 from core.workflow.nodes.base import BaseNodeData
 from core.workflow.nodes.llm.entities import ModelConfig, VisionConfig
+from core.workflow.variables.types import SegmentType
 
 _OLD_BOOL_TYPE_NAME = "bool"
 _OLD_SELECT_TYPE_NAME = "select"

api/core/workflow/nodes/parameter_extractor/exc.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from core.variables.types import SegmentType
+from core.workflow.variables.types import SegmentType
 
 
 class ParameterExtractorNodeError(ValueError):

api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py

@@ -5,7 +5,6 @@ import uuid
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, cast
 
-from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities import ImagePromptMessageContent
@@ -25,14 +24,14 @@ from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
 from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate
 from core.prompt.simple_prompt_transform import ModelMode
 from core.prompt.utils.prompt_message_util import PromptMessageUtil
-from core.variables.types import ArrayValidation, SegmentType
 from core.workflow.enums import NodeType, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
 from core.workflow.file import File
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base import variable_template_parser
 from core.workflow.nodes.base.node import Node
-from core.workflow.nodes.llm import ModelConfig, llm_utils
+from core.workflow.nodes.llm import llm_utils
 from core.workflow.runtime import VariablePool
+from core.workflow.variables.types import ArrayValidation, SegmentType
 from factories.variable_factory import build_segment_with_type
 
 from .entities import ParameterExtractorNodeData
@@ -95,8 +94,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
 
     node_type = NodeType.PARAMETER_EXTRACTOR
 
-    _model_instance: ModelInstance | None = None
-    _model_config: ModelConfigWithCredentialsEntity | None = None
+    _model_instance: ModelInstance
     _credentials_provider: "CredentialsProvider"
     _model_factory: "ModelFactory"
 
@@ -109,6 +107,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         *,
         credentials_provider: "CredentialsProvider",
         model_factory: "ModelFactory",
+        model_instance: ModelInstance,
     ) -> None:
         super().__init__(
             id=id,
@@ -118,6 +117,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         )
         self._credentials_provider = credentials_provider
         self._model_factory = model_factory
+        self._model_instance = model_instance
 
     @classmethod
     def get_default_config(cls, filters: Mapping[str, object] | None = None) -> Mapping[str, object]:
@@ -155,18 +155,14 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
             else []
         )
 
-        model_instance, model_config = self._fetch_model_config(node_data.model)
+        model_instance = self._model_instance
         if not isinstance(model_instance.model_type_instance, LargeLanguageModel):
             raise InvalidModelTypeError("Model is not a Large Language Model")
 
-        llm_model = model_instance.model_type_instance
-        model_schema = llm_model.get_model_schema(
-            model=model_config.model,
-            credentials=model_config.credentials,
-        )
-        if not model_schema:
-            raise ModelSchemaNotFoundError("Model schema not found")
-
+        try:
+            model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
+        except ValueError as exc:
+            raise ModelSchemaNotFoundError("Model schema not found") from exc
         # fetch memory
         memory = llm_utils.fetch_memory(
             variable_pool=variable_pool,
@@ -184,7 +180,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
                 node_data=node_data,
                 query=query,
                 variable_pool=self.graph_runtime_state.variable_pool,
-                model_config=model_config,
+                model_instance=model_instance,
                 memory=memory,
                 files=files,
                 vision_detail=node_data.vision.configs.detail,
@@ -195,7 +191,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
                 data=node_data,
                 query=query,
                 variable_pool=self.graph_runtime_state.variable_pool,
-                model_config=model_config,
+                model_instance=model_instance,
                 memory=memory,
                 files=files,
                 vision_detail=node_data.vision.configs.detail,
@@ -211,24 +207,23 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         }
 
         process_data = {
-            "model_mode": model_config.mode,
+            "model_mode": node_data.model.mode,
             "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
-                model_mode=model_config.mode, prompt_messages=prompt_messages
+                model_mode=node_data.model.mode, prompt_messages=prompt_messages
             ),
             "usage": None,
             "function": {} if not prompt_message_tools else jsonable_encoder(prompt_message_tools[0]),
             "tool_call": None,
-            "model_provider": model_config.provider,
-            "model_name": model_config.model,
+            "model_provider": model_instance.provider,
+            "model_name": model_instance.model_name,
         }
 
         try:
             text, usage, tool_call = self._invoke(
-                node_data_model=node_data.model,
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 tools=prompt_message_tools,
-                stop=model_config.stop,
+                stop=model_instance.stop,
             )
             process_data["usage"] = jsonable_encoder(usage)
             process_data["tool_call"] = jsonable_encoder(tool_call)
@@ -290,17 +285,16 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
 
     def _invoke(
         self,
-        node_data_model: ModelConfig,
         model_instance: ModelInstance,
         prompt_messages: list[PromptMessage],
         tools: list[PromptMessageTool],
-        stop: list[str],
+        stop: Sequence[str],
     ) -> tuple[str, LLMUsage, AssistantPromptMessage.ToolCall | None]:
         invoke_result = model_instance.invoke_llm(
             prompt_messages=prompt_messages,
-            model_parameters=node_data_model.completion_params,
+            model_parameters=dict(model_instance.parameters),
             tools=tools,
-            stop=stop,
+            stop=list(stop),
             stream=False,
             user=self.user_id,
         )
@@ -324,7 +318,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         node_data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_instance: ModelInstance,
         memory: TokenBufferMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -337,7 +331,13 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         )
 
         prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
-        rest_token = self._calculate_rest_token(node_data, query, variable_pool, model_config, "")
+        rest_token = self._calculate_rest_token(
+            node_data=node_data,
+            query=query,
+            variable_pool=variable_pool,
+            model_instance=model_instance,
+            context="",
+        )
         prompt_template = self._get_function_calling_prompt_template(
             node_data, query, variable_pool, memory, rest_token
         )
@@ -349,7 +349,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
             context="",
             memory_config=node_data.memory,
             memory=None,
-            model_config=model_config,
+            model_instance=model_instance,
             image_detail_config=vision_detail,
         )
 
@@ -406,7 +406,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_instance: ModelInstance,
         memory: TokenBufferMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -421,7 +421,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
                 node_data=data,
                 query=query,
                 variable_pool=variable_pool,
-                model_config=model_config,
+                model_instance=model_instance,
                 memory=memory,
                 files=files,
                 vision_detail=vision_detail,
@@ -431,7 +431,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
                 node_data=data,
                 query=query,
                 variable_pool=variable_pool,
-                model_config=model_config,
+                model_instance=model_instance,
                 memory=memory,
                 files=files,
                 vision_detail=vision_detail,
@@ -444,7 +444,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         node_data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_instance: ModelInstance,
         memory: TokenBufferMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -454,7 +454,11 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         """
         prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
         rest_token = self._calculate_rest_token(
-            node_data=node_data, query=query, variable_pool=variable_pool, model_config=model_config, context=""
+            node_data=node_data,
+            query=query,
+            variable_pool=variable_pool,
+            model_instance=model_instance,
+            context="",
         )
         prompt_template = self._get_prompt_engineering_prompt_template(
             node_data=node_data, query=query, variable_pool=variable_pool, memory=memory, max_token_limit=rest_token
@@ -467,7 +471,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
             context="",
             memory_config=node_data.memory,
             memory=memory,
-            model_config=model_config,
+            model_instance=model_instance,
             image_detail_config=vision_detail,
         )
 
@@ -478,7 +482,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         node_data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_instance: ModelInstance,
         memory: TokenBufferMemory | None,
         files: Sequence[File],
         vision_detail: ImagePromptMessageContent.DETAIL | None = None,
@@ -488,7 +492,11 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         """
         prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
         rest_token = self._calculate_rest_token(
-            node_data=node_data, query=query, variable_pool=variable_pool, model_config=model_config, context=""
+            node_data=node_data,
+            query=query,
+            variable_pool=variable_pool,
+            model_instance=model_instance,
+            context="",
         )
         prompt_template = self._get_prompt_engineering_prompt_template(
             node_data=node_data,
@@ -508,7 +516,7 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
             context="",
             memory_config=node_data.memory,
             memory=None,
-            model_config=model_config,
+            model_instance=model_instance,
             image_detail_config=vision_detail,
         )
 
@@ -769,21 +777,16 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
         node_data: ParameterExtractorNodeData,
         query: str,
         variable_pool: VariablePool,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_instance: ModelInstance,
         context: str | None,
     ) -> int:
+        try:
+            model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
+        except ValueError as exc:
+            raise ModelSchemaNotFoundError("Model schema not found") from exc
         prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
 
-        model_instance, model_config = self._fetch_model_config(node_data.model)
-        if not isinstance(model_instance.model_type_instance, LargeLanguageModel):
-            raise InvalidModelTypeError("Model is not a Large Language Model")
-
-        llm_model = model_instance.model_type_instance
-        model_schema = llm_model.get_model_schema(model_config.model, model_config.credentials)
-        if not model_schema:
-            raise ModelSchemaNotFoundError("Model schema not found")
-
-        if set(model_schema.features or []) & {ModelFeature.MULTI_TOOL_CALL, ModelFeature.MULTI_TOOL_CALL}:
+        if set(model_schema.features or []) & {ModelFeature.TOOL_CALL, ModelFeature.MULTI_TOOL_CALL}:
             prompt_template = self._get_function_calling_prompt_template(node_data, query, variable_pool, None, 2000)
         else:
             prompt_template = self._get_prompt_engineering_prompt_template(node_data, query, variable_pool, None, 2000)
@@ -796,27 +799,28 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
             context=context,
             memory_config=node_data.memory,
             memory=None,
-            model_config=model_config,
+            model_instance=model_instance,
         )
         rest_tokens = 2000
 
-        model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
+        model_context_tokens = model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
         if model_context_tokens:
-            model_type_instance = model_config.provider_model_bundle.model_type_instance
-            model_type_instance = cast(LargeLanguageModel, model_type_instance)
-
+            model_type_instance = cast(LargeLanguageModel, model_instance.model_type_instance)
             curr_message_tokens = (
-                model_type_instance.get_num_tokens(model_config.model, model_config.credentials, prompt_messages) + 1000
+                model_type_instance.get_num_tokens(
+                    model_instance.model_name, model_instance.credentials, prompt_messages
+                )
+                + 1000
             )  # add 1000 to ensure tool call messages
 
             max_tokens = 0
-            for parameter_rule in model_config.model_schema.parameter_rules:
+            for parameter_rule in 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 "")
+                        model_instance.parameters.get(parameter_rule.name)
+                        or model_instance.parameters.get(parameter_rule.use_template or "")
                     ) or 0
 
             rest_tokens = model_context_tokens - max_tokens - curr_message_tokens
@@ -824,21 +828,6 @@ class ParameterExtractorNode(Node[ParameterExtractorNodeData]):
 
         return rest_tokens
 
-    def _fetch_model_config(
-        self, node_data_model: ModelConfig
-    ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]:
-        """
-        Fetch model config.
-        """
-        if not self._model_instance or not self._model_config:
-            self._model_instance, self._model_config = llm_utils.fetch_model_config(
-                node_data_model=node_data_model,
-                credentials_provider=self._credentials_provider,
-                model_factory=self._model_factory,
-            )
-
-        return self._model_instance, self._model_config
-
     @classmethod
     def _extract_variable_selector_to_variable_mapping(
         cls,

api/core/workflow/nodes/question_classifier/question_classifier_node.py

@@ -3,12 +3,10 @@ import re
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
-from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance
 from core.model_runtime.entities import LLMUsage, ModelPropertyKey, PromptMessageRole
 from core.model_runtime.utils.encoders import jsonable_encoder
-from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
 from core.prompt.simple_prompt_transform import ModelMode
 from core.prompt.utils.prompt_message_util import PromptMessageUtil
 from core.workflow.entities import GraphInitParams
@@ -22,7 +20,12 @@ from core.workflow.node_events import ModelInvokeCompletedEvent, NodeRunResult
 from core.workflow.nodes.base.entities import VariableSelector
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.base.variable_template_parser import VariableTemplateParser
-from core.workflow.nodes.llm import LLMNode, LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate, llm_utils
+from core.workflow.nodes.llm import (
+    LLMNode,
+    LLMNodeChatModelMessage,
+    LLMNodeCompletionModelPromptTemplate,
+    llm_utils,
+)
 from core.workflow.nodes.llm.file_saver import FileSaverImpl, LLMFileSaver
 from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory
 from libs.json_in_md_parser import parse_and_check_json_markdown
@@ -52,6 +55,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
     _llm_file_saver: LLMFileSaver
     _credentials_provider: "CredentialsProvider"
     _model_factory: "ModelFactory"
+    _model_instance: ModelInstance
 
     def __init__(
         self,
@@ -62,6 +66,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         *,
         credentials_provider: "CredentialsProvider",
         model_factory: "ModelFactory",
+        model_instance: ModelInstance,
         llm_file_saver: LLMFileSaver | None = None,
     ):
         super().__init__(
@@ -75,6 +80,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
 
         self._credentials_provider = credentials_provider
         self._model_factory = model_factory
+        self._model_instance = model_instance
 
         if llm_file_saver is None:
             llm_file_saver = FileSaverImpl(
@@ -95,18 +101,8 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         variable = variable_pool.get(node_data.query_variable_selector) if node_data.query_variable_selector else None
         query = variable.value if variable else None
         variables = {"query": query}
-        # fetch model config
-        model_instance, model_config = llm_utils.fetch_model_config(
-            node_data_model=node_data.model,
-            credentials_provider=self._credentials_provider,
-            model_factory=self._model_factory,
-        )
-        model_schema = model_instance.model_type_instance.get_model_schema(
-            model_instance.model_name,
-            model_instance.credentials,
-        )
-        if not model_schema:
-            raise ValueError(f"Model schema not found for {model_instance.model_name}")
+        # fetch model instance
+        model_instance = self._model_instance
         # fetch memory
         memory = llm_utils.fetch_memory(
             variable_pool=variable_pool,
@@ -131,7 +127,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         rest_token = self._calculate_rest_token(
             node_data=node_data,
             query=query or "",
-            model_config=model_config,
+            model_instance=model_instance,
             context="",
         )
         prompt_template = self._get_prompt_template(
@@ -149,9 +145,7 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
             sys_query="",
             memory=memory,
             model_instance=model_instance,
-            model_schema=model_schema,
-            model_parameters=node_data.model.completion_params,
-            stop=model_config.stop,
+            stop=model_instance.stop,
             sys_files=files,
             vision_enabled=node_data.vision.enabled,
             vision_detail=node_data.vision.configs.detail,
@@ -166,7 +160,6 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         try:
             # handle invoke result
             generator = LLMNode.invoke_llm(
-                node_data_model=node_data.model,
                 model_instance=model_instance,
                 prompt_messages=prompt_messages,
                 stop=stop,
@@ -205,14 +198,14 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
                     category_name = classes_map[category_id_result]
                     category_id = category_id_result
             process_data = {
-                "model_mode": model_config.mode,
+                "model_mode": node_data.model.mode,
                 "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
-                    model_mode=model_config.mode, prompt_messages=prompt_messages
+                    model_mode=node_data.model.mode, prompt_messages=prompt_messages
                 ),
                 "usage": jsonable_encoder(usage),
                 "finish_reason": finish_reason,
-                "model_provider": model_config.provider,
-                "model_name": model_config.model,
+                "model_provider": model_instance.provider,
+                "model_name": model_instance.model_name,
             }
             outputs = {
                 "class_name": category_name,
@@ -285,39 +278,40 @@ class QuestionClassifierNode(Node[QuestionClassifierNodeData]):
         self,
         node_data: QuestionClassifierNodeData,
         query: str,
-        model_config: ModelConfigWithCredentialsEntity,
+        model_instance: ModelInstance,
         context: str | None,
     ) -> int:
-        prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
+        model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
+
         prompt_template = self._get_prompt_template(node_data, query, None, 2000)
-        prompt_messages = prompt_transform.get_prompt(
+        prompt_messages, _ = LLMNode.fetch_prompt_messages(
             prompt_template=prompt_template,
-            inputs={},
-            query="",
-            files=[],
+            sys_query="",
+            sys_files=[],
             context=context,
-            memory_config=node_data.memory,
             memory=None,
-            model_config=model_config,
+            model_instance=model_instance,
+            stop=model_instance.stop,
+            memory_config=node_data.memory,
+            vision_enabled=False,
+            vision_detail=node_data.vision.configs.detail,
+            variable_pool=self.graph_runtime_state.variable_pool,
+            jinja2_variables=[],
         )
         rest_tokens = 2000
 
-        model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
+        model_context_tokens = model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
         if model_context_tokens:
-            model_instance = ModelInstance(
-                provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
-            )
-
             curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
 
             max_tokens = 0
-            for parameter_rule in model_config.model_schema.parameter_rules:
+            for parameter_rule in 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 "")
+                        model_instance.parameters.get(parameter_rule.name)
+                        or model_instance.parameters.get(parameter_rule.use_template or "")
                     ) or 0
 
             rest_tokens = model_context_tokens - max_tokens - curr_message_tokens

api/core/workflow/nodes/tool/tool_node.py

@@ -11,8 +11,6 @@ from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
 from core.tools.errors import ToolInvokeError
 from core.tools.tool_engine import ToolEngine
 from core.tools.utils.message_transformer import ToolFileMessageTransformer
-from core.variables.segments import ArrayAnySegment, ArrayFileSegment
-from core.variables.variables import ArrayAnyVariable
 from core.workflow.enums import (
     NodeType,
     SystemVariableKey,
@@ -23,6 +21,8 @@ from core.workflow.file import File, FileTransferMethod
 from core.workflow.node_events import NodeEventBase, NodeRunResult, StreamChunkEvent, StreamCompletedEvent
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.base.variable_template_parser import VariableTemplateParser
+from core.workflow.variables.segments import ArrayAnySegment, ArrayFileSegment
+from core.workflow.variables.variables import ArrayAnyVariable
 from extensions.ext_database import db
 from factories import file_factory
 from models import ToolFile

api/core/workflow/nodes/trigger_webhook/node.py

@@ -2,14 +2,14 @@ import logging
 from collections.abc import Mapping
 from typing import Any
 
-from core.variables.types import SegmentType
-from core.variables.variables import FileVariable
 from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionStatus
 from core.workflow.enums import NodeExecutionType, NodeType
 from core.workflow.file import FileTransferMethod
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
+from core.workflow.variables.types import SegmentType
+from core.workflow.variables.variables import FileVariable
 from factories import file_factory
 from factories.variable_factory import build_segment_with_type
 

api/core/workflow/nodes/variable_aggregator/entities.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel
 
-from core.variables.types import SegmentType
 from core.workflow.nodes.base import BaseNodeData
+from core.workflow.variables.types import SegmentType
 
 
 class AdvancedSettings(BaseModel):

api/core/workflow/nodes/variable_aggregator/variable_aggregator_node.py

@@ -1,10 +1,10 @@
 from collections.abc import Mapping
 
-from core.variables.segments import Segment
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.variable_aggregator.entities import VariableAggregatorNodeData
+from core.workflow.variables.segments import Segment
 
 
 class VariableAggregatorNode(Node[VariableAggregatorNodeData]):

api/core/workflow/nodes/variable_assigner/common/helpers.py

@@ -3,9 +3,9 @@ from typing import Any, TypeVar
 
 from pydantic import BaseModel
 
-from core.variables import Segment
-from core.variables.consts import SELECTORS_LENGTH
-from core.variables.types import SegmentType
+from core.workflow.variables import Segment
+from core.workflow.variables.consts import SELECTORS_LENGTH
+from core.workflow.variables.types import SegmentType
 
 # Use double underscore (`__`) prefix for internal variables
 # to minimize risk of collision with user-defined variable names.

api/core/workflow/nodes/variable_assigner/v1/node.py

@@ -1,7 +1,6 @@
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
-from core.variables import SegmentType, VariableBase
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.entities import GraphInitParams
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
@@ -9,6 +8,7 @@ from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.variable_assigner.common import helpers as common_helpers
 from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError
+from core.workflow.variables import SegmentType, VariableBase
 
 from .node_data import VariableAssignerData, WriteMode
 

api/core/workflow/nodes/variable_assigner/v2/helpers.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from core.variables import SegmentType
+from core.workflow.variables import SegmentType
 
 from .enums import Operation
 

api/core/workflow/nodes/variable_assigner/v2/node.py

@@ -2,14 +2,14 @@ import json
 from collections.abc import Mapping, MutableMapping, Sequence
 from typing import TYPE_CHECKING, Any
 
-from core.variables import SegmentType, VariableBase
-from core.variables.consts import SELECTORS_LENGTH
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
 from core.workflow.node_events import NodeRunResult
 from core.workflow.nodes.base.node import Node
 from core.workflow.nodes.variable_assigner.common import helpers as common_helpers
 from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError
+from core.workflow.variables import SegmentType, VariableBase
+from core.workflow.variables.consts import SELECTORS_LENGTH
 
 from . import helpers
 from .entities import VariableAssignerNodeData, VariableOperationItem

api/core/workflow/runtime/graph_runtime_state_protocol.py

@@ -2,8 +2,8 @@ from collections.abc import Mapping, Sequence
 from typing import Any, Protocol
 
 from core.model_runtime.entities.llm_entities import LLMUsage
-from core.variables.segments import Segment
 from core.workflow.system_variable import SystemVariableReadOnlyView
+from core.workflow.variables.segments import Segment
 
 
 class ReadOnlyVariablePool(Protocol):

api/core/workflow/runtime/read_only_wrappers.py

@@ -5,8 +5,8 @@ from copy import deepcopy
 from typing import Any
 
 from core.model_runtime.entities.llm_entities import LLMUsage
-from core.variables.segments import Segment
 from core.workflow.system_variable import SystemVariableReadOnlyView
+from core.workflow.variables.segments import Segment
 
 from .graph_runtime_state import GraphRuntimeState
 from .variable_pool import VariablePool

api/core/workflow/runtime/variable_pool.py

@@ -8,10 +8,6 @@ from typing import Annotated, Any, Union, cast
 
 from pydantic import BaseModel, Field
 
-from core.variables import Segment, SegmentGroup, VariableBase
-from core.variables.consts import SELECTORS_LENGTH
-from core.variables.segments import FileSegment, ObjectSegment
-from core.variables.variables import RAGPipelineVariableInput, Variable
 from core.workflow.constants import (
     CONVERSATION_VARIABLE_NODE_ID,
     ENVIRONMENT_VARIABLE_NODE_ID,
@@ -20,6 +16,10 @@ from core.workflow.constants import (
 )
 from core.workflow.file import File, FileAttribute, file_manager
 from core.workflow.system_variable import SystemVariable
+from core.workflow.variables import Segment, SegmentGroup, VariableBase
+from core.workflow.variables.consts import SELECTORS_LENGTH
+from core.workflow.variables.segments import FileSegment, ObjectSegment
+from core.workflow.variables.variables import RAGPipelineVariableInput, Variable
 from factories import variable_factory
 
 VariableValue = Union[str, int, float, dict[str, object], list[object], File]

api/core/workflow/utils/condition/processor.py

@@ -2,10 +2,10 @@ import json
 from collections.abc import Mapping, Sequence
 from typing import Literal, NamedTuple
 
-from core.variables import ArrayFileSegment
-from core.variables.segments import ArrayBooleanSegment, BooleanSegment
 from core.workflow.file import FileAttribute, file_manager
 from core.workflow.runtime import VariablePool
+from core.workflow.variables import ArrayFileSegment
+from core.workflow.variables.segments import ArrayBooleanSegment, BooleanSegment
 
 from .entities import Condition, SubCondition, SupportedComparisonOperator
 

api/core/workflow/variable_loader.py

@@ -2,9 +2,9 @@ import abc
 from collections.abc import Mapping, Sequence
 from typing import Any, Protocol
 
-from core.variables import VariableBase
-from core.variables.consts import SELECTORS_LENGTH
 from core.workflow.runtime import VariablePool
+from core.workflow.variables import VariableBase
+from core.workflow.variables.consts import SELECTORS_LENGTH
 
 
 class VariableLoader(Protocol):

api/core/workflow/variables/variables.py

@@ -4,8 +4,6 @@ from uuid import uuid4
 
 from pydantic import BaseModel, Discriminator, Field, Tag
 
-from core.helper import encrypter
-
 from .segments import (
     ArrayAnySegment,
     ArrayBooleanSegment,
@@ -27,6 +25,14 @@ from .segments import (
 from .types import SegmentType
 
 
+def _obfuscated_token(token: str) -> str:
+    if not token:
+        return token
+    if len(token) <= 8:
+        return "*" * 20
+    return token[:6] + "*" * 12 + token[-2:]
+
+
 class VariableBase(Segment):
     """
     A variable is a segment that has a name.
@@ -86,7 +92,7 @@ class SecretVariable(StringVariable):
 
     @property
     def log(self) -> str:
-        return encrypter.obfuscated_token(self.value)
+        return _obfuscated_token(self.value)
 
 
 class NoneVariable(NoneSegment, VariableBase):

api/core/workflow/workflow_type_encoder.py

@@ -4,8 +4,8 @@ from typing import Any, overload
 
 from pydantic import BaseModel
 
-from core.variables import Segment
 from core.workflow.file.models import File
+from core.workflow.variables import Segment
 
 
 class WorkflowRuntimeTypeConverter:

api/extensions/otel/parser/base.py

@@ -9,11 +9,11 @@ from opentelemetry.trace import Span
 from opentelemetry.trace.status import Status, StatusCode
 from pydantic import BaseModel
 
-from core.variables import Segment
 from core.workflow.enums import NodeType
 from core.workflow.file.models import File
 from core.workflow.graph_events import GraphNodeEventBase
 from core.workflow.nodes.base.node import Node
+from core.workflow.variables import Segment
 from extensions.otel.semconv.gen_ai import ChainAttributes, GenAIAttributes
 
 

api/extensions/otel/parser/retrieval.py

@@ -8,9 +8,9 @@ from typing import Any
 
 from opentelemetry.trace import Span
 
-from core.variables import Segment
 from core.workflow.graph_events import GraphNodeEventBase
 from core.workflow.nodes.base.node import Node
+from core.workflow.variables import Segment
 from extensions.otel.parser.base import DefaultNodeOTelParser, safe_json_dumps
 from extensions.otel.semconv.gen_ai import RetrieverAttributes
 

api/factories/variable_factory.py

@@ -3,8 +3,13 @@ from typing import Any, cast
 from uuid import uuid4
 
 from configs import dify_config
-from core.variables.exc import VariableError
-from core.variables.segments import (
+from core.workflow.constants import (
+    CONVERSATION_VARIABLE_NODE_ID,
+    ENVIRONMENT_VARIABLE_NODE_ID,
+)
+from core.workflow.file import File
+from core.workflow.variables.exc import VariableError
+from core.workflow.variables.segments import (
     ArrayAnySegment,
     ArrayBooleanSegment,
     ArrayFileSegment,
@@ -21,8 +26,8 @@ from core.variables.segments import (
     Segment,
     StringSegment,
 )
-from core.variables.types import SegmentType
-from core.variables.variables import (
+from core.workflow.variables.types import SegmentType
+from core.workflow.variables.variables import (
     ArrayAnyVariable,
     ArrayBooleanVariable,
     ArrayFileVariable,
@@ -39,11 +44,6 @@ from core.variables.variables import (
     StringVariable,
     VariableBase,
 )
-from core.workflow.constants import (
-    CONVERSATION_VARIABLE_NODE_ID,
-    ENVIRONMENT_VARIABLE_NODE_ID,
-)
-from core.workflow.file import File
 
 
 class UnsupportedSegmentTypeError(Exception):

api/fields/_value_type_serializer.py

@@ -1,7 +1,7 @@
 from typing import TypedDict
 
-from core.variables.segments import Segment
-from core.variables.types import SegmentType
+from core.workflow.variables.segments import Segment
+from core.workflow.variables.types import SegmentType
 
 
 class _VarTypedDict(TypedDict, total=False):

api/fields/workflow_fields.py

@@ -1,7 +1,7 @@
 from flask_restx import fields
 
 from core.helper import encrypter
-from core.variables import SecretVariable, SegmentType, VariableBase
+from core.workflow.variables import SecretVariable, SegmentType, VariableBase
 from fields.member_fields import simple_account_fields
 from libs.helper import TimestampField
 

api/models/workflow.py

@@ -22,8 +22,6 @@ from sqlalchemy import (
 from sqlalchemy.orm import Mapped, declared_attr, mapped_column
 from typing_extensions import deprecated
 
-from core.variables import utils as variable_utils
-from core.variables.variables import FloatVariable, IntegerVariable, StringVariable
 from core.workflow.constants import (
     CONVERSATION_VARIABLE_NODE_ID,
     SYSTEM_VARIABLE_NODE_ID,
@@ -33,6 +31,8 @@ from core.workflow.entities.pause_reason import HumanInputRequired, PauseReason,
 from core.workflow.enums import NodeType, WorkflowExecutionStatus
 from core.workflow.file.constants import maybe_file_object
 from core.workflow.file.models import File
+from core.workflow.variables import utils as variable_utils
+from core.workflow.variables.variables import FloatVariable, IntegerVariable, StringVariable
 from extensions.ext_storage import Storage
 from factories.variable_factory import TypeMismatchError, build_segment_with_type
 from libs.datetime_utils import naive_utc_now
@@ -46,7 +46,7 @@ if TYPE_CHECKING:
 
 from constants import DEFAULT_FILE_NUMBER_LIMITS, HIDDEN_VALUE
 from core.helper import encrypter
-from core.variables import SecretVariable, Segment, SegmentType, VariableBase
+from core.workflow.variables import SecretVariable, Segment, SegmentType, VariableBase
 from factories import variable_factory
 from libs import helper
 

api/services/conversation_service.py

@@ -10,7 +10,7 @@ from configs import dify_config
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.db.session_factory import session_factory
 from core.llm_generator.llm_generator import LLMGenerator
-from core.variables.types import SegmentType
+from core.workflow.variables.types import SegmentType
 from extensions.ext_database import db
 from factories import variable_factory
 from libs.datetime_utils import naive_utc_now

api/services/conversation_variable_updater.py

@@ -1,7 +1,7 @@
 from sqlalchemy import select
 from sqlalchemy.orm import Session, sessionmaker
 
-from core.variables.variables import VariableBase
+from core.workflow.variables.variables import VariableBase
 from models import ConversationVariable
 
 

api/services/rag_pipeline/rag_pipeline.py

@@ -36,7 +36,6 @@ from core.rag.entities.event import (
 )
 from core.repositories.factory import DifyCoreRepositoryFactory
 from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository
-from core.variables.variables import VariableBase
 from core.workflow.entities.workflow_node_execution import (
     WorkflowNodeExecution,
     WorkflowNodeExecutionStatus,
@@ -52,6 +51,7 @@ from core.workflow.nodes.node_mapping import LATEST_VERSION, NODE_TYPE_CLASSES_M
 from core.workflow.repositories.workflow_node_execution_repository import OrderConfig
 from core.workflow.runtime import VariablePool
 from core.workflow.system_variable import SystemVariable
+from core.workflow.variables.variables import VariableBase
 from core.workflow.workflow_entry import WorkflowEntry
 from extensions.ext_database import db
 from libs.infinite_scroll_pagination import InfiniteScrollPagination

api/services/trigger/webhook_service.py

@@ -16,9 +16,9 @@ from werkzeug.exceptions import RequestEntityTooLarge
 from configs import dify_config
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.tools.tool_file_manager import ToolFileManager
-from core.variables.types import SegmentType
 from core.workflow.enums import NodeType
 from core.workflow.file.models import FileTransferMethod
+from core.workflow.variables.types import SegmentType
 from enums.quota_type import QuotaType
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client

api/services/variable_truncator.py

@@ -6,7 +6,9 @@ from collections.abc import Mapping
 from typing import Any, Generic, TypeAlias, TypeVar, overload
 
 from configs import dify_config
-from core.variables.segments import (
+from core.workflow.file.models import File
+from core.workflow.nodes.variable_assigner.common.helpers import UpdatedVariable
+from core.workflow.variables.segments import (
     ArrayFileSegment,
     ArraySegment,
     BooleanSegment,
@@ -18,9 +20,7 @@ from core.variables.segments import (
     Segment,
     StringSegment,
 )
-from core.variables.utils import dumps_with_segments
-from core.workflow.file.models import File
-from core.workflow.nodes.variable_assigner.common.helpers import UpdatedVariable
+from core.workflow.variables.utils import dumps_with_segments
 
 _MAX_DEPTH = 100
 

api/services/workflow_draft_variable_service.py

@@ -14,20 +14,20 @@ from sqlalchemy.sql.expression import and_, or_
 
 from configs import dify_config
 from core.app.entities.app_invoke_entities import InvokeFrom
-from core.variables import Segment, StringSegment, VariableBase
-from core.variables.consts import SELECTORS_LENGTH
-from core.variables.segments import (
-    ArrayFileSegment,
-    FileSegment,
-)
-from core.variables.types import SegmentType
-from core.variables.utils import dumps_with_segments
 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.file.models import File
 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 core.workflow.variables import Segment, StringSegment, VariableBase
+from core.workflow.variables.consts import SELECTORS_LENGTH
+from core.workflow.variables.segments import (
+    ArrayFileSegment,
+    FileSegment,
+)
+from core.workflow.variables.types import SegmentType
+from core.workflow.variables.utils import dumps_with_segments
 from extensions.ext_storage import storage
 from factories.file_factory import StorageKeyLoader
 from factories.variable_factory import build_segment, segment_to_variable

api/services/workflow_service.py

@@ -15,8 +15,6 @@ from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.repositories import DifyCoreRepositoryFactory
 from core.repositories.human_input_repository import HumanInputFormRepositoryImpl
-from core.variables import VariableBase
-from core.variables.variables import Variable
 from core.workflow.entities import GraphInitParams, WorkflowNodeExecution
 from core.workflow.entities.pause_reason import HumanInputRequired
 from core.workflow.enums import ErrorStrategy, WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
@@ -41,6 +39,8 @@ from core.workflow.repositories.human_input_form_repository import FormCreatePar
 from core.workflow.runtime import GraphRuntimeState, VariablePool
 from core.workflow.system_variable import SystemVariable
 from core.workflow.variable_loader import load_into_variable_pool
+from core.workflow.variables import VariableBase
+from core.workflow.variables.variables import Variable
 from core.workflow.workflow_entry import WorkflowEntry
 from enums.cloud_plan import CloudPlan
 from events.app_event import app_draft_workflow_was_synced, app_published_workflow_was_updated

api/tests/integration_tests/libs/test_api_token_cache_integration.py

@@ -360,7 +360,7 @@ class TestEndToEndCacheFlow:
 class TestRedisFailover:
     """Test behavior when Redis is unavailable."""
 
-    @patch("services.api_token_service.redis_client")
+    @patch("services.api_token_service.redis_client", autospec=True)
     def test_graceful_degradation_when_redis_fails(self, mock_redis):
         """Test system degrades gracefully when Redis is unavailable."""
         from redis import RedisError

api/tests/integration_tests/services/test_workflow_draft_variable_service.py

@@ -6,11 +6,11 @@ import pytest
 from sqlalchemy import delete
 from sqlalchemy.orm import Session
 
-from core.variables.segments import StringSegment
-from core.variables.types import SegmentType
-from core.variables.variables import StringVariable
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
 from core.workflow.nodes import NodeType
+from core.workflow.variables.segments import StringSegment
+from core.workflow.variables.types import SegmentType
+from core.workflow.variables.variables import StringVariable
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from factories.variable_factory import build_segment

api/tests/integration_tests/tasks/test_remove_app_and_related_data_task.py

@@ -5,7 +5,7 @@ import pytest
 from sqlalchemy import delete
 
 from core.db.session_factory import session_factory
-from core.variables.segments import StringSegment
+from core.workflow.variables.segments import StringSegment
 from models import Tenant
 from models.enums import CreatorUserRole
 from models.model import App, UploadFile
@@ -191,7 +191,7 @@ class TestDeleteDraftVariablesWithOffloadIntegration:
     @pytest.fixture
     def setup_offload_test_data(self, app_and_tenant):
         tenant, app = app_and_tenant
-        from core.variables.types import SegmentType
+        from core.workflow.variables.types import SegmentType
         from libs.datetime_utils import naive_utc_now
 
         with session_factory.create_session() as session:
@@ -422,7 +422,7 @@ class TestDeleteDraftVariablesSessionCommit:
     @pytest.fixture
     def setup_offload_test_data(self, app_and_tenant):
         """Create test data with offload files for session commit tests."""
-        from core.variables.types import SegmentType
+        from core.workflow.variables.types import SegmentType
         from libs.datetime_utils import naive_utc_now
 
         tenant, app = app_and_tenant

api/tests/integration_tests/vdb/opensearch/test_opensearch.py

@@ -41,17 +41,15 @@ class TestOpenSearchConfig:
         assert params["connection_class"].__name__ == "Urllib3HttpConnection"
         assert params["http_auth"] == ("admin", "password")
 
-    @patch("boto3.Session")
-    @patch("core.rag.datasource.vdb.opensearch.opensearch_vector.Urllib3AWSV4SignerAuth")
+    @patch("boto3.Session", autospec=True)
+    @patch("core.rag.datasource.vdb.opensearch.opensearch_vector.Urllib3AWSV4SignerAuth", autospec=True)
     def test_to_opensearch_params_with_aws_managed_iam(
         self, mock_aws_signer_auth: MagicMock, mock_boto_session: MagicMock
     ):
         mock_credentials = MagicMock()
         mock_boto_session.return_value.get_credentials.return_value = mock_credentials
 
-        mock_auth_instance = MagicMock()
-        mock_aws_signer_auth.return_value = mock_auth_instance
-
+        mock_auth_instance = mock_aws_signer_auth.return_value
         aws_region = "ap-southeast-2"
         aws_service = "aoss"
         host = f"aoss-endpoint.{aws_region}.aoss.amazonaws.com"
@@ -157,7 +155,7 @@ class TestOpenSearchVector:
         doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
         embedding = [0.1] * 128
 
-        with patch("opensearchpy.helpers.bulk") as mock_bulk:
+        with patch("opensearchpy.helpers.bulk", autospec=True) as mock_bulk:
             mock_bulk.return_value = ([], [])
             self.vector.add_texts([doc], [embedding])
 
@@ -171,7 +169,7 @@ class TestOpenSearchVector:
         doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
         embedding = [0.1] * 128
 
-        with patch("opensearchpy.helpers.bulk") as mock_bulk:
+        with patch("opensearchpy.helpers.bulk", autospec=True) as mock_bulk:
             mock_bulk.return_value = ([], [])
             self.vector.add_texts([doc], [embedding])
 

api/tests/integration_tests/workflow/nodes/__mock/model.py

@@ -48,3 +48,19 @@ def get_mocked_fetch_model_config(
     )
 
     return MagicMock(return_value=(model_instance, model_config))
+
+
+def get_mocked_fetch_model_instance(
+    provider: str,
+    model: str,
+    mode: str,
+    credentials: dict,
+):
+    mock_fetch_model_config = get_mocked_fetch_model_config(
+        provider=provider,
+        model=model,
+        mode=mode,
+        credentials=credentials,
+    )
+    model_instance, _ = mock_fetch_model_config()
+    return MagicMock(return_value=model_instance)

api/tests/integration_tests/workflow/nodes/test_llm.py

@@ -5,13 +5,13 @@ from collections.abc import Generator
 from unittest.mock import MagicMock, patch
 
 from core.app.entities.app_invoke_entities import InvokeFrom
-from core.app.workflow.node_factory import DifyNodeFactory
 from core.llm_generator.output_parser.structured_output import _parse_structured_output
+from core.model_manager import ModelInstance
 from core.workflow.entities import GraphInitParams
 from core.workflow.enums import WorkflowNodeExecutionStatus
-from core.workflow.graph import Graph
 from core.workflow.node_events import StreamCompletedEvent
 from core.workflow.nodes.llm.node import LLMNode
+from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory
 from core.workflow.runtime import GraphRuntimeState, VariablePool
 from core.workflow.system_variable import SystemVariable
 from extensions.ext_database import db
@@ -67,21 +67,14 @@ def init_llm_node(config: dict) -> LLMNode:
 
     graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
 
-    # Create node factory
-    node_factory = DifyNodeFactory(
-        graph_init_params=init_params,
-        graph_runtime_state=graph_runtime_state,
-    )
-
-    graph = Graph.init(graph_config=graph_config, node_factory=node_factory)
-
     node = LLMNode(
         id=str(uuid.uuid4()),
         config=config,
         graph_init_params=init_params,
         graph_runtime_state=graph_runtime_state,
-        credentials_provider=MagicMock(),
-        model_factory=MagicMock(),
+        credentials_provider=MagicMock(spec=CredentialsProvider),
+        model_factory=MagicMock(spec=ModelFactory),
+        model_instance=MagicMock(spec=ModelInstance),
     )
 
     return node
@@ -116,8 +109,7 @@ def test_execute_llm():
 
     db.session.close = MagicMock()
 
-    # Mock the _fetch_model_config to avoid database calls
-    def mock_fetch_model_config(*_args, **_kwargs):
+    def build_mock_model_instance() -> MagicMock:
         from decimal import Decimal
         from unittest.mock import MagicMock
 
@@ -125,7 +117,20 @@ def test_execute_llm():
         from core.model_runtime.entities.message_entities import AssistantPromptMessage
 
         # Create mock model instance
-        mock_model_instance = MagicMock()
+        mock_model_instance = MagicMock(spec=ModelInstance)
+        mock_model_instance.provider = "openai"
+        mock_model_instance.model_name = "gpt-3.5-turbo"
+        mock_model_instance.credentials = {}
+        mock_model_instance.parameters = {}
+        mock_model_instance.stop = []
+        mock_model_instance.model_type_instance = MagicMock()
+        mock_model_instance.model_type_instance.get_model_schema.return_value = MagicMock(
+            model_properties={},
+            parameter_rules=[],
+            features=[],
+        )
+        mock_model_instance.provider_model_bundle = MagicMock()
+        mock_model_instance.provider_model_bundle.configuration.using_provider_type = "custom"
         mock_usage = LLMUsage(
             prompt_tokens=30,
             prompt_unit_price=Decimal("0.001"),
@@ -149,14 +154,7 @@ def test_execute_llm():
         )
         mock_model_instance.invoke_llm.return_value = mock_llm_result
 
-        # Create mock model config
-        mock_model_config = MagicMock()
-        mock_model_config.mode = "chat"
-        mock_model_config.provider = "openai"
-        mock_model_config.model = "gpt-3.5-turbo"
-        mock_model_config.parameters = {}
-
-        return mock_model_instance, mock_model_config
+        return mock_model_instance
 
     # Mock fetch_prompt_messages to avoid database calls
     def mock_fetch_prompt_messages_1(**_kwargs):
@@ -167,10 +165,9 @@ def test_execute_llm():
             UserPromptMessage(content="what's the weather today?"),
         ], []
 
-    with (
-        patch.object(LLMNode, "_fetch_model_config", mock_fetch_model_config),
-        patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_1),
-    ):
+    node._model_instance = build_mock_model_instance()
+
+    with patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_1):
         # execute node
         result = node._run()
         assert isinstance(result, Generator)
@@ -228,8 +225,7 @@ def test_execute_llm_with_jinja2():
     # Mock db.session.close()
     db.session.close = MagicMock()
 
-    # Mock the _fetch_model_config method
-    def mock_fetch_model_config(*_args, **_kwargs):
+    def build_mock_model_instance() -> MagicMock:
         from decimal import Decimal
         from unittest.mock import MagicMock
 
@@ -237,7 +233,20 @@ def test_execute_llm_with_jinja2():
         from core.model_runtime.entities.message_entities import AssistantPromptMessage
 
         # Create mock model instance
-        mock_model_instance = MagicMock()
+        mock_model_instance = MagicMock(spec=ModelInstance)
+        mock_model_instance.provider = "openai"
+        mock_model_instance.model_name = "gpt-3.5-turbo"
+        mock_model_instance.credentials = {}
+        mock_model_instance.parameters = {}
+        mock_model_instance.stop = []
+        mock_model_instance.model_type_instance = MagicMock()
+        mock_model_instance.model_type_instance.get_model_schema.return_value = MagicMock(
+            model_properties={},
+            parameter_rules=[],
+            features=[],
+        )
+        mock_model_instance.provider_model_bundle = MagicMock()
+        mock_model_instance.provider_model_bundle.configuration.using_provider_type = "custom"
         mock_usage = LLMUsage(
             prompt_tokens=30,
             prompt_unit_price=Decimal("0.001"),
@@ -261,14 +270,7 @@ def test_execute_llm_with_jinja2():
         )
         mock_model_instance.invoke_llm.return_value = mock_llm_result
 
-        # Create mock model config
-        mock_model_config = MagicMock()
-        mock_model_config.mode = "chat"
-        mock_model_config.provider = "openai"
-        mock_model_config.model = "gpt-3.5-turbo"
-        mock_model_config.parameters = {}
-
-        return mock_model_instance, mock_model_config
+        return mock_model_instance
 
     # Mock fetch_prompt_messages to avoid database calls
     def mock_fetch_prompt_messages_2(**_kwargs):
@@ -279,10 +281,9 @@ def test_execute_llm_with_jinja2():
             UserPromptMessage(content="what's the weather today?"),
         ], []
 
-    with (
-        patch.object(LLMNode, "_fetch_model_config", mock_fetch_model_config),
-        patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_2),
-    ):
+    node._model_instance = build_mock_model_instance()
+
+    with patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_2):
         # execute node
         result = node._run()
 

api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py

@@ -4,18 +4,17 @@ import uuid
 from unittest.mock import MagicMock
 
 from core.app.entities.app_invoke_entities import InvokeFrom
-from core.app.workflow.node_factory import DifyNodeFactory
+from core.model_manager import ModelInstance
 from core.model_runtime.entities import AssistantPromptMessage
 from core.workflow.entities import GraphInitParams
 from core.workflow.enums import WorkflowNodeExecutionStatus
-from core.workflow.graph import Graph
 from core.workflow.nodes.llm.protocols import CredentialsProvider, ModelFactory
 from core.workflow.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode
 from core.workflow.runtime import GraphRuntimeState, VariablePool
 from core.workflow.system_variable import SystemVariable
 from extensions.ext_database import db
 from models.enums import UserFrom
-from tests.integration_tests.workflow.nodes.__mock.model import get_mocked_fetch_model_config
+from tests.integration_tests.workflow.nodes.__mock.model import get_mocked_fetch_model_instance
 
 """FOR MOCK FIXTURES, DO NOT REMOVE"""
 from tests.integration_tests.model_runtime.__mock.plugin_daemon import setup_model_mock
@@ -72,14 +71,6 @@ def init_parameter_extractor_node(config: dict):
 
     graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
 
-    # Create node factory
-    node_factory = DifyNodeFactory(
-        graph_init_params=init_params,
-        graph_runtime_state=graph_runtime_state,
-    )
-
-    graph = Graph.init(graph_config=graph_config, node_factory=node_factory)
-
     node = ParameterExtractorNode(
         id=str(uuid.uuid4()),
         config=config,
@@ -87,6 +78,7 @@ def init_parameter_extractor_node(config: dict):
         graph_runtime_state=graph_runtime_state,
         credentials_provider=MagicMock(spec=CredentialsProvider),
         model_factory=MagicMock(spec=ModelFactory),
+        model_instance=MagicMock(spec=ModelInstance),
     )
     return node
 
@@ -116,12 +108,12 @@ def test_function_calling_parameter_extractor(setup_model_mock):
         }
     )
 
-    node._fetch_model_config = get_mocked_fetch_model_config(
+    node._model_instance = get_mocked_fetch_model_instance(
         provider="langgenius/openai/openai",
         model="gpt-3.5-turbo",
         mode="chat",
         credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
-    )
+    )()
     db.session.close = MagicMock()
 
     result = node._run()
@@ -157,12 +149,12 @@ def test_instructions(setup_model_mock):
         },
     )
 
-    node._fetch_model_config = get_mocked_fetch_model_config(
+    node._model_instance = get_mocked_fetch_model_instance(
         provider="langgenius/openai/openai",
         model="gpt-3.5-turbo",
         mode="chat",
         credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
-    )
+    )()
     db.session.close = MagicMock()
 
     result = node._run()
@@ -207,12 +199,12 @@ def test_chat_parameter_extractor(setup_model_mock):
         },
     )
 
-    node._fetch_model_config = get_mocked_fetch_model_config(
+    node._model_instance = get_mocked_fetch_model_instance(
         provider="langgenius/openai/openai",
         model="gpt-3.5-turbo",
         mode="chat",
         credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
-    )
+    )()
     db.session.close = MagicMock()
 
     result = node._run()
@@ -258,12 +250,12 @@ def test_completion_parameter_extractor(setup_model_mock):
         },
     )
 
-    node._fetch_model_config = get_mocked_fetch_model_config(
+    node._model_instance = get_mocked_fetch_model_instance(
         provider="langgenius/openai/openai",
         model="gpt-3.5-turbo-instruct",
         mode="completion",
         credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
-    )
+    )()
     db.session.close = MagicMock()
 
     result = node._run()
@@ -383,12 +375,12 @@ def test_chat_parameter_extractor_with_memory(setup_model_mock, monkeypatch):
         },
     )
 
-    node._fetch_model_config = get_mocked_fetch_model_config(
+    node._model_instance = get_mocked_fetch_model_instance(
         provider="langgenius/openai/openai",
         model="gpt-3.5-turbo",
         mode="chat",
         credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
-    )
+    )()
     # Test the mock before running the actual test
     monkeypatch.setattr("core.workflow.nodes.llm.llm_utils.fetch_memory", get_mocked_fetch_memory("customized memory"))
     db.session.close = MagicMock()

api/tests/test_containers_integration_tests/services/test_agent_service.py

@@ -19,14 +19,14 @@ class TestAgentService:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("services.agent_service.PluginAgentClient") as mock_plugin_agent_client,
-            patch("services.agent_service.ToolManager") as mock_tool_manager,
-            patch("services.agent_service.AgentConfigManager") as mock_agent_config_manager,
+            patch("services.agent_service.PluginAgentClient", autospec=True) as mock_plugin_agent_client,
+            patch("services.agent_service.ToolManager", autospec=True) as mock_tool_manager,
+            patch("services.agent_service.AgentConfigManager", autospec=True) as mock_agent_config_manager,
             patch("services.agent_service.current_user", create_autospec(Account, instance=True)) as mock_current_user,
-            patch("services.app_service.FeatureService") as mock_feature_service,
-            patch("services.app_service.EnterpriseService") as mock_enterprise_service,
-            patch("services.app_service.ModelManager") as mock_model_manager,
-            patch("services.account_service.FeatureService") as mock_account_feature_service,
+            patch("services.app_service.FeatureService", autospec=True) as mock_feature_service,
+            patch("services.app_service.EnterpriseService", autospec=True) as mock_enterprise_service,
+            patch("services.app_service.ModelManager", autospec=True) as mock_model_manager,
+            patch("services.account_service.FeatureService", autospec=True) as mock_account_feature_service,
         ):
             # Setup default mock returns for agent service
             mock_plugin_agent_client_instance = mock_plugin_agent_client.return_value

api/tests/test_containers_integration_tests/services/test_app_generate_service.py

@@ -18,18 +18,22 @@ class TestAppGenerateService:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("services.billing_service.BillingService") as mock_billing_service,
-            patch("services.app_generate_service.WorkflowService") as mock_workflow_service,
-            patch("services.app_generate_service.RateLimit") as mock_rate_limit,
-            patch("services.app_generate_service.CompletionAppGenerator") as mock_completion_generator,
-            patch("services.app_generate_service.ChatAppGenerator") as mock_chat_generator,
-            patch("services.app_generate_service.AgentChatAppGenerator") as mock_agent_chat_generator,
-            patch("services.app_generate_service.AdvancedChatAppGenerator") as mock_advanced_chat_generator,
-            patch("services.app_generate_service.WorkflowAppGenerator") as mock_workflow_generator,
-            patch("services.app_generate_service.MessageBasedAppGenerator") as mock_message_based_generator,
-            patch("services.account_service.FeatureService") as mock_account_feature_service,
-            patch("services.app_generate_service.dify_config") as mock_dify_config,
-            patch("configs.dify_config") as mock_global_dify_config,
+            patch("services.billing_service.BillingService", autospec=True) as mock_billing_service,
+            patch("services.app_generate_service.WorkflowService", autospec=True) as mock_workflow_service,
+            patch("services.app_generate_service.RateLimit", autospec=True) as mock_rate_limit,
+            patch("services.app_generate_service.CompletionAppGenerator", autospec=True) as mock_completion_generator,
+            patch("services.app_generate_service.ChatAppGenerator", autospec=True) as mock_chat_generator,
+            patch("services.app_generate_service.AgentChatAppGenerator", autospec=True) as mock_agent_chat_generator,
+            patch(
+                "services.app_generate_service.AdvancedChatAppGenerator", autospec=True
+            ) as mock_advanced_chat_generator,
+            patch("services.app_generate_service.WorkflowAppGenerator", autospec=True) as mock_workflow_generator,
+            patch(
+                "services.app_generate_service.MessageBasedAppGenerator", autospec=True
+            ) as mock_message_based_generator,
+            patch("services.account_service.FeatureService", autospec=True) as mock_account_feature_service,
+            patch("services.app_generate_service.dify_config", autospec=True) as mock_dify_config,
+            patch("configs.dify_config", autospec=True) as mock_global_dify_config,
         ):
             # Setup default mock returns for billing service
             mock_billing_service.update_tenant_feature_plan_usage.return_value = {
@@ -983,7 +987,7 @@ class TestAppGenerateService:
         }
 
         # Execute the method under test
-        with patch("services.app_generate_service.AppExecutionParams") as mock_exec_params:
+        with patch("services.app_generate_service.AppExecutionParams", autospec=True) as mock_exec_params:
             mock_payload = MagicMock()
             mock_payload.workflow_run_id = fake.uuid4()
             mock_payload.model_dump_json.return_value = "{}"

api/tests/test_containers_integration_tests/services/test_model_load_balancing_service.py

@@ -17,10 +17,12 @@ class TestModelLoadBalancingService:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("services.model_load_balancing_service.ProviderManager") as mock_provider_manager,
-            patch("services.model_load_balancing_service.LBModelManager") as mock_lb_model_manager,
-            patch("services.model_load_balancing_service.ModelProviderFactory") as mock_model_provider_factory,
-            patch("services.model_load_balancing_service.encrypter") as mock_encrypter,
+            patch("services.model_load_balancing_service.ProviderManager", autospec=True) as mock_provider_manager,
+            patch("services.model_load_balancing_service.LBModelManager", autospec=True) as mock_lb_model_manager,
+            patch(
+                "services.model_load_balancing_service.ModelProviderFactory", autospec=True
+            ) as mock_model_provider_factory,
+            patch("services.model_load_balancing_service.encrypter", autospec=True) as mock_encrypter,
         ):
             # Setup default mock returns
             mock_provider_manager_instance = mock_provider_manager.return_value

api/tests/test_containers_integration_tests/services/test_model_provider_service.py

@@ -17,8 +17,8 @@ class TestModelProviderService:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("services.model_provider_service.ProviderManager") as mock_provider_manager,
-            patch("services.model_provider_service.ModelProviderFactory") as mock_model_provider_factory,
+            patch("services.model_provider_service.ProviderManager", autospec=True) as mock_provider_manager,
+            patch("services.model_provider_service.ModelProviderFactory", autospec=True) as mock_model_provider_factory,
         ):
             # Setup default mock returns
             mock_provider_manager.return_value.get_configurations.return_value = MagicMock()
@@ -526,7 +526,9 @@ class TestModelProviderService:
 
         # Act: Execute the method under test
         service = ModelProviderService()
-        with patch.object(service, "get_provider_credential", return_value=expected_credentials) as mock_method:
+        with patch.object(
+            service, "get_provider_credential", return_value=expected_credentials, autospec=True
+        ) as mock_method:
             result = service.get_provider_credential(tenant.id, "openai")
 
             # Assert: Verify the expected outcomes
@@ -854,7 +856,9 @@ class TestModelProviderService:
 
         # Act: Execute the method under test
         service = ModelProviderService()
-        with patch.object(service, "get_model_credential", return_value=expected_credentials) as mock_method:
+        with patch.object(
+            service, "get_model_credential", return_value=expected_credentials, autospec=True
+        ) as mock_method:
             result = service.get_model_credential(tenant.id, "openai", "llm", "gpt-4", None)
 
             # Assert: Verify the expected outcomes

api/tests/test_containers_integration_tests/services/test_webhook_service.py

@@ -22,16 +22,13 @@ class TestWebhookService:
     def mock_external_dependencies(self):
         """Mock external service dependencies."""
         with (
-            patch("services.trigger.webhook_service.AsyncWorkflowService") as mock_async_service,
-            patch("services.trigger.webhook_service.ToolFileManager") as mock_tool_file_manager,
-            patch("services.trigger.webhook_service.file_factory") as mock_file_factory,
-            patch("services.account_service.FeatureService") as mock_feature_service,
+            patch("services.trigger.webhook_service.AsyncWorkflowService", autospec=True) as mock_async_service,
+            patch("services.trigger.webhook_service.ToolFileManager", autospec=True) as mock_tool_file_manager,
+            patch("services.trigger.webhook_service.file_factory", autospec=True) as mock_file_factory,
+            patch("services.account_service.FeatureService", autospec=True) as mock_feature_service,
         ):
             # Mock ToolFileManager
-            mock_tool_file_instance = MagicMock()
-            mock_tool_file_manager.return_value = mock_tool_file_instance
-
-            # Mock file creation
+            mock_tool_file_instance = mock_tool_file_manager.return_value  # Mock file creation
             mock_tool_file = MagicMock()
             mock_tool_file.id = "test_file_id"
             mock_tool_file_instance.create_file_by_raw.return_value = mock_tool_file
@@ -435,12 +432,12 @@ class TestWebhookService:
 
         with flask_app_with_containers.app_context():
             # Mock tenant owner lookup to return the test account
-            with patch("services.trigger.webhook_service.select") as mock_select:
+            with patch("services.trigger.webhook_service.select", autospec=True) as mock_select:
                 mock_query = MagicMock()
                 mock_select.return_value.join.return_value.where.return_value = mock_query
 
                 # Mock the session to return our test account
-                with patch("services.trigger.webhook_service.Session") as mock_session:
+                with patch("services.trigger.webhook_service.Session", autospec=True) as mock_session:
                     mock_session_instance = MagicMock()
                     mock_session.return_value.__enter__.return_value = mock_session_instance
                     mock_session_instance.scalar.return_value = test_data["account"]
@@ -462,7 +459,7 @@ class TestWebhookService:
         with flask_app_with_containers.app_context():
             # Mock EndUserService to raise an exception
             with patch(
-                "services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type"
+                "services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type", autospec=True
             ) as mock_end_user:
                 mock_end_user.side_effect = ValueError("Failed to create end user")
 

api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py

@@ -1,8 +1,8 @@
 import pytest
 from faker import Faker
 
-from core.variables.segments import StringSegment
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID
+from core.workflow.variables.segments import StringSegment
 from models import App, Workflow
 from models.enums import DraftVariableType
 from models.workflow import WorkflowDraftVariable
@@ -467,7 +467,7 @@ class TestWorkflowDraftVariableService:
         fake = Faker()
         app = self._create_test_app(db_session_with_containers, mock_external_service_dependencies, fake=fake)
         workflow = self._create_test_workflow(db_session_with_containers, app, fake=fake)
-        from core.variables.variables import StringVariable
+        from core.workflow.variables.variables import StringVariable
 
         conv_var = StringVariable(
             id=fake.uuid4(),
@@ -650,7 +650,7 @@ class TestWorkflowDraftVariableService:
         fake = Faker()
         app = self._create_test_app(db_session_with_containers, mock_external_service_dependencies, fake=fake)
         workflow = self._create_test_workflow(db_session_with_containers, app, fake=fake)
-        from core.variables.variables import StringVariable
+        from core.workflow.variables.variables import StringVariable
 
         conv_var1 = StringVariable(
             id=fake.uuid4(),

api/tests/test_containers_integration_tests/services/test_workflow_service.py

@@ -764,7 +764,7 @@ class TestWorkflowService:
         # Act - Mock current_user context and pass session
         from unittest.mock import patch
 
-        with patch("flask_login.utils._get_user", return_value=account):
+        with patch("flask_login.utils._get_user", return_value=account, autospec=True):
             result = workflow_service.publish_workflow(
                 session=db_session_with_containers, app_model=app, account=account
             )
@@ -1391,10 +1391,21 @@ class TestWorkflowService:
 
         workflow_service = WorkflowService()
 
+        from unittest.mock import patch
+
+        from core.app.workflow.node_factory import DifyNodeFactory
+        from core.model_manager import ModelInstance
+
         # Act
-        result = workflow_service.run_free_workflow_node(
-            node_data=node_data, tenant_id=tenant_id, user_id=user_id, node_id=node_id, user_inputs=user_inputs
-        )
+        with patch.object(
+            DifyNodeFactory,
+            "_build_model_instance_for_llm_node",
+            return_value=MagicMock(spec=ModelInstance),
+            autospec=True,
+        ):
+            result = workflow_service.run_free_workflow_node(
+                node_data=node_data, tenant_id=tenant_id, user_id=user_id, node_id=node_id, user_inputs=user_inputs
+            )
 
         # Assert
         assert result is not None

api/tests/test_containers_integration_tests/tasks/test_add_document_to_index_task.py

@@ -18,7 +18,9 @@ class TestAddDocumentToIndexTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.add_document_to_index_task.IndexProcessorFactory") as mock_index_processor_factory,
+            patch(
+                "tasks.add_document_to_index_task.IndexProcessorFactory", autospec=True
+            ) as mock_index_processor_factory,
         ):
             # Setup mock index processor
             mock_processor = MagicMock()
@@ -378,7 +380,7 @@ class TestAddDocumentToIndexTask:
         redis_client.set(indexing_cache_key, "processing", ex=300)
 
         # Mock the get_child_chunks method for each segment
-        with patch.object(DocumentSegment, "get_child_chunks") as mock_get_child_chunks:
+        with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks:
             # Setup mock to return child chunks for each segment
             mock_child_chunks = []
             for i in range(2):  # Each segment has 2 child chunks

api/tests/test_containers_integration_tests/tasks/test_batch_create_segment_to_index_task.py

@@ -51,9 +51,9 @@ class TestBatchCreateSegmentToIndexTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.batch_create_segment_to_index_task.storage") as mock_storage,
-            patch("tasks.batch_create_segment_to_index_task.ModelManager") as mock_model_manager,
-            patch("tasks.batch_create_segment_to_index_task.VectorService") as mock_vector_service,
+            patch("tasks.batch_create_segment_to_index_task.storage", autospec=True) as mock_storage,
+            patch("tasks.batch_create_segment_to_index_task.ModelManager", autospec=True) as mock_model_manager,
+            patch("tasks.batch_create_segment_to_index_task.VectorService", autospec=True) as mock_vector_service,
         ):
             # Setup default mock returns
             mock_storage.download.return_value = None

api/tests/test_containers_integration_tests/tasks/test_clean_dataset_task.py

@@ -63,8 +63,8 @@ class TestCleanDatasetTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.clean_dataset_task.storage") as mock_storage,
-            patch("tasks.clean_dataset_task.IndexProcessorFactory") as mock_index_processor_factory,
+            patch("tasks.clean_dataset_task.storage", autospec=True) as mock_storage,
+            patch("tasks.clean_dataset_task.IndexProcessorFactory", autospec=True) as mock_index_processor_factory,
         ):
             # Setup default mock returns
             mock_storage.delete.return_value = None
@@ -597,7 +597,7 @@ class TestCleanDatasetTask:
         db_session_with_containers.commit()
 
         # Mock the get_image_upload_file_ids function to return our image file IDs
-        with patch("tasks.clean_dataset_task.get_image_upload_file_ids") as mock_get_image_ids:
+        with patch("tasks.clean_dataset_task.get_image_upload_file_ids", autospec=True) as mock_get_image_ids:
             mock_get_image_ids.return_value = [f.id for f in image_files]
 
             # Execute the task

api/tests/test_containers_integration_tests/tasks/test_create_segment_to_index_task.py

@@ -41,7 +41,7 @@ class TestCreateSegmentToIndexTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.create_segment_to_index_task.IndexProcessorFactory") as mock_factory,
+            patch("tasks.create_segment_to_index_task.IndexProcessorFactory", autospec=True) as mock_factory,
         ):
             # Setup default mock returns
             mock_processor = MagicMock()
@@ -708,7 +708,7 @@ class TestCreateSegmentToIndexTask:
         redis_client.set(cache_key, "processing", ex=300)
 
         # Mock Redis to raise exception in finally block
-        with patch.object(redis_client, "delete", side_effect=Exception("Redis connection failed")):
+        with patch.object(redis_client, "delete", side_effect=Exception("Redis connection failed"), autospec=True):
             # Act: Execute the task - Redis failure should not prevent completion
             with pytest.raises(Exception) as exc_info:
                 create_segment_to_index_task(segment.id)

api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py

@@ -37,7 +37,7 @@ class _TrackedSessionContext:
             self._closed_sessions.append(self._session)
             return original_close(*args, **kwargs)
 
-        self._close_patcher = patch.object(self._session, "close", side_effect=_tracked_close)
+        self._close_patcher = patch.object(self._session, "close", side_effect=_tracked_close, autospec=True)
         self._close_patcher.start()
         return self._session
 
@@ -69,7 +69,9 @@ def session_close_tracker():
         original_context_manager = original_create_session(*args, **kwargs)
         return _TrackedSessionContext(original_context_manager, opened_sessions, closed_sessions)
 
-    with patch.object(task_module.session_factory, "create_session", side_effect=_tracked_create_session):
+    with patch.object(
+        task_module.session_factory, "create_session", side_effect=_tracked_create_session, autospec=True
+    ):
         yield {"opened_sessions": opened_sessions, "closed_sessions": closed_sessions}
 
 
@@ -77,13 +79,11 @@ def session_close_tracker():
 def patched_external_dependencies():
     """Patch non-DB collaborators while keeping database behavior real."""
     with (
-        patch("tasks.document_indexing_task.IndexingRunner") as mock_indexing_runner,
-        patch("tasks.document_indexing_task.FeatureService") as mock_feature_service,
-        patch("tasks.document_indexing_task.generate_summary_index_task") as mock_summary_task,
+        patch("tasks.document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner,
+        patch("tasks.document_indexing_task.FeatureService", autospec=True) as mock_feature_service,
+        patch("tasks.document_indexing_task.generate_summary_index_task", autospec=True) as mock_summary_task,
     ):
-        mock_runner_instance = MagicMock()
-        mock_indexing_runner.return_value = mock_runner_instance
-
+        mock_runner_instance = mock_indexing_runner.return_value
         mock_features = MagicMock()
         mock_features.billing.enabled = False
         mock_features.billing.subscription.plan = CloudPlan.PROFESSIONAL
@@ -307,9 +307,17 @@ class TestDatasetIndexingTaskIntegration:
 
         # Act
         with (
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[next_task]),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time") as set_waiting_spy,
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
+                return_value=[next_task],
+                autospec=True,
+            ),
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True
+            ) as set_waiting_spy,
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
+            ) as delete_key_spy,
         ):
             _document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
 
@@ -336,8 +344,10 @@ class TestDatasetIndexingTaskIntegration:
 
         # Act
         with (
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[]),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
+            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[], autospec=True),
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
+            ) as delete_key_spy,
         ):
             _document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
 
@@ -426,9 +436,13 @@ class TestDatasetIndexingTaskIntegration:
 
         # Act
         with (
-            patch("tasks.document_indexing_task._document_indexing", side_effect=Exception("failed")),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[next_task]),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time"),
+            patch("tasks.document_indexing_task._document_indexing", side_effect=Exception("failed"), autospec=True),
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
+                return_value=[next_task],
+                autospec=True,
+            ),
+            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True),
         ):
             _document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
 
@@ -511,8 +525,11 @@ class TestDatasetIndexingTaskIntegration:
             patch(
                 "tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
                 return_value=pending_tasks[:concurrency_limit],
+                autospec=True,
             ),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time") as set_waiting_spy,
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True
+            ) as set_waiting_spy,
         ):
             _document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
 
@@ -538,8 +555,12 @@ class TestDatasetIndexingTaskIntegration:
         # Act
         with (
             patch("tasks.document_indexing_task.dify_config.TENANT_ISOLATED_TASK_CONCURRENCY", 3),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=ordered_tasks),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time"),
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks",
+                return_value=ordered_tasks,
+                autospec=True,
+            ),
+            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.set_task_waiting_time", autospec=True),
         ):
             _document_indexing_with_tenant_queue(dataset.tenant_id, dataset.id, document_ids, task_dispatch_spy)
 
@@ -578,8 +599,10 @@ class TestDatasetIndexingTaskIntegration:
 
         # Act
         with (
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[]),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
+            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[], autospec=True),
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
+            ) as delete_key_spy,
         ):
             normal_document_indexing_task(dataset.tenant_id, dataset.id, document_ids)
 
@@ -599,8 +622,10 @@ class TestDatasetIndexingTaskIntegration:
 
         # Act
         with (
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[]),
-            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key") as delete_key_spy,
+            patch("tasks.document_indexing_task.TenantIsolatedTaskQueue.pull_tasks", return_value=[], autospec=True),
+            patch(
+                "tasks.document_indexing_task.TenantIsolatedTaskQueue.delete_task_key", autospec=True
+            ) as delete_key_spy,
         ):
             priority_document_indexing_task(dataset.tenant_id, dataset.id, document_ids)
 

api/tests/test_containers_integration_tests/tasks/test_delete_segment_from_index_task.py

@@ -216,7 +216,7 @@ class TestDeleteSegmentFromIndexTask:
         db_session_with_containers.commit()
         return segments
 
-    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
+    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
     def test_delete_segment_from_index_task_success(self, mock_index_processor_factory, db_session_with_containers):
         """
         Test successful segment deletion from index with comprehensive verification.
@@ -399,7 +399,7 @@ class TestDeleteSegmentFromIndexTask:
         # Verify the task completed without exceptions
         assert result is None  # Task should return None when indexing is not completed
 
-    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
+    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
     def test_delete_segment_from_index_task_index_processor_clean(
         self, mock_index_processor_factory, db_session_with_containers
     ):
@@ -457,7 +457,7 @@ class TestDeleteSegmentFromIndexTask:
             mock_index_processor_factory.reset_mock()
             mock_processor.reset_mock()
 
-    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
+    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
     def test_delete_segment_from_index_task_exception_handling(
         self, mock_index_processor_factory, db_session_with_containers
     ):
@@ -501,7 +501,7 @@ class TestDeleteSegmentFromIndexTask:
         assert call_args[1]["with_keywords"] is True
         assert call_args[1]["delete_child_chunks"] is True
 
-    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
+    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
     def test_delete_segment_from_index_task_empty_index_node_ids(
         self, mock_index_processor_factory, db_session_with_containers
     ):
@@ -543,7 +543,7 @@ class TestDeleteSegmentFromIndexTask:
         assert call_args[1]["with_keywords"] is True
         assert call_args[1]["delete_child_chunks"] is True
 
-    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory")
+    @patch("tasks.delete_segment_from_index_task.IndexProcessorFactory", autospec=True)
     def test_delete_segment_from_index_task_large_index_node_ids(
         self, mock_index_processor_factory, db_session_with_containers
     ):

api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py

@@ -32,14 +32,11 @@ class TestDocumentIndexingTasks:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.document_indexing_task.IndexingRunner") as mock_indexing_runner,
-            patch("tasks.document_indexing_task.FeatureService") as mock_feature_service,
+            patch("tasks.document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner,
+            patch("tasks.document_indexing_task.FeatureService", autospec=True) as mock_feature_service,
         ):
             # Setup mock indexing runner
-            mock_runner_instance = MagicMock()
-            mock_indexing_runner.return_value = mock_runner_instance
-
-            # Setup mock feature service
+            mock_runner_instance = mock_indexing_runner.return_value  # Setup mock feature service
             mock_features = MagicMock()
             mock_features.billing.enabled = False
             mock_feature_service.get_features.return_value = mock_features

api/tests/test_containers_integration_tests/tasks/test_document_indexing_update_task.py

@@ -16,15 +16,13 @@ class TestDocumentIndexingUpdateTask:
         - IndexingRunner.run([...])
         """
         with (
-            patch("tasks.document_indexing_update_task.IndexProcessorFactory") as mock_factory,
-            patch("tasks.document_indexing_update_task.IndexingRunner") as mock_runner,
+            patch("tasks.document_indexing_update_task.IndexProcessorFactory", autospec=True) as mock_factory,
+            patch("tasks.document_indexing_update_task.IndexingRunner", autospec=True) as mock_runner,
         ):
             processor_instance = MagicMock()
             mock_factory.return_value.init_index_processor.return_value = processor_instance
 
-            runner_instance = MagicMock()
-            mock_runner.return_value = runner_instance
-
+            runner_instance = mock_runner.return_value
             yield {
                 "factory": mock_factory,
                 "processor": processor_instance,

api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py

@@ -31,15 +31,14 @@ class TestDuplicateDocumentIndexingTasks:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.duplicate_document_indexing_task.IndexingRunner") as mock_indexing_runner,
-            patch("tasks.duplicate_document_indexing_task.FeatureService") as mock_feature_service,
-            patch("tasks.duplicate_document_indexing_task.IndexProcessorFactory") as mock_index_processor_factory,
+            patch("tasks.duplicate_document_indexing_task.IndexingRunner", autospec=True) as mock_indexing_runner,
+            patch("tasks.duplicate_document_indexing_task.FeatureService", autospec=True) as mock_feature_service,
+            patch(
+                "tasks.duplicate_document_indexing_task.IndexProcessorFactory", autospec=True
+            ) as mock_index_processor_factory,
         ):
             # Setup mock indexing runner
-            mock_runner_instance = MagicMock()
-            mock_indexing_runner.return_value = mock_runner_instance
-
-            # Setup mock feature service
+            mock_runner_instance = mock_indexing_runner.return_value  # Setup mock feature service
             mock_features = MagicMock()
             mock_features.billing.enabled = False
             mock_feature_service.get_features.return_value = mock_features
@@ -650,7 +649,7 @@ class TestDuplicateDocumentIndexingTasks:
             updated_document = db_session_with_containers.query(Document).where(Document.id == doc_id).first()
             assert updated_document.indexing_status == "parsing"
 
-    @patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue")
+    @patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue", autospec=True)
     def test_normal_duplicate_document_indexing_task_with_tenant_queue(
         self, mock_queue_class, db_session_with_containers, mock_external_service_dependencies
     ):
@@ -693,7 +692,7 @@ class TestDuplicateDocumentIndexingTasks:
             updated_document = db_session_with_containers.query(Document).where(Document.id == doc_id).first()
             assert updated_document.indexing_status == "parsing"
 
-    @patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue")
+    @patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue", autospec=True)
     def test_priority_duplicate_document_indexing_task_with_tenant_queue(
         self, mock_queue_class, db_session_with_containers, mock_external_service_dependencies
     ):
@@ -737,7 +736,7 @@ class TestDuplicateDocumentIndexingTasks:
             updated_document = db_session_with_containers.query(Document).where(Document.id == doc_id).first()
             assert updated_document.indexing_status == "parsing"
 
-    @patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue")
+    @patch("tasks.duplicate_document_indexing_task.TenantIsolatedTaskQueue", autospec=True)
     def test_tenant_queue_wrapper_processes_next_tasks(
         self, mock_queue_class, db_session_with_containers, mock_external_service_dependencies
     ):

api/tests/test_containers_integration_tests/tasks/test_enable_segments_to_index_task.py

@@ -18,7 +18,9 @@ class TestEnableSegmentsToIndexTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.enable_segments_to_index_task.IndexProcessorFactory") as mock_index_processor_factory,
+            patch(
+                "tasks.enable_segments_to_index_task.IndexProcessorFactory", autospec=True
+            ) as mock_index_processor_factory,
         ):
             # Setup mock index processor
             mock_processor = MagicMock()
@@ -370,7 +372,7 @@ class TestEnableSegmentsToIndexTask:
             redis_client.set(indexing_cache_key, "processing", ex=300)
 
         # Mock the get_child_chunks method for each segment
-        with patch.object(DocumentSegment, "get_child_chunks") as mock_get_child_chunks:
+        with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks:
             # Setup mock to return child chunks for each segment
             mock_child_chunks = []
             for i in range(2):  # Each segment has 2 child chunks

api/tests/test_containers_integration_tests/tasks/test_mail_account_deletion_task.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 from faker import Faker
@@ -16,16 +16,14 @@ class TestMailAccountDeletionTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.mail_account_deletion_task.mail") as mock_mail,
-            patch("tasks.mail_account_deletion_task.get_email_i18n_service") as mock_get_email_service,
+            patch("tasks.mail_account_deletion_task.mail", autospec=True) as mock_mail,
+            patch("tasks.mail_account_deletion_task.get_email_i18n_service", autospec=True) as mock_get_email_service,
         ):
             # Setup mock mail service
             mock_mail.is_inited.return_value = True
 
             # Setup mock email service
-            mock_email_service = MagicMock()
-            mock_get_email_service.return_value = mock_email_service
-
+            mock_email_service = mock_get_email_service.return_value
             yield {
                 "mail": mock_mail,
                 "get_email_service": mock_get_email_service,

api/tests/test_containers_integration_tests/tasks/test_mail_change_mail_task.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 from faker import Faker
@@ -15,16 +15,14 @@ class TestMailChangeMailTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.mail_change_mail_task.mail") as mock_mail,
-            patch("tasks.mail_change_mail_task.get_email_i18n_service") as mock_get_email_i18n_service,
+            patch("tasks.mail_change_mail_task.mail", autospec=True) as mock_mail,
+            patch("tasks.mail_change_mail_task.get_email_i18n_service", autospec=True) as mock_get_email_i18n_service,
         ):
             # Setup mock mail service
             mock_mail.is_inited.return_value = True
 
             # Setup mock email i18n service
-            mock_email_service = MagicMock()
-            mock_get_email_i18n_service.return_value = mock_email_service
-
+            mock_email_service = mock_get_email_i18n_service.return_value
             yield {
                 "mail": mock_mail,
                 "email_i18n_service": mock_email_service,

api/tests/test_containers_integration_tests/tasks/test_mail_email_code_login_task.py

@@ -53,8 +53,8 @@ class TestSendEmailCodeLoginMailTask:
     def mock_external_service_dependencies(self):
         """Mock setup for external service dependencies."""
         with (
-            patch("tasks.mail_email_code_login.mail") as mock_mail,
-            patch("tasks.mail_email_code_login.get_email_i18n_service") as mock_email_service,
+            patch("tasks.mail_email_code_login.mail", autospec=True) as mock_mail,
+            patch("tasks.mail_email_code_login.get_email_i18n_service", autospec=True) as mock_email_service,
         ):
             # Setup default mock returns
             mock_mail.is_inited.return_value = True
@@ -573,7 +573,7 @@ class TestSendEmailCodeLoginMailTask:
             mock_email_service_instance.send_email.side_effect = exception
 
             # Mock logging to capture error messages
-            with patch("tasks.mail_email_code_login.logger") as mock_logger:
+            with patch("tasks.mail_email_code_login.logger", autospec=True) as mock_logger:
                 # Act: Execute the task - it should handle the exception gracefully
                 send_email_code_login_mail_task(
                     language=test_language,