Feature/pkgsplit (#83)

* Starting to spawn base package
* More package hacking
* Bedrock and VertexAI
* Parquet split
* Updated templates
* Utils

trustgraph-flow/trustgraph/metering/__init__.py

@@ -0,0 +1,3 @@
+
+from . counter import *
+

trustgraph-flow/trustgraph/metering/__main__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . counter import run
+
+if __name__ == '__main__':
+    run()
+

trustgraph-flow/trustgraph/metering/counter.py

@@ -0,0 +1,75 @@
+"""
+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
+from .. log_level import LogLevel
+from .. base import Consumer
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = text_completion_response_queue
+default_subscriber = module
+
+
+class Processor(Consumer):
+
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+
+        super(Processor, self).__init__(
+            **params | {
+                "input_queue": input_queue,
+                "subscriber": subscriber,
+                "input_schema": TextCompletionResponse,
+            }
+        )
+
+    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"]
+
+        print(f"Handling response {id}...", flush=True)
+
+        num_in = v.in_token
+        num_out = v.out_token
+
+        model_input_price, model_output_price = self.get_prices(price_list, modelname)
+
+        if model_input_price == None:
+            cost_per_call = f"Model Not Found in Price list"
+        else:
+            cost_in = num_in * model_input_price
+            cost_out = num_out * model_output_price
+            cost_per_call = round(cost_in + cost_out, 6)
+
+        print(f"Input Tokens: {num_in}", flush=True)
+        print(f"Output Tokens: {num_out}", flush=True)
+        print(f"Cost for call: ${cost_per_call}", flush=True)
+
+    @staticmethod
+    def add_args(parser):
+
+        Consumer.add_args(
+            parser, default_input_queue, default_subscriber,
+        )
+
+def run():
+
+    Processor.start(module, __doc__)

trustgraph-flow/trustgraph/metering/pricelist.py

@@ -0,0 +1,104 @@
+price_list = {
+  "price_list": [
+    {
+      "model_name": "mistral.mistral-large-2407-v1:0",
+      "input_price": 0.000004,
+      "output_price": 0.000012
+    },
+    {
+      "model_name": "meta.llama3-1-405b-instruct-v1:0",
+      "input_price": 0.00000532,
+      "output_price": 0.000016
+    },
+    {
+      "model_name": "mistral.mixtral-8x7b-instruct-v0:1",
+      "input_price": 0.00000045,
+      "output_price": 0.0000007
+    },
+    {
+      "model_name": "meta.llama3-1-70b-instruct-v1:0",
+      "input_price": 0.00000099,
+      "output_price": 0.00000099
+    },
+    {
+      "model_name": "meta.llama3-1-8b-instruct-v1:0",
+      "input_price": 0.00000022,
+      "output_price": 0.00000022
+    },
+    {
+      "model_name": "anthropic.claude-3-haiku-20240307-v1:0",
+      "input_price": 0.00000025,
+      "output_price": 0.00000125
+    },
+    {
+      "model_name": "anthropic.claude-3-5-sonnet-20240620-v1:0",
+      "input_price": 0.000003,
+      "output_price": 0.000015
+    },
+    {
+      "model_name": "cohere.command-r-plus-v1:0",
+      "input_price": 0.0000030,
+      "output_price": 0.0000150
+    },
+    {
+      "model_name": "ollama",
+      "input_price": 0,
+      "output_price": 0
+    },
+    {
+      "model_name": "claude-3-haiku-20240307",
+      "input_price": 0.00000025,
+      "output_price": 0.00000125
+    },
+    {
+      "model_name": "claude-3-5-sonnet-20240620",
+      "input_price": 0.000003,
+      "output_price": 0.000015
+    },
+    {
+      "model_name": "claude-3-opus-20240229",
+      "input_price": 0.000015,
+      "output_price": 0.000075
+    },
+    {
+      "model_name": "claude-3-sonnet-20240229",
+      "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
+    },
+    {
+      "model_name": "llama.cpp",
+      "input_price": 0,
+      "output_price": 0
+    },
+    {
+      "model_name": "gpt-4o",
+      "input_price": 0.000005,
+      "output_price": 0.000015
+    },
+    {
+      "model_name": "gpt-4o-2024-08-06",
+      "input_price": 0.0000025,
+      "output_price": 0.000010
+    },
+    {
+      "model_name": "gpt-4o-2024-05-13",
+      "input_price": 0.000005,
+      "output_price": 0.000015
+    },
+    {
+      "model_name": "gpt-4o-mini",
+      "input_price": 0.00000015,
+      "output_price": 0.0000006
+    },
+  ]
+}