diff --git a/README.md b/README.md index 89f5943..274364b 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,8 @@ pip install -r requirements.txt # LOCAL_DB_HOST, LOCAL_DB_PORT, LOCAL_DB_NAME, LOCAL_DB_USER, LOCAL_DB_PASSWORD # REDIS_LOCAL_HOST, REDIS_LOCAL_PORT # JWT_SECRET, JWT_ALGORITHM -# GEMINI_API_KEY +# GEMINI_API_KEY (단일 키, 하위 호환) +# GEMINI_API_KEYS=key1,key2,key3 (다중 키 round-robin, 설정 시 KEYS 우선) # REDIS_LOCAL_HOST, REDIS_LOCAL_PORT (일일 사용 제한; local) # AI_PIN_GENERATION_* / ISSUE_PIN_CREATE_* / ISSUE_PIN_EDIT_* (rate limit, docs/issue-pin-create-and-reliability.md §9) # AWS_ACCESS_KEY, AWS_SECRET_KEY, AWS_BUCKET (이미지 업로드 시) diff --git a/app/core/config.py b/app/core/config.py index d89dbd8..80e0ef3 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -94,6 +94,11 @@ class Settings(BaseSettings): # Gemini/Vector DB gemini_api_key: SecretStr | None = Field(default=None, alias="GEMINI_API_KEY") + gemini_api_keys: str = Field( + default="", + alias="GEMINI_API_KEYS", + description="콤마 구분 다중 키 (설정 시 GEMINI_API_KEY보다 우선)", + ) # 이슈 핀 신뢰도 VLM. 3.1-pro는 503·지연이 잦아 기본은 flash 계열. gemini_vlm_model: str = Field( default="gemini-2.5-flash", @@ -669,6 +674,17 @@ def async_database_url(self) -> str: db_password=conn.password, ) + def resolved_gemini_api_keys(self) -> tuple[str, ...]: + from app.services.internal.ai.gemini_retry import parse_gemini_model_list + + if self.gemini_api_keys.strip(): + raw = list(parse_gemini_model_list(self.gemini_api_keys)) + elif self.gemini_api_key is not None: + raw = [self.gemini_api_key.get_secret_value()] + else: + raw = [] + return tuple(dict.fromkeys(k for k in raw if k)) + @lru_cache def get_settings() -> Settings: diff --git a/app/core/deps.py b/app/core/deps.py index 4bf6826..c73a368 100644 --- a/app/core/deps.py +++ b/app/core/deps.py @@ -48,6 +48,7 @@ build_vlm_service, require_gemini_api_key, ) +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool, get_gemini_key_pool from app.services.internal.complaint_wiring import build_complaint_email_service from app.services.internal.geo.ImageExifLocationResolveService import ImageExifLocationResolveService from app.services.internal.geo.ImageMultipartGeoService import ImageMultipartGeoService @@ -148,6 +149,14 @@ def get_vector_store_service(request: Request) -> VectorStoreService: VectorStoreServiceDep = Annotated[VectorStoreService, Depends(get_vector_store_service)] +def get_gemini_key_pool_dep(request: Request) -> GeminiKeyPool | None: + pool = getattr(request.app.state, "gemini_key_pool", None) + return pool or get_gemini_key_pool() + + +GeminiKeyPoolDep = Annotated[GeminiKeyPool | None, Depends(get_gemini_key_pool_dep)] + + def get_pin_repo(session: DbSessionDep) -> PinRepo: return PinRepo(session) diff --git a/app/main.py b/app/main.py index 51dbfd7..2e9c329 100644 --- a/app/main.py +++ b/app/main.py @@ -21,6 +21,7 @@ from app.services.internal.ContestPinSchedulerService import ContestPinSchedulerService from app.services.internal.PolicyPinSchedulerService import PolicyPinSchedulerService from app.services.internal.complaint_wiring import build_complaint_email_service +from app.services.internal.ai.gemini_key_pool import init_gemini_key_pool from app.services.VectorStoreService import VectorStoreService from app.services.vector_domains import build_hnsw_kwargs, build_vector_domain_configs from app.schemas.IssueDTO import ( @@ -89,14 +90,23 @@ async def lifespan(app: FastAPI): exc, ) - gemini_api_key_secret = settings.gemini_api_key - if gemini_api_key_secret is None: - logger.warning("GEMINI_API_KEY is not configured. VectorStoreService initialization skipped.") + gemini_api_keys = settings.resolved_gemini_api_keys() + gemini_key_pool = init_gemini_key_pool() + if gemini_key_pool is not None: + app.state.gemini_key_pool = gemini_key_pool + logger.info( + "Gemini key pool initialized: %d key(s), rotation=%s", + gemini_key_pool.size, + "on" if gemini_key_pool.enabled else "off", + ) else: + logger.warning("GEMINI_API_KEY is not configured. VectorStoreService initialization skipped.") + + if gemini_api_keys: try: domain_configs = build_vector_domain_configs(settings) app.state.vector_store_service = VectorStoreService( - api_key=gemini_api_key_secret.get_secret_value(), + api_key=gemini_api_keys[0], table_name=settings.vector_table_name, default_embedding_model=settings.gemini_embedding_model, default_embed_dim=settings.vector_embed_dim, @@ -105,6 +115,7 @@ async def lifespan(app: FastAPI): text_search_config=settings.vector_text_search_config, embedding_batch_size_override=settings.gemini_embedding_batch_size, hnsw_kwargs=build_hnsw_kwargs(settings), + key_pool=gemini_key_pool, ) if settings.vector_dim_check: dimension_checks = ( @@ -149,12 +160,12 @@ async def lifespan(app: FastAPI): complaint_scheduler = None if ( - gemini_api_key_secret is not None + gemini_api_keys and getattr(app.state, "vector_store_service", None) is not None ): try: complaint_email_service = build_complaint_email_service( - api_key=gemini_api_key_secret.get_secret_value(), + api_key=gemini_api_keys[0], vector_store_service=app.state.vector_store_service, ) complaint_scheduler = ComplaintPetitionSchedulerService( @@ -167,7 +178,7 @@ async def lifespan(app: FastAPI): logger.warning("Complaint scheduler initialization failed: %s", exc) policy_pin_scheduler = None - if settings.policy_news_service_key is not None and gemini_api_key_secret is not None: + if settings.policy_news_service_key is not None and gemini_api_keys: try: policy_pin_scheduler = PolicyPinSchedulerService(s3_util=app.state.s3_util) policy_pin_scheduler.start() @@ -176,7 +187,7 @@ async def lifespan(app: FastAPI): logger.warning("Policy pin scheduler initialization failed: %s", exc) contest_pin_scheduler = None - if gemini_api_key_secret is not None: + if gemini_api_keys: try: contest_pin_scheduler = ContestPinSchedulerService(s3_util=app.state.s3_util) contest_pin_scheduler.start() diff --git a/app/services/RagRerankService.py b/app/services/RagRerankService.py index d4701ea..5010611 100644 --- a/app/services/RagRerankService.py +++ b/app/services/RagRerankService.py @@ -8,6 +8,9 @@ import numpy as np from llama_index.embeddings.google_genai import GoogleGenAIEmbedding +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool +from app.services.internal.ai.rotating_gemini_embedding import RotatingGoogleGenAIEmbedding + logger = logging.getLogger(__name__) @@ -26,16 +29,26 @@ def __init__( embedding_model: str, embed_dim: int, embedding_batch_size: int | None = None, + key_pool: GeminiKeyPool | None = None, ) -> None: batch_size = embedding_batch_size if batch_size is None: batch_size = 1 if embedding_model.strip().endswith("embedding-2") else 10 - self._embed_model = GoogleGenAIEmbedding( - model_name=embedding_model, - api_key=api_key, - embedding_config={"output_dimensionality": embed_dim}, - embed_batch_size=max(1, batch_size), - ) + resolved_batch_size = max(1, batch_size) + if key_pool is not None and key_pool.enabled: + self._embed_model = RotatingGoogleGenAIEmbedding( + key_pool=key_pool, + model_name=embedding_model, + embed_dim=embed_dim, + embed_batch_size=resolved_batch_size, + ) + else: + self._embed_model = GoogleGenAIEmbedding( + model_name=embedding_model, + api_key=api_key, + embedding_config={"output_dimensionality": embed_dim}, + embed_batch_size=resolved_batch_size, + ) async def rerank( self, diff --git a/app/services/VectorStoreService.py b/app/services/VectorStoreService.py index 6825670..36206e6 100644 --- a/app/services/VectorStoreService.py +++ b/app/services/VectorStoreService.py @@ -15,6 +15,8 @@ from app.core.database import async_database_url as default_async_database_url from app.core.database import sync_database_url as default_sync_database_url +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool +from app.services.internal.ai.rotating_gemini_embedding import RotatingGoogleGenAIEmbedding from app.services.vector_domains import DomainVectorConfig, VectorDomain logger = logging.getLogger(__name__) @@ -47,6 +49,7 @@ def __init__( text_search_config: str = "simple", embedding_batch_size_override: int | None = None, hnsw_kwargs: dict[str, Any] | None = None, + key_pool: GeminiKeyPool | None = None, ) -> None: self._sync_database_url = sync_database_url self._async_database_url = async_database_url @@ -56,6 +59,7 @@ def __init__( self._hybrid_search = hybrid_search self._text_search_config = text_search_config self._api_key = api_key + self._key_pool = key_pool self._bundles: dict[str, _VectorIndexBundle] = {} self._domain_configs = domain_configs or {} self._embed_models: dict[tuple[str, int], BaseEmbedding] = {} @@ -106,12 +110,20 @@ def _get_or_create_embed_model(self, model_name: str, embed_dim: int) -> BaseEmb embed_batch_size = max(1, int(self._embedding_batch_size_override)) else: embed_batch_size = 1 if model_name.strip().endswith("embedding-2") else 10 - embed_model = GoogleGenAIEmbedding( - model_name=model_name, - api_key=self._api_key, - embedding_config={"output_dimensionality": embed_dim}, - embed_batch_size=embed_batch_size, - ) + if self._key_pool is not None and self._key_pool.enabled: + embed_model = RotatingGoogleGenAIEmbedding( + key_pool=self._key_pool, + model_name=model_name, + embed_dim=embed_dim, + embed_batch_size=embed_batch_size, + ) + else: + embed_model = GoogleGenAIEmbedding( + model_name=model_name, + api_key=self._api_key, + embedding_config={"output_dimensionality": embed_dim}, + embed_batch_size=embed_batch_size, + ) self._embed_models[model_key] = embed_model return embed_model diff --git a/app/services/internal/ai/ComplaintEmailLLMService.py b/app/services/internal/ai/ComplaintEmailLLMService.py index bc8061e..e9f29ec 100644 --- a/app/services/internal/ai/ComplaintEmailLLMService.py +++ b/app/services/internal/ai/ComplaintEmailLLMService.py @@ -15,6 +15,7 @@ ComplaintEmailOpinionRenderer, OpinionAttachmentImage, ) +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool from app.services.internal.ai.gemini_retry import generate_content_with_retry from app.services.prompts.complaint_email_opinion import complaint_opinion_prompt @@ -25,6 +26,7 @@ class ComplaintEmailLLMService: api_key: str model_name: str = "gemini-2.5-flash" + key_pool: GeminiKeyPool | None = None client: genai.Client = field(init=False, repr=False) _opinion_renderer: ComplaintEmailOpinionRenderer = field( default_factory=ComplaintEmailOpinionRenderer, @@ -75,6 +77,7 @@ async def _generate_json_text(self, prompt: str) -> str: contents=text, config=config, log_prefix="ComplaintEmailLLM", + key_pool=self.key_pool, ) except genai_errors.APIError as exc: raise BusinessException( diff --git a/app/services/internal/ai/ComplaintEmailVLMService.py b/app/services/internal/ai/ComplaintEmailVLMService.py index 9fd86a8..f8bff32 100644 --- a/app/services/internal/ai/ComplaintEmailVLMService.py +++ b/app/services/internal/ai/ComplaintEmailVLMService.py @@ -17,6 +17,7 @@ ComplaintEmailVlmInput, ) from app.schemas.IssueDTO import ImageWithLocation +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool from app.services.internal.ai.gemini_retry import generate_content_with_retry from app.services.internal.ai.VLMService import resolve_upload_image_mime from app.services.prompts.complaint_email_vlm import ( @@ -111,12 +112,14 @@ def __init__( catalog: type[ComplaintEmailVlmCatalog] = ComplaintEmailVlmCatalog, prompt_builder: ComplaintEmailVlmPromptBuilder | None = None, result_processor: ComplaintEmailVlmResultProcessor | None = None, + key_pool: GeminiKeyPool | None = None, ) -> None: self._api_key = api_key self._model = model self._catalog = catalog self._prompt_builder = prompt_builder or ComplaintEmailVlmPromptBuilder(catalog) self._processor = result_processor or ComplaintEmailVlmResultProcessor(catalog) + self._key_pool = key_pool self._client: genai.Client | None = None if api_key: self._client = genai.Client(api_key=api_key) @@ -199,6 +202,7 @@ async def _call_vlm(self, prompt: str, images: list[ImageWithLocation]) -> str: contents=parts, config=config, log_prefix="ComplaintEmailVLM", + key_pool=self._key_pool, ) except genai_errors.APIError as exc: raise BusinessException( diff --git a/app/services/internal/ai/IssuePinLLMService.py b/app/services/internal/ai/IssuePinLLMService.py index 13c6d12..9806322 100644 --- a/app/services/internal/ai/IssuePinLLMService.py +++ b/app/services/internal/ai/IssuePinLLMService.py @@ -11,6 +11,7 @@ from app.core.codes import ErrorCode from app.core.exceptions import BusinessException, raise_business_exception +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool from app.services.internal.ai.gemini_retry import generate_content_with_retry from app.services.prompts.issue_pin import ISSUE_PIN_OUTPUT_SCHEMA @@ -23,6 +24,7 @@ class IssuePinLLMService: api_key: str model_name: str = "gemini-2.5-flash" + key_pool: GeminiKeyPool | None = None client: genai.Client = field(init=False, repr=False) fallback_models: tuple[str, ...] = ("gemini-2.5-flash-lite", "gemini-2.5-flash") @@ -42,6 +44,7 @@ async def _generate_with_retry( contents=contents, config=config, log_prefix="Pin", + key_pool=self.key_pool, ) async def generate_pin_copy(self, *, prompt: str) -> dict[str, str]: diff --git a/app/services/internal/ai/IssueRagPlannerService.py b/app/services/internal/ai/IssueRagPlannerService.py index 87d184a..41a2695 100644 --- a/app/services/internal/ai/IssueRagPlannerService.py +++ b/app/services/internal/ai/IssueRagPlannerService.py @@ -9,6 +9,7 @@ from google import genai from google.genai import types +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool from app.services.internal.ai.gemini_retry import generate_content_with_retry logger = logging.getLogger(__name__) @@ -74,6 +75,7 @@ def _normalize_string_list(value: object) -> list[str]: class IssueRagPlannerService: api_key: str model_name: str = "gemini-2.5-flash-lite" + key_pool: GeminiKeyPool | None = None client: genai.Client = field(init=False, repr=False) fallback_models: tuple[str, ...] = ("gemini-2.5-flash",) @@ -95,6 +97,7 @@ async def _generate_with_retry( config=config, log_prefix="Planner", log_context=log_context, + key_pool=self.key_pool, ) async def rewrite_queries( diff --git a/app/services/internal/ai/PolicyCardnewsImageService.py b/app/services/internal/ai/PolicyCardnewsImageService.py index 6244250..42e24c9 100644 --- a/app/services/internal/ai/PolicyCardnewsImageService.py +++ b/app/services/internal/ai/PolicyCardnewsImageService.py @@ -5,9 +5,11 @@ from dataclasses import dataclass, field from google import genai +from google.genai import errors as genai_errors from google.genai import types from app.core.config import settings +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool, get_gemini_key_pool from app.services.internal.ai.gemini_retry import ( is_retryable_gemini_error, parse_gemini_model_list, @@ -25,6 +27,7 @@ class PolicyCardnewsImageService: # 정책 카드뉴스 슬라이드 이미지 생성 (텍스트 모델과 분리) api_key: str model_name: str = "gemini-2.5-flash-image" + key_pool: GeminiKeyPool | None = None client: genai.Client = field(init=False, repr=False) fallback_models: tuple[str, ...] = ("gemini-3-pro-image-preview",) @@ -33,14 +36,15 @@ def __post_init__(self) -> None: @classmethod def from_settings(cls) -> PolicyCardnewsImageService: - secret = settings.gemini_api_key - if secret is None: + keys = settings.resolved_gemini_api_keys() + if not keys: raise RuntimeError("GEMINI_API_KEY가 설정되어 있지 않습니다.") fallbacks = parse_gemini_model_list(settings.gemini_cardnews_image_fallback_models) return cls( - api_key=secret.get_secret_value(), + api_key=keys[0], model_name=settings.gemini_cardnews_image_model.strip(), fallback_models=fallbacks, + key_pool=get_gemini_key_pool(), ) async def generate_slide_image_bytes(self, *, prompt: str) -> bytes: @@ -93,36 +97,86 @@ async def generate_slide_image_bytes(self, *, prompt: str) -> bytes: ) from last_error async def _generate_with_gemini_image(self, *, model: str, prompt: str) -> bytes: - response = await self.client.aio.models.generate_content( - model=model, - contents=prompt, - config=types.GenerateContentConfig( - response_modalities=[types.Modality.IMAGE], - ), + pool = self.key_pool + config = types.GenerateContentConfig( + response_modalities=[types.Modality.IMAGE], ) - return self._extract_image_bytes(response) + + async def invoke(client: genai.Client) -> bytes: + response = await client.aio.models.generate_content( + model=model, + contents=prompt, + config=config, + ) + return self._extract_image_bytes(response) + + if pool is not None and pool.enabled: + key_index, _, client = await pool.acquire() + try: + return await invoke(client) + except (genai_errors.ServerError, genai_errors.APIError) as exc: + if not is_retryable_gemini_error(exc): + raise + last_error: Exception | None = exc + for failover_idx in pool.failover_indices(key_index): + try: + return await invoke(pool.client_at(failover_idx)) + except (genai_errors.ServerError, genai_errors.APIError) as failover_exc: + last_error = failover_exc + if not is_retryable_gemini_error(failover_exc): + raise + if last_error is not None: + raise last_error + raise + + return await invoke(self.client) async def _generate_with_imagen(self, *, model: str, prompt: str) -> bytes: - response = await self.client.aio.models.generate_images( - model=model, - prompt=prompt, - config=types.GenerateImagesConfig( - number_of_images=1, - aspect_ratio="3:4", - output_mime_type="image/png", - ), + pool = self.key_pool + image_config = types.GenerateImagesConfig( + number_of_images=1, + aspect_ratio="3:4", + output_mime_type="image/png", ) - generated = response.generated_images or [] - if not generated: - raise RuntimeError("Imagen 응답에 generated_images가 없음") - data = generated[0].image.image_bytes - if not data: - raise RuntimeError("Imagen image_bytes가 비어 있음") - if isinstance(data, str): - import base64 - - return base64.b64decode(data) - return bytes(data) + + async def invoke(client: genai.Client) -> bytes: + response = await client.aio.models.generate_images( + model=model, + prompt=prompt, + config=image_config, + ) + generated = response.generated_images or [] + if not generated: + raise RuntimeError("Imagen 응답에 generated_images가 없음") + data = generated[0].image.image_bytes + if not data: + raise RuntimeError("Imagen image_bytes가 비어 있음") + if isinstance(data, str): + import base64 + + return base64.b64decode(data) + return bytes(data) + + if pool is not None and pool.enabled: + key_index, _, client = await pool.acquire() + try: + return await invoke(client) + except (genai_errors.ServerError, genai_errors.APIError) as exc: + if not is_retryable_gemini_error(exc): + raise + last_error: Exception | None = exc + for failover_idx in pool.failover_indices(key_index): + try: + return await invoke(pool.client_at(failover_idx)) + except (genai_errors.ServerError, genai_errors.APIError) as failover_exc: + last_error = failover_exc + if not is_retryable_gemini_error(failover_exc): + raise + if last_error is not None: + raise last_error + raise + + return await invoke(self.client) @staticmethod def _extract_image_bytes(response: object) -> bytes: diff --git a/app/services/internal/ai/VLMService.py b/app/services/internal/ai/VLMService.py index c30ed5d..c6194fa 100644 --- a/app/services/internal/ai/VLMService.py +++ b/app/services/internal/ai/VLMService.py @@ -21,6 +21,7 @@ VLM_PRIVACY_NOTES, build_vlm_prompt, ) +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool from app.services.internal.ai.gemini_retry import generate_content_with_retry from app.services.prompts.confidence_basis import CONFIDENCE_BASIS_ARRAY_SCHEMA from app.services.prompts.issue_reliability_text import ( @@ -257,6 +258,7 @@ class VLMService: api_key: str model_name: str = "gemini-3.1-pro-preview" fallback_models: tuple[str, ...] = ("gemini-2.5-flash", "gemini-2.5-pro") + key_pool: GeminiKeyPool | None = None client: genai.Client = field(init=False, repr=False) def __post_init__(self) -> None: @@ -280,6 +282,7 @@ async def _generate_with_retry( max_attempts_per_model=max_attempts_per_model or 5, log_prefix="VLM", log_context=log_context, + key_pool=self.key_pool, ) async def analyze_image( diff --git a/app/services/internal/ai/gemini_factory.py b/app/services/internal/ai/gemini_factory.py index 00f03d2..edf28e7 100644 --- a/app/services/internal/ai/gemini_factory.py +++ b/app/services/internal/ai/gemini_factory.py @@ -9,21 +9,23 @@ from app.services.internal.ai.IssuePinLLMService import IssuePinLLMService from app.services.internal.ai.IssueRagPlannerService import IssueRagPlannerService from app.services.internal.ai.VLMService import VLMService +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool, get_gemini_key_pool from app.services.internal.ai.gemini_retry import parse_gemini_model_list -def _gemini_api_key_or_none() -> str | None: - secret = settings.gemini_api_key - if secret is None: - return None - return secret.get_secret_value() +def _resolved_keys() -> tuple[str, ...]: + return settings.resolved_gemini_api_keys() def require_gemini_api_key() -> str: - api_key = _gemini_api_key_or_none() - if api_key is None: + keys = _resolved_keys() + if not keys: raise_business_exception(ErrorCode.VLM_NOT_CONFIGURED) - return api_key + return keys[0] + + +def _key_pool() -> GeminiKeyPool | None: + return get_gemini_key_pool() def build_vlm_service(*, api_key: str | None = None) -> VLMService: @@ -32,19 +34,21 @@ def build_vlm_service(*, api_key: str | None = None) -> VLMService: api_key=key, model_name=settings.gemini_vlm_model, fallback_models=parse_gemini_model_list(settings.gemini_vlm_fallback_models), + key_pool=_key_pool(), ) def build_issue_pin_llm_service(*, model: str | None = None) -> IssuePinLLMService: - api_key = _gemini_api_key_or_none() - if api_key is None: + keys = _resolved_keys() + if not keys: raise RuntimeError("GEMINI_API_KEY가 설정되어 있지 않습니다.") model_name = (model or settings.gemini_pin_text_model).strip() fallbacks = parse_gemini_model_list(settings.gemini_pin_text_fallback_models) return IssuePinLLMService( - api_key=api_key, + api_key=keys[0], model_name=model_name, fallback_models=fallbacks, + key_pool=_key_pool(), ) @@ -54,6 +58,7 @@ def build_issue_rag_planner_service(*, api_key: str | None = None) -> IssueRagPl api_key=key, model_name=settings.gemini_rag_planner_model, fallback_models=parse_gemini_model_list(settings.gemini_rag_planner_fallback_models), + key_pool=_key_pool(), ) @@ -62,6 +67,7 @@ def build_complaint_email_vlm_service(*, api_key: str | None = None) -> Complain return ComplaintEmailVlmService( api_key=key, model=settings.gemini_vlm_model, + key_pool=_key_pool(), ) @@ -70,6 +76,7 @@ def build_complaint_email_llm_service(*, api_key: str | None = None) -> Complain return ComplaintEmailLLMService( api_key=key, model_name=settings.gemini_pin_text_model, + key_pool=_key_pool(), ) @@ -80,4 +87,5 @@ def build_rag_rerank_service(*, api_key: str | None = None) -> RagRerankService: embedding_model=settings.gemini_embedding_model, embed_dim=settings.vector_embed_dim, embedding_batch_size=settings.gemini_embedding_batch_size, + key_pool=_key_pool(), ) diff --git a/app/services/internal/ai/gemini_key_pool.py b/app/services/internal/ai/gemini_key_pool.py new file mode 100644 index 0000000..8bce008 --- /dev/null +++ b/app/services/internal/ai/gemini_key_pool.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +from google import genai + +from app.core.config import settings + +if TYPE_CHECKING: + pass + +_pool: GeminiKeyPool | None = None + + +def mask_api_key(api_key: str) -> str: + key = (api_key or "").strip() + if len(key) <= 8: + return "***" + return f"...{key[-4:]}" + + +class GeminiKeyPool: + def __init__(self, api_keys: tuple[str, ...]) -> None: + if not api_keys: + raise ValueError("GeminiKeyPool requires at least one API key") + self._api_keys = api_keys + self._index = 0 + self._lock = threading.Lock() + self._clients: dict[int, genai.Client] = {} + + @classmethod + def from_settings(cls) -> GeminiKeyPool: + keys = settings.resolved_gemini_api_keys() + if not keys: + raise ValueError("No Gemini API keys configured") + return cls(keys) + + @property + def size(self) -> int: + return len(self._api_keys) + + @property + def enabled(self) -> bool: + return self.size >= 2 + + def primary_api_key(self) -> str: + return self._api_keys[0] + + def api_key_at(self, index: int) -> str: + if index < 0 or index >= self.size: + raise IndexError(f"Gemini key index out of range: {index}") + return self._api_keys[index] + + def _advance_index(self) -> int: + with self._lock: + idx = self._index + self._index = (idx + 1) % self.size + return idx + + def client_at(self, index: int) -> genai.Client: + if index < 0 or index >= self.size: + raise IndexError(f"Gemini key index out of range: {index}") + with self._lock: + client = self._clients.get(index) + if client is None: + client = genai.Client(api_key=self._api_keys[index]) + self._clients[index] = client + return client + + async def acquire(self) -> tuple[int, str, genai.Client]: + idx = self._advance_index() + api_key = self._api_keys[idx] + return idx, api_key, self.client_at(idx) + + def acquire_sync(self) -> tuple[int, str, genai.Client]: + idx = self._advance_index() + api_key = self._api_keys[idx] + return idx, api_key, self.client_at(idx) + + def failover_indices(self, start_index: int) -> list[int]: + if self.size <= 1: + return [] + return [(start_index + offset) % self.size for offset in range(1, self.size)] + + +def init_gemini_key_pool() -> GeminiKeyPool | None: + global _pool + keys = settings.resolved_gemini_api_keys() + if not keys: + _pool = None + return None + _pool = GeminiKeyPool(keys) + return _pool + + +def get_gemini_key_pool() -> GeminiKeyPool | None: + global _pool + if _pool is None: + return init_gemini_key_pool() + return _pool diff --git a/app/services/internal/ai/gemini_retry.py b/app/services/internal/ai/gemini_retry.py index f803c94..f0e3639 100644 --- a/app/services/internal/ai/gemini_retry.py +++ b/app/services/internal/ai/gemini_retry.py @@ -2,10 +2,13 @@ import asyncio import logging -from typing import Any +from typing import TYPE_CHECKING, Any from google.genai import errors as genai_errors +if TYPE_CHECKING: + from app.services.internal.ai.gemini_key_pool import GeminiKeyPool + logger = logging.getLogger(__name__) DEFAULT_MAX_ATTEMPTS_PER_MODEL = 5 @@ -16,8 +19,18 @@ def parse_gemini_model_list(raw: str) -> tuple[str, ...]: return tuple(part.strip() for part in raw.split(",") if part.strip()) -def is_retryable_gemini_error(exc: Exception) -> bool: +def _gemini_error_status(exc: Exception) -> int | None: status_code = getattr(exc, "status_code", None) + if isinstance(status_code, int): + return status_code + code = getattr(exc, "code", None) + if isinstance(code, int): + return code + return None + + +def is_retryable_gemini_error(exc: Exception) -> bool: + status_code = _gemini_error_status(exc) if status_code in {429, 500, 502, 503, 504}: return True text = str(exc).lower() @@ -34,7 +47,7 @@ def is_retryable_gemini_error(exc: Exception) -> bool: def should_skip_gemini_model(exc: Exception) -> bool: - status_code = getattr(exc, "status_code", None) + status_code = _gemini_error_status(exc) if status_code in {400, 404}: return True text = str(exc).lower() @@ -47,6 +60,77 @@ def _format_log_context(log_context: str | None) -> str: return f" [{log_context}]" +async def _generate_once( + client: Any, + *, + model: str, + contents: Any, + config: Any | None, +) -> Any: + kwargs: dict[str, Any] = {"model": model, "contents": contents} + if config is not None: + kwargs["config"] = config + return await client.aio.models.generate_content(**kwargs) + + +async def _generate_with_key_failover( + *, + key_pool: GeminiKeyPool, + start_index: int, + start_client: Any, + model: str, + contents: Any, + config: Any | None, + log_prefix: str, + ctx: str, + attempt_no: int, + max_attempts_per_model: int, +) -> Any: + clients_to_try: list[tuple[int, Any]] = [(start_index, start_client)] + for failover_idx in key_pool.failover_indices(start_index): + clients_to_try.append((failover_idx, key_pool.client_at(failover_idx))) + + last_error: Exception | None = None + for key_idx, try_client in clients_to_try: + try: + result = await _generate_once( + try_client, + model=model, + contents=contents, + config=config, + ) + if key_idx != start_index: + logger.warning( + "%s key failover success%s: model=%s key_index=%d/%d attempt=%d/%d", + log_prefix, + ctx, + model, + key_idx + 1, + key_pool.size, + attempt_no, + max_attempts_per_model, + ) + return result + except (genai_errors.ServerError, genai_errors.APIError) as exc: + last_error = exc + if not is_retryable_gemini_error(exc): + raise + if key_idx != clients_to_try[-1][0]: + logger.warning( + "%s key failover retry%s: model=%s key_index=%d/%d status=%s err=%s", + log_prefix, + ctx, + model, + key_idx + 1, + key_pool.size, + getattr(exc, "status_code", None) or _gemini_error_status(exc), + exc, + ) + if last_error is not None: + raise last_error + raise RuntimeError(f"{log_prefix} generate_content 호출에 실패했습니다.") + + async def generate_content_with_retry( client: Any, *, @@ -58,20 +142,23 @@ async def generate_content_with_retry( base_delay_seconds: float = DEFAULT_BASE_DELAY_SECONDS, log_prefix: str = "Gemini", log_context: str | None = None, + key_pool: GeminiKeyPool | None = None, ) -> Any: model_candidates = [model_name] model_candidates.extend(m for m in fallback_models if m != model_name) last_error: Exception | None = None ctx = _format_log_context(log_context) + pool_active = key_pool is not None and key_pool.enabled logger.info( - "%s call start%s: primary=%s fallbacks=%s chain=%s max_attempts_per_model=%d", + "%s call start%s: primary=%s fallbacks=%s chain=%s max_attempts_per_model=%d key_pool=%s", log_prefix, ctx, model_name, list(fallback_models), model_candidates, max_attempts_per_model, + f"{key_pool.size}keys" if pool_active else "off", ) for model_index, model in enumerate(model_candidates): @@ -88,26 +175,58 @@ async def generate_content_with_retry( for attempt in range(max_attempts_per_model): attempt_no = attempt + 1 - logger.info( - "%s attempt start%s: model=%s attempt=%d/%d", - log_prefix, - ctx, - model, - attempt_no, - max_attempts_per_model, - ) + active_client = client + key_index: int | None = None + if pool_active and key_pool is not None: + key_index, _, active_client = await key_pool.acquire() + logger.info( + "%s attempt start%s: model=%s attempt=%d/%d key_index=%d/%d", + log_prefix, + ctx, + model, + attempt_no, + max_attempts_per_model, + key_index + 1, + key_pool.size, + ) + else: + logger.info( + "%s attempt start%s: model=%s attempt=%d/%d", + log_prefix, + ctx, + model, + attempt_no, + max_attempts_per_model, + ) try: - kwargs: dict[str, Any] = {"model": model, "contents": contents} - if config is not None: - kwargs["config"] = config - result = await client.aio.models.generate_content(**kwargs) + if pool_active and key_pool is not None and key_index is not None: + result = await _generate_with_key_failover( + key_pool=key_pool, + start_index=key_index, + start_client=active_client, + model=model, + contents=contents, + config=config, + log_prefix=log_prefix, + ctx=ctx, + attempt_no=attempt_no, + max_attempts_per_model=max_attempts_per_model, + ) + else: + result = await _generate_once( + active_client, + model=model, + contents=contents, + config=config, + ) logger.info( - "%s success%s: model=%s attempt=%d/%d", + "%s success%s: model=%s attempt=%d/%d%s", log_prefix, ctx, model, attempt_no, max_attempts_per_model, + f" key_index={key_index + 1}/{key_pool.size}" if pool_active and key_index is not None and key_pool else "", ) return result except (genai_errors.ServerError, genai_errors.APIError) as exc: diff --git a/app/services/internal/ai/rotating_gemini_embedding.py b/app/services/internal/ai/rotating_gemini_embedding.py new file mode 100644 index 0000000..3fabd77 --- /dev/null +++ b/app/services/internal/ai/rotating_gemini_embedding.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import threading +from typing import Any + +from llama_index.core.base.embeddings.base import BaseEmbedding +from llama_index.embeddings.google_genai import GoogleGenAIEmbedding + +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool +from app.services.internal.ai.gemini_retry import is_retryable_gemini_error + + +class RotatingGoogleGenAIEmbedding(BaseEmbedding): + """GeminiKeyPool round-robin으로 GoogleGenAIEmbedding 호출을 분산.""" + + def __init__( + self, + *, + key_pool: GeminiKeyPool, + model_name: str, + embed_dim: int, + embed_batch_size: int, + ) -> None: + super().__init__( + model_name=model_name, + embed_batch_size=max(1, embed_batch_size), + ) + self._key_pool = key_pool + self._embed_dim = embed_dim + self._embed_batch_size = max(1, embed_batch_size) + self._models_by_index: dict[int, GoogleGenAIEmbedding] = {} + self._lock = threading.Lock() + + def _embed_model_at(self, index: int) -> GoogleGenAIEmbedding: + with self._lock: + model = self._models_by_index.get(index) + if model is None: + model = GoogleGenAIEmbedding( + model_name=self.model_name, + api_key=self._key_pool.api_key_at(index), + embedding_config={"output_dimensionality": self._embed_dim}, + embed_batch_size=self._embed_batch_size, + ) + self._models_by_index[index] = model + return model + + def _execute_with_failover(self, func_name: str, *args: Any, **kwargs: Any) -> Any: + idx, _, _ = self._key_pool.acquire_sync() + try: + return getattr(self._embed_model_at(idx), func_name)(*args, **kwargs) + except Exception as exc: + if not is_retryable_gemini_error(exc): + raise + last_error: Exception | None = exc + for failover_idx in self._key_pool.failover_indices(idx): + try: + return getattr(self._embed_model_at(failover_idx), func_name)(*args, **kwargs) + except Exception as failover_exc: + last_error = failover_exc + if not is_retryable_gemini_error(failover_exc): + raise + if last_error is not None: + raise last_error + raise + + async def _aexecute_with_failover(self, func_name: str, *args: Any, **kwargs: Any) -> Any: + idx, _, _ = await self._key_pool.acquire() + try: + return await getattr(self._embed_model_at(idx), func_name)(*args, **kwargs) + except Exception as exc: + if not is_retryable_gemini_error(exc): + raise + last_error: Exception | None = exc + for failover_idx in self._key_pool.failover_indices(idx): + try: + return await getattr(self._embed_model_at(failover_idx), func_name)(*args, **kwargs) + except Exception as failover_exc: + last_error = failover_exc + if not is_retryable_gemini_error(failover_exc): + raise + if last_error is not None: + raise last_error + raise + + def _get_text_embedding(self, text: str) -> list[float]: + return self._execute_with_failover("get_text_embedding", text) + + async def _aget_text_embedding(self, text: str) -> list[float]: + return await self._aexecute_with_failover("aget_text_embedding", text) + + def _get_text_embeddings(self, texts: list[str]) -> list[list[float]]: + return self._execute_with_failover("get_text_embedding_batch", texts) + + async def _aget_text_embeddings(self, texts: list[str]) -> list[list[float]]: + return await self._aexecute_with_failover("aget_text_embedding_batch", texts) + + def _get_query_embedding(self, query: str) -> list[float]: + return self._execute_with_failover("get_query_embedding", query) + + async def _aget_query_embedding(self, query: str) -> list[float]: + return await self._aexecute_with_failover("aget_query_embedding", query) + + @classmethod + def class_name(cls) -> str: + return "RotatingGoogleGenAIEmbedding" + + def _get_text_embedding_safe(self, text: str) -> list[float]: + return self._get_text_embedding(text) + + async def _aget_text_embedding_safe(self, text: str) -> list[float]: + return await self._aget_text_embedding(text) + + def _get_text_embeddings_safe(self, texts: list[str]) -> list[list[float]]: + return self._get_text_embeddings(texts) + + async def _aget_text_embeddings_safe(self, texts: list[str]) -> list[list[float]]: + return await self._aget_text_embeddings(texts) + + def _get_query_embedding_safe(self, query: str) -> list[float]: + return self._get_query_embedding(query) + + async def _aget_query_embedding_safe(self, query: str) -> list[float]: + return await self._aget_query_embedding(query) diff --git a/tests/test_gemini_key_pool.py b/tests/test_gemini_key_pool.py new file mode 100644 index 0000000..25131e7 --- /dev/null +++ b/tests/test_gemini_key_pool.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +from google.genai import errors as genai_errors + +from app.core.config import Settings +from app.services.internal.ai.gemini_key_pool import GeminiKeyPool, init_gemini_key_pool +from app.services.internal.ai.gemini_retry import generate_content_with_retry +from app.services.internal.ai.rotating_gemini_embedding import RotatingGoogleGenAIEmbedding + + +class ResolvedGeminiApiKeysTest(unittest.TestCase): + def test_keys_only(self) -> None: + settings = Settings( + GEMINI_API_KEYS="key1,key2,key3", + ) + self.assertEqual(settings.resolved_gemini_api_keys(), ("key1", "key2", "key3")) + + def test_single_key_fallback(self) -> None: + settings = Settings( + GEMINI_API_KEY="only", + ) + self.assertEqual(settings.resolved_gemini_api_keys(), ("only",)) + + def test_keys_take_priority_over_single_key(self) -> None: + settings = Settings( + GEMINI_API_KEY="ignored", + GEMINI_API_KEYS="a,b", + ) + self.assertEqual(settings.resolved_gemini_api_keys(), ("a", "b")) + + def test_dedupe_preserves_order(self) -> None: + settings = Settings( + GEMINI_API_KEYS="a,b,a,c,b", + ) + self.assertEqual(settings.resolved_gemini_api_keys(), ("a", "b", "c")) + + +class GeminiKeyPoolTest(unittest.IsolatedAsyncioTestCase): + async def test_round_robin_three_keys(self) -> None: + pool = GeminiKeyPool(("k1", "k2", "k3")) + self.assertTrue(pool.enabled) + indices = [] + for _ in range(5): + idx, key, _client = await pool.acquire() + indices.append(idx) + self.assertEqual(key, ("k1", "k2", "k3")[idx]) + self.assertEqual(indices, [0, 1, 2, 0, 1]) + + async def test_single_key_not_enabled(self) -> None: + pool = GeminiKeyPool(("only",)) + self.assertFalse(pool.enabled) + idx, key, _ = await pool.acquire() + self.assertEqual((idx, key), (0, "only")) + + def test_failover_indices(self) -> None: + pool = GeminiKeyPool(("a", "b", "c")) + self.assertEqual(pool.failover_indices(0), [1, 2]) + self.assertEqual(pool.failover_indices(2), [0, 1]) + + @patch("app.services.internal.ai.gemini_key_pool.settings") + def test_init_from_settings_empty(self, mock_settings: MagicMock) -> None: + mock_settings.resolved_gemini_api_keys.return_value = () + self.assertIsNone(init_gemini_key_pool()) + + +class GeminiRetryKeyRotationTest(unittest.IsolatedAsyncioTestCase): + async def test_failover_uses_second_client_on_429(self) -> None: + pool = GeminiKeyPool(("key-a", "key-b")) + client_a = MagicMock() + client_b = MagicMock() + pool._clients[0] = client_a + pool._clients[1] = client_b + + err_429 = genai_errors.ClientError(429, {"error": {"message": "rate"}}) + ok_response = MagicMock() + + client_a.aio.models.generate_content = AsyncMock(side_effect=err_429) + client_b.aio.models.generate_content = AsyncMock(return_value=ok_response) + + result = await generate_content_with_retry( + client_a, + model_name="gemini-2.5-flash", + contents="hello", + max_attempts_per_model=1, + key_pool=pool, + log_prefix="Test", + ) + self.assertIs(result, ok_response) + client_b.aio.models.generate_content.assert_awaited_once() + + async def test_failover_skips_second_client_on_400(self) -> None: + pool = GeminiKeyPool(("key-a", "key-b")) + client_a = MagicMock() + client_b = MagicMock() + pool._clients[0] = client_a + pool._clients[1] = client_b + + err_400 = genai_errors.ClientError(400, {"error": {"message": "bad request"}}) + client_a.aio.models.generate_content = AsyncMock(side_effect=err_400) + client_b.aio.models.generate_content = AsyncMock() + + with self.assertRaises(genai_errors.ClientError): + await generate_content_with_retry( + client_a, + model_name="gemini-2.5-flash", + contents="hello", + max_attempts_per_model=1, + key_pool=pool, + log_prefix="Test", + ) + client_a.aio.models.generate_content.assert_awaited_once() + client_b.aio.models.generate_content.assert_not_awaited() + + +class RotatingEmbeddingFailoverTest(unittest.TestCase): + @patch("app.services.internal.ai.rotating_gemini_embedding.GoogleGenAIEmbedding") + def test_failover_uses_second_key_on_429(self, mock_embedding_cls: MagicMock) -> None: + pool = GeminiKeyPool(("key-a", "key-b")) + model_a = MagicMock() + model_b = MagicMock() + err_429 = genai_errors.ClientError(429, {"error": {"message": "rate"}}) + model_a.get_text_embedding.side_effect = err_429 + model_b.get_text_embedding.return_value = [0.1, 0.2] + mock_embedding_cls.side_effect = [model_a, model_b] + + embedder = RotatingGoogleGenAIEmbedding( + key_pool=pool, + model_name="text-embedding-004", + embed_dim=768, + embed_batch_size=10, + ) + result = embedder._get_text_embedding("hello") + self.assertEqual(result, [0.1, 0.2]) + model_a.get_text_embedding.assert_called_once_with("hello") + model_b.get_text_embedding.assert_called_once_with("hello") + + @patch("app.services.internal.ai.rotating_gemini_embedding.GoogleGenAIEmbedding") + def test_non_retryable_error_does_not_failover(self, mock_embedding_cls: MagicMock) -> None: + pool = GeminiKeyPool(("key-a", "key-b")) + model_a = MagicMock() + err_400 = genai_errors.ClientError(400, {"error": {"message": "bad request"}}) + model_a.get_text_embedding.side_effect = err_400 + mock_embedding_cls.return_value = model_a + + embedder = RotatingGoogleGenAIEmbedding( + key_pool=pool, + model_name="text-embedding-004", + embed_dim=768, + embed_batch_size=10, + ) + with self.assertRaises(genai_errors.ClientError): + embedder._get_text_embedding("hello") + model_a.get_text_embedding.assert_called_once_with("hello") + mock_embedding_cls.assert_called_once() + + @patch("app.services.internal.ai.rotating_gemini_embedding.GoogleGenAIEmbedding") + def test_query_embedding_uses_get_query_embedding(self, mock_embedding_cls: MagicMock) -> None: + pool = GeminiKeyPool(("key-a",)) + model = MagicMock() + model.get_query_embedding.return_value = [0.3, 0.4] + mock_embedding_cls.return_value = model + + embedder = RotatingGoogleGenAIEmbedding( + key_pool=pool, + model_name="text-embedding-004", + embed_dim=768, + embed_batch_size=10, + ) + result = embedder._get_query_embedding("search query") + self.assertEqual(result, [0.3, 0.4]) + model.get_query_embedding.assert_called_once_with("search query") + model.get_text_embedding.assert_not_called()