WIP: P4

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

@@ -3,7 +3,7 @@ import time
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, NamedTuple, NewType, Union, final
+from typing import Any, NewType, Union
 
 from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom, WorkflowAppGenerateEntity
 from core.app.entities.queue_entities import (

api/core/app/apps/message_based_app_generator.py

@@ -2,12 +2,9 @@ import json
 import logging
 import time
 import uuid
-from collections.abc import Generator
-from typing import Any, Mapping, Union, cast
+from collections.abc import Generator, Mapping
+from typing import Any, Union, cast
 
-from libs.broadcast_channel.channel import Subscription, Topic
-from libs.broadcast_channel.exc import SubscriptionClosedError
-from libs.broadcast_channel.redis.channel import BroadcastChannel as RedisBroadcastChannel
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
@@ -34,6 +31,9 @@ from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBa
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
+from libs.broadcast_channel.channel import Topic
+from libs.broadcast_channel.exc import SubscriptionClosedError
+from libs.broadcast_channel.redis.channel import BroadcastChannel as RedisBroadcastChannel
 from libs.datetime_utils import naive_utc_now
 from models import Account
 from models.enums import CreatorUserRole

api/core/app/apps/message_generator.py

@@ -0,0 +1,59 @@
+import json
+import time
+from collections.abc import Generator, Mapping
+from typing import Any
+
+from core.app.entities.task_entities import (
+    StreamEvent,
+)
+from extensions.ext_redis import redis_client
+from libs.broadcast_channel.channel import Topic
+from libs.broadcast_channel.exc import SubscriptionClosedError
+from libs.broadcast_channel.redis.channel import BroadcastChannel as RedisBroadcastChannel
+from models.model import AppMode
+
+
+class MessageGenerator:
+    @staticmethod
+    def _make_channel_key(app_mode: AppMode, workflow_run_id: str):
+        return f"channel:{app_mode}:{str(workflow_run_id)}"
+
+    @classmethod
+    def get_response_topic(cls, app_mode: AppMode, workflow_run_id: str) -> Topic:
+        key = cls._make_channel_key(app_mode, workflow_run_id)
+        channel = RedisBroadcastChannel(redis_client)
+        topic = channel.topic(key)
+        return topic
+
+    @classmethod
+    def retrieve_events(
+        cls, app_mode: AppMode, workflow_run_id: str, idle_timeout=300
+    ) -> Generator[Mapping | str, None, None]:
+        topic = cls.get_response_topic(app_mode, workflow_run_id)
+        return _topic_msg_generator(topic, idle_timeout)
+
+
+def _topic_msg_generator(topic: Topic, idle_timeout: float) -> Generator[Mapping[str, Any], None, None]:
+    last_msg_time = time.time()
+    with topic.subscribe() as sub:
+        while True:
+            try:
+                msg = sub.receive()
+            except SubscriptionClosedError:
+                return
+            if msg is None:
+                current_time = time.time()
+                if current_time - last_msg_time > idle_timeout:
+                    return
+                # skip the `None` message
+                continue
+
+            last_msg_time = time.time()
+            event = json.loads(msg)
+            yield event
+            if not isinstance(event, dict):
+                continue
+
+            event_type = event.get("event")
+            if event_type in (StreamEvent.WORKFLOW_FINISHED, StreamEvent.WORKFLOW_PAUSED):
+                return

api/core/app/entities/app_invoke_entities.py

@@ -1,6 +1,6 @@
 from collections.abc import Mapping, Sequence
 from enum import StrEnum
-from typing import TYPE_CHECKING, Any, Literal, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
 

api/core/repositories/human_input_reposotiry.py

@@ -1,12 +1,10 @@
-import abc
 import dataclasses
 import json
-import uuid
-from collections.abc import Sequence
-from typing import Any, Mapping
+from collections.abc import Mapping, Sequence
+from typing import Any
 
 from sqlalchemy import Engine, select
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import sessionmaker
 
 from core.workflow.nodes.human_input.entities import (
     DeliveryChannelConfig,