AgentEvalTool/backend/agenteval/channels/base.py

"""Abstract base class for message channels."""

import json
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Any, Awaitable, Callable, Optional


@dataclass
class ChannelHealth:
    ok: bool
    message: str = ""


@dataclass
class SendResult:
    ok: bool
    question_msg_id: Optional[str] = None
    raw_response: Optional[dict[str, Any]] = None
    error: Optional[str] = None


@dataclass
class Reply:
    question_msg_id: str
    content: Any
    sender_name: Optional[str] = None
    msg_time: Optional[datetime] = None
    raw_message: Optional[dict[str, Any]] = None


class ExchangeStatus(str, Enum):
    """Outcome categories for one complete channel exchange."""

    SUCCESS = "success"
    SEND_FAILED = "send_failed"
    REPLY_TIMEOUT = "reply_timeout"
    POLL_FAILED = "poll_failed"


def normalize_reply_text(content: Any) -> str:
    """Convert a channel reply payload into stable plain text.

    Adapters may return a plain string, a Tutu-style ``msgBody`` object, or a
    provider-specific mapping.  The exchange interface exposes only text so
    callers do not need to learn each provider's response shape.
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        body = content.get("msgBody")
        if isinstance(body, dict):
            for key in ("content", "text", "message"):
                value = body.get(key)
                if isinstance(value, str) and value:
                    return value
        elif isinstance(body, str) and body:
            return body

        for key in ("content", "text", "message"):
            value = content.get(key)
            if isinstance(value, str) and value:
                return value
        return json.dumps(content, ensure_ascii=False)
    return str(content)


@dataclass(frozen=True)
class ExchangeOutcome:
    """Result of a complete send-and-reply exchange.

    ``diagnostic`` is intentionally opaque to callers.  Adapters may retain
    provider-specific data there without expanding the public interface.
    """

    status: ExchangeStatus
    correlation_id: Optional[str] = None
    reply_text: Optional[str] = None
    latency_ms: Optional[int] = None
    reason: Optional[str] = None
    diagnostic: Any = None

    def __post_init__(self) -> None:
        if self.status is ExchangeStatus.SUCCESS:
            if not self.correlation_id:
                raise ValueError("a successful exchange requires a correlation_id")
            if self.reply_text is None:
                raise ValueError("a successful exchange requires reply_text")
            if self.latency_ms is None or self.latency_ms < 0:
                raise ValueError("a successful exchange requires non-negative latency_ms")

    @property
    def ok(self) -> bool:
        return self.status is ExchangeStatus.SUCCESS

    @property
    def expected_failure(self) -> bool:
        return self.status in {
            ExchangeStatus.SEND_FAILED,
            ExchangeStatus.REPLY_TIMEOUT,
            ExchangeStatus.POLL_FAILED,
        }

    @classmethod
    def succeeded(
        cls,
        *,
        correlation_id: str,
        reply: Any,
        latency_ms: int,
        diagnostic: Any = None,
    ) -> "ExchangeOutcome":
        return cls(
            status=ExchangeStatus.SUCCESS,
            correlation_id=correlation_id,
            reply_text=normalize_reply_text(reply),
            latency_ms=latency_ms,
            diagnostic=diagnostic,
        )

    @classmethod
    def send_failed(cls, reason: str, *, diagnostic: Any = None) -> "ExchangeOutcome":
        return cls(status=ExchangeStatus.SEND_FAILED, reason=reason, diagnostic=diagnostic)

    @classmethod
    def reply_timeout(
        cls,
        *,
        correlation_id: Optional[str] = None,
        latency_ms: Optional[int] = None,
        diagnostic: Any = None,
    ) -> "ExchangeOutcome":
        return cls(
            status=ExchangeStatus.REPLY_TIMEOUT,
            correlation_id=correlation_id,
            latency_ms=latency_ms,
            reason="reply timeout",
            diagnostic=diagnostic,
        )

    @classmethod
    def poll_failed(
        cls,
        reason: str,
        *,
        correlation_id: Optional[str] = None,
        latency_ms: Optional[int] = None,
        diagnostic: Any = None,
    ) -> "ExchangeOutcome":
        return cls(
            status=ExchangeStatus.POLL_FAILED,
            correlation_id=correlation_id,
            latency_ms=latency_ms,
            reason=reason,
            diagnostic=diagnostic,
        )


class EvalChannel(ABC):
    """Abstract message channel used to interact with an evaluation target.

    All methods are async so push-based channels (WebSocket, SSE) and
    poll-based channels (REST) share the same interface.
    """

    @abstractmethod
    async def health_check(self) -> ChannelHealth:
        """Verify the channel can reach the target."""
        ...

    async def exchange(
        self,
        content: str,
        *,
        timeout: float = 30.0,
        poll_interval: float = 1.0,
        on_sent: Optional[Callable[[SendResult], Awaitable[None]]] = None,
        **kwargs: Any,
    ) -> ExchangeOutcome:
        """Complete one send-and-reply exchange through this channel.

        ``on_sent`` runs after the adapter confirms a successful send and
        before reply polling begins.  A hook failure is deliberately allowed
        to propagate so callers cannot observe a reply without a durable sent
        record.
        """
        started_at = time.monotonic()
        try:
            send_result = await self._send(content, **kwargs)
        except Exception as exc:
            return ExchangeOutcome.send_failed(str(exc))
        if not send_result.ok:
            return ExchangeOutcome.send_failed(send_result.error or "channel send failed")

        correlation_id = send_result.question_msg_id
        if not correlation_id:
            raise ValueError("a successful send must provide question_msg_id")

        if on_sent is not None:
            await on_sent(send_result)

        try:
            reply = await self._poll_reply(
                correlation_id,
                timeout=timeout,
                poll_interval=poll_interval,
            )
        except Exception as exc:
            return ExchangeOutcome.poll_failed(
                str(exc),
                correlation_id=correlation_id,
                latency_ms=_elapsed_ms(started_at),
            )

        latency_ms = _elapsed_ms(started_at)
        if reply is None:
            return ExchangeOutcome.reply_timeout(
                correlation_id=correlation_id,
                latency_ms=latency_ms,
            )

        return ExchangeOutcome.succeeded(
            correlation_id=correlation_id,
            reply=reply.content,
            latency_ms=latency_ms,
            diagnostic={"send": send_result.raw_response, "reply": reply.raw_message},
        )

    @abstractmethod
    async def _send(self, content: str, **kwargs: Any) -> SendResult:
        """Transport primitive used by :meth:`exchange`."""
        ...

    @abstractmethod
    async def _poll_reply(
        self,
        question_msg_id: str,
        timeout: float = 30.0,
        poll_interval: float = 1.0,
    ) -> Optional[Reply]:
        """Transport primitive used by :meth:`exchange`."""
        ...


def _elapsed_ms(started_at: float) -> int:
    return max(0, int((time.monotonic() - started_at) * 1000))