app/adapters/gemini_api.py

"""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)