Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Model cost calculator fails with NPE for some adapters #201 #252

Merged
merged 2 commits into from
Feb 29, 2024
Merged

fix: Model cost calculator fails with NPE for some adapters #201 #252

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5f073f1
fix: Model cost calculator fails with NPE for some adapters #201
astsiapanay Feb 29, 2024
1effcdb
fix: unit tests
astsiapanay Feb 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 122 additions & 1 deletion src/main/java/com/epam/aidial/core/util/ModelCostCalculator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,20 @@
import com.epam.aidial.core.ProxyContext;
import com.epam.aidial.core.config.Deployment;
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.token.TokenUsage;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.netty.buffer.ByteBufInputStream;
import io.vertx.core.buffer.Buffer;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;

import java.io.InputStream;
import java.math.BigDecimal;
import java.util.Scanner;

@Slf4j
@UtilityClass
Expand All @@ -25,7 +33,12 @@ public static BigDecimal calculate(ProxyContext context) {
return null;
}

return calculate(context.getTokenUsage(), pricing.getPrompt(), pricing.getCompletion());
return switch (pricing.getUnit()) {
case "token" -> calculate(context.getTokenUsage(), pricing.getPrompt(), pricing.getCompletion());
case "char_without_whitespace" ->
calculate(model.getType(), context.getRequestBody(), context.getResponseBody(), pricing.getPrompt(), pricing.getCompletion());
default -> null;
};
}

private static BigDecimal calculate(TokenUsage tokenUsage, String promptRate, String completionRate) {
Expand All @@ -47,4 +60,112 @@ private static BigDecimal calculate(TokenUsage tokenUsage, String promptRate, St
return cost;
}

private static BigDecimal calculate(ModelType modelType, Buffer requestBody, Buffer responseBody, String promptRate, String completionRate) {
RequestLengthResult requestLengthResult = getRequestContentLength(modelType, requestBody);
int responseLength = getResponseContentLength(modelType, responseBody, requestLengthResult.stream());
BigDecimal cost = null;
if (promptRate != null) {
cost = new BigDecimal(requestLengthResult.length()).multiply(new BigDecimal(promptRate));
}
if (completionRate != null) {
BigDecimal completionCost = new BigDecimal(responseLength).multiply(new BigDecimal(completionRate));
if (cost == null) {
cost = completionCost;
} else {
cost = cost.add(completionCost);
}
}
return cost;
}

private static int getResponseContentLength(ModelType modelType, Buffer responseBody, boolean isStreamingResponse) {
if (modelType == ModelType.EMBEDDING) {
return 0;
}
if (isStreamingResponse) {
try (Scanner scanner = new Scanner(new ByteBufInputStream(responseBody.getByteBuf()))) {
scanner.useDelimiter("\n*data: *");
int len = 0;
while (scanner.hasNext()) {
String chunk = scanner.next();
if (chunk.startsWith("[DONE]")) {
break;
}
ObjectNode tree = (ObjectNode) ProxyUtil.MAPPER.readTree(chunk);
ArrayNode choices = (ArrayNode) tree.get("choices");
if (choices == null) {
// skip error message
continue;
}
JsonNode contentNode = choices.get(0).get("delta").get("content");
if (contentNode != null) {
len += getLengthWithoutWhitespace(contentNode.textValue());
}
}
return len;
} catch (Throwable e) {
throw new RuntimeException(e);
}
} else {
try (InputStream stream = new ByteBufInputStream(responseBody.getByteBuf())) {
ObjectNode tree = (ObjectNode) ProxyUtil.MAPPER.readTree(stream);
ArrayNode choices = (ArrayNode) tree.get("choices");
if (choices == null) {
// skip error message
return 0;
}
JsonNode contentNode = choices.get(0).get("message").get("content");
return getLengthWithoutWhitespace(contentNode.textValue());
} catch (Throwable e) {
throw new RuntimeException(e);
}
}
}

private static RequestLengthResult getRequestContentLength(ModelType modelType, Buffer requestBody) {
try (InputStream stream = new ByteBufInputStream(requestBody.getByteBuf())) {
int len;
ObjectNode tree = (ObjectNode) ProxyUtil.MAPPER.readTree(stream);
if (modelType == ModelType.CHAT) {
ArrayNode messages = (ArrayNode) tree.get("messages");
len = 0;
for (int i = 0; i < messages.size(); i++) {
JsonNode message = messages.get(i);
len += getLengthWithoutWhitespace(message.get("content").textValue());
}
return new RequestLengthResult(len, tree.get("stream").asBoolean(false));
} else {
JsonNode input = tree.get("input");
if (input instanceof ArrayNode array) {
len = 0;
for (int i = 0; i < array.size(); i++) {
len += getLengthWithoutWhitespace(array.get(i).textValue());
}
} else {
len = getLengthWithoutWhitespace(input.textValue());
}
}
return new RequestLengthResult(len, false);
} catch (Throwable e) {
throw new RuntimeException(e);
}
}

private static int getLengthWithoutWhitespace(String s) {
if (s == null) {
return 0;
}
int len = 0;
for (int i = 0; i < s.length(); i++) {
if (s.charAt(i) != ' ') {
len++;
}
}
return len;
}

private record RequestLengthResult(int length, boolean stream) {

}

}
Loading
Loading