"""Gemini 适配器入口:拼装上游请求、调用、超时、错误处理。""" from __future__ import annotations import json import re from typing import Any, AsyncIterator, Optional import httpx from ..config import ( GEMINI_API_CLIENT, GEMINI_API_VERSION, GEMINI_BASE_URL, GEMINI_DEFAULT_EMBEDDINGS_MODEL, get_settings, ) from ..errors import HttpError, error_code_for_status, is_retryable_status from ..http_client import get_http_client from ..utils.sse import parse_sse_stream from .gemini.config import build_generation_config from .gemini.messages import convert_messages_to_contents from .gemini.response import ( convert_embeddings_response, convert_models_response, convert_response, ) from .gemini.stream import stream_openai_chunks from .gemini.tools import convert_tools def _build_headers(api_key: str) -> dict[str, str]: return { "x-goog-api-client": GEMINI_API_CLIENT, "x-goog-api-key": api_key, "Content-Type": "application/json", } def _strip_model_suffix(model: str) -> tuple[str, bool, bool]: """剥离模型名后缀(:search / -search-preview),返回 (clean_model, use_search, _removed)。""" use_search = False m = model if m.endswith(":search"): use_search = True m = m[: -len(":search")] elif m.endswith("-search-preview"): use_search = True m = m[: -len("-search-preview")] return m, use_search, False def _build_body( *, messages: list[dict], body: dict, model: str, ) -> dict[str, Any]: """构造 Gemini generateContent body。""" contents, system_text = convert_messages_to_contents(messages) gemini_body: dict[str, Any] = {"contents": contents} if system_text: gemini_body["system_instruction"] = {"parts": [{"text": system_text}]} # 处理 google extra_body extra_google = None extra_body = body.get("extra_body") or body.get("google") if isinstance(extra_body, dict): extra_google = extra_body.get("google") if "google" in extra_body else extra_body gen_config, safety_settings = build_generation_config(body, extra_google) if gen_config: gemini_body["generationConfig"] = gen_config if safety_settings: gemini_body["safetySettings"] = safety_settings # cached content cached = body.get("_cached_content") or (extra_google or {}).get("cached_content") if extra_google else None if cached: gemini_body["cachedContent"] = cached # tools tools = convert_tools(body) # 模型后缀 :search -> 自动追加 googleSearch 工具 clean_model, use_search, _ = _strip_model_suffix(model) if use_search: search_tool = {"google_search": {}} if False else {"googleSearch": {}} if tools: tools.append(search_tool) else: tools = [search_tool] if tools: gemini_body["tools"] = tools return gemini_body, clean_model async def chat_completions( *, api_key: str, model: str, messages: list[dict], body: dict, stream: bool = False, ) -> Any: """调用 Gemini,返回 OpenAI 格式结果或 chunk 流。""" gemini_body, clean_model = _build_body(messages=messages, body=body, model=model) if stream: return _stream_chat(api_key, clean_model, gemini_body, body) url = ( f"{GEMINI_BASE_URL}/{GEMINI_API_VERSION}/models/{clean_model}:generateContent" ) client = get_http_client() try: resp = await client.post( url, headers=_build_headers(api_key), json=gemini_body, ) except httpx.TimeoutException as e: raise HttpError( f"Gemini upstream timeout: {e}", status=504, code="upstream_timeout", ) from e except httpx.HTTPError as e: raise HttpError( f"Gemini upstream error: {e}", status=502, code="bad_gateway", ) from e if resp.status_code >= 400: await _raise_upstream_error(resp) try: data = resp.json() except Exception as e: raise HttpError( f"Gemini upstream returned non-JSON: {resp.text[:200]}", status=502, code="bad_gateway", ) from e return convert_response(data, model=clean_model) async def _stream_chat( api_key: str, clean_model: str, gemini_body: dict, body: dict, ) -> AsyncIterator[dict]: """流式调用 Gemini,yield OpenAI chunk dict。""" url = ( f"{GEMINI_BASE_URL}/{GEMINI_API_VERSION}/models/{clean_model}:streamGenerateContent" "?alt=sse" ) client = get_http_client() include_usage = False so = body.get("stream_options") if isinstance(so, dict) and so.get("include_usage"): include_usage = True try: async with client.stream( "POST", url, headers=_build_headers(api_key), json=gemini_body, ) as resp: if resp.status_code >= 400: text = await resp.aread() await _raise_upstream_error_from_raw(resp.status_code, text) sse_iter = parse_sse_stream(resp.aiter_bytes()) async for chunk in stream_openai_chunks( sse_iter, model=clean_model, include_usage=include_usage ): yield chunk except httpx.TimeoutException as e: raise HttpError( f"Gemini stream timeout: {e}", status=504, code="upstream_timeout", ) from e except HttpError: raise except httpx.HTTPError as e: raise HttpError( f"Gemini stream error: {e}", status=502, code="bad_gateway", ) from e async def embeddings( *, api_key: str, model: str, input_data: Any, ) -> dict: """调用 Gemini batchEmbedContents,返回 OpenAI embeddings 格式。""" # 归一化 input 为 list if isinstance(input_data, str): inputs = [input_data] elif isinstance(input_data, list): inputs = [str(x) for x in input_data] else: inputs = [str(input_data)] clean_model = model or GEMINI_DEFAULT_EMBEDDINGS_MODEL url = f"{GEMINI_BASE_URL}/{GEMINI_API_VERSION}/models/{clean_model}:batchEmbedContents" body = {"requests": [{"model": f"models/{clean_model}", "content": {"parts": [{"text": t}]}} for t in inputs]} client = get_http_client() try: resp = await client.post(url, headers=_build_headers(api_key), json=body) except httpx.TimeoutException as e: raise HttpError(f"Gemini embeddings timeout: {e}", status=504, code="upstream_timeout") from e except httpx.HTTPError as e: raise HttpError(f"Gemini embeddings error: {e}", status=502, code="bad_gateway") from e if resp.status_code >= 400: await _raise_upstream_error(resp) try: data = resp.json() except Exception as e: raise HttpError("Gemini embeddings non-JSON", status=502, code="bad_gateway") from e return convert_embeddings_response(data, model=clean_model) async def list_models(*, api_key: str) -> dict: url = f"{GEMINI_BASE_URL}/{GEMINI_API_VERSION}/models" client = get_http_client() try: resp = await client.get(url, headers=_build_headers(api_key), params={"pageSize": "200"}) except httpx.HTTPError as e: raise HttpError(f"Gemini list models error: {e}", status=502, code="bad_gateway") from e if resp.status_code >= 400: await _raise_upstream_error(resp) try: data = resp.json() except Exception as e: raise HttpError("Gemini list models non-JSON", status=502, code="bad_gateway") from e return convert_models_response(data) async def list_models_raw(*, api_key: str) -> list[dict]: """原始模型列表(供后台拉取模型用)。""" url = f"{GEMINI_BASE_URL}/{GEMINI_API_VERSION}/models" client = get_http_client() try: resp = await client.get(url, headers=_build_headers(api_key), params={"pageSize": "200"}) except httpx.HTTPError as e: raise HttpError(f"Gemini list models error: {e}", status=502, code="bad_gateway") from e if resp.status_code >= 400: await _raise_upstream_error(resp) try: data = resp.json() except Exception as e: raise HttpError("Gemini list models non-JSON", status=502, code="bad_gateway") from e out = [] for m in data.get("models") or []: name = (m.get("name") or "").removeprefix("models/") if name: out.append({"id": name, "name": m.get("displayName") or name}) return out async def _raise_upstream_error(resp: httpx.Response) -> None: text = resp.text try: body = resp.json() except Exception: body = None upstream_err = None if isinstance(body, dict): upstream_err = body.get("error") if isinstance(body.get("error"), dict) else body code = (upstream_err or {}).get("code") if isinstance(upstream_err, dict) else None code = code or error_code_for_status(resp.status_code, "upstream_error") message = (upstream_err or {}).get("message") if isinstance(upstream_err, dict) else None message = message or text or f"Upstream error ({resp.status_code})" raise HttpError( message, status=resp.status_code, code=code, upstream=upstream_err, ) async def _raise_upstream_error_from_raw(status: int, raw: bytes) -> None: text = raw.decode("utf-8", errors="replace") try: body = json.loads(text) except Exception: body = None upstream_err = None if isinstance(body, dict): upstream_err = body.get("error") if isinstance(body.get("error"), dict) else body code = (upstream_err or {}).get("code") if isinstance(upstream_err, dict) else None code = code or error_code_for_status(status, "upstream_error") message = (upstream_err or {}).get("message") if isinstance(upstream_err, dict) else None message = message or text or f"Upstream error ({status})" raise HttpError(message, status=status, code=code, upstream=upstream_err)