Messenger Interface rework

chatsky/messengers/common/interface.py

@@ -11,25 +11,36 @@
 import logging
 from pathlib import Path
 from tempfile import gettempdir
-from typing import Optional, Any, List, Tuple, Hashable, TYPE_CHECKING, Type
+
+from typing import Optional, Any, Hashable, TYPE_CHECKING, Type
+from pydantic import BaseModel
 
 if TYPE_CHECKING:
     from chatsky.script import Context, Message
     from chatsky.pipeline.types import PipelineRunnerFunction
     from chatsky.messengers.common.types import PollingInterfaceLoopFunction
     from chatsky.script.core.message import Attachment
+    from chatsky.pipeline.pipeline.pipeline import Pipeline
 
 logger = logging.getLogger(__name__)
 
 
-class MessengerInterface(abc.ABC):
+class MessengerInterface(abc.ABC, BaseModel):
     """
     Class that represents a message interface used for communication between pipeline and users.
     It is responsible for connection between user and pipeline, as well as for request-response transactions.
     """
 
+    running: bool = True
+    """Shows whether the interface is still accepting new requests."""
+    finished_working: bool = False
+    """Shows whether the interface has finished processing all of the requests received."""
+
     @abc.abstractmethod
-    async def connect(self, pipeline_runner: PipelineRunnerFunction):
+    async def connect(
+        self,
+        pipeline_runner: PipelineRunnerFunction,
+    ):
         """
         Method invoked when message interface is instantiated and connection is established.
         May be used for sending an introduction message or displaying general bot information.
@@ -39,6 +50,15 @@ async def connect(self, pipeline_runner: PipelineRunnerFunction):
         """
         raise NotImplementedError
 
+    async def cleanup(self):
+        """
+        A placeholder method for any cleanup code you want to be
+        called before shutting down the program.
+        You can redefine this method in your class.
+        Note you also need to call cleanup() of the parent class.
+        """
+        pass
+
 
 class MessengerInterfaceWithAttachments(MessengerInterface, abc.ABC):
     """
@@ -93,75 +113,135 @@ class PollingMessengerInterface(MessengerInterface):
     """
     Polling message interface runs in a loop, constantly asking users for a new input.
     """
+    number_of_workers: int = 2
+    _request_queue = asyncio.Queue()
+    _worker_tasks = []
 
     @abc.abstractmethod
-    def _request(self) -> List[Tuple[Message, Hashable]]:
+    async def _respond(self, ctx_id, last_response):
         """
-        Method used for sending users request for their input.
+        Method used for sending users responses for their last input.
 
-        :return: A list of tuples: user inputs and context ids (any user ids) associated with the inputs.
+        :param ctx_id: Context id, specifies the user id. Without multiple messenger interfaces it's basically a
+         redundant parameter, because this function is just a more complex `print(last_response)`. (Change before merge)
+        :param last_response: Latest response from the pipeline which should be relayed to the specified user.
         """
         raise NotImplementedError
 
-    @abc.abstractmethod
-    def _respond(self, responses: List[Context]):
-        """
-        Method used for sending users responses for their last input.
+    async def _process_request(self, ctx_id: Any, update: Message, pipeline_runner: PipelineRunnerFunction):
+        """Process a new update for ctx."""
+        context = await pipeline_runner(update, ctx_id)
+        await self._respond(ctx_id, context.last_response)
 
-        :param responses: A list of contexts, representing dialogs with the users;
-            `last_response`, `id` and some dialog info can be extracted from there.
+    async def _worker_job(self, pipeline_runner: PipelineRunnerFunction, worker_timeout: float):
         """
-        raise NotImplementedError
+        Obtain Lock over the current context,
+        Process the update and send it.
+        """
+        request = await self._request_queue.get()
+        if request is not None:
+            (ctx_id, update) = request
+            async with self.pipeline.context_lock[ctx_id]:  # get exclusive access to this context among interfaces
+                await asyncio.wait_for(
+                    self._process_request(ctx_id, update, pipeline_runner),
+                    timeout=worker_timeout,
+                )
+            return False
+        else:
+            return True
 
-    def _on_exception(self, e: BaseException):
+    # This worker doesn't save the request and basically deletes it from the queue in case it can't process it.
+    # An option to save the request may be fitting? Maybe with an amount of retries.
+    async def _worker(self, pipeline_runner: PipelineRunnerFunction, worker_timeout: float):
+        while self.running or not self._request_queue.empty():
+            try:
+                no_more_jobs = self._worker_job(pipeline_runner, worker_timeout)
+                if no_more_jobs:
+                    logger.info("Worker finished working - all remaining requests have been processed.")
+                    # Polling_loop should give the required data on whether the stop signal was sent or if
+                    # the loop() function gave 'False'.
+                    break
+            except TimeoutError:
+                logger.info("worker couldn't process request in time. A request *may* have been lost.")
+
+    @abc.abstractmethod
+    async def _get_updates(self) -> list[tuple[Any, Message]]:
         """
-        Method that is called on polling cycle exceptions, in some cases it should show users the exception.
-        By default, it logs all exit exceptions to `info` log and all non-exit exceptions to `error`.
+        Obtain updates from another server
 
-        :param e: The exception.
+        Example:
+            self.bot.request_updates()
         """
-        if isinstance(e, Exception):
-            logger.error(f"Exception in {type(self).__name__} loop!", exc_info=e)
-        else:
-            logger.info(f"{type(self).__name__} has stopped polling.")
+
+    async def _polling_job(self, poll_timeout: float):
+        try:
+            received_updates = await asyncio.wait_for(self._get_updates(), timeout=poll_timeout)
+            if received_updates is not None:
+                for update in received_updates:
+                    await self._request_queue.put(update)
+        except TimeoutError:
+            logger.debug("polling_job failed - timed out")
 
     async def _polling_loop(
         self,
-        pipeline_runner: PipelineRunnerFunction,
+        loop: PollingInterfaceLoopFunction = lambda: True,
+        poll_timeout: float = None,
         timeout: float = 0,
     ):
-        """
-        Method running the request - response cycle once.
-        """
-        user_updates = self._request()
-        responses = [await pipeline_runner(request, ctx_id) for request, ctx_id in user_updates]
-        self._respond(responses)
-        await asyncio.sleep(timeout)
+        try:
+            while loop() and self.running:
+                await asyncio.shield(self._polling_job(poll_timeout))  # shield from cancellation
+                await asyncio.sleep(timeout)
+        finally:
+            self.running = False
+            # If loop() is somehow True after being False once, this logging will be wrong.
+            # But no user would want to break their own logging, right?
+            if loop() is False:
+                logger.info("polling_loop stopped working - the loop() condition was False")
+            else:
+                logger.info("polling_loop stopped working - the stop signal was received.")
+            # If there are no more jobs/stop signal received, a special 'None' request is
+            # sent to the queue (one for each worker), they shut down the workers.
+            # In case of more workers than two, change the number of 'None' requests to the new number of workers.
+            for i in range(self.number_of_workers):
+                self._request_queue.put_nowait(None)
 
     async def connect(
         self,
         pipeline_runner: PipelineRunnerFunction,
         loop: PollingInterfaceLoopFunction = lambda: True,
+        poll_timeout: float = None,
+        worker_timeout: float = None,
         timeout: float = 0,
     ):
+        # Saving strong references to workers, so that they can be cleaned up properly.
+        # shield() creates a task just like create_task() according to docs.
+        # But for safety we have two task wrappers, I guess.
+        for i in range(self.number_of_workers):
+            task = asyncio.create_task(asyncio.shield(self._worker(pipeline_runner, worker_timeout)))
+            self._worker_tasks.append(task)
+        print("worker tasks:", self._worker_tasks)
+        await self._polling_loop(loop=loop, poll_timeout=poll_timeout, timeout=timeout)
+
+    # Maybe "worker_cleanup" instead of this function name?
+    async def cleanup(self):
+        """
+        Blocks until all workers are done.
         """
-        Method, running a request - response cycle in a loop.
-        The looping behavior is determined by `loop` and `timeout`,
-        for most cases the loop itself shouldn't be overridden.
+        await super().cleanup()
+        await asyncio.wait(self._worker_tasks)
 
-        :param pipeline_runner: A function that should process user request and return context;
-            usually it's a :py:meth:`~chatsky.pipeline.pipeline.pipeline.Pipeline._run_pipeline` function.
-        :param loop: a function that determines whether polling should be continued;
-            called in each cycle, should return `True` to continue polling or `False` to stop.
-        :param timeout: a time interval between polls (in seconds).
+    def _on_exception(self, e: BaseException):
         """
-        while loop():
-            try:
-                await self._polling_loop(pipeline_runner, timeout)
+        Method that is called on polling cycle exceptions, in some cases it should show users the exception.
+        By default, it logs all exit exceptions to `info` log and all non-exit exceptions to `error`.
 
-            except BaseException as e:
-                self._on_exception(e)
-                break
+        :param e: The exception.
+        """
+        if isinstance(e, Exception):
+            logger.error(f"Exception in {type(self).__name__} loop!", exc_info=e)
+        else:
+            logger.info(f"{type(self).__name__} has stopped polling.")
 
 
 class CallbackMessengerInterface(MessengerInterface):
@@ -171,8 +251,12 @@ class CallbackMessengerInterface(MessengerInterface):
 
     def __init__(self) -> None:
         self._pipeline_runner: Optional[PipelineRunnerFunction] = None
+        super().__init__()
 
-    async def connect(self, pipeline_runner: PipelineRunnerFunction):
+    async def connect(
+        self,
+        pipeline_runner: PipelineRunnerFunction,
+    ):
         self._pipeline_runner = pipeline_runner
 
     async def on_request_async(

chatsky/messengers/console.py

@@ -1,8 +1,6 @@
 from typing import Any, Hashable, List, Optional, TextIO, Tuple
 from uuid import uuid4
 from chatsky.messengers.common.interface import PollingMessengerInterface
-from chatsky.pipeline.types import PipelineRunnerFunction
-from chatsky.script.core.context import Context
 from chatsky.script.core.message import Message
 
 
@@ -12,9 +10,6 @@ class CLIMessengerInterface(PollingMessengerInterface):
     This message interface can maintain dialog with one user at a time only.
     """
 
-    supported_request_attachment_types = set()
-    supported_response_attachment_types = set()
-
     def __init__(
         self,
         intro: Optional[str] = None,
@@ -29,13 +24,13 @@ def __init__(
         self._prompt_response: str = prompt_response
         self._descriptor: Optional[TextIO] = out_descriptor
 
-    def _request(self) -> List[Tuple[Message, Any]]:
-        return [(Message(input(self._prompt_request)), self._ctx_id)]
+    async def _get_updates(self) -> List[Tuple[Any, Message]]:
+        return [(self._ctx_id, Message(input(self._prompt_request)))]
 
-    def _respond(self, responses: List[Context]):
-        print(f"{self._prompt_response}{responses[0].last_response.text}", file=self._descriptor)
+    async def _respond(self, ctx_id, last_response: Message):
+        print(f"{self._prompt_response}{last_response.text}", file=self._descriptor)
 
-    async def connect(self, pipeline_runner: PipelineRunnerFunction, **kwargs):
+    async def connect(self, *args, **kwargs):
         """
         The CLIProvider generates new dialog id used to user identification on each `connect` call.
 
@@ -46,4 +41,4 @@ async def connect(self, pipeline_runner: PipelineRunnerFunction, **kwargs):
         self._ctx_id = uuid4()
         if self._intro is not None:
             print(self._intro)
-        await super().connect(pipeline_runner, **kwargs)
+        await super().connect(*args, **kwargs)

chatsky/messengers/telegram/abstract.py

@@ -41,8 +41,7 @@
         Update,
         Message as TelegramMessage,
     )
-    from telegram.ext import Application, ExtBot, MessageHandler, CallbackQueryHandler
-    from telegram.ext.filters import ALL
+    from telegram.ext import Application, ExtBot
 
     telegram_available = True
 except ImportError:
@@ -93,8 +92,6 @@ def __init__(self, token: str, attachments_directory: Optional[Path] = None) ->
             raise ImportError("`python-telegram-bot` package is missing.\nTry to run `pip install chatsky[telegram]`.")
 
         self.application = Application.builder().token(token).build()
-        self.application.add_handler(MessageHandler(ALL, self.on_message))
-        self.application.add_handler(CallbackQueryHandler(self.on_callback))
 
     async def get_attachment_bytes(self, source: str) -> bytes:
         file = await self.application.bot.get_file(source)

chatsky/messengers/telegram/interface.py

@@ -8,17 +8,29 @@
 from pathlib import Path
 from typing import Any, Optional
 
+from chatsky.script import Message
+
+from chatsky.messengers.common import PollingMessengerInterface
+
 from chatsky.pipeline.types import PipelineRunnerFunction
 
 from .abstract import _AbstractTelegramInterface
 
+try:
+    from telegram.ext import MessageHandler, CallbackQueryHandler
+    from telegram.ext.filters import ALL
+
+    telegram_available = True
+except ImportError:
+    telegram_available = False
+
 try:
     from telegram import Update
 except ImportError:
     Update = Any
 
 
-class LongpollingInterface(_AbstractTelegramInterface):
+class LongpollingInterface(_AbstractTelegramInterface, PollingMessengerInterface):
     """
     Telegram messenger interface, that requests Telegram API in a loop.
 
@@ -35,11 +47,22 @@ def __init__(
         self.interval = interval
         self.timeout = timeout
 
-    async def connect(self, pipeline_runner: PipelineRunnerFunction, *args, **kwargs):
-        await super().connect(pipeline_runner, *args, **kwargs)
-        self.application.run_polling(
+    async def _get_updates(self) -> list[tuple[Any, Message]]:
+        updates = self.application.bot.get_updates(
             poll_interval=self.interval, timeout=self.timeout, allowed_updates=Update.ALL_TYPES
         )
+        parsed_updates = []
+        for update in updates:
+            data_available = update.message is not None or update.callback_query is not None
+            if update.effective_chat is not None and data_available:
+                message = self.extract_message_from_telegram(update)
+                message.original_message = update
+                parsed_updates.append((update.effective_chat.id, message))
+        return parsed_updates
+
+    async def _respond(self, ctx_id, last_response):
+        if last_response is not None:
+            await self.cast_message_to_telegram_and_send(self.application.bot, ctx_id, last_response)
 
 
 class WebhookInterface(_AbstractTelegramInterface):
@@ -59,6 +82,8 @@ def __init__(
         super().__init__(token, attachments_directory)
         self.listen = host
         self.port = port
+        self.application.add_handler(MessageHandler(ALL, self.on_message))
+        self.application.add_handler(CallbackQueryHandler(self.on_callback))
 
     async def connect(self, pipeline_runner: PipelineRunnerFunction, *args, **kwargs):
         await super().connect(pipeline_runner, *args, **kwargs)