diff --git a/letta/llm_api/google_ai_client.py b/letta/llm_api/google_ai_client.py index 653cea5b..788994c9 100644 --- a/letta/llm_api/google_ai_client.py +++ b/letta/llm_api/google_ai_client.py @@ -49,11 +49,14 @@ def get_gemini_endpoint_and_headers( return url, headers -def google_ai_check_valid_api_key(api_key: str): +async def google_ai_check_valid_api_key_async(api_key: str): + """ + Async version to check if Google AI API key is valid without blocking the event loop. + """ client = genai.Client(api_key=api_key) # use the count token endpoint for a cheap model - as of 5/7/2025 this is slightly faster than fetching the list of models try: - client.models.count_tokens( + await client.aio.models.count_tokens( model=GOOGLE_MODEL_FOR_API_KEY_CHECK, contents="", ) diff --git a/letta/schemas/providers/google_gemini.py b/letta/schemas/providers/google_gemini.py index 161ac568..1261e668 100644 --- a/letta/schemas/providers/google_gemini.py +++ b/letta/schemas/providers/google_gemini.py @@ -21,10 +21,10 @@ class GoogleAIProvider(Provider): base_url: str = "https://generativelanguage.googleapis.com" async def check_api_key(self): - from letta.llm_api.google_ai_client import google_ai_check_valid_api_key + from letta.llm_api.google_ai_client import google_ai_check_valid_api_key_async api_key = self.get_api_key_secret().get_plaintext() - google_ai_check_valid_api_key(api_key) + await google_ai_check_valid_api_key_async(api_key) async def list_llm_models_async(self): from letta.llm_api.google_ai_client import google_ai_get_model_list_async