Retry on VertexAI's InternalServerError

google · Nov 14, 2024 · 1c2d235 · 1c2d235
1 parent 60f3db1
commit 1c2d235
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llm_toolkit/models.py b/llm_toolkit/models.py
@@ -31,7 +31,8 @@
 import openai
 import tiktoken
 import vertexai
-from google.api_core.exceptions import (GoogleAPICallError, InvalidArgument,
+from google.api_core.exceptions import (GoogleAPICallError,
+                                        InternalServerError, InvalidArgument,
                                         ResourceExhausted)
 from vertexai import generative_models
 from vertexai.preview.generative_models import ChatSession, GenerativeModel
@@ -650,6 +651,7 @@ def get_chat_client(self, model: GenerativeModel) -> Any:
           InvalidArgument,
           ValueError,  # TODO(dongge): Handle RECITATION specifically.
           IndexError,  # A known error from vertexai.
+          InternalServerError,  # A known error from vertexai.
       ],
       other_exceptions={ResourceExhausted: 100})
   def _do_generate(self, client: ChatSession, prompt: str,