88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
"""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
|
|
|
|
async def _post(self, config: ModelRuntimeConfig, payload: dict[str, Any]) -> dict[str, Any]:
|
|
adapter = self._adapter(config)
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout, transport=self.transport) as 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:
|
|
adapter = self._adapter(config)
|
|
try:
|
|
payload = adapter.chat_payload(config.model_name, messages, temperature)
|
|
return adapter.parse_chat(await self._post(config, payload))
|
|
except ProtocolAdapterError as exc:
|
|
raise ModelGatewayError(str(exc)) from exc
|
|
|
|
async def embed(self, config: ModelRuntimeConfig, inputs: str | list[str]) -> list[list[float]]:
|
|
adapter = self._adapter(config)
|
|
try:
|
|
payload = adapter.embedding_payload(config.model_name, inputs)
|
|
return adapter.parse_embeddings(await self._post(config, payload))
|
|
except ProtocolAdapterError as exc:
|
|
raise ModelGatewayError(str(exc)) from exc
|
|
|
|
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 "连接成功"
|