AgentEvalTool/backend/agenteval/model_gateway.py

"""Shared multi-protocol model transport used by evaluation features."""

from typing import Any

import httpx

from agenteval.model_protocols import ProtocolAdapterError, get_protocol_adapter
from agenteval.services.model_configs import ModelRuntimeConfig


class ModelGatewayError(RuntimeError):
    pass


class ModelGateway:
    def __init__(self, timeout: float = 60.0, transport: httpx.AsyncBaseTransport | None = None):
        self.timeout = timeout
        self.transport = transport
        # 评测侧 LLM 调用的累计 token 用量（引擎结束时写入 run summary）
        self.total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        self._client: httpx.AsyncClient | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        # 单实例复用客户端，省去每次 LLM 调用的 TCP/TLS 握手
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(timeout=self.timeout, transport=self.transport)
        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

    async def _post(self, config: ModelRuntimeConfig, payload: dict[str, Any]) -> dict[str, Any]:
        adapter = self._adapter(config)
        try:
            client = await self._get_client()
            response = await client.post(
                config.endpoint_url,
                headers=adapter.headers(config.api_key),
                json=payload,
            )
            response.raise_for_status()
            data = response.json()
            if not isinstance(data, dict):
                raise ModelGatewayError("模型接口返回格式不是 JSON 对象")
            return data
        except ModelGatewayError:
            raise
        except httpx.HTTPStatusError as exc:
            detail = exc.response.text[:300]
            raise ModelGatewayError(f"模型接口返回 HTTP {exc.response.status_code}: {detail}") from exc
        except Exception as exc:
            raise ModelGatewayError(f"模型接口调用失败: {exc}") from exc

    @staticmethod
    def _adapter(config: ModelRuntimeConfig):
        try:
            return get_protocol_adapter(config.provider)
        except ProtocolAdapterError as exc:
            raise ModelGatewayError(str(exc)) from exc

    async def chat(
        self,
        config: ModelRuntimeConfig,
        messages: list[dict[str, str]],
        temperature: float = 0.2,
    ) -> str:
        content, _ = await self.chat_with_usage(config, messages, temperature)
        return content

    async def chat_with_usage(
        self,
        config: ModelRuntimeConfig,
        messages: list[dict[str, str]],
        temperature: float = 0.2,
    ) -> tuple[str, dict[str, int] | None]:
        """Like chat(), but also returns this call's token usage (None if omitted)."""
        adapter = self._adapter(config)
        try:
            payload = adapter.chat_payload(config.model_name, messages, temperature)
            data = await self._post(config, payload)
            usage = adapter.parse_usage(data)
            self._accumulate(usage)
            return adapter.parse_chat(data), usage
        except ProtocolAdapterError as exc:
            raise ModelGatewayError(str(exc)) from exc

    async def embed(self, config: ModelRuntimeConfig, inputs: str | list[str]) -> list[list[float]]:
        vectors, _ = await self.embed_with_usage(config, inputs)
        return vectors

    async def embed_with_usage(
        self, config: ModelRuntimeConfig, inputs: str | list[str]
    ) -> tuple[list[list[float]], dict[str, int] | None]:
        adapter = self._adapter(config)
        try:
            payload = adapter.embedding_payload(config.model_name, inputs)
            data = await self._post(config, payload)
            usage = adapter.parse_usage(data)
            self._accumulate(usage)
            return adapter.parse_embeddings(data), usage
        except ProtocolAdapterError as exc:
            raise ModelGatewayError(str(exc)) from exc

    def _accumulate(self, usage: dict[str, int] | None) -> None:
        if not usage:
            return
        for key in self.total_usage:
            self.total_usage[key] += int(usage.get(key) or 0)

    async def moderate(self, config: ModelRuntimeConfig, text: str) -> dict[str, Any]:
        adapter = self._adapter(config)
        try:
            payload = adapter.moderation_payload(config.model_name, text)
            return adapter.parse_moderation(await self._post(config, payload))
        except ProtocolAdapterError as exc:
            raise ModelGatewayError(str(exc)) from exc

    async def test_connection(self, config: ModelRuntimeConfig) -> str:
        if config.capability.value == "chat":
            await self.chat(config, [{"role": "user", "content": "回复 OK"}], temperature=0)
        elif config.capability.value == "embedding":
            vectors = await self.embed(config, "connection test")
            if not vectors or not vectors[0]:
                raise ModelGatewayError("Embedding 接口未返回向量")
        else:
            await self.moderate(config, "connection test")
        return "连接成功"