From 420e0282de1856f86cd8c0a574b8df6552898d82 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 25 Oct 2024 00:00:52 +0100 Subject: [PATCH] Catch rate limit --- .../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 @@ Input is prompt, output is response. 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 @@ class Processor(ConsumerProducer): # 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)