feat: Gemini 다중 API 키 round-robin 로테이션 및 429 failover#129
Merged
Conversation
GEMINI_API_KEYS(콤마 구분)로 여러 키를 설정하고 GeminiKeyPool로 generate_content·embedding 호출을 분산한다. 429 발생 시 다른 키로 즉시 재시도하며, GEMINI_API_KEY 단일 설정은 하위 호환을 유지한다.
7 tasks
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a Gemini API key pool (GeminiKeyPool) to support round-robin rotation and failover across multiple API keys, enhancing the reliability of various AI services (VLM, LLM, embedding, and image generation). Feedback on these changes highlights several critical improvements: ensuring that failover in gemini_retry.py is only triggered for retryable errors to prevent wasting API calls, adding thread-safety and failover capabilities to the rotating embedding service, securing thread-safety when initializing clients in the key pool, and refactoring duplicate failover logic in the cardnews image service into a shared helper.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Collaborator
Author
|
/gemini review |
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a Gemini API key pool (GeminiKeyPool) to support round-robin rotation and failover across multiple API keys, updating various AI services and embeddings to utilize this pool. The review feedback highlights two key improvement opportunities: first, correcting RotatingGoogleGenAIEmbedding to use query-specific embedding methods instead of document embedding methods to avoid degrading RAG retrieval performance; second, refactoring duplicated failover logic in PolicyCardnewsImageService into a reusable helper method to improve maintainability.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Collaborator
Author
|
/gemini review |
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a round-robin Gemini API key rotation and failover mechanism via a new GeminiKeyPool and RotatingGoogleGenAIEmbedding. This allows distributing API requests across multiple keys configured via GEMINI_API_KEYS and automatically failing over to other keys when retryable errors (like 429 rate limits) occur. The feedback suggests several improvements: securing the new gemini_api_keys configuration with SecretStr to prevent credential exposure, calling protected embedding methods in RotatingGoogleGenAIEmbedding to avoid duplicate LlamaIndex callback triggers, updating the corresponding unit tests, and refactoring duplicate failover logic in PolicyCardnewsImageService into a shared helper method.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+97
to
+101
| gemini_api_keys: str = Field( | ||
| default="", | ||
| alias="GEMINI_API_KEYS", | ||
| description="콤마 구분 다중 키 (설정 시 GEMINI_API_KEY보다 우선)", | ||
| ) |
Contributor
There was a problem hiding this comment.
API 키와 같은 민감한 정보는 로그 출력이나 직렬화 시 노출되는 것을 방지하기 위해 SecretStr 타입을 사용하는 것이 안전합니다.
Suggested change
| gemini_api_keys: str = Field( | |
| default="", | |
| alias="GEMINI_API_KEYS", | |
| description="콤마 구분 다중 키 (설정 시 GEMINI_API_KEY보다 우선)", | |
| ) | |
| gemini_api_keys: SecretStr = Field( | |
| default=SecretStr(""), | |
| alias="GEMINI_API_KEYS", | |
| description="콤마 구분 다중 키 (설정 시 GEMINI_API_KEY보다 우선)", | |
| ) |
Comment on lines
+677
to
+686
| 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)) |
Contributor
There was a problem hiding this comment.
gemini_api_keys가 SecretStr로 변경됨에 따라, 값을 읽어올 때 get_secret_value()를 사용하도록 수정합니다.
Suggested change
| 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)) | |
| def resolved_gemini_api_keys(self) -> tuple[str, ...]: | |
| from app.services.internal.ai.gemini_retry import parse_gemini_model_list | |
| api_keys_str = self.gemini_api_keys.get_secret_value() | |
| if api_keys_str.strip(): | |
| raw = list(parse_gemini_model_list(api_keys_str)) | |
| 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)) |
Comment on lines
+85
to
+101
| 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) |
Contributor
There was a problem hiding this comment.
GoogleGenAIEmbedding (또는 BaseEmbedding)의 공개 메서드(get_text_embedding, get_text_embedding_batch 등)를 호출하면 LlamaIndex 내부의 콜백 매니저가 중복으로 트리거되어 불필요한 로깅이나 트레이싱 스팬이 중복 생성될 수 있습니다. 내부 구현 메서드인 protected 메서드(_get_text_embedding, _get_text_embeddings 등)를 호출하도록 변경하는 것이 좋습니다.
Suggested change
| 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) | |
| 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_embeddings", texts) | |
| async def _aget_text_embeddings(self, texts: list[str]) -> list[list[float]]: | |
| return await self._aexecute_with_failover("_aget_text_embeddings", 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) |
Comment on lines
+118
to
+175
| 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() |
Contributor
There was a problem hiding this comment.
RotatingGoogleGenAIEmbedding이 내부적으로 protected 메서드를 호출하도록 변경됨에 따라, 테스트 코드의 모킹 대상 메서드도 _get_text_embedding 및 _get_query_embedding으로 수정합니다.
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()
Comment on lines
99
to
+179
| 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) |
Contributor
There was a problem hiding this comment.
_generate_with_gemini_image와 _generate_with_imagen 메서드에 중복되어 있는 복잡한 failover 로직을 공통 헬퍼 메서드로 추출하여 코드 중복을 제거하고 가독성을 향상시킬 수 있습니다.
async def _execute_with_failover(self, invoke_fn) -> bytes:
pool = self.key_pool
if pool is not None and pool.enabled:
key_index, _, client = await pool.acquire()
try:
return await invoke_fn(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_fn(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_fn(self.client)
async def _generate_with_gemini_image(self, *, model: str, prompt: str) -> bytes:
config = types.GenerateContentConfig(
response_modalities=[types.Modality.IMAGE],
)
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)
return await self._execute_with_failover(invoke)
async def _generate_with_imagen(self, *, model: str, prompt: str) -> bytes:
image_config = types.GenerateImagesConfig(
number_of_images=1,
aspect_ratio="3:4",
output_mime_type="image/png",
)
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)
return await self._execute_with_failover(invoke)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🔗 Related Issue
✨ 작업 개요
Gemini API 다중 키 round-robin 로테이션을 도입해 rate limit 부담을 여러 계정에 분산합니다.
GEMINI_API_KEYS(콤마 구분)로 N개 키 설정,GEMINI_API_KEY단일 키 하위 호환GeminiKeyPool: 호출마다 키 순환, 429 시 다른 키로 즉시 failovergenerate_content(LLM/VLM) 및 embedding(VectorStore, RAG rerank) 경로에 공통 적용주요 변경 파일
app/services/internal/ai/gemini_key_pool.py(신규)app/services/internal/ai/rotating_gemini_embedding.py(신규)app/services/internal/ai/gemini_retry.pyapp/services/internal/ai/gemini_factory.pyapp/main.py,app/core/config.py,app/core/deps.pyVectorStoreService,RagRerankService, LLM/VLM 서비스,PolicyCardnewsImageServicetests/test_gemini_key_pool.py설정 예시