57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
"""Google Gemini generateContent REST protocol."""
|
|
|
|
from typing import Any
|
|
|
|
from agenteval.models import ModelCapability, ModelProtocol
|
|
|
|
from .base import ModelProtocolAdapter, ProtocolAdapterError
|
|
|
|
|
|
class GoogleGeminiAdapter(ModelProtocolAdapter):
|
|
protocol = ModelProtocol.GOOGLE_GEMINI
|
|
supported_capabilities = frozenset({ModelCapability.CHAT})
|
|
|
|
def headers(self, api_key: str | None) -> dict[str, str]:
|
|
headers = super().headers(api_key)
|
|
if api_key:
|
|
headers["x-goog-api-key"] = api_key
|
|
return headers
|
|
|
|
def chat_payload(
|
|
self,
|
|
model_name: str | None,
|
|
messages: list[dict[str, str]],
|
|
temperature: float,
|
|
) -> dict[str, Any]:
|
|
self.require_model(model_name)
|
|
system_parts: list[dict[str, str]] = []
|
|
contents: list[dict[str, Any]] = []
|
|
for message in messages:
|
|
role = message.get("role", "user")
|
|
part = {"text": message.get("content", "")}
|
|
if role in {"system", "developer"}:
|
|
system_parts.append(part)
|
|
else:
|
|
contents.append({"role": "model" if role == "assistant" else "user", "parts": [part]})
|
|
payload: dict[str, Any] = {
|
|
"contents": contents,
|
|
"generationConfig": {"temperature": temperature},
|
|
}
|
|
if system_parts:
|
|
payload["systemInstruction"] = {"parts": system_parts}
|
|
return payload
|
|
|
|
def parse_chat(self, data: dict[str, Any]) -> str:
|
|
try:
|
|
parts = data["candidates"][0]["content"]["parts"]
|
|
except (KeyError, IndexError, TypeError) as exc:
|
|
raise ProtocolAdapterError("Gemini 接口返回格式不正确") from exc
|
|
content = "\n".join(
|
|
str(part.get("text", ""))
|
|
for part in parts
|
|
if isinstance(part, dict) and part.get("text")
|
|
)
|
|
if not content:
|
|
raise ProtocolAdapterError("模型接口返回内容为空")
|
|
return content
|