Implement custom bedrock json parser for v1 & v2 by yiyuan-he · Pull Request #25 · mxiamxia/opentelemetry-java-instrumentation

yiyuan-he · 2024-12-11T18:32:11Z

Description of changes:

Implement custom bedrock json parser for our Java v1 & v2 SDK auto-instrumentation using only native language features to avoid introducing external dependencies.

TODO:

~~Address PR comments from @mxiamxia~~ Pending further review.
~~Implementation parity between v2 and v1~~ Done.
~~Add unit tests for custom parser~~ Tentatively done. May need to add more test cases
~~Update behavior to not add attribute when value is null. Currently it assigns "null" as attribute value.~~
~~Add support for new Amazon Nova Models.~~ Done for both v1 & v2.
~~Implement custom parser for Java v1.~~ Updated this PR with v1 changes.
~~Investigate why only Mistral contract test case is failing.~~ There is minor rounding difference when using custom parser. Had to update the assertion for gen_ai.usage.input_tokens from 15 to 16.

Test plan:

Added unit tests for various edge cases.
Ran the existing contract tests in our ADOT package and verified the test cases for all 6 models are passing.

Verified generated spans for 6 models that we currently support. Also added some logging statements to measure the processing time of the custom parser.

Amazon Nova

Amazon Titan

Anthropic Claude

AI21 Jamba

Mistral

Meta Llama

Cohere Command R

mxiamxia · 2024-12-12T00:46:27Z

+    public JsonParser(String json) {
+      this.json = json.trim();
+      this.position = 0;
+    }


I think we should avoid to create new JsonParser instances with JsonString text from constructor. Can we update public Map<String, Object> parse() to static and have it take the JsongString text as input so we can call JsonParser.parse(json) directly in the caller method?

ok. I found JsonParser needs to maintain position state in each instance for each data processing. is it necessary?

As we discussed, suggest changing JsonParser to something like JSONObject

The JsonParser needs to maintain some kind of position state since we are parsing the string character by character.

However, this state does not need to be exposed as public. I think we can hide this state from the user by making it private instead.

mxiamxia · 2024-12-12T00:54:10Z

+      BedrockJsonParser.JsonParser parser = new BedrockJsonParser.JsonParser(jsonString);
+      Map<String, Object> jsonBody = parser.parse();


BedrockJsonParser.JsonObject jsonObject = new BedrockJsonParser.JsonObject(jsonString);

Will implement something similar with this:

// Parse the LLM JSON string into a Map LlmJson llmJson = BedrockJsonParser.parse(jsonString);

mxiamxia · 2024-12-12T01:07:44Z

+      return null;
+    }
+
+    private static Object resolvePath(Map<String, Object> json, String path) {


Suggest changing ot the custom JsonObject as input instead of generic Map which will restrict to pass in the correct/expect data while calling this method

mxiamxia · 2024-12-12T03:25:47Z

+      return result.toString();
+    }
+
+    private Object readValue() {


I found an open-source mini-json parser without requiring Dep, which has very similar Impl like yours which is good. :) Could double check test a few things? Will you impl be able to handle the attribute value with special chars such as \n?

https://github.com/ralfstx/minimal-json/blob/master/com.eclipsesource.json/src/main/java/com/eclipsesource/json/JsonParser.java#L159

Could you add the unit tests for the json utils? We test against different bedrock response format pattern and edge cases

Added unit tests to cover special chars and unicode that may appear in the input/output prompts. Also added some test cases to cover empty objects and arrays.

yiyuan-he · 2024-12-13T18:50:24Z

    try {
-      JsonNode jsonBody;
+      // Extract JSON string from target if it is a Bedrock Runtime JSON blob
+      String jsonString;
      if (target instanceof SdkBytes) {
-        String jsonString = ((SdkBytes) target).asUtf8String();
-        jsonBody = objectMapper.readTree(jsonString);
+        jsonString = ((SdkBytes) target).asUtf8String();
      } else {
        if (target != null) {
          return target.toString();
        }
        return null;
      }

+      // Parse the LLM JSON string into a Map
+      LlmJson llmJson = BedrockJsonParser.parse(jsonString);
+
+      // Use attribute name to extract the corresponding value
      switch (attributeName) {
        case "gen_ai.request.max_tokens":
-          return getMaxTokens(jsonBody);
+          return llmJson.getMaxTokens();
        case "gen_ai.request.temperature":
-          return getTemperature(jsonBody);
+          return llmJson.getTemperature();
        case "gen_ai.request.top_p":
-          return getTopP(jsonBody);
+          return llmJson.getTopP();
        case "gen_ai.response.finish_reasons":
-          return getFinishReasons(jsonBody);
+          return llmJson.getFinishReasons();
        case "gen_ai.usage.input_tokens":
-          return getInputTokens(jsonBody);
+          return llmJson.getInputTokens();
        case "gen_ai.usage.output_tokens":
-          return getOutputTokens(jsonBody);
+          return llmJson.getOutputTokens();
        default:
          return null;
      }
-    } catch (JsonProcessingException e) {
+    } catch (RuntimeException e) {
      return null;
    }


This logic should swallow all Exceptions raised by the custom parser so that we do not propagate it to the application level

mxiamxia

LGTM!

yiyuan-he · 2024-12-17T17:46:27Z

Closing this PR since #27 was merged to our new v1_33_6_patch branch.

set up custom bedrock json parser

a09688c

yiyuan-he changed the title ~~set up custom bedrock json parser~~ Set up custom bedrock json parser for Java v2 Dec 11, 2024

yiyuan-he changed the title ~~Set up custom bedrock json parser for Java v2~~ POC: Set up custom bedrock json parser for Java v2 Dec 11, 2024

apply formatting

815b8cd

yiyuan-he changed the title ~~POC: Set up custom bedrock json parser for Java v2~~ PoC: Set up custom bedrock json parser for Java v2 Dec 11, 2024

yiyuanh and others added 2 commits December 11, 2024 22:29

set up custom parser in v1

da8e4fa

apply formatting

1c4fe0e

yiyuan-he changed the title ~~PoC: Set up custom bedrock json parser for Java v2~~ PoC: Set up custom bedrock json parser for v1 & v2 Dec 11, 2024

yiyuan-he changed the base branch from v1_32_1_dev to v1_33_6_patch December 11, 2024 22:55

yiyuan-he changed the base branch from v1_33_6_patch to v1_32_1_dev December 11, 2024 22:56

yiyuanh added 2 commits December 11, 2024 15:46

add support for amazon nova models in v2

a593ff8

add support for amazon nova models in v1

6eda62e

mxiamxia reviewed Dec 12, 2024

View reviewed changes

yiyuanh and others added 4 commits December 12, 2024 11:55

refactor for more type safety

16236c7

properly handle null attribute values

c384cb2

restrict public path resolver to only use llmjson type

1d1699d

handle escape chars and add unit tests in java v2

64da476

yiyuan-he changed the title ~~PoC: Set up custom bedrock json parser for v1 & v2~~ Implement custom bedrock json parser for v1 & v2 Dec 13, 2024

yiyuan-he commented Dec 13, 2024

View reviewed changes

yiyuanh added 4 commits December 13, 2024 12:06

refactor to decouple business logic

4ed3868

v1 implementation parity

7bf71bb

add groovy unit tests

cf2d606

apply formatting

41c6b4d

yiyuan-he changed the base branch from v1_32_1_dev to v1_33_6_patch December 16, 2024 21:55

yiyuan-he changed the base branch from v1_33_6_patch to v1_32_1_dev December 16, 2024 21:56

mxiamxia approved these changes Dec 16, 2024

View reviewed changes

yiyuan-he mentioned this pull request Dec 17, 2024

Merge custom parser to v1_33_6 branch #27

Merged

yiyuan-he closed this Dec 17, 2024

yiyuan-he mentioned this pull request Dec 20, 2024

Merge recent gen ai changes to v2_10_0 branch #29

Merged

		BedrockJsonParser.JsonParser parser = new BedrockJsonParser.JsonParser(jsonString);
		Map<String, Object> jsonBody = parser.parse();

Conversation

yiyuan-he commented Dec 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mxiamxia left a comment

Choose a reason for hiding this comment

Uh oh!

yiyuan-he commented Dec 17, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yiyuan-he commented Dec 11, 2024 •

edited

Loading