honor $ in usage string

wso2 · Sep 27, 2024 · 07a0e09 · 07a0e09
1 parent 44efbda
commit 07a0e09
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/....wso2.apk.enforcer/src/main/java/org/wso2/apk/enforcer/grpc/ExternalProcessorService.java b/....wso2.apk.enforcer/src/main/java/org/wso2/apk/enforcer/grpc/ExternalProcessorService.java
@@ -44,6 +44,7 @@
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -309,6 +310,9 @@ private static Usage extractUsageFromBody(String body, String completionTokenPat
             JsonNode rootNode = mapper.readTree(body);
             // Extract prompt token count
             String[] keysForPromtTokens = promptTokenPath.split("\\.");
+            if (keysForPromtTokens.length > 0 && "$".equals(keysForPromtTokens[0])) {
+                keysForPromtTokens = Arrays.copyOfRange(keysForPromtTokens, 1, keysForPromtTokens.length);
+            }
             JsonNode currentNodeForPromtToken = null;
             if (rootNode.has(keysForPromtTokens[0])) {
                 currentNodeForPromtToken = rootNode.get(keysForPromtTokens[0]);