From db4e1df0fbae574f4a3d0807e5c5fef5b0e91b50 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Fri, 25 Oct 2024 00:01:40 +0100 Subject: [PATCH] Catch rate limit (#124) --- .../trustgraph/model/text_completion/googleaistudio/llm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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)