AgentEvalTool/backend/agenteval/channels/tutu.py

"""Tutu API message channel implementation."""

import asyncio
import json
import time
from typing import Any, Optional

import httpx

from agenteval.channels.base import ChannelHealth, EvalChannel, Reply, SendResult


class TutuApiChannel(EvalChannel):
    """Message channel backed by the Tutu chat API.

    Uses a single ``httpx.AsyncClient`` per instance to reuse TCP connections
    across the many small poll requests during an evaluation run.
    """

    def __init__(self, config: dict[str, Any]):
        self.base_url = config["base_url"].rstrip("/")
        self.token = config["token"]
        self.tenant = config["tenant"]
        self.chat_channel_id = config["chat_channel_id"]
        self.chat_contact_id = config["chat_contact_id"]
        self.sender_type = config.get("sender_type", "WORK_WE_CUSTOMER")
        self.contact_type = config.get("contact_type", "EXTERNAL")
        self._client: Optional[httpx.AsyncClient] = None

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(timeout=15.0)
        return self._client

    async def close(self) -> None:
        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    def _build_url(self, path: str) -> str:
        return f"{self.base_url}/api/{self.tenant}/{path}"

    def _build_headers(self, accept: str = "*/*") -> dict[str, str]:
        return {
            "accept": accept,
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }

    async def health_check(self) -> ChannelHealth:
        """Send a lightweight request to verify connectivity."""
        client = await self._get_client()
        try:
            url = self._build_url("v1/chat/message")
            params = {
                "chatChannelId": self.chat_channel_id,
                "chatContactId": self.chat_contact_id,
                "page": 0,
                "size": 1,
            }
            resp = await client.get(url, headers=self._build_headers(), params=params)
            if resp.status_code == 200:
                return ChannelHealth(ok=True, message="通道正常")
            return ChannelHealth(ok=False, message=f"HTTP {resp.status_code}: {resp.text[:200]}")
        except Exception as exc:
            return ChannelHealth(ok=False, message=f"请求异常: {exc}")

    async def send(self, content: str, **kwargs: Any) -> SendResult:
        """Send a text message to the configured chat contact."""
        payload = {
            "chatChannelId": self.chat_channel_id,
            "chatContactType": self.contact_type,
            "chatContactId": self.chat_contact_id,
            "msgType": kwargs.get("msg_type", "text"),
            "msgBody": json.dumps({"content": content}, ensure_ascii=False),
            "actualSenderType": self.sender_type,
            "sender": {"type": self.contact_type},
        }
        client = await self._get_client()
        try:
            url = self._build_url("v1/chat/message/sendMsg")
            resp = await client.post(url, headers=self._build_headers(), json=payload)
            if resp.status_code != 200:
                return SendResult(ok=False, error=f"HTTP {resp.status_code}: {resp.text[:500]}")

            data = resp.json()
            # The reply references the sent message via metadata.questionMsgId == sent msgId.
            question_msg_id = data.get("msgId")
            return SendResult(ok=True, question_msg_id=question_msg_id, raw_response=data)
        except Exception as exc:
            return SendResult(ok=False, error=f"发送异常: {exc}")

    async def poll_reply(
        self,
        question_msg_id: str,
        timeout: float = 30.0,
        poll_interval: float = 1.0,
    ) -> Optional[Reply]:
        """Poll chat history until a reply matching the questionMsgId arrives."""
        deadline = time.time() + timeout
        seen_ids: set[str] = set()
        client = await self._get_client()

        while time.time() < deadline:
            try:
                url = self._build_url("v1/chat/message")
                params = {
                    "chatChannelId": self.chat_channel_id,
                    "chatContactId": self.chat_contact_id,
                    "page": 0,
                    "size": 20,
                }
                resp = await client.get(url, headers=self._build_headers(), params=params)
                if resp.status_code == 200:
                    data = resp.json()
                    records = data.get("data", []) if isinstance(data, dict) else []
                    for msg in records:
                        msg_id = msg.get("id") or msg.get("msgId")
                        if not msg_id or msg_id in seen_ids:
                            continue
                        seen_ids.add(msg_id)

                        meta = msg.get("metadata", {})
                        if meta.get("questionMsgId") == question_msg_id:
                            return Reply(
                                question_msg_id=question_msg_id,
                                content=msg.get("msgBody"),
                                sender_name=msg.get("senderName") or msg.get("actualSenderName"),
                                msg_time=msg.get("msgTime"),
                                raw_message=msg,
                            )
            except Exception:
                pass

            await asyncio.sleep(poll_interval)

        return None