From 1294826aff7b4339b45135f881ec84feee16b077 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sat, 17 Aug 2024 18:41:25 +0100 Subject: [PATCH] Rate limit exception handling framework, but no much implementation --- trustgraph/base/consumer.py | 9 +++++++++ trustgraph/model/text_completion/azure/llm.py | 5 +++++ trustgraph/model/text_completion/bedrock/llm.py | 3 +++ trustgraph/model/text_completion/claude/llm.py | 2 ++ trustgraph/model/text_completion/cohere/llm.py | 17 +++++++++-------- trustgraph/model/text_completion/ollama/llm.py | 2 ++ trustgraph/model/text_completion/openai/llm.py | 2 ++ .../model/text_completion/vertexai/llm.py | 6 +++--- 8 files changed, 35 insertions(+), 11 deletions(-) diff --git a/trustgraph/base/consumer.py b/trustgraph/base/consumer.py index ec20349b..d5b9f133 100644 --- a/trustgraph/base/consumer.py +++ b/trustgraph/base/consumer.py @@ -1,8 +1,10 @@ from pulsar.schema import JsonSchema from prometheus_client import start_http_server, Histogram, Info, Counter +import time from . base_processor import BaseProcessor +from .. exceptions import TooManyRequests class Consumer(BaseProcessor): @@ -59,6 +61,13 @@ class Consumer(BaseProcessor): __class__.processing_metric.labels(status="success").inc() + except TooManyRequests: + self.consumer.negative_acknowledge(msg) + print("TooManyRequests: will retry") + __class__.processing_metric.labels(status="rate-limit").inc() + time.sleep(15) + continue + except Exception as e: print("Exception:", e, flush=True) diff --git a/trustgraph/model/text_completion/azure/llm.py b/trustgraph/model/text_completion/azure/llm.py index e3e7a559..71338455 100755 --- a/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph/model/text_completion/azure/llm.py @@ -12,6 +12,7 @@ from .... schema import text_completion_request_queue from .... schema import text_completion_response_queue from .... log_level import LogLevel from .... base import ConsumerProducer +from .... exceptions import TooManyRequests module = ".".join(__name__.split(".")[1:-1]) @@ -76,6 +77,10 @@ class Processor(ConsumerProducer): } resp = requests.post(url, data=body, headers=headers) + + if resp.status_code == 429: + raise TooManyRequests() + result = resp.json() message_content = result['choices'][0]['message']['content'] diff --git a/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph/model/text_completion/bedrock/llm.py index 8a158183..11485af0 100755 --- a/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph/model/text_completion/bedrock/llm.py @@ -13,6 +13,7 @@ from .... schema import text_completion_request_queue from .... schema import text_completion_response_queue from .... log_level import LogLevel from .... base import ConsumerProducer +from .... exceptions import TooManyRequests module = ".".join(__name__.split(".")[1:-1]) @@ -121,6 +122,8 @@ class Processor(ConsumerProducer): accept = 'application/json' contentType = 'application/json' + # FIXME: Consider catching request limits and raise TooManyRequests + # See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html response = self.bedrock.invoke_model(body=promptbody, modelId=self.model, accept=accept, contentType=contentType) # Mistral Response Structure diff --git a/trustgraph/model/text_completion/claude/llm.py b/trustgraph/model/text_completion/claude/llm.py index f5665201..29e5277d 100755 --- a/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph/model/text_completion/claude/llm.py @@ -57,6 +57,8 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) prompt = v.prompt + + # FIXME: Rate limits? response = message = self.claude.messages.create( model=self.model, max_tokens=1000, diff --git a/trustgraph/model/text_completion/cohere/llm.py b/trustgraph/model/text_completion/cohere/llm.py index ce1aa364..e52c2889 100755 --- a/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph/model/text_completion/cohere/llm.py @@ -59,15 +59,16 @@ class Processor(ConsumerProducer): prompt = v.prompt + # FIXME: Deal with rate limits? output = self.cohere.chat( - model=self.model, - message=prompt, - preamble = "You are a helpful AI-assistant.", - temperature=0.0, - chat_history=[], - prompt_truncation='auto', - connectors=[] - ) + model=self.model, + message=prompt, + preamble = "You are a helpful AI-assistant.", + temperature=0.0, + chat_history=[], + prompt_truncation='auto', + connectors=[] + ) resp = output.text print(resp, flush=True) diff --git a/trustgraph/model/text_completion/ollama/llm.py b/trustgraph/model/text_completion/ollama/llm.py index 90cd9a53..900140a2 100755 --- a/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph/model/text_completion/ollama/llm.py @@ -65,6 +65,8 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) prompt = v.prompt + + # FIXME: Rate limits? response = self.llm.invoke(prompt) print("Send response...", flush=True) diff --git a/trustgraph/model/text_completion/openai/llm.py b/trustgraph/model/text_completion/openai/llm.py index 44226bbd..25b84070 100755 --- a/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph/model/text_completion/openai/llm.py @@ -57,6 +57,8 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) prompt = v.prompt + + # FIXME: Rate limits resp = self.openai.chat.completions.create( model=self.model, messages=[ diff --git a/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph/model/text_completion/vertexai/llm.py index 4861dd8e..3da0c041 100755 --- a/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph/model/text_completion/vertexai/llm.py @@ -26,6 +26,7 @@ from .... schema import text_completion_request_queue from .... schema import text_completion_response_queue from .... log_level import LogLevel from .... base import ConsumerProducer +from .... exceptions import TooManyRequests module = ".".join(__name__.split(".")[1:-1]) @@ -139,9 +140,8 @@ class Processor(ConsumerProducer): except google.api_core.exceptions.ResourceExhausted: - print("429, resource busy, sleeping", flush=True) - time.sleep(15) - self.consumer.negative_acknowledge(msg) + # 429 / rate limits case + raise TooManyRequests # Let other exceptions fall through