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

Added AzureAI support

Browse source
This commit is contained in:
JackColquitt 2024-09-28 16:55:43 -07:00
parent d389c4f45e
commit 8825ba530a
1 changed files with 20 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

26
trustgraph/model/text_completion/azure/llm.py
View file

@ -108,9 +108,7 @@ class Processor(ConsumerProducer):
result = resp.json()
message_content = result['choices'][0]['message']['content']
return message_content
return result
def handle(self, msg):
@ -132,9 +130,17 @@ class Processor(ConsumerProducer):
with __class__.text_completion_metric.time():
response = self.call_llm(prompt)
resp = response['choices'][0]['message']['content']
inputtokens = response['usage']['prompt_tokens']
outputtokens = response['usage']['completion_tokens']
print(resp, flush=True)
print(f"Input Tokens: {inputtokens}", flush=True)
print(f"Output Tokens: {outputtokens}", flush=True)
print("Send response...", flush=True)
r = TextCompletionResponse(response=response, error=None)
r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model="AzureAI")
self.producer.send(r, properties={"id": id})
except TooManyRequests:
@ -145,7 +151,11 @@ class Processor(ConsumerProducer):
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})
@ -162,7 +172,11 @@ class Processor(ConsumerProducer):
error=Error(
type = "llm-error",
message = str(e),
)
),
response=None,
in_token=None,
out_token=None,
model=None,
)
self.producer.send(r, properties={"id": id})