136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
"""Google Gemini 适配器（ARCH §4.2/§4.4）。
|
||
|
||
- 能力：原生结构化输出（`response_mime_type=application/json` + `response_schema`）、思考。
|
||
前缀缓存走 Gemini 隐式/显式缓存，机制差异大，原型保守标 `prefix_cache=False`。
|
||
- 客户端注入（`GeminiClient` Protocol，`google.genai.Client().aio` 满足）——测试注入替身，
|
||
绝不联网；真实 SDK（`google-genai`）仅在 apps/api 构造客户端时需要，本模块不硬 import。
|
||
- 瞬时故障翻译为 `TransientProviderError`，交网关退避/回退。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import AsyncIterator
|
||
from typing import Any, Protocol
|
||
|
||
from ..errors import TransientProviderError
|
||
from ..types import LlmRequest
|
||
from .base import (
|
||
Capabilities,
|
||
ProviderResult,
|
||
ProviderUsage,
|
||
StreamChunk,
|
||
is_transient_by_name,
|
||
provider_error_summary,
|
||
)
|
||
|
||
_TRANSIENT_NAMES = frozenset(
|
||
{
|
||
"ResourceExhausted",
|
||
"ServiceUnavailable",
|
||
"DeadlineExceeded",
|
||
"InternalServerError",
|
||
"ServerError",
|
||
"APIError",
|
||
}
|
||
)
|
||
|
||
|
||
class GeminiModels(Protocol):
|
||
async def generate_content(self, **kwargs: Any) -> Any: ...
|
||
|
||
async def generate_content_stream(self, **kwargs: Any) -> Any: ...
|
||
|
||
|
||
class GeminiAio(Protocol):
|
||
models: GeminiModels
|
||
|
||
|
||
class GeminiClient(Protocol):
|
||
"""`google.genai.Client()` 满足此 Protocol（用其 `.aio.models`）。"""
|
||
|
||
aio: GeminiAio
|
||
|
||
|
||
def _is_transient(exc: Exception) -> bool:
|
||
# Gemini SDK 部分错误用 `exc.code` 携带状态码 → extra_codes=True 一并检查。
|
||
return is_transient_by_name(exc, _TRANSIENT_NAMES, extra_codes=True)
|
||
|
||
|
||
def _contents(req: LlmRequest) -> str:
|
||
if isinstance(req.input, str):
|
||
return req.input
|
||
return "\n\n".join(b.text for b in req.input)
|
||
|
||
|
||
def _system_text(req: LlmRequest) -> str:
|
||
return "\n\n".join(b.text for b in req.system)
|
||
|
||
|
||
def _usage_from(meta: Any) -> ProviderUsage:
|
||
if meta is None:
|
||
return ProviderUsage(input_tokens=0, output_tokens=0)
|
||
return ProviderUsage(
|
||
input_tokens=int(getattr(meta, "prompt_token_count", 0) or 0),
|
||
output_tokens=int(getattr(meta, "candidates_token_count", 0) or 0),
|
||
cache_read_tokens=int(getattr(meta, "cached_content_token_count", 0) or 0),
|
||
)
|
||
|
||
|
||
def _config(req: LlmRequest) -> dict[str, Any]:
|
||
cfg: dict[str, Any] = {}
|
||
system = _system_text(req)
|
||
if system:
|
||
cfg["system_instruction"] = system
|
||
if req.max_tokens is not None:
|
||
cfg["max_output_tokens"] = req.max_tokens
|
||
if req.output_schema is not None:
|
||
cfg["response_mime_type"] = "application/json"
|
||
cfg["response_schema"] = req.output_schema
|
||
return cfg
|
||
|
||
|
||
class GeminiAdapter:
|
||
def __init__(self, provider: str, client: GeminiClient) -> None:
|
||
self.provider = provider
|
||
self._client = client
|
||
|
||
def capabilities(self) -> Capabilities:
|
||
return Capabilities(structured_output=True, prefix_cache=False, thinking=True)
|
||
|
||
async def complete(self, req: LlmRequest, model: str) -> ProviderResult:
|
||
try:
|
||
resp = await self._client.aio.models.generate_content(
|
||
model=model, contents=_contents(req), config=_config(req)
|
||
)
|
||
except Exception as exc:
|
||
if _is_transient(exc):
|
||
raise TransientProviderError(
|
||
provider_error_summary(exc), provider=self.provider
|
||
) from exc
|
||
raise
|
||
text = getattr(resp, "text", "") or ""
|
||
usage = _usage_from(getattr(resp, "usage_metadata", None))
|
||
if req.output_schema is not None:
|
||
parsed = req.output_schema.model_validate_json(text)
|
||
return ProviderResult(text=text, usage=usage, parsed=parsed)
|
||
return ProviderResult(text=text, usage=usage)
|
||
|
||
async def stream(self, req: LlmRequest, model: str) -> AsyncIterator[StreamChunk]:
|
||
try:
|
||
stream = await self._client.aio.models.generate_content_stream(
|
||
model=model, contents=_contents(req), config=_config(req)
|
||
)
|
||
async for chunk in stream:
|
||
text = getattr(chunk, "text", "") or ""
|
||
if text:
|
||
yield StreamChunk(text=text)
|
||
meta = getattr(chunk, "usage_metadata", None)
|
||
if meta is not None:
|
||
yield StreamChunk(usage=_usage_from(meta))
|
||
except Exception as exc:
|
||
if _is_transient(exc):
|
||
raise TransientProviderError(
|
||
provider_error_summary(exc), provider=self.provider
|
||
) from exc
|
||
raise
|