11from __future__ import annotations
22
3+ import asyncio
34import hashlib
5+ import random
46import struct
57from typing import Protocol
68
9+ import structlog
10+
11+ log = structlog .get_logger (__name__ )
12+
713
814class EmbeddingError (Exception ):
915 def __init__ (self , * , code : str , message : str , retriable : bool ) -> None :
@@ -13,6 +19,16 @@ def __init__(self, *, code: str, message: str, retriable: bool) -> None:
1319 self .retriable = retriable
1420
1521
22+ # Gemini 가 한도 초과(429)일 때만 백오프 재시도한다. 다른 오류(인증·잘못된 입력 등)는
23+ # 재시도해도 동일하므로 즉시 실패시킨다. SDK ClientError 는 .code(HTTP)·.status 를 노출한다.
24+ def _is_rate_limited (exc : Exception ) -> bool :
25+ if getattr (exc , "code" , None ) == 429 :
26+ return True
27+ if getattr (exc , "status" , None ) == "RESOURCE_EXHAUSTED" :
28+ return True
29+ return "RESOURCE_EXHAUSTED" in str (exc )
30+
31+
1632# 구현체는 바꿔서 사용할 수 있음
1733class EmbeddingProvider (Protocol ):
1834 @property
@@ -61,7 +77,20 @@ def _embed_one(self, text: str) -> list[float]:
6177# Gemini Embedding 을 사용합니다.
6278# 이건 충대키로 안되니 키 발급 필요함
6379class GeminiEmbeddingProvider :
64- def __init__ (self , * , api_key : str , model : str , dim : int ) -> None :
80+ # 한 요청에 너무 많은 청크를 담으면 분당 토큰 한도(TPM)에 걸려 429 가 난다.
81+ # batch_size 로 쪼개 순차 호출하고, 429 는 지수 백오프로 재시도한다.
82+ _MAX_BACKOFF_SEC = 30.0
83+
84+ def __init__ (
85+ self ,
86+ * ,
87+ api_key : str ,
88+ model : str ,
89+ dim : int ,
90+ batch_size : int = 32 ,
91+ max_retries : int = 5 ,
92+ retry_base_delay_sec : float = 2.0 ,
93+ ) -> None :
6594 if not api_key :
6695 raise ValueError ("GEMINI_API_KEY 누락 — provider=gemini 사용 불가" )
6796 if dim <= 0 :
@@ -71,6 +100,9 @@ def __init__(self, *, api_key: str, model: str, dim: int) -> None:
71100 self ._client = genai .Client (api_key = api_key )
72101 self ._model = model
73102 self ._dim = dim
103+ self ._batch_size = max (1 , batch_size )
104+ self ._max_retries = max (0 , max_retries )
105+ self ._retry_base_delay = max (0.0 , retry_base_delay_sec )
74106
75107 @property
76108 def dim (self ) -> int :
@@ -87,25 +119,50 @@ async def embed(
87119 return []
88120 from google .genai import types as genai_types
89121
90- try :
91- resp = await self ._client .aio .models .embed_content (
92- model = self ._model ,
93- contents = texts ,
94- config = genai_types .EmbedContentConfig (
95- # 인덱싱은 RETRIEVAL_DOCUMENT, 검색 쿼리는 RETRIEVAL_QUERY 로
96- # 분리해야 Gemini embedding 의 코사인 정합도가 최적화된다.
97- task_type = task_type ,
98- output_dimensionality = self ._dim ,
99- ),
100- )
101- except Exception as exc :
102- raise EmbeddingError (
103- code = "GEMINI_FAILED" ,
104- message = f"Gemini embedding 호출 실패: { exc } ,
105- retriable = True ,
106- ) from exc
107-
108- return [list (e .values ) for e in resp .embeddings ]
122+ config = genai_types .EmbedContentConfig (
123+ # 인덱싱은 RETRIEVAL_DOCUMENT, 검색 쿼리는 RETRIEVAL_QUERY 로
124+ # 분리해야 Gemini embedding 의 코사인 정합도가 최적화된다.
125+ task_type = task_type ,
126+ output_dimensionality = self ._dim ,
127+ )
128+
129+ vectors : list [list [float ]] = []
130+ for start in range (0 , len (texts ), self ._batch_size ):
131+ batch = texts [start : start + self ._batch_size ]
132+ resp = await self ._embed_batch_with_retry (batch , config )
133+ vectors .extend (list (e .values ) for e in resp .embeddings )
134+ return vectors
135+
136+ async def _embed_batch_with_retry (self , batch : list [str ], config : object ) -> object :
137+ attempt = 0
138+ while True :
139+ try :
140+ return await self ._client .aio .models .embed_content (
141+ model = self ._model ,
142+ contents = batch ,
143+ config = config ,
144+ )
145+ except Exception as exc :
146+ rate_limited = _is_rate_limited (exc )
147+ if rate_limited and attempt < self ._max_retries :
148+ delay = min (
149+ self ._retry_base_delay * (2 ** attempt ), self ._MAX_BACKOFF_SEC
150+ )
151+ delay += random .uniform (0.0 , self ._retry_base_delay * 0.1 )
152+ log .warning (
153+ "embed.gemini.rate_limited" ,
154+ attempt = attempt + 1 ,
155+ max_retries = self ._max_retries ,
156+ delay_sec = round (delay , 2 ),
157+ )
158+ await asyncio .sleep (delay )
159+ attempt += 1
160+ continue
161+ raise EmbeddingError (
162+ code = "GEMINI_RATE_LIMITED" if rate_limited else "GEMINI_FAILED" ,
163+ message = f"Gemini embedding 호출 실패: { exc } ,
164+ retriable = True ,
165+ ) from exc
109166
110167
111168def build_embedding_provider (
@@ -114,11 +171,21 @@ def build_embedding_provider(
114171 dim : int ,
115172 model : str ,
116173 gemini_api_key : str = "" ,
174+ batch_size : int = 32 ,
175+ max_retries : int = 5 ,
176+ retry_base_delay_sec : float = 2.0 ,
117177) -> EmbeddingProvider :
118178 if provider == "mock" :
119179 return MockEmbeddingProvider (dim = dim , model = model )
120180 if provider == "gemini" :
121- return GeminiEmbeddingProvider (api_key = gemini_api_key , model = model , dim = dim )
181+ return GeminiEmbeddingProvider (
182+ api_key = gemini_api_key ,
183+ model = model ,
184+ dim = dim ,
185+ batch_size = batch_size ,
186+ max_retries = max_retries ,
187+ retry_base_delay_sec = retry_base_delay_sec ,
188+ )
122189 if provider == "openai" :
123190 raise NotImplementedError ("openai embedding provider 미구현 — 후속 PR에서 추가" )
124191 if provider == "ollama" :
0 commit comments