diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index 345a5441..01dc9736 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -6,6 +6,7 @@ import google.generativeai as genai from google.generativeai.types import HarmCategory, HarmBlockThreshold +from google.api_core.exceptions import ResourceExhausted from prometheus_client import Histogram import os @@ -139,7 +140,7 @@ def handle(self, msg): # FIXME: Wrong exception, don't know what this LLM throws # for a rate limit - except TooManyRequests: + except ResourceExhausted as e: print("Send rate limit response...", flush=True)