"""Gemini 이미지 생성 클라이언트 — urllib 기반 구현 + 키 라운드로빈 + 자동 로깅.""" import base64 import json import time import socket import urllib.error import urllib.request from typing import Dict, List, Optional, Tuple GEMINI_API_URL_TEMPLATE = ( "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}" ) from app.core.config import settings as _settings from app.modules.llm.gemini_key_pool import get_next_key from app.modules.llm.llm_logger import log_llm_call class ModerationError(Exception): """Raised when Gemini blocks content due to safety/moderation policy.""" def __init__(self, block_reason: str, block_categories: list, raw_response: dict): self.block_reason = block_reason self.block_categories = block_categories self.raw_response = raw_response super().__init__(f"Content moderation blocked: {block_reason}") class GeminiImageClient: """Gemini image generation client using the REST API + 키 라운드로빈.""" def __init__(self, api_key: str = None, model: str = None) -> None: self._fixed_api_key = api_key # 명시 전달 시 고정, None이면 라운드로빈 self._model = model or _settings.gemini_image_model self._ctx: Dict = {} def set_context(self, **kwargs) -> "GeminiImageClient": """로깅 컨텍스트 설정.""" self._ctx.update(kwargs) return self def _get_api_key(self) -> str: if self._fixed_api_key: return self._fixed_api_key return get_next_key() def generate_image( self, prompt: str, reference_images: Optional[List[bytes]] = None, aspect_ratio: str = "16:9", labeled_references: Optional[List[Tuple[str, bytes]]] = None, ) -> Tuple[bytes, int]: """Generate an image from a text prompt with optional reference images. Returns: Tuple of (PNG image bytes, response_time_ms). """ parts: list = [] # 참조 이미지 ID 수집 (로깅용) ref_image_ids = [] if labeled_references: parts.append({"text": prompt}) for i, (label, img_bytes) in enumerate(labeled_references, 1): parts.append({"text": f"Reference image {i}:"}) parts.append({ "inline_data": { "mime_type": "image/png", "data": base64.b64encode(img_bytes).decode("ascii"), } }) # label에서 image ID 추출 시도 (없으면 무시) if label: ref_image_ids.append(label) elif reference_images: parts.append({"text": prompt}) for img_bytes in reference_images: parts.append({ "inline_data": { "mime_type": "image/png", "data": base64.b64encode(img_bytes).decode("ascii"), } }) else: parts.append({"text": prompt}) body: Dict = { "contents": [{"parts": parts}], "generationConfig": { "responseModalities": ["TEXT", "IMAGE"], "imageConfig": { "aspectRatio": aspect_ratio, "imageSize": "2K", }, }, } api_key = self._get_api_key() url = GEMINI_API_URL_TEMPLATE.format(model=self._model, api_key=api_key) req = urllib.request.Request( url, data=json.dumps(body).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) start_time = time.monotonic() last_error: Exception | None = None timeout = _settings.llm_timeout_image_gen max_retries = _settings.llm_max_retries for attempt in range(1, max_retries + 2): try: with urllib.request.urlopen(req, timeout=timeout) as response: payload = json.loads(response.read().decode("utf-8")) break except urllib.error.HTTPError as exc: error_text = exc.read().decode("utf-8", errors="replace") last_error = RuntimeError(f"Gemini image API error {exc.code}: {error_text}") if exc.code in {429, 500, 502, 503, 504} and attempt <= max_retries: # 429면 다음 키로 재시도 if exc.code == 429 and not self._fixed_api_key: api_key = self._get_api_key() url = GEMINI_API_URL_TEMPLATE.format(model=self._model, api_key=api_key) req = urllib.request.Request( url, data=json.dumps(body).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) time.sleep(2 * attempt) continue elapsed_ms = int((time.monotonic() - start_time) * 1000) log_llm_call( model_name=self._model, user_prompt=prompt, status="error", error_message=str(last_error), duration_ms=elapsed_ms, reference_image_ids=ref_image_ids, **self._ctx, ) raise RuntimeError(f"Gemini image API error {exc.code}: {error_text}") from exc except (urllib.error.URLError, socket.timeout) as exc: last_error = exc if attempt <= max_retries: time.sleep(2 * attempt) continue elapsed_ms = int((time.monotonic() - start_time) * 1000) log_llm_call( model_name=self._model, user_prompt=prompt, status="error", error_message=str(last_error), duration_ms=elapsed_ms, reference_image_ids=ref_image_ids, **self._ctx, ) raise RuntimeError(f"Gemini image API request failed after retries: {exc}") from exc else: elapsed_ms = int((time.monotonic() - start_time) * 1000) log_llm_call( model_name=self._model, user_prompt=prompt, status="error", error_message=str(last_error), duration_ms=elapsed_ms, reference_image_ids=ref_image_ids, **self._ctx, ) raise RuntimeError(f"Gemini image API request failed: {last_error}") elapsed_ms = int((time.monotonic() - start_time) * 1000) # Check for moderation blocks try: self._check_moderation_block(payload) except ModerationError as exc: log_llm_call( model_name=self._model, user_prompt=prompt, status="error", error_message=str(exc), duration_ms=elapsed_ms, reference_image_ids=ref_image_ids, **self._ctx, ) raise # Extract image bytes candidates = payload.get("candidates", []) for candidate in candidates: content = candidate.get("content", {}) for part in content.get("parts", []): inline_data = part.get("inlineData") or part.get("inline_data") if isinstance(inline_data, dict): b64 = inline_data.get("data") if b64: # 성공 로깅 (출력은 이미지이므로 텍스트 없음) log_llm_call( model_name=self._model, user_prompt=prompt, output_text="[image generated]", status="success", duration_ms=elapsed_ms, reference_image_ids=ref_image_ids, **self._ctx, ) return base64.b64decode(b64), elapsed_ms log_llm_call( model_name=self._model, user_prompt=prompt, status="error", error_message="No image parts in response", duration_ms=elapsed_ms, reference_image_ids=ref_image_ids, **self._ctx, ) raise RuntimeError(f"Gemini image API returned no image parts: {json.dumps(payload)[:500]}") @staticmethod def _check_moderation_block(payload: dict) -> None: """Detect moderation blocks in the Gemini response.""" prompt_feedback = payload.get("promptFeedback", {}) block_reason = prompt_feedback.get("blockReason") if block_reason: categories = [ r.get("category", "UNKNOWN") for r in prompt_feedback.get("safetyRatings", []) if r.get("probability", "NEGLIGIBLE") not in ("NEGLIGIBLE", "LOW") ] raise ModerationError( block_reason=block_reason, block_categories=categories, raw_response=payload, ) candidates = payload.get("candidates", []) if candidates: finish_reason = candidates[0].get("finishReason", "") if finish_reason in ("SAFETY", "IMAGE_SAFETY"): safety_ratings = candidates[0].get("safetyRatings", []) categories = [ r.get("category", "UNKNOWN") for r in safety_ratings if r.get("probability", "NEGLIGIBLE") not in ("NEGLIGIBLE", "LOW") ] raise ModerationError( block_reason="SAFETY", block_categories=categories, raw_response=payload, ) if not candidates: safety_ratings = prompt_feedback.get("safetyRatings", []) high_risk = [ r.get("category", "UNKNOWN") for r in safety_ratings if r.get("probability", "NEGLIGIBLE") not in ("NEGLIGIBLE", "LOW") ] if high_risk: raise ModerationError( block_reason="SAFETY", block_categories=high_risk, raw_response=payload, )