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

Token counting working for Bedrock

Browse source
This commit is contained in:
JackColquitt 2024-09-27 17:25:23 -07:00 • committed by Cyber MacGeddon
parent 74fcba1b3e
commit 8cbe1a27a7
2 changed files with 16 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
trustgraph/metering/counter.py
View file

@ -42,8 +42,8 @@ class Processor(Consumer):
num_in = v.in_token
num_out = v.out_token
print(f"Input Tokens: {num_in}")
print(f"Output Tokens: {num_out}")
print(f"Input Tokens: {num_in}", flush=True)
print(f"Output Tokens: {num_out}", flush=True)
@staticmethod
def add_args(parser):

16
trustgraph/model/text_completion/bedrock/llm.py
View file

@ -209,14 +209,22 @@ class Processor(ConsumerProducer):
# Use Mistral as default
else:
response_body = json.loads(response.get("body").read())
outputtext = response_body['outputs'][0]['text']
outputtext = response_body['outputs'][0]['text']
metadata = response['ResponseMetadata']['HTTPHeaders']
inputtokens = int(metadata['x-amzn-bedrock-input-token-count'])
outputtokens = int(metadata['x-amzn-bedrock-output-token-count'])
print(outputtext, flush=True)
print(f"Input Tokens: {inputtokens}", flush=True)
print(f"Output Tokens: {outputtokens}", flush=True)
print("Send response...", flush=True)
r = TextCompletionResponse(
error=None,
response=outputtext
response=outputtext,
in_token=inputtokens,
out_token=outputtokens,
)
self.send(r, properties={"id": id})
@ -236,6 +244,8 @@ class Processor(ConsumerProducer):
message = str(e),
),
response=None,
in_token=None,
out_token=None,
)
self.producer.send(r, properties={"id": id})
@ -254,6 +264,8 @@ class Processor(ConsumerProducer):
message = str(e),
),
response=None,
in_token=None,
out_token=None,
)
self.consumer.acknowledge(msg)