dialogflow.py

"""Core dialogflow handling things"""

from __future__ import annotations
import asyncio
from asyncio.queues import QueueEmpty
from intent_deployer.types import IntentDeployException

from asyncio import Queue
import logging
from dataclasses import dataclass, field
from typing import Any, Callable, ClassVar, Coroutine, Optional
from weakref import WeakValueDictionary

from pydantic.main import BaseModel
from starlette.background import BackgroundTasks
from webhook.types import DialogflowMessage, DialogflowMessageText, DialogflowRequest

from intent_deployer.compile import (
    Context,
    register_dialogflow_handler,
    register_intent_executor,
)
from intent_deployer.notifiers import Notifier, NullNotifier, get_notifier

log: logging.Logger = logging.getLogger(__name__)


@dataclass
class DialogflowContext(Context):
    session: str
    notifier_meta: Optional[tuple[str, dict]]
    fulfillment_messages: list[DialogflowMessage] = field(default_factory=list)
    should_inhibit_fulfillment_message: bool = field(default=False)

    _queue: Queue[Any] = field(default_factory=Queue)

    instances: ClassVar[
        WeakValueDictionary[str, DialogflowContext]
    ] = WeakValueDictionary()

    def reset_messages(self):
        # @simmsb: might be better to just collect response messages instead of
        # storing them in the context, so that we don't have to do this. should
        # be fine for now though.
        self.fulfillment_messages = []
        self.should_inhibit_fulfillment_message = False

    def add_message(self, *msgs: str):
        for msg in msgs:
            self.fulfillment_messages.append(DialogflowMessageText(text=[msg]))

    def inhibit_fulfillment_message(self):
        self.should_inhibit_fulfillment_message = True

    def get_notifier(self) -> Optional[Notifier]:
        if self.notifier_meta is None:
            return None

        return get_notifier(*self.notifier_meta)

    @classmethod
    def get_or_create(cls, background_tasks, session, *args, **kwargs):
        if (instance := cls.instances.get(session)) is not None:
            # replace the background tasks instance though
            instance.background_tasks = background_tasks
            return instance

        instance = cls(background_tasks, None, session, *args, **kwargs)
        cls.instances[session] = instance
        return instance

    @classmethod
    def from_request(
        cls, background_tasks: BackgroundTasks, req: DialogflowRequest
    ) -> DialogflowContext:
        notifier_meta = None
        if (src := req.original_detect_intent_request) is not None:
            assert src.payload is not None
            notifier_meta = (src.source, src.payload)

        return cls.get_or_create(background_tasks, req.session, notifier_meta)

    def push_response(self, resp: Any):
        self._queue.put_nowait(resp)

    async def wait_for_response(self) -> Any:
        return await self._queue.get()

    async def wait_for_confirmation(self) -> ConfirmationIntent:
        """Wait for the user to use the 'confirm' intent

        Basically if the user enters 'yes', 'okay', 'ok', 'confirm', 'no',
        'cancel' this will return with the answer.

        """
        while not isinstance(rsp := await self.wait_for_response(), ConfirmationIntent):
            pass

        return rsp

    async def confirm_if_needed(
        self, msg: str, on_confirm: Callable[[], Coroutine[None, None, None]]
    ):
        """Send a message back to the user, then execute `on_confirm` if the
        user confirms.
        """
        try:
            # make sure the queue is empty, :S
            while True:
                self._queue.get_nowait()
        except QueueEmpty:
            pass

        self.add_message(msg, "Please confirm")

        async def inner():
            rsp = await self.wait_for_confirmation()

            notifier = self.get_notifier() or NullNotifier()

            if rsp.confirmed:
                try:
                    await on_confirm()
                    await notifier.send("Intent deployed")
                except IntentDeployException as e:
                    log.exception(e)
                    self.push_status("failure", str(e))
                    await notifier.send(f"Could not deploy that intent: {e}")
                except Exception as e:
                    log.exception(e)
                    self.push_status("failure", str(e))
                    await notifier.send(f"Could not deploy that intent, internal error")
            else:
                self.push_status("cancelled")
                await notifier.send("Cancelled")

        self.background_tasks.add_task(inner)


@dataclass
class ConfirmationIntent:
    confirmed: bool


@register_dialogflow_handler("confirmation")
def handle_confirmation(params: dict[str, Any]) -> ConfirmationIntent:
    return ConfirmationIntent(params["confirm"] == "yes")


@register_intent_executor(ConfirmationIntent)
async def handle_confirmation_intent(ctx: Context, intent: ConfirmationIntent):
    if isinstance(ctx, DialogflowContext):
        ctx.inhibit_fulfillment_message()
        ctx.push_response(intent)
    else:
        log.warning(
            "Received a confirmation event but with a non-dialogflow context, ignoring"
        )