feat: added pricing params to the listing (#45) (#34)

epam · Nov 22, 2023 · 6208c6d · 6208c6d
1 parent e0a9549
commit 6208c6d
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 7 deletions.
diff --git a/src/main/java/com/epam/aidial/core/config/Model.java b/src/main/java/com/epam/aidial/core/config/Model.java
@@ -12,5 +12,6 @@
 public class Model extends Deployment {
     private ModelType type;
     private TokenLimits limits;
+    private Pricing pricing;
     private List<Upstream> upstreams = List.of();
-}
+}
diff --git a/src/main/java/com/epam/aidial/core/config/Pricing.java b/src/main/java/com/epam/aidial/core/config/Pricing.java
@@ -0,0 +1,10 @@
+package com.epam.aidial.core.config;
+
+import lombok.Data;
+
+@Data
+public class Pricing {
+    private String unit;
+    private Double prompt;
+    private Double completion;
+}
diff --git a/src/main/java/com/epam/aidial/core/controller/ModelController.java b/src/main/java/com/epam/aidial/core/controller/ModelController.java
@@ -4,9 +4,11 @@
 import com.epam.aidial.core.config.Config;
 import com.epam.aidial.core.config.Model;
 import com.epam.aidial.core.config.ModelType;
+import com.epam.aidial.core.config.Pricing;
 import com.epam.aidial.core.config.TokenLimits;
 import com.epam.aidial.core.data.ListData;
 import com.epam.aidial.core.data.ModelData;
+import com.epam.aidial.core.data.PricingData;
 import com.epam.aidial.core.data.TokenLimitsData;
 import com.epam.aidial.core.util.HttpStatus;
 import io.vertx.core.Future;
@@ -72,16 +74,29 @@ private static ModelData createModel(Model model) {
         data.setLimits(createLimits(model.getLimits()));
         data.setInputAttachmentTypes(model.getInputAttachmentTypes());
         data.setMaxInputAttachments(model.getMaxInputAttachments());
+        data.setPricing(createPricing(model.getPricing()));
         return data;
     }
 
     private static TokenLimitsData createLimits(TokenLimits limits) {
+        if (limits == null) {
+            return null;
+        }
         TokenLimitsData data = new TokenLimitsData();
-        if (limits != null) {
-            data.setMaxPromptTokens(limits.getMaxPromptTokens());
-            data.setMaxCompletionTokens(limits.getMaxCompletionTokens());
-            data.setMaxTotalTokens(limits.getMaxTotalTokens());
+        data.setMaxPromptTokens(limits.getMaxPromptTokens());
+        data.setMaxCompletionTokens(limits.getMaxCompletionTokens());
+        data.setMaxTotalTokens(limits.getMaxTotalTokens());
+        return data;
+    }
+
+    private static PricingData createPricing(Pricing pricing) {
+        if (pricing == null) {
+            return null;
         }
+        PricingData data = new PricingData();
+        data.setUnit(pricing.getUnit());
+        data.setPrompt(pricing.getPrompt());
+        data.setCompletion(pricing.getCompletion());
         return data;
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/data/ModelData.java b/src/main/java/com/epam/aidial/core/data/ModelData.java
@@ -14,7 +14,8 @@ public class ModelData extends DeploymentData {
 
     private String lifecycleStatus = "generally-available";
     private CapabilitiesData capabilities = new CapabilitiesData();
-    private TokenLimitsData limits = new TokenLimitsData();
+    private TokenLimitsData limits;
+    private PricingData pricing;
 
     {
         setObject("model");

diff --git a/src/main/java/com/epam/aidial/core/data/PricingData.java b/src/main/java/com/epam/aidial/core/data/PricingData.java
@@ -0,0 +1,15 @@
+package com.epam.aidial.core.data;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.databind.PropertyNamingStrategies;
+import com.fasterxml.jackson.databind.annotation.JsonNaming;
+import lombok.Data;
+
+@Data
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
+public class PricingData {
+    private String unit;
+    private Double prompt;
+    private Double completion;
+}
diff --git a/src/main/resources/aidial.config.json b/src/main/resources/aidial.config.json
@@ -54,6 +54,11 @@
       ],
       "limits": {
         "maxTotalTokens": 4096
+      },
+      "pricing": {
+        "unit": "token",
+        "prompt": 0.00005,
+        "completion": 0.0001
       }
     },
     "embedding-ada": {
@@ -64,11 +69,20 @@
       ],
       "limits": {
         "maxTotalTokens": 8192
+      },
+      "pricing": {
+        "unit": "token",
+        "prompt": 0.00005
       }
     },
     "exotic-model": {
       "type": "chat",
-      "endpoint" : "http://localhost:7001/openai/deployments/exotic-model/chat/completions"
+      "endpoint" : "http://localhost:7001/openai/deployments/exotic-model/chat/completions",
+      "pricing": {
+        "unit": "char_without_whitespace",
+        "prompt": 0.00005,
+        "completion": 0.0001
+      }
     }
   },
   "keys": {