diff --git a/trustgraph/metering/pricelist.py b/trustgraph/metering/pricelist.py index da54b44f..41c1adb2 100644 --- a/trustgraph/metering/pricelist.py +++ b/trustgraph/metering/pricelist.py @@ -65,5 +65,15 @@ price_list = { "input_price": 0.000003, "output_price": 0.000015 }, + { + "model_name": "command-r-08-202", + "input_price": 0.0000025, + "output_price": 0.000010 + }, + { + "model_name": "c4ai-aya-23-8b", + "input_price": 0, + "output_price": 0 + }, ] } \ No newline at end of file diff --git a/trustgraph/model/text_completion/cohere/llm.py b/trustgraph/model/text_completion/cohere/llm.py index af55fd2c..75ef8ae4 100755 --- a/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph/model/text_completion/cohere/llm.py @@ -91,10 +91,15 @@ class Processor(ConsumerProducer): ) resp = output.text + inputtokens = output.meta.billed_units.input_tokens + outputtokens = output.meta.billed_units.output_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=resp, error=None) + 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) @@ -111,6 +116,9 @@ class Processor(ConsumerProducer): message = str(e), ), response=None, + in_token=None, + out_token=None, + model=None, ) self.producer.send(r, properties={"id": id}) @@ -129,6 +137,9 @@ class Processor(ConsumerProducer): message = str(e), ), response=None, + in_token=None, + out_token=None, + model=None, ) self.producer.send(r, properties={"id": id})