Refactor rate limit handling (#280)

* - Refactored retry for rate limits into the base class
- ConsumerProducer is derived from Consumer to simplify code
- Added rate_limit_count metrics for rate limit events

* Add rate limit events to VertexAI and Google AI Studio

* Added Grafana rate limit dashboard

* Add rate limit handling to all LLMs

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)