Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions ufo/llm/gemini.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,8 @@ def chat_completion(
f"Retrying in {sleep_time:.2f}s..."
)
time.sleep(sleep_time)
if attempt == self.max_retry - 1:
raise

return self.get_text_from_all_candidates(response), cost

Expand Down Expand Up @@ -234,12 +236,21 @@ def get_text_from_all_candidates(

return all_texts

@functools.lru_cache()
@staticmethod
def get_gemini_client(api_key: str) -> genai.Client:
"""
Create a Gemini client using the provided API key.
:param api_key: The API key for authentication.
:return: A Gemini client instance.
"""
return genai.Client(api_key=api_key)

# Module-level cache for Gemini clients
_gemini_client_cache: dict = {}


def _get_gemini_client(api_key: str) -> genai.Client:
"""
Create or retrieve a cached Gemini client using the provided API key.
:param api_key: The API key for authentication.
:return: A Gemini client instance.
"""
if api_key not in _gemini_client_cache:
_gemini_client_cache[api_key] = genai.Client(api_key=api_key)
return _gemini_client_cache[api_key]


# Bind as static method for backward compatibility
GeminiService.get_gemini_client = staticmethod(_get_gemini_client)