From 3f7a4745cbf25dfca7696ebd752556795c75850d Mon Sep 17 00:00:00 2001 From: selnem Date: Tue, 16 Jun 2026 18:03:37 +0900 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20Gemini=20=EB=8B=A4=EC=A4=91=20API?= =?UTF-8?q?=20=ED=82=A4=20round-robin=20=EB=A1=9C=ED=85=8C=EC=9D=B4?= =?UTF-8?q?=EC=85=98=20=EB=B0=8F=20429=20failover=20GEMINI=5FAPI=5FKEYS(?= =?UTF-8?q?=EC=BD=A4=EB=A7=88=20=EA=B5=AC=EB=B6=84)=EB=A1=9C=20=EC=97=AC?= =?UTF-8?q?=EB=9F=AC=20=ED=82=A4=EB=A5=BC=20=EC=84=A4=EC=A0=95=ED=95=98?= =?UTF-8?q?=EA=B3=A0=20GeminiKeyPool=EB=A1=9C=20generate=5Fcontent=C2=B7em?= =?UTF-8?q?bedding=20=ED=98=B8=EC=B6=9C=EC=9D=84=20=EB=B6=84=EC=82=B0?= =?UTF-8?q?=ED=95=9C=EB=8B=A4.=20429=20=EB=B0=9C=EC=83=9D=20=EC=8B=9C=20?= =?UTF-8?q?=EB=8B=A4=EB=A5=B8=20=ED=82=A4=EB=A1=9C=20=EC=A6=89=EC=8B=9C=20?= =?UTF-8?q?=EC=9E=AC=EC=8B=9C=EB=8F=84=ED=95=98=EB=A9=B0,=20GEMINI=5FAPI?= =?UTF-8?q?=5FKEY=20=EB=8B=A8=EC=9D=BC=20=EC=84=A4=EC=A0=95=EC=9D=80=20?= =?UTF-8?q?=ED=95=98=EC=9C=84=20=ED=98=B8=ED=99=98=EC=9D=84=20=EC=9C=A0?= =?UTF-8?q?=EC=A7=80=ED=95=9C=EB=8B=A4.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 3 +- app/core/config.py | 16 ++ app/core/deps.py | 9 ++ app/main.py | 27 +++- app/services/RagRerankService.py | 25 +++- app/services/VectorStoreService.py | 24 ++- .../internal/ai/ComplaintEmailLLMService.py | 3 + .../internal/ai/ComplaintEmailVLMService.py | 4 + .../internal/ai/IssuePinLLMService.py | 3 + .../internal/ai/IssueRagPlannerService.py | 3 + .../internal/ai/PolicyCardnewsImageService.py | 112 ++++++++++---- app/services/internal/ai/VLMService.py | 3 + app/services/internal/ai/gemini_factory.py | 30 ++-- app/services/internal/ai/gemini_key_pool.py | 100 +++++++++++++ app/services/internal/ai/gemini_retry.py | 137 ++++++++++++++++-- .../internal/ai/rotating_gemini_embedding.py | 85 +++++++++++ tests/test_gemini_key_pool.py | 91 ++++++++++++ 17 files changed, 599 insertions(+), 76 deletions(-) create mode 100644 app/services/internal/ai/gemini_key_pool.py create mode 100644 app/services/internal/ai/rotating_gemini_embedding.py create mode 100644 tests/test_gemini_key_pool.py 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..34284ac --- /dev/null +++ b/app/services/internal/ai/gemini_key_pool.py @@ -0,0 +1,100 @@ +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}") + 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..916743d 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 @@ -47,6 +50,75 @@ 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 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), + 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 +130,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 +163,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..4de2526 --- /dev/null +++ b/app/services/internal/ai/rotating_gemini_embedding.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +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 + + +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] = {} + + def _embed_model_at(self, index: int) -> GoogleGenAIEmbedding: + 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 _get_text_embedding(self, text: str) -> list[float]: + idx, _, _ = self._key_pool.acquire_sync() + return self._embed_model_at(idx).get_text_embedding(text) + + async def _aget_text_embedding(self, text: str) -> list[float]: + idx, _, _ = await self._key_pool.acquire() + return await self._embed_model_at(idx).aget_text_embedding(text) + + def _get_text_embeddings(self, texts: list[str]) -> list[list[float]]: + idx, _, _ = self._key_pool.acquire_sync() + return self._embed_model_at(idx).get_text_embedding_batch(texts) + + async def _aget_text_embeddings(self, texts: list[str]) -> list[list[float]]: + idx, _, _ = await self._key_pool.acquire() + return await self._embed_model_at(idx).aget_text_embedding_batch(texts) + + def _get_query_embedding(self, query: str) -> list[float]: + return self._get_text_embedding(query) + + async def _aget_query_embedding(self, query: str) -> list[float]: + return await self._aget_text_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..b0ff9ed --- /dev/null +++ b/tests/test_gemini_key_pool.py @@ -0,0 +1,91 @@ +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 + + +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() From a798876e75ec870c67c642130ebbd16e30042af4 Mon Sep 17 00:00:00 2001 From: selnem Date: Tue, 16 Jun 2026 18:22:22 +0900 Subject: [PATCH 2/3] =?UTF-8?q?fix:=20non-retryable=20=EC=98=A4=EB=A5=98?= =?UTF-8?q?=20=EC=A6=89=EC=8B=9C=20=EC=A4=91=EB=8B=A8=ED=95=98=EA=B3=A0=20?= =?UTF-8?q?embedding/key=20pool=20failover=20thread-safety=20=EA=B0=9C?= =?UTF-8?q?=EC=84=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/services/internal/ai/gemini_key_pool.py | 11 +-- app/services/internal/ai/gemini_retry.py | 18 ++++- .../internal/ai/rotating_gemini_embedding.py | 74 ++++++++++++++----- tests/test_gemini_key_pool.py | 66 +++++++++++++++++ 4 files changed, 143 insertions(+), 26 deletions(-) diff --git a/app/services/internal/ai/gemini_key_pool.py b/app/services/internal/ai/gemini_key_pool.py index 34284ac..8bce008 100644 --- a/app/services/internal/ai/gemini_key_pool.py +++ b/app/services/internal/ai/gemini_key_pool.py @@ -61,11 +61,12 @@ def _advance_index(self) -> int: 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}") - client = self._clients.get(index) - if client is None: - client = genai.Client(api_key=self._api_keys[index]) - self._clients[index] = client - return client + 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() diff --git a/app/services/internal/ai/gemini_retry.py b/app/services/internal/ai/gemini_retry.py index 916743d..f0e3639 100644 --- a/app/services/internal/ai/gemini_retry.py +++ b/app/services/internal/ai/gemini_retry.py @@ -19,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() @@ -37,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() @@ -103,6 +113,8 @@ async def _generate_with_key_failover( 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", @@ -111,7 +123,7 @@ async def _generate_with_key_failover( model, key_idx + 1, key_pool.size, - getattr(exc, "status_code", None), + getattr(exc, "status_code", None) or _gemini_error_status(exc), exc, ) if last_error is not None: diff --git a/app/services/internal/ai/rotating_gemini_embedding.py b/app/services/internal/ai/rotating_gemini_embedding.py index 4de2526..bfbaa7f 100644 --- a/app/services/internal/ai/rotating_gemini_embedding.py +++ b/app/services/internal/ai/rotating_gemini_embedding.py @@ -1,11 +1,13 @@ 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): @@ -27,34 +29,70 @@ def __init__( 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: - 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 + 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]: - idx, _, _ = self._key_pool.acquire_sync() - return self._embed_model_at(idx).get_text_embedding(text) + return self._execute_with_failover("get_text_embedding", text) async def _aget_text_embedding(self, text: str) -> list[float]: - idx, _, _ = await self._key_pool.acquire() - return await self._embed_model_at(idx).aget_text_embedding(text) + return await self._aexecute_with_failover("aget_text_embedding", text) def _get_text_embeddings(self, texts: list[str]) -> list[list[float]]: - idx, _, _ = self._key_pool.acquire_sync() - return self._embed_model_at(idx).get_text_embedding_batch(texts) + return self._execute_with_failover("get_text_embedding_batch", texts) async def _aget_text_embeddings(self, texts: list[str]) -> list[list[float]]: - idx, _, _ = await self._key_pool.acquire() - return await self._embed_model_at(idx).aget_text_embedding_batch(texts) + return await self._aexecute_with_failover("aget_text_embedding_batch", texts) def _get_query_embedding(self, query: str) -> list[float]: return self._get_text_embedding(query) diff --git a/tests/test_gemini_key_pool.py b/tests/test_gemini_key_pool.py index b0ff9ed..d48aa77 100644 --- a/tests/test_gemini_key_pool.py +++ b/tests/test_gemini_key_pool.py @@ -8,6 +8,7 @@ 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): @@ -89,3 +90,68 @@ async def test_failover_uses_second_client_on_429(self) -> None: ) 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() From 4302e45aadd6b7fd14290086d2544ca205190cde Mon Sep 17 00:00:00 2001 From: selnem Date: Tue, 16 Jun 2026 18:36:23 +0900 Subject: [PATCH 3/3] =?UTF-8?q?fix:=20rotating=20=EC=9E=84=EB=B2=A0?= =?UTF-8?q?=EB=94=A9=20=EC=BF=BC=EB=A6=AC=20=EA=B2=BD=EB=A1=9C=EB=A5=BC=20?= =?UTF-8?q?get=5Fquery=5Fembedding=EC=9C=BC=EB=A1=9C=20=EB=B6=84=EB=A6=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/ai/rotating_gemini_embedding.py | 4 ++-- tests/test_gemini_key_pool.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/app/services/internal/ai/rotating_gemini_embedding.py b/app/services/internal/ai/rotating_gemini_embedding.py index bfbaa7f..3fabd77 100644 --- a/app/services/internal/ai/rotating_gemini_embedding.py +++ b/app/services/internal/ai/rotating_gemini_embedding.py @@ -95,10 +95,10 @@ 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._get_text_embedding(query) + return self._execute_with_failover("get_query_embedding", query) async def _aget_query_embedding(self, query: str) -> list[float]: - return await self._aget_text_embedding(query) + return await self._aexecute_with_failover("aget_query_embedding", query) @classmethod def class_name(cls) -> str: diff --git a/tests/test_gemini_key_pool.py b/tests/test_gemini_key_pool.py index d48aa77..25131e7 100644 --- a/tests/test_gemini_key_pool.py +++ b/tests/test_gemini_key_pool.py @@ -155,3 +155,21 @@ def test_non_retryable_error_does_not_failover(self, mock_embedding_cls: MagicMo 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()