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 Cohere API support

Browse source
This commit is contained in:
JackColquitt 2024-09-28 16:05:32 -07:00
parent 75d0a90b4f
commit 843dee5c90
2 changed files with 22 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

10
trustgraph/metering/pricelist.py
View file

@ -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
},
]
}

13
trustgraph/model/text_completion/cohere/llm.py
View file

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