94 lines
3.7 KiB
Python
94 lines
3.7 KiB
Python
"""Shared OpenAI-compatible model transport used by evaluation features."""
|
|
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from agenteval.services.model_configs import ModelRuntimeConfig
|
|
from agenteval.utils.llm import extract_content_from_llm_response
|
|
|
|
|
|
class ModelGatewayError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class ModelGateway:
|
|
def __init__(self, timeout: float = 60.0):
|
|
self.timeout = timeout
|
|
|
|
@staticmethod
|
|
def _headers(config: ModelRuntimeConfig) -> dict[str, str]:
|
|
headers = {"Content-Type": "application/json"}
|
|
if config.api_key:
|
|
headers["Authorization"] = f"Bearer {config.api_key}"
|
|
return headers
|
|
|
|
async def _post(self, config: ModelRuntimeConfig, payload: dict[str, Any]) -> dict[str, Any]:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
response = await client.post(
|
|
config.endpoint_url,
|
|
headers=self._headers(config),
|
|
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
|
|
|
|
async def chat(
|
|
self,
|
|
config: ModelRuntimeConfig,
|
|
messages: list[dict[str, str]],
|
|
temperature: float = 0.2,
|
|
) -> str:
|
|
payload = {
|
|
"model": config.model_name,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
}
|
|
content = extract_content_from_llm_response(await self._post(config, payload))
|
|
if not content:
|
|
raise ModelGatewayError("模型接口返回内容为空")
|
|
return content
|
|
|
|
async def embed(self, config: ModelRuntimeConfig, inputs: str | list[str]) -> list[list[float]]:
|
|
data = await self._post(config, {"model": config.model_name, "input": inputs})
|
|
try:
|
|
rows = sorted(data["data"], key=lambda item: item.get("index", 0))
|
|
return [row["embedding"] for row in rows]
|
|
except (KeyError, TypeError) as exc:
|
|
raise ModelGatewayError("Embedding 接口返回格式不正确") from exc
|
|
|
|
async def moderate(self, config: ModelRuntimeConfig, text: str) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {"input": text}
|
|
if config.model_name:
|
|
payload["model"] = config.model_name
|
|
data = await self._post(config, payload)
|
|
try:
|
|
result = data["results"][0]
|
|
except (KeyError, IndexError, TypeError) as exc:
|
|
raise ModelGatewayError("Moderation 接口返回格式不正确") from exc
|
|
if not isinstance(result, dict):
|
|
raise ModelGatewayError("Moderation 接口返回格式不正确")
|
|
return result
|
|
|
|
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 "连接成功"
|