restrict public path resolver to only use llmjson type

mxiamxia · Dec 12, 2024 · 1d1699d · 1d1699d
1 parent c384cb2
commit 1d1699d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/...library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/BedrockJsonParser.java b/...library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/BedrockJsonParser.java
@@ -166,9 +166,9 @@ private JsonPathResolver() {
       throw new UnsupportedOperationException("Utility class");
     }
 
-    public static Object resolvePath(Map<String, Object> json, String... paths) {
+    public static Object resolvePath(LlmJson llmJson, String... paths) {
       for (String path : paths) {
-        Object value = resolvePath(json, path);
+        Object value = resolvePath(llmJson.getJsonBody(), path);
         if (value != null) {
           return value;
         }

diff --git a/...s-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/LlmJson.java b/...s-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/LlmJson.java
@@ -17,12 +17,16 @@ public LlmJson(Map<String, Object> jsonBody) {
     this.jsonBody = jsonBody;
   }
 
+  public Map<String, Object> getJsonBody() {
+    return jsonBody;
+  }
+
   @Nullable
   private String approximateTokenCount(String... textPaths) {
     return Arrays.stream(textPaths)
         .map(
             path -> {
-              Object value = BedrockJsonParser.JsonPathResolver.resolvePath(jsonBody, path);
+              Object value = BedrockJsonParser.JsonPathResolver.resolvePath(this, path);
               if (value instanceof String) {
                 int tokenEstimate = (int) Math.ceil(((String) value).length() / 6.0);
                 return Integer.toString(tokenEstimate);
@@ -47,7 +51,7 @@ private String approximateTokenCount(String... textPaths) {
   public String getMaxTokens() {
     Object value =
         BedrockJsonParser.JsonPathResolver.resolvePath(
-            jsonBody,
+            this,
             "/max_tokens",
             "/max_gen_len",
             "/textGenerationConfig/maxTokenCount",
@@ -68,7 +72,7 @@ public String getMaxTokens() {
   public String getTemperature() {
     Object value =
         BedrockJsonParser.JsonPathResolver.resolvePath(
-            jsonBody,
+            this,
             "/temperature",
             "/textGenerationConfig/temperature",
             "inferenceConfig/temperature");
@@ -88,7 +92,7 @@ public String getTemperature() {
   public String getTopP() {
     Object value =
         BedrockJsonParser.JsonPathResolver.resolvePath(
-            jsonBody, "/top_p", "/p", "/textGenerationConfig/topP", "inferenceConfig/top_p");
+            this, "/top_p", "/p", "/textGenerationConfig/topP", "inferenceConfig/top_p");
     return value != null ? String.valueOf(value) : null;
   }
 
@@ -105,7 +109,7 @@ public String getTopP() {
   public String getFinishReasons() {
     Object value =
         BedrockJsonParser.JsonPathResolver.resolvePath(
-            jsonBody,
+            this,
             "/stopReason",
             "/finish_reason",
             "/stop_reason",
@@ -131,7 +135,7 @@ public String getInputTokens() {
     // Try direct tokens counts first
     Object directCount =
         BedrockJsonParser.JsonPathResolver.resolvePath(
-            jsonBody,
+            this,
             "/inputTextTokenCount",
             "/prompt_token_count",
             "/usage/input_tokens",
@@ -162,7 +166,7 @@ public String getOutputTokens() {
     // Try direct token counts first
     Object directCount =
         BedrockJsonParser.JsonPathResolver.resolvePath(
-            jsonBody,
+            this,
             "/generation_token_count",
             "/results/0/tokenCount",
             "/usage/output_tokens",