Merge pull request #329 from airtai/limit-event-handlers

Limit event handlers

docs/docs/SUMMARY.md

@@ -80,6 +80,7 @@ search:
                 - [FunctionCallExecution](api/fastagency/base/FunctionCallExecution.md)
                 - [IOMessage](api/fastagency/base/IOMessage.md)
                 - [IOMessageVisitor](api/fastagency/base/IOMessageVisitor.md)
+                - [KeepAlive](api/fastagency/base/KeepAlive.md)
                 - [MultipleChoice](api/fastagency/base/MultipleChoice.md)
                 - [ProviderProtocol](api/fastagency/base/ProviderProtocol.md)
                 - [Runnable](api/fastagency/base/Runnable.md)
@@ -169,6 +170,7 @@ search:
                         - [handle_message](api/fastagency/ui/mesop/message/handle_message.md)
                         - [message_box](api/fastagency/ui/mesop/message/message_box.md)
                     - send_prompt
+                        - [get_more_messages](api/fastagency/ui/mesop/send_prompt/get_more_messages.md)
                         - [send_prompt_to_autogen](api/fastagency/ui/mesop/send_prompt/send_prompt_to_autogen.md)
                         - [send_user_feedback_to_autogen](api/fastagency/ui/mesop/send_prompt/send_user_feedback_to_autogen.md)
                     - styles
@@ -178,6 +180,9 @@ search:
                         - [MesopMultipleChoiceInnerStyles](api/fastagency/ui/mesop/styles/MesopMultipleChoiceInnerStyles.md)
                         - [MesopSingleChoiceInnerStyles](api/fastagency/ui/mesop/styles/MesopSingleChoiceInnerStyles.md)
                         - [MesopTextInputInnerStyles](api/fastagency/ui/mesop/styles/MesopTextInputInnerStyles.md)
+                    - timer
+                        - [configure_static_file_serving](api/fastagency/ui/mesop/timer/configure_static_file_serving.md)
+                        - [wakeup_component](api/fastagency/ui/mesop/timer/wakeup_component.md)
     - [CLI](cli/cli.md)
 - [Contributing](contributing/index.md)
     - [Development](contributing/CONTRIBUTING.md)

docs/docs/en/api/fastagency/base/KeepAlive.md

@@ -0,0 +1,11 @@
+---
+# 0.5 - API
+# 2 - Release
+# 3 - Contributing
+# 5 - Template Page
+# 10 - Default
+search:
+  boost: 0.5
+---
+
+::: fastagency.base.KeepAlive

docs/docs/en/api/fastagency/ui/mesop/send_prompt/get_more_messages.md

@@ -0,0 +1,11 @@
+---
+# 0.5 - API
+# 2 - Release
+# 3 - Contributing
+# 5 - Template Page
+# 10 - Default
+search:
+  boost: 0.5
+---
+
+::: fastagency.ui.mesop.send_prompt.get_more_messages

docs/docs/en/api/fastagency/ui/mesop/timer/configure_static_file_serving.md

@@ -0,0 +1,11 @@
+---
+# 0.5 - API
+# 2 - Release
+# 3 - Contributing
+# 5 - Template Page
+# 10 - Default
+search:
+  boost: 0.5
+---
+
+::: fastagency.ui.mesop.timer.configure_static_file_serving

docs/docs/en/api/fastagency/ui/mesop/timer/wakeup_component.md

@@ -0,0 +1,11 @@
+---
+# 0.5 - API
+# 2 - Release
+# 3 - Contributing
+# 5 - Template Page
+# 10 - Default
+search:
+  boost: 0.5
+---
+
+::: fastagency.ui.mesop.timer.wakeup_component

fastagency/base.py

@@ -29,6 +29,7 @@
     "MultipleChoice",
     "Runnable",
     "SuggestedFunctionCall",
+    "KeepAlive",
     "SystemMessage",
     "TextInput",
     "TextMessage",
@@ -45,6 +46,7 @@
     "text_input",
     "multiple_choice",
     "system_message",
+    "keep_alive",
     "workflow_completed",
     "error",
 ]
@@ -77,6 +79,7 @@ def _get_message_class(type: Optional[MessageType]) -> "Type[IOMessage]":
             "function_call_execution": FunctionCallExecution,
             "text_input": TextInput,
             "multiple_choice": MultipleChoice,
+            "keep_alive": KeepAlive,
             "system_message": SystemMessage,
             "workflow_completed": WorkflowCompleted,
             "error": Error,
@@ -171,6 +174,10 @@ class Error(IOMessage):
     long: Optional[str] = None
 
 
+@dataclass
+class KeepAlive(IOMessage): ...
+
+
 class IOMessageVisitor(ABC):
     def visit(self, message: IOMessage) -> Optional[str]:
         method_name = f"visit_{message.type}"
@@ -202,6 +209,9 @@ def visit_multiple_choice(self, message: MultipleChoice) -> Optional[str]:
     def visit_system_message(self, message: SystemMessage) -> Optional[str]:
         return self.visit_default(message)
 
+    def visit_keep_alive(self, message: KeepAlive) -> Optional[str]:
+        return self.visit_default(message)
+
     def visit_workflow_completed(self, message: WorkflowCompleted) -> Optional[str]:
         return self.visit_default(message)
 

fastagency/ui/mesop/base.py

@@ -1,4 +1,5 @@
 import threading
+import time
 from collections.abc import Generator, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -16,6 +17,7 @@
     AskingMessage,
     IOMessage,
     IOMessageVisitor,
+    KeepAlive,
     MultipleChoice,
     ProviderProtocol,
     Runnable,
@@ -25,6 +27,7 @@
 )
 from ...logging import get_logger
 from .styles import MesopHomePageStyles
+from .timer import configure_static_file_serving
 
 logger = get_logger(__name__)
 
@@ -50,13 +53,15 @@ def __init__(
         *,
         security_policy: Optional[me.SecurityPolicy] = None,
         styles: Optional[MesopHomePageStyles] = None,
+        keep_alive: Optional[bool] = False,
     ) -> None:
         """Initialize the console UI object.
 
         Args:
             super_conversation (Optional[MesopUI], optional): The super conversation. Defaults to None.
             security_policy (Optional[me.SecurityPolicy], optional): The security policy. Defaults to None.
             styles (Optional[MesopHomePageStyles], optional): The styles. Defaults to None.
+            keep_alive (Optional[bool]): If keep alive messages should be inserted, defaults to False`
         """
         logger.info(f"Initializing MesopUI: {self}")
         try:
@@ -66,9 +71,13 @@ def __init__(
             self._in_queue: Optional[Queue[str]] = None
             self._out_queue: Optional[Queue[MesopMessage]] = None
 
+            self._keep_me_alive = keep_alive
+            self._keep_alive_thread: Optional[threading.Thread] = None
             if super_conversation is None:
                 self._in_queue = Queue()
                 self._out_queue = Queue()
+                self.keep_me_alive()
+
             MesopUI.register(self)
 
             if MesopUI._me is None:
@@ -84,6 +93,23 @@ def __init__(
 
     _registry: ClassVar[dict[str, "MesopUI"]] = {}
 
+    def keep_me_alive(self) -> None:
+        def keep_alive_worker() -> None:
+            while self._keep_me_alive:
+                time.sleep(3)
+                if self._out_queue:
+                    msg = KeepAlive()
+                    mesop_msg = self._mesop_message(msg)
+                    logger.info(f"putting keepalive {msg.uuid}")
+                    self._out_queue.put(mesop_msg)
+
+        if self._keep_me_alive and self._keep_alive_thread is None:
+            self._keep_alive_thread = threading.Thread(target=keep_alive_worker)
+            self._keep_alive_thread.start()
+
+    def do_not_keep_me_alive(self) -> None:
+        self._keep_me_alive = False
+
     @classmethod
     def get_created_instance(cls) -> "MesopUI":
         created_instance = cls._created_instance
@@ -215,7 +241,7 @@ def create_subconversation(self) -> "MesopUI":
         return sub_conversation
 
     def _is_stream_braker(self, message: IOMessage) -> bool:
-        return isinstance(message, (AskingMessage, WorkflowCompleted))
+        return isinstance(message, (AskingMessage, WorkflowCompleted, KeepAlive))
 
     def respond(self, message: str) -> None:
         self.in_queue.put(message)
@@ -246,6 +272,9 @@ def handle_wsgi(
         MesopUI._created_instance = self
         MesopUI._app = app
 
+        if configure_static_file_serving is None:  # pragme: no cover
+            logger.error("configure_static_file_serving is None")
+
         if MesopUI._me is None:
             logger.error("MesopUI._me is None")
             raise RuntimeError("MesopUI._me is None")
@@ -276,7 +305,6 @@ def conversation_worker(ui: MesopUI, subconversation: MesopUI) -> None:
                 ui=subconversation,  # type: ignore[arg-type]
                 initial_message=initial_message,
             )
-
             ui.process_message(
                 IOMessage.create(
                     sender="user",
@@ -288,7 +316,6 @@ def conversation_worker(ui: MesopUI, subconversation: MesopUI) -> None:
                     },
                 )
             )
-
             ui.process_message(
                 IOMessage.create(
                     sender="user",
@@ -317,9 +344,10 @@ def conversation_worker(ui: MesopUI, subconversation: MesopUI) -> None:
                     result="Exception raised",
                 )
             )
-            return
+        finally:
+            ui.do_not_keep_me_alive()
 
-    ui = MesopUI()
+    ui = MesopUI(keep_alive=True)
     subconversation = ui.create_subconversation()
     thread = threading.Thread(target=conversation_worker, args=(ui, subconversation))
     thread.start()

fastagency/ui/mesop/main.py

@@ -22,7 +22,10 @@
 
 
 DEFAULT_SECURITY_POLICY = me.SecurityPolicy(
-    allowed_iframe_parents=["https://fastagency.ai"]
+    allowed_script_srcs=[
+        "https://cdn.jsdelivr.net",
+    ],
+    allowed_iframe_parents=["https://fastagency.ai"],
 )
 
 

fastagency/ui/mesop/message.py

@@ -4,6 +4,7 @@
 from uuid import uuid4
 
 import mesop as me
+import mesop.labs as mel
 
 from fastagency.helpers import jsonify_string
 
@@ -13,6 +14,7 @@
     FunctionCallExecution,
     IOMessage,
     IOMessageVisitor,
+    KeepAlive,
     MultipleChoice,
     SuggestedFunctionCall,
     SystemMessage,
@@ -24,8 +26,9 @@
 from .base import MesopMessage
 from .components.inputs import input_text
 from .data_model import Conversation, ConversationMessage, State
-from .send_prompt import send_user_feedback_to_autogen
+from .send_prompt import get_more_messages, send_user_feedback_to_autogen
 from .styles import MesopHomePageStyles, MesopMessageStyles
+from .timer import wakeup_component
 
 logger = get_logger(__name__)
 
@@ -193,14 +196,23 @@ def visit_system_message(self, message: SystemMessage) -> None:
             if "heading" in message.message and "body" in message.message
             else json.dumps(message.message)
         )
-
         self.visit_default(
             message,
             content=content,
             style=self._styles.message.system,
             scrollable=True,
         )
 
+    def visit_keep_alive(self, message: KeepAlive) -> None:
+        def on_wakeup(e: mel.WebEvent) -> Iterator[None]:
+            logger.info("waking up, after the keep alive")
+            self._conversation_message.feedback_completed = True
+            yield from consume_responses(get_more_messages())
+
+        with me.box():
+            if not (self._readonly or self._conversation_message.feedback_completed):
+                wakeup_component(on_wakeup=on_wakeup)
+
     def visit_suggested_function_call(self, message: SuggestedFunctionCall) -> None:
         content = f"""**function_name**: `{message.function_name}`<br>
 **call_id**: `{message.call_id}`<br>

fastagency/ui/mesop/send_prompt.py

@@ -22,3 +22,10 @@ def send_user_feedback_to_autogen(user_response: str) -> Iterable[MesopMessage]:
     mesop_io = MesopUI.get_conversation(mesop_id)
     mesop_io.respond(user_response)
     return mesop_io.get_message_stream()
+
+
+def get_more_messages() -> Iterable[MesopMessage]:
+    state = me.state(State)
+    mesop_id = state.conversation.fastagency
+    mesop_io = MesopUI.get_conversation(mesop_id)
+    return mesop_io.get_message_stream()

fastagency/ui/mesop/timer.py

@@ -0,0 +1,71 @@
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+import mesop.labs as mel
+import mesop.server.static_file_serving
+import mesop.server.wsgi_app
+from flask import Flask, Response
+from mesop.server.static_file_serving import (
+    WEB_COMPONENTS_PATH_SEGMENT,
+    noop,
+    send_file_compressed,
+)
+from mesop.server.static_file_serving import (
+    configure_static_file_serving as configure_static_file_serving_original,
+)
+
+from ...logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def configure_static_file_serving(
+    app: Flask,
+    static_file_runfiles_base: str,
+    livereload_script_url: Optional[str] = None,
+    preprocess_request: Callable[[], None] = noop,
+    disable_gzip_cache: bool = False,
+    default_allowed_iframe_parents: str = "'self'",
+) -> None:
+    logger.info("Configuring static file serving with patched method")
+
+    configure_static_file_serving_original(
+        app=app,
+        static_file_runfiles_base=static_file_runfiles_base,
+        livereload_script_url=livereload_script_url,
+        preprocess_request=preprocess_request,
+        disable_gzip_cache=disable_gzip_cache,
+        default_allowed_iframe_parents=default_allowed_iframe_parents,
+    )
+
+    @app.route(f"/{WEB_COMPONENTS_PATH_SEGMENT}/__fast_agency_internal__/<path:path>")
+    def serve_web_components_fast_agency(path: str) -> Response:
+        logger.info(f"Serve web components fast agency: {path}")
+
+        root = Path(__file__).parents[3].resolve()
+        serving_path = f"{root}/{path}"
+
+        return send_file_compressed(  # type: ignore[no-any-return]
+            serving_path,
+            disable_gzip_cache=disable_gzip_cache,
+        )
+
+
+logger.info("Patching static file serving in Mesop")
+mesop.server.wsgi_app.configure_static_file_serving = configure_static_file_serving
+
+
+@mel.web_component(path="/__fast_agency_internal__/javascript/wakeup_component.js")  # type: ignore[misc]
+def wakeup_component(
+    *,
+    on_wakeup: Callable[[mel.WebEvent], Any],
+    key: Optional[str] = None,
+) -> Any:
+    return mel.insert_web_component(
+        name="wakeup-component",
+        key=key,
+        events={
+            "wakeupEvent": on_wakeup,
+        },
+        properties={},
+    )