From e80f5a4d6952585776188102f1d6a235a834a7c8 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sun, 26 Jan 2025 17:59:07 +0000 Subject: [PATCH] Add rate limit events to VertexAI and Google AI Studio --- .../text_completion/googleaistudio/llm.py | 38 ++++++++----------- .../model/text_completion/vertexai/llm.py | 20 +++------- 2 files changed, 21 insertions(+), 37 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index a249998d..5d5b23a0 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -88,7 +88,8 @@ class Processor(ConsumerProducer): HarmCategory.HARM_CATEGORY_HARASSMENT: block_level, HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: block_level, HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: block_level, - # There is a documentation conflict on whether or not CIVIC_INTEGRITY is a valid category + # There is a documentation conflict on whether or not + # CIVIC_INTEGRITY is a valid category # HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level, } @@ -122,8 +123,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits? - with __class__.text_completion_metric.time(): chat_session = self.llm.start_chat( @@ -140,35 +139,30 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) + r = TextCompletionResponse( + response=resp, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model + ) self.send(r, properties={"id": id}) print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit except ResourceExhausted as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the default handler + raise TooManyRequests() except Exception as e: - print(f"Exception: {e}") + # Apart from rate limits, treat all exceptions as unrecoverable + + print(type(e), flush=True) + print(f"Exception: {e}", flush=True) print("Send error response...", flush=True) diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index cb817836..524316f2 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -178,25 +178,15 @@ class Processor(ConsumerProducer): except google.api_core.exceptions.ResourceExhausted as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the default handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True)