From d884a4d9f334054332d4c49b227d039c19d668f4 Mon Sep 17 00:00:00 2001
From: Uladzislau Vishneuski <Vladisavvv@gmail.com>
Date: Tue, 21 Nov 2023 03:39:37 +0100
Subject: [PATCH] feat: calculate price based on deployment if it is possible
 (#11) (#13)

---
 aidial_analytics_realtime/analytics.py |  2 +-
 aidial_analytics_realtime/rates.py     | 16 +++++--
 tests/test_rates.py                    | 60 +++++++++++++++++++++++---
 3 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/aidial_analytics_realtime/analytics.py b/aidial_analytics_realtime/analytics.py
index 5941390..f4613f3 100644
--- a/aidial_analytics_realtime/analytics.py
+++ b/aidial_analytics_realtime/analytics.py
@@ -91,7 +91,7 @@ def make_point(
             )
 
     price = rates_calculator.calculate_price(
-        model, request_content, response_content, usage
+        deployment, model, request_content, response_content, usage
     )
 
     point = (
diff --git a/aidial_analytics_realtime/rates.py b/aidial_analytics_realtime/rates.py
index e855f93..f8c88ff 100644
--- a/aidial_analytics_realtime/rates.py
+++ b/aidial_analytics_realtime/rates.py
@@ -63,15 +63,25 @@ def __init__(self, rates_str: str | None = None):
         assert rates_str is not None
         self.rates = parse_raw_as(Rates, rates_str)
 
+    def get_rate(self, deployment: str, model: str):
+        deployment_rate = self.rates.get(deployment)
+
+        if deployment_rate is not None:
+            return deployment_rate
+        else:
+            return self.rates.get(model)
+
     def calculate_price(
         self,
+        deployment: str,
         model: str,
         request_content: str,
         response_content: str,
         usage: dict | None,
     ) -> Decimal:
-        model_rate = self.rates.get(model)
-        if not model_rate:
+        rate = self.get_rate(deployment, model)
+
+        if not rate:
             return Decimal(0)
 
-        return model_rate.calculate(request_content, response_content, usage)
+        return rate.calculate(request_content, response_content, usage)
diff --git a/tests/test_rates.py b/tests/test_rates.py
index c3ad46b..620fa67 100644
--- a/tests/test_rates.py
+++ b/tests/test_rates.py
@@ -1,6 +1,8 @@
 import json
 from decimal import Decimal
 
+import pytest
+
 from aidial_analytics_realtime.rates import RatesCalculator
 
 MODEL_RATES = json.dumps(
@@ -10,6 +12,16 @@
             "prompt_price": "0.00003",
             "completion_price": "0.00006",
         },
+        "gpt-35-turbo": {
+            "unit": "token",
+            "prompt_price": "0.0000015",
+            "completion_price": "0.000002",
+        },
+        "chat-bison": {
+            "unit": "char_without_whitespace",
+            "prompt_price": "0.00005",
+            "completion_price": "0.00005",
+        },
         "chat-bison@001": {
             "unit": "char_without_whitespace",
             "prompt_price": "0.0000005",
@@ -19,15 +31,49 @@
 )
 
 
-def test_token_rates():
+TOKEN_RATES_TESTDATA = [
+    ("gpt-4", "gpt-4", "0.00012"),
+    ("gpt-4", "gpt-35-turbo", "0.00012"),
+    ("gpt-4", "non-existent", "0.00012"),
+    ("gpt-4", "gpt-4", "0.00012"),
+    ("gpt-35-turbo", "gpt-4", "0.0000050"),
+    ("non-existent", "gpt-4", "0.00012"),
+    ("gpt-35-turbo", "gpt-4", "0.0000050"),
+    ("gpt-35-turbo", "gpt-35-turbo", "0.0000050"),
+    ("non-existent", "non-existent", "0"),
+]
+
+
+@pytest.mark.parametrize("deployment, model, price", TOKEN_RATES_TESTDATA)
+def test_token_rates(deployment: str, model: str, price: str):
     rates = RatesCalculator(MODEL_RATES)
-    price = rates.calculate_price(
-        "gpt-4", "hello", "hello", {"prompt_tokens": 2, "completion_tokens": 1}
+    calculated_price = rates.calculate_price(
+        deployment,
+        model,
+        "hello",
+        "hello",
+        {"prompt_tokens": 2, "completion_tokens": 1},
     )
-    assert price == Decimal("0.00012")
+    assert calculated_price == Decimal(price)
+
 
+CHAR_RATES_TESTDATA = [
+    ("chat-bison@001", "chat-bison@001", "0.0000035"),
+    ("chat-bison@001", "chat-bison", "0.0000035"),
+    ("chat-bison@001", "non-existent", "0.0000035"),
+    ("chat-bison@001", "chat-bison@001", "0.0000035"),
+    ("chat-bison", "chat-bison@001", "0.00035"),
+    ("non-existent", "chat-bison@001", "0.0000035"),
+    ("chat-bison", "chat-bison@001", "0.00035"),
+    ("chat-bison", "chat-bison", "0.00035"),
+    ("non-existent", "non-existent", "0"),
+]
 
-def test_char_rates():
+
+@pytest.mark.parametrize("deployment, model, price", CHAR_RATES_TESTDATA)
+def test_char_rates(deployment: str, model: str, price: str):
     rates = RatesCalculator(MODEL_RATES)
-    price = rates.calculate_price("chat-bison@001", "hi", "hello", None)
-    assert price == Decimal("0.0000035")
+    calculated_price = rates.calculate_price(
+        deployment, model, "hi", "hello", None
+    )
+    assert calculated_price == Decimal(price)