Price calc using price list

trustgraph/metering/counter.py

@@ -3,6 +3,7 @@ Simple token counter for each LLM response.
 """
 
 from prometheus_client import Histogram, Info
+from . pricelist import price_list
 
 from .. schema import TextCompletionResponse, Error
 from .. schema import text_completion_response_queue
@@ -30,9 +31,16 @@ class Processor(Consumer):
             }
         )
 
+    def get_prices(self, prices, modelname):
+        for model in prices["price_list"]:
+            if model["model_name"] == modelname:
+                return model["input_price"], model["output_price"]
+        return None, None  # Return None if model is not found
+
     def handle(self, msg):
 
         v = msg.value()
+        modelname = v.model
 
         # Sender-produced ID
         id = msg.properties()["id"]
@@ -42,8 +50,14 @@ class Processor(Consumer):
         num_in = v.in_token
         num_out = v.out_token
 
+        model_input_price, model_output_price = self.get_prices(price_list, modelname)
+        cost_in = num_in * model_input_price
+        cost_out = num_out * model_output_price
+        cost_per_call = cost_in + cost_out
+
         print(f"Input Tokens: {num_in}", flush=True)
         print(f"Output Tokens: {num_out}", flush=True)
+        print(f"Cost for call: ${cost_per_call:.6f}", flush=True)
 
     @staticmethod
     def add_args(parser):

trustgraph/metering/pricelist.py

@@ -0,0 +1,29 @@
+price_list = {
+  "price_list": [
+    {
+      "model_name": "GPT-3.5-Turbo",
+      "input_price": 0.0015,
+      "output_price": 0.002
+    },
+    {
+      "model_name": "GPT-4",
+      "input_price": 0.03,
+      "output_price": 0.06
+    },
+    {
+      "model_name": "mistral.mixtral-8x7b-instruct-v0:1",
+      "input_price": 0.00000045,
+      "output_price": 0.0000007
+    },
+    {
+      "model_name": "DALL-E 2",
+      "input_price": 0.016,
+      "output_price": 0.016
+    },
+    {
+      "model_name": "Jurassic-1 Jumbo",
+      "input_price": 0.01,
+      "output_price": 0.01
+    }
+  ]
+}

trustgraph/model/text_completion/bedrock/llm.py

@@ -225,6 +225,7 @@ class Processor(ConsumerProducer):
                 response=outputtext,
                 in_token=inputtokens,
                 out_token=outputtokens,
+                model=str(self.model),
             )
 
             self.send(r, properties={"id": id})
@@ -246,6 +247,7 @@ class Processor(ConsumerProducer):
                 response=None,
                 in_token=None,
                 out_token=None,
+                model=None,
             )
 
             self.producer.send(r, properties={"id": id})
@@ -266,6 +268,7 @@ class Processor(ConsumerProducer):
                 response=None,
                 in_token=None,
                 out_token=None,
+                model=None,
             )
 
             self.consumer.acknowledge(msg)

trustgraph/schema/models.py

@@ -16,6 +16,7 @@ class TextCompletionResponse(Record):
     response = String()
     in_token = Integer()
     out_token = Integer()
+    model = String()
 
 text_completion_request_queue = topic(
     'text-completion', kind='non-persistent', namespace='request'