AI_design_agent/scr/gemini_client_request.py

import base64
import io
import json
import os
from typing import List, Optional
from ..logger_setup import logger
from ..upload_tos import upload_image
import requests
from PIL import Image


DEFAULT_API_BASE_URL = "https://api.openaius.com"

# GOOGLE_API_KEY
class GeminiAPIError(RuntimeError):
    pass


def _plugin_dir() -> str:
    return os.path.dirname(__file__)



def _get_api_key() -> str:
    # 1) Environment variables override
    key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
    if key:
        return key


    raise GeminiAPIError(
        "Google Gemini API key not found. Provide it via environment variable 'GOOGLE_API_KEY' (or 'GEMINI_API_KEY'),"
        " config file 'gemini_config.json' (field 'api_key'), or a 'gemini_api_key.txt' file in the plugin directory."
    )


def _get_base_url() -> str:
    # 2) Base URL via env
    base = os.environ.get("GOOGLE_API_BASE_URL") or os.environ.get("GEMINI_API_BASE_URL")
    if base:
        return base
    return DEFAULT_API_BASE_URL


def _build_full_endpoint(model: str) -> str:

    base = _get_base_url().rstrip("/")
    return f"{base}/v1beta/models/{model}:generateContent"


def _apply_auth(headers: dict, params: dict, api_key: str) -> None:
    """Apply authentication to headers or query according to env/config.

    Priority:
    - Env vars (header)
    - Config (auth_header_name/auth_header_value_template)
    - Query param (env GEMINI_QUERY_PARAM_NAME) or config (query_param_name)
    - Default query param 'key'
    """
    cfg = {}

    # Extra headers from config
    extra_headers = cfg.get("extra_headers") or {}
    if isinstance(extra_headers, dict):
        headers.update({str(k): str(v) for k, v in extra_headers.items()})

    # Auth via header (env)
    env_hdr_name = os.environ.get("GEMINI_AUTH_HEADER_NAME")
    env_hdr_value = os.environ.get("GEMINI_AUTH_HEADER_VALUE")
    if env_hdr_name and env_hdr_value:
        headers[str(env_hdr_name)] = env_hdr_value.format(api_key=api_key)
        return

    # Auth via header (config)
    hdr_name = cfg.get("auth_header_name")
    hdr_value_tmpl = cfg.get("auth_header_value_template")
    if hdr_name and hdr_value_tmpl:
        headers[str(hdr_name)] = str(hdr_value_tmpl).format(api_key=api_key)
        return

    # Auth via query param (env or config)
    query_param = os.environ.get("GEMINI_QUERY_PARAM_NAME") or cfg.get("query_param_name") or "key"
    params[str(query_param)] = api_key


def _encode_image_to_base64(img: Image.Image) -> dict:
    from utils.image_io import pil_to_png_bytes

    png_bytes = pil_to_png_bytes(img)
    return {
        "mime_type": "image/png",
        "data": base64.b64encode(png_bytes).decode("utf-8"),
    }


def _build_payload(prompt: str, images: List[Image.Image], seed: Optional[int] = None) -> dict:
    parts: List[dict] = [{"text": prompt}]
    for img in images:
        parts.append({"inline_data": _encode_image_to_base64(img)})
    generation_config = {
        "response_mime_type": "image/png",
    }
    if seed is not None:
        # Include commonly seen keys for better proxy compatibility
        try:
            generation_config["seed"] = int(seed)
            generation_config["random_seed"] = int(seed)
        except Exception:
            pass

    payload = {
        "contents": [
            {
                "role": "user",
                "parts": parts,
            }
        ],
        # Match requirement: request both IMAGE and TEXT modalities, but we expect image primary
        "generationConfig": generation_config,
        "responseModalities": ["IMAGE", "TEXT"],
    }
    return payload


def _extract_image_bytes_from_response(resp_json: dict) -> bytes:
    # Typical success path: candidates[0].content.parts[*].inline_data.data (base64)
    candidates = resp_json.get("candidates") or []
    if not candidates:
        # Check for promptFeedback block
        feedback = resp_json.get("promptFeedback") or {}
        block_reason = feedback.get("blockReason") or feedback.get("block_reason")
        if block_reason:
            raise GeminiAPIError(f"Gemini blocked the request: {block_reason}")
        raise GeminiAPIError("Gemini API returned no candidates.")

    # inspect first candidate with inline image data
    for cand in candidates:
        content = cand.get("content") or {}
        parts = content.get("parts") or []
        for part in parts:
            inline = part.get("inline_data") or part.get("inlineData")
            if inline and isinstance(inline, dict):
                mime = inline.get("mime_type") or inline.get("mimeType") or ""
                data_b64 = inline.get("data")
                if data_b64 and mime.startswith("image/"):
                    return base64.b64decode(data_b64)

    # Fallback: check ground-level parts
    contents = resp_json.get("contents") or []
    for content in contents:
        for part in content.get("parts", []):
            inline = part.get("inline_data") or part.get("inlineData")
            if inline and isinstance(inline, dict):
                mime = inline.get("mime_type") or inline.get("mimeType") or ""
                data_b64 = inline.get("data")
                if data_b64 and mime.startswith("image/"):
                    return base64.b64decode(data_b64)

    # If we reach here, extract helpful diagnostics
    finish_reason = candidates[0].get("finishReason") if candidates else None
    raise GeminiAPIError(
        f"No image found in response. finishReason={finish_reason}, raw={json.dumps(resp_json)[:800]}"
    )


def call_gemini_generate_image(
    prompt: str,
    images: List[Image.Image],
    model: str = "gemini-2.5-flash-image",
    api_key: Optional[str] = None,
    timeout: float = 60.0,
    seed: Optional[int] = None,
) -> bytes:
    """Call Gemini API to generate an image response from prompt + images.

    Returns PNG bytes on success, raises GeminiAPIError on failure.
    """
    key = api_key or _get_api_key()
    url = _build_full_endpoint(model)
    print(url)
    headers = {
        "Content-Type": "application/json",
    }
    params = {}
    _apply_auth(headers, params, key)
    payload = _build_payload(prompt, images, seed=seed)

    try:
        resp = requests.post(url, headers=headers, params=params, json=payload, timeout=timeout)
    except requests.RequestException as ex:
        raise GeminiAPIError(f"Network error calling Gemini API: {ex}")

    if resp.status_code != 200:
        # Try to parse error payload
        try:
            err_json = resp.json()
        except Exception:
            err_json = None
        if err_json:
            # Common format: {"error": {"code":..., "message":..., "status":...}}
            err = err_json.get("error") or {}
            message = err.get("message") or json.dumps(err_json)
            raise GeminiAPIError(f"Gemini API error {resp.status_code}: {message}")
        raise GeminiAPIError(f"Gemini API error {resp.status_code}: {resp.text}")

    try:
        resp_json = resp.json()
    except ValueError as ex:
        raise GeminiAPIError(f"Failed to parse Gemini response JSON: {ex}")

    img_bytes=_extract_image_bytes_from_response(resp_json)
    img = Image.open(io.BytesIO(img_bytes))
    img_url=upload_image(img)
    return img_url


def call_gemini_generate_image_from_images(
    prompt: str,
    image_urls: List[str],
    model: str = "gemini-2.5-flash-image",
    resolution: str = "1K",
    api_key: Optional[str] = None,
    timeout: float = 60.0,
    seed: Optional[int] = None,
) -> Optional[Image.Image]:
    """
    从图片URL列表生成图片
    
    Args:
        prompt: 提示词
        image_urls: 图片URL列表
        model: 使用的模型
        resolution: 分辨率 ("1K" 或 "2K")
        api_key: API密钥（可选）
        timeout: 超时时间
        seed: 随机种子（可选）
    
    Returns:
        PIL Image对象，失败返回None
    """
    try:
        # 下载图片URL并转换为PIL Image
        images = []
        for image_url in image_urls:
            try:
                response = requests.get(image_url, timeout=30)
                response.raise_for_status()
                img = Image.open(io.BytesIO(response.content)).convert("RGB")
                images.append(img)
            except Exception as e:
                logger.error(f"下载图片失败 {image_url}: {e}")
                return None
        
        if not images:
            logger.error("没有成功下载任何图片")
            return None
        
        # 调用生成函数
        key = api_key or _get_api_key()
        url = _build_full_endpoint(model)
        headers = {
            "Content-Type": "application/json",
        }
        params = {}
        _apply_auth(headers, params, key)
        payload = _build_payload(prompt, images, seed=seed)

        try:
            resp = requests.post(url, headers=headers, params=params, json=payload, timeout=timeout)
        except requests.RequestException as ex:
            logger.error(f"网络错误: {ex}")
            return None

        if resp.status_code != 200:
            try:
                err_json = resp.json()
                err = err_json.get("error") or {}
                message = err.get("message") or json.dumps(err_json)
                logger.error(f"Gemini API错误 {resp.status_code}: {message}")
            except Exception:
                logger.error(f"Gemini API错误 {resp.status_code}: {resp.text}")
            return None

        try:
            resp_json = resp.json()
        except ValueError as ex:
            logger.error(f"解析响应JSON失败: {ex}")
            return None

        img_bytes = _extract_image_bytes_from_response(resp_json)
        img = Image.open(io.BytesIO(img_bytes))
        return img
        
    except Exception as e:
        logger.error(f"生成图片时出错: {e}")
        return None






if __name__=="__main__":
    image=r"D:\ppt\企业微信截图_1760499008629.png"
    image=[Image.open(image)]
    res=call_gemini_generate_image("把这条上衣变成红色",image,api_key="sk-qnsfJw0vsAitlnrXcOeBYrLbTv9LXfsN1m3jIUfMJagan5IR")
    with open("output.png", "wb") as f:
        f.write(res)