devplacepy/devplacepy/services/devii/interaction/broker.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import logging
import time
from typing import Any, Awaitable, Callable, Optional

import uuid_utils

from .capabilities import (
    CHANNEL_CLI,
    CHANNEL_SITE_CHAT,
    CHANNEL_TELEGRAM,
    ChannelContext,
    channel_id_for_session,
    context_for,
)
from .markdown import project_markdown, project_plain_menu
from .parse import parse_plain_reply
from .schema import (
    InteractionRequest,
    InteractionResult,
    count_input_fields,
    validate_prompt_args,
)

logger = logging.getLogger("devii.interaction")

EmitFn = Callable[[dict[str, Any]], Awaitable[None]]
WaitFn = Callable[[str, dict[str, Any], float], Awaitable[Any]]
TelegramPresentFn = Callable[[dict[str, Any]], Awaitable[Any]]
PlainPresentFn = Callable[[str], Awaitable[str]]


class InteractionBroker:
    def __init__(
        self,
        session_channel: str = "main",
        *,
        owner_kind: str = "guest",
        owner_id: str = "",
        emit: Optional[EmitFn] = None,
        wait_site: Optional[WaitFn] = None,
        present_telegram: Optional[TelegramPresentFn] = None,
        present_plain: Optional[PlainPresentFn] = None,
    ) -> None:
        self._session_channel = session_channel
        self._channel_id = channel_id_for_session(session_channel)
        self._owner_kind = owner_kind
        self._owner_id = owner_id
        self._emit = emit
        self._wait_site = wait_site
        self._present_telegram = present_telegram
        self._present_plain = present_plain
        self._open: dict[str, dict[str, Any]] = {}
        self._lock = asyncio.Lock()
        self._single_flight = True

    @property
    def channel_id(self) -> str:
        return self._channel_id

    def set_channel_id(self, channel_id: str) -> None:
        self._channel_id = channel_id

    def open_id(self) -> str | None:
        if not self._open:
            return None
        return next(iter(self._open))

    def open_request(self) -> InteractionRequest | None:
        open_id = self.open_id()
        if not open_id:
            return None
        payload = self._open.get(open_id) or {}
        request = payload.get("request_model")
        if isinstance(request, InteractionRequest):
            return request
        raw = payload.get("request")
        if isinstance(raw, dict):
            try:
                return InteractionRequest.model_validate(raw)
            except Exception:
                return None
        return None

    def answer_text(self, text: str) -> dict[str, Any] | None:
        open_id = self.open_id()
        if not open_id:
            return None
        request = self.open_request()
        if request is None:
            return {
                "handled": True,
                "status": "error",
                "error": "Open interaction has no parseable schema.",
                "interaction_id": open_id,
            }
        status, values, error = parse_plain_reply(request, text)
        if status == "error":
            return {
                "handled": True,
                "status": "error",
                "error": error or "Could not parse reply.",
                "interaction_id": open_id,
            }
        result = {
            "status": status,
            "interaction_id": open_id,
            "values": values if status == "submitted" else {},
            "meta": {
                "adapter": self._channel_id,
                "degraded": True,
                "via": "text",
            },
            "error": error,
        }
        return {
            "handled": True,
            "status": status,
            "interaction_id": open_id,
            "result": result,
        }

    def channel_context(self) -> ChannelContext:
        owner_kind = getattr(self, "_owner_kind", "guest")
        owner_id = getattr(self, "_owner_id", "")
        return context_for(
            self._channel_id,
            open_interaction_id=self.open_id(),
            owner_kind=owner_kind,
            owner_id=owner_id,
        )

    def bind(
        self,
        *,
        emit: Optional[EmitFn] = None,
        wait_site: Optional[WaitFn] = None,
        present_telegram: Optional[TelegramPresentFn] = None,
        present_plain: Optional[PlainPresentFn] = None,
    ) -> None:
        if emit is not None:
            self._emit = emit
        if wait_site is not None:
            self._wait_site = wait_site
        if present_telegram is not None:
            self._present_telegram = present_telegram
        if present_plain is not None:
            self._present_plain = present_plain

    async def prompt(self, arguments: dict[str, Any]) -> InteractionResult:
        request = validate_prompt_args(arguments)
        async with self._lock:
            if self._single_flight and self._open:
                for open_id in list(self._open):
                    await self._finish(
                        open_id,
                        status="superseded",
                        values={},
                        error=None,
                    )
            interaction_id = request.id or f"ix-{uuid_utils.uuid7().hex[:8]}"
            if not interaction_id.startswith("ix-") and not request.id:
                interaction_id = f"ix-{interaction_id}"
            started = time.monotonic()
            payload = {
                "interaction_id": interaction_id,
                "request": request.model_dump(),
                "request_model": request,
                "markdown": project_markdown(request, interaction_id),
                "plain": project_plain_menu(request),
                "channel_id": self._channel_id,
                "started": started,
                "future": asyncio.get_event_loop().create_future(),
            }
            self._open[interaction_id] = payload
            logger.info(
                "interaction open id=%s channel=%s widgets=%d",
                interaction_id,
                self._channel_id,
                count_input_fields(request.widgets),
            )

        try:
            result = await self._present(interaction_id, request, payload)
        except Exception as exc:
            logger.exception("interaction failed id=%s", interaction_id)
            result = InteractionResult(
                interaction_id=interaction_id,
                status="error",
                channel_id=self._channel_id,
                values={},
                meta={"adapter": self._channel_id, "degraded": True},
                error=str(exc),
            )
            await self._resolve_local(interaction_id, result)
            return result

        latency_ms = int((time.monotonic() - started) * 1000)
        if "latency_ms" not in result.meta:
            result.meta["latency_ms"] = latency_ms
        if "adapter" not in result.meta:
            result.meta["adapter"] = self._channel_id
        self._open.pop(interaction_id, None)
        logger.info(
            "interaction closed id=%s status=%s latency_ms=%s",
            interaction_id,
            result.status,
            result.meta.get("latency_ms"),
        )
        return result

    async def cancel(
        self, interaction_id: str = "", reason: str = "cancelled"
    ) -> InteractionResult:
        target = interaction_id or self.open_id()
        if not target or target not in self._open:
            return InteractionResult(
                interaction_id=target or "",
                status="error",
                channel_id=self._channel_id,
                error="No open interaction to cancel.",
            )
        status = reason if reason in (
            "cancelled",
            "timeout",
            "superseded",
            "error",
            "channel_changed",
        ) else "cancelled"
        result = InteractionResult(
            interaction_id=target,
            status=status,  # type: ignore[arg-type]
            channel_id=self._channel_id,
            values={},
            meta={"adapter": self._channel_id, "reason": reason},
        )
        await self._resolve_local(target, result)
        return result

    def resolve(
        self,
        interaction_id: str,
        *,
        status: str,
        values: dict[str, Any] | None = None,
        error: str | None = None,
        meta: dict[str, Any] | None = None,
    ) -> bool:
        payload = self._open.get(interaction_id)
        if payload is None:
            return False
        result = InteractionResult(
            interaction_id=interaction_id,
            status=status,  # type: ignore[arg-type]
            channel_id=self._channel_id,
            values=values or {},
            meta=meta or {"adapter": self._channel_id},
            error=error,
        )
        future = payload.get("future")
        if future is not None and not future.done():
            future.set_result(result)
        return True

    async def _present(
        self,
        interaction_id: str,
        request: InteractionRequest,
        payload: dict[str, Any],
    ) -> InteractionResult:
        timeout = float(request.timeout_sec or 0)
        if self._wait_site is not None and self._channel_id in (
            CHANNEL_SITE_CHAT,
            CHANNEL_TELEGRAM,
        ):
            return await self._present_site(interaction_id, request, payload, timeout)
        if self._channel_id == CHANNEL_TELEGRAM:
            return await self._present_telegram_path(
                interaction_id, request, payload, timeout
            )
        return await self._present_plain_path(
            interaction_id, request, payload, timeout
        )

    async def _present_site(
        self,
        interaction_id: str,
        request: InteractionRequest,
        payload: dict[str, Any],
        timeout: float,
    ) -> InteractionResult:
        if self._wait_site is None:
            return await self._present_plain_path(
                interaction_id, request, payload, timeout, degraded=True
            )
        frame = {
            "interaction_id": interaction_id,
            "channel": self._channel_id,
            "title": request.title,
            "description": request.description,
            "widgets": [w.model_dump() for w in request.widgets],
            "submit_label": request.submit_label,
            "cancel_label": request.cancel_label,
            "cancelable": request.cancelable,
            "timeout_sec": request.timeout_sec,
            "markdown": payload["markdown"],
            "plain": payload["plain"],
            "request": request.model_dump(),
        }
        try:
            raw = await self._wait_site(interaction_id, frame, timeout)
        except asyncio.TimeoutError:
            return InteractionResult(
                interaction_id=interaction_id,
                status="timeout",
                channel_id=self._channel_id,
                meta={
                    "adapter": self._channel_id,
                    "degraded": self._channel_id != CHANNEL_SITE_CHAT,
                },
            )
        return self._normalize_site_result(
            interaction_id,
            raw,
            adapter=self._channel_id,
            degraded=self._channel_id != CHANNEL_SITE_CHAT,
        )

    async def _present_telegram_path(
        self,
        interaction_id: str,
        request: InteractionRequest,
        payload: dict[str, Any],
        timeout: float,
    ) -> InteractionResult:
        if self._present_telegram is not None:
            try:
                raw = await self._present_telegram(
                    {
                        "interaction_id": interaction_id,
                        "request": request.model_dump(),
                        "markdown": payload["markdown"],
                        "plain": payload["plain"],
                        "timeout_sec": request.timeout_sec,
                    }
                )
                return self._normalize_site_result(
                    interaction_id, raw, adapter=CHANNEL_TELEGRAM, degraded=True
                )
            except asyncio.TimeoutError:
                return InteractionResult(
                    interaction_id=interaction_id,
                    status="timeout",
                    channel_id=self._channel_id,
                    meta={"adapter": CHANNEL_TELEGRAM, "degraded": True},
                )
            except Exception as exc:
                logger.exception("telegram adapter failed, falling back to plain")
                payload["plain_error"] = str(exc)
        return await self._present_plain_path(
            interaction_id, request, payload, timeout, degraded=True
        )

    async def _present_plain_path(
        self,
        interaction_id: str,
        request: InteractionRequest,
        payload: dict[str, Any],
        timeout: float,
        *,
        degraded: bool = False,
    ) -> InteractionResult:
        menu = payload["plain"]
        if self._present_plain is not None:
            try:
                if timeout > 0:
                    reply = await asyncio.wait_for(
                        self._present_plain(menu), timeout=timeout
                    )
                else:
                    reply = await self._present_plain(menu)
            except asyncio.TimeoutError:
                return InteractionResult(
                    interaction_id=interaction_id,
                    status="timeout",
                    channel_id=self._channel_id,
                    meta={
                        "adapter": self._channel_id or CHANNEL_CLI,
                        "degraded": True,
                    },
                )
            status, values, error = parse_plain_reply(request, reply)
            return InteractionResult(
                interaction_id=interaction_id,
                status=status,  # type: ignore[arg-type]
                channel_id=self._channel_id,
                values=values,
                meta={
                    "adapter": self._channel_id or CHANNEL_CLI,
                    "degraded": True,
                },
                error=error,
            )
        future: asyncio.Future = payload["future"]
        if self._emit is not None:
            await self._emit(
                {
                    "type": "interaction",
                    "id": interaction_id,
                    "mode": "plain",
                    "text": menu,
                    "markdown": payload["markdown"],
                }
            )
        try:
            if timeout > 0:
                result = await asyncio.wait_for(future, timeout=timeout)
            else:
                result = await future
            if isinstance(result, InteractionResult):
                result.meta.setdefault("degraded", degraded or True)
                return result
            return self._normalize_site_result(
                interaction_id, result, adapter=self._channel_id, degraded=True
            )
        except asyncio.TimeoutError:
            return InteractionResult(
                interaction_id=interaction_id,
                status="timeout",
                channel_id=self._channel_id,
                meta={"adapter": self._channel_id, "degraded": True},
            )

    def _normalize_site_result(
        self,
        interaction_id: str,
        raw: Any,
        *,
        adapter: str = CHANNEL_SITE_CHAT,
        degraded: bool = False,
    ) -> InteractionResult:
        if isinstance(raw, InteractionResult):
            return raw
        if not isinstance(raw, dict):
            return InteractionResult(
                interaction_id=interaction_id,
                status="error",
                channel_id=self._channel_id,
                error="Invalid interaction result payload.",
                meta={"adapter": adapter, "degraded": degraded},
            )
        if raw.get("error") and not raw.get("status"):
            return InteractionResult(
                interaction_id=interaction_id,
                status="error",
                channel_id=self._channel_id,
                error=str(raw.get("error")),
                meta={"adapter": adapter, "degraded": degraded},
            )
        status = str(raw.get("status") or "submitted")
        if status not in (
            "submitted",
            "cancelled",
            "timeout",
            "superseded",
            "error",
            "channel_changed",
        ):
            status = "error"
        values = raw.get("values") if isinstance(raw.get("values"), dict) else {}
        meta = raw.get("meta") if isinstance(raw.get("meta"), dict) else {}
        meta.setdefault("adapter", adapter)
        meta.setdefault("degraded", degraded)
        return InteractionResult(
            interaction_id=str(raw.get("interaction_id") or interaction_id),
            status=status,  # type: ignore[arg-type]
            channel_id=self._channel_id,
            values=values,
            meta=meta,
            error=str(raw["error"]) if raw.get("error") else None,
        )

    async def _finish(
        self,
        interaction_id: str,
        *,
        status: str,
        values: dict[str, Any],
        error: str | None,
    ) -> None:
        result = InteractionResult(
            interaction_id=interaction_id,
            status=status,  # type: ignore[arg-type]
            channel_id=self._channel_id,
            values=values,
            meta={"adapter": self._channel_id},
            error=error,
        )
        await self._resolve_local(interaction_id, result)

    async def _resolve_local(
        self, interaction_id: str, result: InteractionResult
    ) -> None:
        payload = self._open.pop(interaction_id, None)
        if payload is None:
            return
        future = payload.get("future")
        if future is not None and not future.done():
            future.set_result(result)