62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
"""OpenAI-compatible chat, embedding, and moderation protocol."""
|
|
|
|
from typing import Any
|
|
|
|
from agenteval.models import ModelCapability, ModelProtocol
|
|
from agenteval.utils.llm import extract_content_from_llm_response
|
|
|
|
from .base import ModelProtocolAdapter, ProtocolAdapterError
|
|
|
|
|
|
class OpenAICompatibleAdapter(ModelProtocolAdapter):
|
|
protocol = ModelProtocol.OPENAI_COMPATIBLE
|
|
supported_capabilities = frozenset(ModelCapability)
|
|
|
|
def headers(self, api_key: str | None) -> dict[str, str]:
|
|
headers = super().headers(api_key)
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
return headers
|
|
|
|
def chat_payload(
|
|
self,
|
|
model_name: str | None,
|
|
messages: list[dict[str, str]],
|
|
temperature: float,
|
|
) -> dict[str, Any]:
|
|
return {"model": self.require_model(model_name), "messages": messages, "temperature": temperature}
|
|
|
|
def parse_chat(self, data: dict[str, Any]) -> str:
|
|
content = extract_content_from_llm_response(data)
|
|
if not content:
|
|
raise ProtocolAdapterError("模型接口返回内容为空")
|
|
return content
|
|
|
|
def embedding_payload(self, model_name: str | None, inputs: str | list[str]) -> dict[str, Any]:
|
|
return {"model": self.require_model(model_name), "input": inputs}
|
|
|
|
def parse_embeddings(self, data: dict[str, Any]) -> list[list[float]]:
|
|
try:
|
|
rows = sorted(data["data"], key=lambda item: item.get("index", 0))
|
|
vectors = [row["embedding"] for row in rows]
|
|
except (KeyError, TypeError) as exc:
|
|
raise ProtocolAdapterError("Embedding 接口返回格式不正确") from exc
|
|
if not all(isinstance(vector, list) for vector in vectors):
|
|
raise ProtocolAdapterError("Embedding 接口返回格式不正确")
|
|
return vectors
|
|
|
|
def moderation_payload(self, model_name: str | None, text: str) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {"input": text}
|
|
if model_name:
|
|
payload["model"] = model_name
|
|
return payload
|
|
|
|
def parse_moderation(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
try:
|
|
result = data["results"][0]
|
|
except (KeyError, IndexError, TypeError) as exc:
|
|
raise ProtocolAdapterError("Moderation 接口返回格式不正确") from exc
|
|
if not isinstance(result, dict):
|
|
raise ProtocolAdapterError("Moderation 接口返回格式不正确")
|
|
return result
|