apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 01:01:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Add rate limit events to VertexAI and Google AI Studio

Browse source
This commit is contained in:
Cyber MacGeddon 2025-01-26 17:59:07 +00:00
parent 24ef170797
commit e80f5a4d69
2 changed files with 21 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

38
trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py
View file

@ -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)

20
trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py
View file

@ -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)