Skip to content

Commit 6455d56

Browse files
committed
Fix Gemini client handling improvements
1 parent 196225e commit 6455d56

File tree

1 file changed

+20
-11
lines changed

1 file changed

+20
-11
lines changed

ufo/llm/gemini.py

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,8 @@ def chat_completion(
101101
"yellow"
102102
)
103103
time.sleep(sleep_time)
104+
if attempt == self.max_retry - 1:
105+
raise
104106

105107
return self.get_text_from_all_candidates(response), cost
106108

@@ -209,14 +211,21 @@ def get_text_from_all_candidates(self, response: GenerateContentResponse) -> Lis
209211

210212
return all_texts
211213

212-
@functools.lru_cache()
213-
@staticmethod
214-
def get_gemini_client(api_key: str) -> genai.Client:
215-
"""
216-
Create a Gemini client using the provided API key.
217-
:param api_key: The API key for authentication.
218-
:return: A Gemini client instance.
219-
"""
220-
return genai.Client(
221-
api_key=api_key
222-
)
214+
215+
# Module-level cache for Gemini clients
216+
_gemini_client_cache: dict = {}
217+
218+
219+
def _get_gemini_client(api_key: str) -> genai.Client:
220+
"""
221+
Create or retrieve a cached Gemini client using the provided API key.
222+
:param api_key: The API key for authentication.
223+
:return: A Gemini client instance.
224+
"""
225+
if api_key not in _gemini_client_cache:
226+
_gemini_client_cache[api_key] = genai.Client(api_key=api_key)
227+
return _gemini_client_cache[api_key]
228+
229+
230+
# Bind as static method for backward compatibility
231+
GeminiService.get_gemini_client = staticmethod(_get_gemini_client)

0 commit comments

Comments
 (0)