490 lines
18 KiB
Python
490 lines
18 KiB
Python
"""Gemini Provider - 原生 Google Gemini API 支持"""
|
||
|
||
import json
|
||
import logging
|
||
import time
|
||
|
||
import httpx
|
||
|
||
from agentkit.core.exceptions import LLMProviderError
|
||
from agentkit.llm.protocol import (
|
||
LLMProvider,
|
||
LLMRequest,
|
||
LLMResponse,
|
||
StreamChunk,
|
||
TokenUsage,
|
||
ToolCall,
|
||
)
|
||
from agentkit.llm.retry import (
|
||
CircuitBreaker,
|
||
CircuitBreakerConfig,
|
||
RetryConfig,
|
||
RetryPolicy,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class _GeminiStreamContext:
|
||
"""Wraps an httpx streaming response context manager for use with retry/circuit breaker."""
|
||
|
||
def __init__(self, response_ctx, response):
|
||
self._response_ctx = response_ctx
|
||
self._response = response
|
||
|
||
async def __aenter__(self):
|
||
return self._response
|
||
|
||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||
return await self._response_ctx.__aexit__(exc_type, exc_val, exc_tb)
|
||
|
||
|
||
class GeminiProvider(LLMProvider):
|
||
"""Google Gemini API 原生 Provider"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str,
|
||
model: str = "gemini-2.0-flash",
|
||
max_output_tokens: int = 4096,
|
||
base_url: str = "https://generativelanguage.googleapis.com",
|
||
timeout: float = 120.0,
|
||
safety_settings: list | None = None,
|
||
retry_config: RetryConfig | None = None,
|
||
circuit_breaker_config: CircuitBreakerConfig | None = None,
|
||
max_connections: int = 100,
|
||
max_keepalive_connections: int = 20,
|
||
keepalive_expiry: float = 30.0,
|
||
):
|
||
self._api_key = api_key
|
||
self._model = model
|
||
self._max_output_tokens = max_output_tokens
|
||
self._base_url = base_url.rstrip("/")
|
||
self._timeout = timeout
|
||
self._safety_settings = safety_settings
|
||
self._limits = httpx.Limits(
|
||
max_connections=max_connections,
|
||
max_keepalive_connections=max_keepalive_connections,
|
||
keepalive_expiry=keepalive_expiry,
|
||
)
|
||
self._client: httpx.AsyncClient | None = None
|
||
self._retry_policy = RetryPolicy(retry_config) if retry_config else None
|
||
self._circuit_breaker = (
|
||
CircuitBreaker(circuit_breaker_config, provider="gemini")
|
||
if circuit_breaker_config
|
||
else None
|
||
)
|
||
|
||
def _get_client(self) -> httpx.AsyncClient:
|
||
"""Lazy client initialization"""
|
||
if self._client is None:
|
||
self._client = httpx.AsyncClient(timeout=self._timeout, limits=self._limits)
|
||
return self._client
|
||
|
||
async def close(self) -> None:
|
||
"""关闭 HTTP 客户端连接池"""
|
||
if self._client is not None:
|
||
await self._client.aclose()
|
||
self._client = None
|
||
|
||
def _convert_messages(
|
||
self, messages: list[dict[str, str]]
|
||
) -> tuple[dict[str, object] | None, list[dict[str, object]]]:
|
||
"""将 OpenAI 风格消息转换为 Gemini 格式
|
||
|
||
Returns:
|
||
(system_instruction, contents)
|
||
"""
|
||
system_instruction: dict[str, object] | None = None
|
||
contents: list[dict[str, object]] = []
|
||
|
||
for msg in messages:
|
||
role = msg.get("role", "")
|
||
content = msg.get("content", "")
|
||
|
||
if role == "system":
|
||
system_instruction = {"parts": [{"text": content}]}
|
||
continue
|
||
|
||
if role == "user":
|
||
# Check if this is a tool result message
|
||
if msg.get("tool_call_id"):
|
||
# Tool response: role="user" with functionResponse part
|
||
tool_name = msg.get("name", "")
|
||
# If name not at top level, try to extract from content
|
||
if not tool_name and isinstance(content, str):
|
||
try:
|
||
parsed = json.loads(content)
|
||
tool_name = parsed.get("name", "")
|
||
except (json.JSONDecodeError, AttributeError):
|
||
pass
|
||
contents.append(
|
||
{
|
||
"role": "user",
|
||
"parts": [
|
||
{
|
||
"functionResponse": {
|
||
"name": tool_name,
|
||
"response": {
|
||
"content": content,
|
||
},
|
||
},
|
||
}
|
||
],
|
||
}
|
||
)
|
||
else:
|
||
contents.append(
|
||
{
|
||
"role": "user",
|
||
"parts": [{"text": content}],
|
||
}
|
||
)
|
||
continue
|
||
|
||
if role == "assistant":
|
||
tool_calls = msg.get("tool_calls")
|
||
if tool_calls:
|
||
parts: list[dict[str, object]] = []
|
||
if content:
|
||
parts.append({"text": content})
|
||
for tc in tool_calls:
|
||
func = tc.get("function", {})
|
||
arguments = func.get("arguments", "{}")
|
||
if isinstance(arguments, str):
|
||
try:
|
||
arguments = json.loads(arguments)
|
||
except json.JSONDecodeError:
|
||
arguments = {"raw": arguments}
|
||
parts.append(
|
||
{
|
||
"functionCall": {
|
||
"name": func.get("name", ""),
|
||
"args": arguments,
|
||
},
|
||
}
|
||
)
|
||
contents.append({"role": "model", "parts": parts})
|
||
else:
|
||
contents.append(
|
||
{
|
||
"role": "model",
|
||
"parts": [{"text": content}],
|
||
}
|
||
)
|
||
continue
|
||
|
||
if role == "tool":
|
||
# OpenAI format: {"role": "tool", "tool_call_id": "...", "content": "..."}
|
||
tool_name = msg.get("name", "")
|
||
tool_content = msg.get("content", "")
|
||
contents.append(
|
||
{
|
||
"role": "user",
|
||
"parts": [
|
||
{
|
||
"functionResponse": {
|
||
"name": tool_name,
|
||
"response": {
|
||
"content": tool_content,
|
||
},
|
||
},
|
||
}
|
||
],
|
||
}
|
||
)
|
||
|
||
return system_instruction, contents
|
||
|
||
def _convert_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]:
|
||
"""将 OpenAI function 格式转换为 Gemini functionDeclarations"""
|
||
declarations = []
|
||
for tool in tools:
|
||
if tool.get("type") == "function":
|
||
func = tool.get("function", {})
|
||
declarations.append(
|
||
{
|
||
"name": func.get("name", ""),
|
||
"description": func.get("description", ""),
|
||
"parameters": func.get("parameters", {"type": "object", "properties": {}}),
|
||
}
|
||
)
|
||
if not declarations:
|
||
return []
|
||
return [{"functionDeclarations": declarations}]
|
||
|
||
def _convert_tool_choice(self, tool_choice: str) -> dict[str, object] | None:
|
||
"""将 OpenAI tool_choice 格式转换为 Gemini toolConfig"""
|
||
if tool_choice == "auto":
|
||
return {"functionCallingConfig": {"mode": "AUTO"}}
|
||
elif tool_choice == "required":
|
||
return {"functionCallingConfig": {"mode": "ANY"}}
|
||
elif tool_choice and tool_choice not in ("none",):
|
||
return {"functionCallingConfig": {"mode": "AUTO"}}
|
||
if tool_choice == "none":
|
||
return {"functionCallingConfig": {"mode": "NONE"}}
|
||
return None
|
||
|
||
def _parse_response(self, data: dict[str, object], model: str) -> LLMResponse:
|
||
"""将 Gemini 响应转换为 LLMResponse"""
|
||
candidates = data.get("candidates", [])
|
||
text_parts: list[str] = []
|
||
tool_calls: list[ToolCall] = []
|
||
tool_call_index = 0
|
||
|
||
if candidates:
|
||
content = candidates[0].get("content", {})
|
||
parts = content.get("parts", [])
|
||
for part in parts:
|
||
if "text" in part:
|
||
text_parts.append(part["text"])
|
||
elif "functionCall" in part:
|
||
fc = part["functionCall"]
|
||
tool_calls.append(
|
||
ToolCall(
|
||
id=f"call_{tool_call_index}",
|
||
name=fc.get("name", ""),
|
||
arguments=fc.get("args", {}),
|
||
)
|
||
)
|
||
tool_call_index += 1
|
||
|
||
usage_metadata = data.get("usageMetadata", {})
|
||
usage = TokenUsage(
|
||
prompt_tokens=usage_metadata.get("promptTokenCount", 0),
|
||
completion_tokens=usage_metadata.get("candidatesTokenCount", 0),
|
||
)
|
||
|
||
return LLMResponse(
|
||
content="".join(text_parts),
|
||
model=data.get("modelVersion", model),
|
||
usage=usage,
|
||
tool_calls=tool_calls,
|
||
)
|
||
|
||
def _handle_error(self, status_code: int, resp_body: bytes) -> None:
|
||
"""处理 Gemini API 错误响应"""
|
||
try:
|
||
error_data = json.loads(resp_body)
|
||
error_info = error_data.get("error", {})
|
||
error_msg = error_info.get("message", f"HTTP {status_code}")
|
||
except (json.JSONDecodeError, AttributeError):
|
||
error_msg = f"HTTP {status_code}"
|
||
|
||
raise LLMProviderError("gemini", f"HTTP {status_code}: {error_msg}")
|
||
|
||
async def chat(self, request: LLMRequest) -> LLMResponse:
|
||
"""发送 chat 请求（带 retry + circuit breaker）"""
|
||
if self._circuit_breaker and self._retry_policy:
|
||
return await self._circuit_breaker.execute(
|
||
self._retry_policy.execute, self._chat_impl, request
|
||
)
|
||
if self._retry_policy:
|
||
return await self._retry_policy.execute(self._chat_impl, request)
|
||
if self._circuit_breaker:
|
||
return await self._circuit_breaker.execute(self._chat_impl, request)
|
||
return await self._chat_impl(request)
|
||
|
||
async def _chat_impl(self, request: LLMRequest) -> LLMResponse:
|
||
client = self._get_client()
|
||
model = request.model or self._model
|
||
url = f"{self._base_url}/v1beta/models/{model}:generateContent?key={self._api_key}"
|
||
|
||
system_instruction, contents = self._convert_messages(request.messages)
|
||
|
||
payload: dict[str, object] = {
|
||
"contents": contents,
|
||
"generationConfig": {
|
||
"temperature": request.temperature,
|
||
"maxOutputTokens": request.max_tokens or self._max_output_tokens,
|
||
},
|
||
}
|
||
|
||
if system_instruction is not None:
|
||
payload["systemInstruction"] = system_instruction
|
||
|
||
if request.tools:
|
||
gemini_tools = self._convert_tools(request.tools)
|
||
if gemini_tools:
|
||
payload["tools"] = gemini_tools
|
||
tool_config = self._convert_tool_choice(request.tool_choice)
|
||
if tool_config is not None:
|
||
payload["toolConfig"] = tool_config
|
||
|
||
if self._safety_settings:
|
||
payload["safetySettings"] = self._safety_settings
|
||
|
||
start = time.monotonic()
|
||
|
||
try:
|
||
resp = await client.post(url, json=payload)
|
||
except httpx.HTTPError as e:
|
||
raise LLMProviderError("gemini", str(e)) from e
|
||
|
||
latency_ms = (time.monotonic() - start) * 1000
|
||
|
||
if resp.status_code != 200:
|
||
self._handle_error(resp.status_code, resp.content)
|
||
|
||
data = resp.json()
|
||
response = self._parse_response(data, model)
|
||
response.latency_ms = latency_ms
|
||
|
||
return response
|
||
|
||
async def chat_stream(self, request: LLMRequest):
|
||
"""Stream chat response using SSE（带 retry + circuit breaker）"""
|
||
if self._circuit_breaker and self._retry_policy:
|
||
ctx = await self._circuit_breaker.execute(
|
||
self._retry_policy.execute, self._open_stream, request
|
||
)
|
||
elif self._retry_policy:
|
||
ctx = await self._retry_policy.execute(self._open_stream, request)
|
||
elif self._circuit_breaker:
|
||
ctx = await self._circuit_breaker.execute(self._open_stream, request)
|
||
else:
|
||
ctx = await self._open_stream(request)
|
||
|
||
async with ctx as response:
|
||
async for chunk in self._iterate_stream(response, request):
|
||
yield chunk
|
||
|
||
async def _open_stream(self, request: LLMRequest):
|
||
"""Open the streaming HTTP connection; returns an async context manager."""
|
||
client = self._get_client()
|
||
model = request.model or self._model
|
||
url = f"{self._base_url}/v1beta/models/{model}:streamGenerateContent?key={self._api_key}&alt=sse"
|
||
|
||
system_instruction, contents = self._convert_messages(request.messages)
|
||
|
||
payload: dict[str, object] = {
|
||
"contents": contents,
|
||
"generationConfig": {
|
||
"temperature": request.temperature,
|
||
"maxOutputTokens": request.max_tokens or self._max_output_tokens,
|
||
},
|
||
}
|
||
|
||
if system_instruction is not None:
|
||
payload["systemInstruction"] = system_instruction
|
||
|
||
if request.tools:
|
||
gemini_tools = self._convert_tools(request.tools)
|
||
if gemini_tools:
|
||
payload["tools"] = gemini_tools
|
||
tool_config = self._convert_tool_choice(request.tool_choice)
|
||
if tool_config is not None:
|
||
payload["toolConfig"] = tool_config
|
||
|
||
if self._safety_settings:
|
||
payload["safetySettings"] = self._safety_settings
|
||
|
||
response_ctx = client.stream("POST", url, json=payload)
|
||
response = await response_ctx.__aenter__()
|
||
|
||
if response.status_code != 200:
|
||
error_body = await response.aread()
|
||
await response_ctx.__aexit__(None, None, None)
|
||
self._handle_error(response.status_code, error_body)
|
||
|
||
return _GeminiStreamContext(response_ctx, response)
|
||
|
||
async def _iterate_stream(self, response, request: LLMRequest):
|
||
"""Iterate over an already-open SSE stream and yield StreamChunks."""
|
||
accumulated_tool_calls: list[dict[str, object]] = []
|
||
model = request.model or self._model
|
||
|
||
async for line in response.aiter_lines():
|
||
line = line.strip()
|
||
if not line or not line.startswith("data: "):
|
||
continue
|
||
|
||
data_str = line[6:]
|
||
try:
|
||
data = json.loads(data_str)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
candidates = data.get("candidates", [])
|
||
if not candidates:
|
||
# Usage-only chunk
|
||
usage_metadata = data.get("usageMetadata")
|
||
if usage_metadata:
|
||
usage = TokenUsage(
|
||
prompt_tokens=usage_metadata.get("promptTokenCount", 0),
|
||
completion_tokens=usage_metadata.get("candidatesTokenCount", 0),
|
||
)
|
||
if accumulated_tool_calls:
|
||
tool_calls = [
|
||
ToolCall(
|
||
id=tc["id"],
|
||
name=tc["name"],
|
||
arguments=tc["arguments"],
|
||
)
|
||
for tc in accumulated_tool_calls
|
||
]
|
||
yield StreamChunk(
|
||
content="",
|
||
model=data.get("modelVersion", model),
|
||
tool_calls=tool_calls,
|
||
usage=usage,
|
||
is_final=True,
|
||
)
|
||
accumulated_tool_calls = []
|
||
else:
|
||
yield StreamChunk(
|
||
content="",
|
||
model=data.get("modelVersion", model),
|
||
usage=usage,
|
||
is_final=True,
|
||
)
|
||
continue
|
||
|
||
content = candidates[0].get("content", {})
|
||
parts = content.get("parts", [])
|
||
|
||
for part in parts:
|
||
if "text" in part:
|
||
text = part["text"]
|
||
if text:
|
||
yield StreamChunk(
|
||
content=text,
|
||
model=data.get("modelVersion", model),
|
||
)
|
||
elif "functionCall" in part:
|
||
fc = part["functionCall"]
|
||
accumulated_tool_calls.append(
|
||
{
|
||
"id": f"call_{len(accumulated_tool_calls)}",
|
||
"name": fc.get("name", ""),
|
||
"arguments": fc.get("args", {}),
|
||
}
|
||
)
|
||
|
||
# Check for finish reason
|
||
finish_reason = candidates[0].get("finishReason", "")
|
||
if finish_reason in ("STOP", "MAX_TOKENS") and accumulated_tool_calls:
|
||
tool_calls = [
|
||
ToolCall(
|
||
id=tc["id"],
|
||
name=tc["name"],
|
||
arguments=tc["arguments"],
|
||
)
|
||
for tc in accumulated_tool_calls
|
||
]
|
||
yield StreamChunk(
|
||
content="",
|
||
model=data.get("modelVersion", model),
|
||
tool_calls=tool_calls,
|
||
is_final=True,
|
||
)
|
||
accumulated_tool_calls = []
|
||
|
||
def get_model_info(self) -> dict[str, object]:
|
||
"""返回 Provider 和模型信息"""
|
||
return {
|
||
"provider": "gemini",
|
||
"model": self._model,
|
||
"max_output_tokens": self._max_output_tokens,
|
||
}
|