feat(api/workflow): Add Conversation.dialogue_count (#7275)

api/contexts/__init__.py

@@ -1,3 +1,7 @@
 from contextvars import ContextVar
 
-tenant_id: ContextVar[str] = ContextVar('tenant_id')
+from core.workflow.entities.variable_pool import VariablePool
+
+tenant_id: ContextVar[str] = ContextVar('tenant_id')
+
+workflow_variable_pool: ContextVar[VariablePool] = ContextVar('workflow_variable_pool')

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

@@ -8,6 +8,8 @@
 
 from flask import Flask, current_app
 from pydantic import ValidationError
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 
 import contexts
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
@@ -18,15 +20,20 @@
 from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedException, PublishFrom
 from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
-from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
+from core.app.entities.app_invoke_entities import (
+    AdvancedChatAppGenerateEntity,
+    InvokeFrom,
+)
 from core.app.entities.task_entities import ChatbotAppBlockingResponse, ChatbotAppStreamResponse
 from core.file.message_file_parser import MessageFileParser
 from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
 from core.ops.ops_trace_manager import TraceQueueManager
+from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.enums import SystemVariable
 from extensions.ext_database import db
 from models.account import Account
 from models.model import App, Conversation, EndUser, Message
-from models.workflow import Workflow
+from models.workflow import ConversationVariable, Workflow
 
 logger = logging.getLogger(__name__)
 
@@ -120,7 +127,7 @@ def generate(
             conversation=conversation,
             stream=stream
         )
-    
+
     def single_iteration_generate(self, app_model: App,
                                   workflow: Workflow,
                                   node_id: str,
@@ -140,10 +147,10 @@ def single_iteration_generate(self, app_model: App,
         """
         if not node_id:
             raise ValueError('node_id is required')
-        
+
         if args.get('inputs') is None:
             raise ValueError('inputs is required')
-        
+
         extras = {
             "auto_generate_conversation_name": False
         }
@@ -209,7 +216,7 @@ def _generate(self, *,
             # update conversation features
             conversation.override_model_configs = workflow.features
             db.session.commit()
-            db.session.refresh(conversation)
+            # db.session.refresh(conversation)
 
         # init queue manager
         queue_manager = MessageBasedAppQueueManager(
@@ -221,15 +228,69 @@ def _generate(self, *,
             message_id=message.id
         )
 
+        # Init conversation variables
+        stmt = select(ConversationVariable).where(
+            ConversationVariable.app_id == conversation.app_id, ConversationVariable.conversation_id == conversation.id
+        )
+        with Session(db.engine) as session:
+            conversation_variables = session.scalars(stmt).all()
+            if not conversation_variables:
+                # Create conversation variables if they don't exist.
+                conversation_variables = [
+                    ConversationVariable.from_variable(
+                        app_id=conversation.app_id, conversation_id=conversation.id, variable=variable
+                    )
+                    for variable in workflow.conversation_variables
+                ]
+                session.add_all(conversation_variables)
+            # Convert database entities to variables.
+            conversation_variables = [item.to_variable() for item in conversation_variables]
+
+            session.commit()
+
+            # Increment dialogue count.
+            conversation.dialogue_count += 1
+
+            conversation_id = conversation.id
+            conversation_dialogue_count = conversation.dialogue_count
+            db.session.commit()
+            db.session.refresh(conversation)
+
+        inputs = application_generate_entity.inputs
+        query = application_generate_entity.query
+        files = application_generate_entity.files
+
+        user_id = None
+        if application_generate_entity.invoke_from in [InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API]:
+            end_user = db.session.query(EndUser).filter(EndUser.id == application_generate_entity.user_id).first()
+            if end_user:
+                user_id = end_user.session_id
+        else:
+            user_id = application_generate_entity.user_id
+
+        # Create a variable pool.
+        system_inputs = {
+            SystemVariable.QUERY: query,
+            SystemVariable.FILES: files,
+            SystemVariable.CONVERSATION_ID: conversation_id,
+            SystemVariable.USER_ID: user_id,
+            SystemVariable.DIALOGUE_COUNT: conversation_dialogue_count,
+        }
+        variable_pool = VariablePool(
+            system_variables=system_inputs,
+            user_inputs=inputs,
+            environment_variables=workflow.environment_variables,
+            conversation_variables=conversation_variables,
+        )
+        contexts.workflow_variable_pool.set(variable_pool)
+
         # new thread
         worker_thread = threading.Thread(target=self._generate_worker, kwargs={
             'flask_app': current_app._get_current_object(),
             'application_generate_entity': application_generate_entity,
             'queue_manager': queue_manager,
-            'conversation_id': conversation.id,
             'message_id': message.id,
-            'user': user,
-            'context': contextvars.copy_context()
+            'context': contextvars.copy_context(),
         })
 
         worker_thread.start()
@@ -242,7 +303,7 @@ def _generate(self, *,
             conversation=conversation,
             message=message,
             user=user,
-            stream=stream
+            stream=stream,
         )
 
         return AdvancedChatAppGenerateResponseConverter.convert(
@@ -253,9 +314,7 @@ def _generate(self, *,
     def _generate_worker(self, flask_app: Flask,
                          application_generate_entity: AdvancedChatAppGenerateEntity,
                          queue_manager: AppQueueManager,
-                         conversation_id: str,
                          message_id: str,
-                         user: Account,
                          context: contextvars.Context) -> None:
         """
         Generate worker in a new thread.
@@ -282,16 +341,14 @@ def _generate_worker(self, flask_app: Flask,
                         user_id=application_generate_entity.user_id
                     )
                 else:
-                    # get conversation and message
-                    conversation = self._get_conversation(conversation_id)
+                    # get message
                     message = self._get_message(message_id)
 
                     # chatbot app
                     runner = AdvancedChatAppRunner()
                     runner.run(
                         application_generate_entity=application_generate_entity,
                         queue_manager=queue_manager,
-                        conversation=conversation,
                         message=message
                     )
             except GenerateTaskStoppedException:
@@ -314,14 +371,17 @@ def _generate_worker(self, flask_app: Flask,
             finally:
                 db.session.close()
 
-    def _handle_advanced_chat_response(self, application_generate_entity: AdvancedChatAppGenerateEntity,
-                                       workflow: Workflow,
-                                       queue_manager: AppQueueManager,
-                                       conversation: Conversation,
-                                       message: Message,
-                                       user: Union[Account, EndUser],
-                                       stream: bool = False) \
-            -> Union[ChatbotAppBlockingResponse, Generator[ChatbotAppStreamResponse, None, None]]:
+    def _handle_advanced_chat_response(
+        self,
+        *,
+        application_generate_entity: AdvancedChatAppGenerateEntity,
+        workflow: Workflow,
+        queue_manager: AppQueueManager,
+        conversation: Conversation,
+        message: Message,
+        user: Union[Account, EndUser],
+        stream: bool = False,
+    ) -> Union[ChatbotAppBlockingResponse, Generator[ChatbotAppStreamResponse, None, None]]:
         """
         Handle response.
         :param application_generate_entity: application generate entity
@@ -341,7 +401,7 @@ def _handle_advanced_chat_response(self, application_generate_entity: AdvancedCh
             conversation=conversation,
             message=message,
             user=user,
-            stream=stream
+            stream=stream,
         )
 
         try:

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

@@ -4,9 +4,6 @@
 from collections.abc import Mapping
 from typing import Any, Optional, cast
 
-from sqlalchemy import select
-from sqlalchemy.orm import Session
-
 from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfig
 from core.app.apps.advanced_chat.workflow_event_trigger_callback import WorkflowEventTriggerCallback
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
@@ -19,13 +16,10 @@
 from core.app.entities.queue_entities import QueueAnnotationReplyEvent, QueueStopEvent, QueueTextChunkEvent
 from core.moderation.base import ModerationException
 from core.workflow.callbacks.base_workflow_callback import WorkflowCallback
-from core.workflow.entities.node_entities import SystemVariable
-from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.nodes.base_node import UserFrom
 from core.workflow.workflow_engine_manager import WorkflowEngineManager
 from extensions.ext_database import db
-from models.model import App, Conversation, EndUser, Message
-from models.workflow import ConversationVariable, Workflow
+from models import App, Message, Workflow
 
 logger = logging.getLogger(__name__)
 
@@ -39,7 +33,6 @@ def run(
         self,
         application_generate_entity: AdvancedChatAppGenerateEntity,
         queue_manager: AppQueueManager,
-        conversation: Conversation,
         message: Message,
     ) -> None:
         """
@@ -63,15 +56,6 @@ def run(
 
         inputs = application_generate_entity.inputs
         query = application_generate_entity.query
-        files = application_generate_entity.files
-
-        user_id = None
-        if application_generate_entity.invoke_from in [InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API]:
-            end_user = db.session.query(EndUser).filter(EndUser.id == application_generate_entity.user_id).first()
-            if end_user:
-                user_id = end_user.session_id
-        else:
-            user_id = application_generate_entity.user_id
 
         # moderation
         if self.handle_input_moderation(
@@ -103,38 +87,6 @@ def run(
         if bool(os.environ.get('DEBUG', 'False').lower() == 'true'):
             workflow_callbacks.append(WorkflowLoggingCallback())
 
-        # Init conversation variables
-        stmt = select(ConversationVariable).where(
-            ConversationVariable.app_id == conversation.app_id, ConversationVariable.conversation_id == conversation.id
-        )
-        with Session(db.engine) as session:
-            conversation_variables = session.scalars(stmt).all()
-            if not conversation_variables:
-                conversation_variables = [
-                    ConversationVariable.from_variable(
-                        app_id=conversation.app_id, conversation_id=conversation.id, variable=variable
-                    )
-                    for variable in workflow.conversation_variables
-                ]
-                session.add_all(conversation_variables)
-                session.commit()
-            # Convert database entities to variables
-            conversation_variables = [item.to_variable() for item in conversation_variables]
-
-        # Create a variable pool.
-        system_inputs = {
-            SystemVariable.QUERY: query,
-            SystemVariable.FILES: files,
-            SystemVariable.CONVERSATION_ID: conversation.id,
-            SystemVariable.USER_ID: user_id,
-        }
-        variable_pool = VariablePool(
-            system_variables=system_inputs,
-            user_inputs=inputs,
-            environment_variables=workflow.environment_variables,
-            conversation_variables=conversation_variables,
-        )
-
         # RUN WORKFLOW
         workflow_engine_manager = WorkflowEngineManager()
         workflow_engine_manager.run_workflow(
@@ -146,7 +98,6 @@ def run(
             invoke_from=application_generate_entity.invoke_from,
             callbacks=workflow_callbacks,
             call_depth=application_generate_entity.call_depth,
-            variable_pool=variable_pool,
         )
 
     def single_iteration_run(
@@ -155,7 +106,7 @@ def single_iteration_run(
         """
         Single iteration run
         """
-        app_record: App = db.session.query(App).filter(App.id == app_id).first()
+        app_record = db.session.query(App).filter(App.id == app_id).first()
         if not app_record:
             raise ValueError('App not found')
 

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

@@ -4,6 +4,7 @@
 from collections.abc import Generator
 from typing import Any, Optional, Union, cast
 
+import contexts
 from constants.tts_auto_play_timeout import TTS_AUTO_PLAY_TIMEOUT, TTS_AUTO_PLAY_YIELD_CPU_TIME
 from core.app.apps.advanced_chat.app_generator_tts_publisher import AppGeneratorTTSPublisher, AudioTrunk
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
@@ -47,7 +48,8 @@
 from core.model_runtime.entities.llm_entities import LLMUsage
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.ops.ops_trace_manager import TraceQueueManager
-from core.workflow.entities.node_entities import NodeType, SystemVariable
+from core.workflow.entities.node_entities import NodeType
+from core.workflow.enums import SystemVariable
 from core.workflow.nodes.answer.answer_node import AnswerNode
 from core.workflow.nodes.answer.entities import TextGenerateRouteChunk, VarGenerateRouteChunk
 from events.message_event import message_was_created
@@ -71,6 +73,7 @@ class AdvancedChatAppGenerateTaskPipeline(BasedGenerateTaskPipeline, WorkflowCyc
     _application_generate_entity: AdvancedChatAppGenerateEntity
     _workflow: Workflow
     _user: Union[Account, EndUser]
+    # Deprecated
     _workflow_system_variables: dict[SystemVariable, Any]
     _iteration_nested_relations: dict[str, list[str]]
 
@@ -81,7 +84,7 @@ def __init__(
             conversation: Conversation,
             message: Message,
             user: Union[Account, EndUser],
-            stream: bool
+            stream: bool,
     ) -> None:
         """
         Initialize AdvancedChatAppGenerateTaskPipeline.
@@ -103,11 +106,12 @@ def __init__(
         self._workflow = workflow
         self._conversation = conversation
         self._message = message
+        # Deprecated
         self._workflow_system_variables = {
             SystemVariable.QUERY: message.query,
             SystemVariable.FILES: application_generate_entity.files,
             SystemVariable.CONVERSATION_ID: conversation.id,
-            SystemVariable.USER_ID: user_id
+            SystemVariable.USER_ID: user_id,
         }
 
         self._task_state = AdvancedChatTaskState(
@@ -613,7 +617,9 @@ def _generate_stream_outputs_when_node_finished(self) -> Optional[Generator]:
 
                 if route_chunk_node_id == 'sys':
                     # system variable
-                    value = self._workflow_system_variables.get(SystemVariable.value_of(value_selector[1]))
+                    value = contexts.workflow_variable_pool.get().get(value_selector)
+                    if value:
+                        value = value.text
                 elif route_chunk_node_id in self._iteration_nested_relations:
                     # it's a iteration variable
                     if not self._iteration_state or route_chunk_node_id not in self._iteration_state.current_iterations: