From eeeb04772027d83332845149a1d1e9691ca381dc Mon Sep 17 00:00:00 2001 From: Min Xia Date: Mon, 24 Jun 2024 22:56:15 -0700 Subject: [PATCH] aws v2 not working --- .../aws-sdk-2.2/library/build.gradle.kts | 1 + .../awssdk/v2_2/FieldMapper.java | 17 ++- .../awssdk/v2_2/Serializer.java | 118 ++++++++++++++++++ 3 files changed, 135 insertions(+), 1 deletion(-) diff --git a/instrumentation/aws-sdk/aws-sdk-2.2/library/build.gradle.kts b/instrumentation/aws-sdk/aws-sdk-2.2/library/build.gradle.kts index 37acd881ee2f..4375c27ac5fd 100644 --- a/instrumentation/aws-sdk/aws-sdk-2.2/library/build.gradle.kts +++ b/instrumentation/aws-sdk/aws-sdk-2.2/library/build.gradle.kts @@ -9,6 +9,7 @@ dependencies { library("software.amazon.awssdk:sqs:2.2.0") library("software.amazon.awssdk:sns:2.2.0") library("software.amazon.awssdk:aws-json-protocol:2.2.0") + implementation("org.json:json:20240303") compileOnly(project(":muzzle")) // For @NoMuzzle testImplementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing")) diff --git a/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapper.java b/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapper.java index 569d0eb5ae18..31b2e0f7e354 100644 --- a/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapper.java +++ b/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapper.java @@ -9,12 +9,16 @@ import java.util.List; import java.util.function.Function; import javax.annotation.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import software.amazon.awssdk.core.SdkRequest; import software.amazon.awssdk.core.SdkResponse; import software.amazon.awssdk.utils.StringUtils; class FieldMapper { + private static final Logger logger = LoggerFactory.getLogger(FieldMapper.class); + private final Serializer serializer; private final MethodHandleFactory methodHandleFactory; @@ -65,8 +69,19 @@ private void mapToAttributes( for (int i = 1; i < path.size() && target != null; i++) { target = next(target, path.get(i)); } + String value; if (target != null) { - String value = serializer.serialize(target); + if (AwsResourceAttributes.isGenAiAttribute(fieldMapping.getAttribute())) { + logger.info("isGenAiAttribute: " + fieldMapping.getAttribute()); + value = serializer.serialize(fieldMapping.getAttribute(), target); + span.setAttribute("gen_ai.system", "AWS Bedrock"); + logger.info("isGenAiAttribute value: " + value); + } else { + logger.info("Not isGenAiAttribute: " + fieldMapping.getAttribute()); + value = serializer.serialize(target); + logger.info("Not isGenAiAttribute value: " + value); + } + if (!StringUtils.isEmpty(value)) { span.setAttribute(fieldMapping.getAttribute(), value); } diff --git a/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/Serializer.java b/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/Serializer.java index 979ecb08e87d..6fbc4bb4afed 100644 --- a/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/Serializer.java +++ b/instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/Serializer.java @@ -12,6 +12,8 @@ import java.util.Optional; import java.util.stream.Collectors; import javax.annotation.Nullable; +import org.json.JSONObject; +import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.SdkPojo; import software.amazon.awssdk.http.ContentStreamProvider; import software.amazon.awssdk.http.SdkHttpFullRequest; @@ -65,4 +67,120 @@ private String serialize(Collection collection) { String serialized = collection.stream().map(this::serialize).collect(Collectors.joining(",")); return (StringUtils.isEmpty(serialized) ? null : "[" + serialized + "]"); } + + @Nullable + String serialize(String attributeName, Object target) { + try { + JSONObject jsonBody; + if (target instanceof SdkBytes) { + jsonBody = new JSONObject(((SdkBytes) target).asUtf8String()); + } else { + if (target != null) { + return target.toString(); + } + return null; + } + switch (attributeName) { + case "gen_ai.response.finish_reason": + return getFinishReason(jsonBody); + case "gen_ai.usage.completion_tokens": + return getOutputTokens(jsonBody); + case "gen_ai.usage.prompt_tokens": + return getInputTokens(jsonBody); + case "gen_ai.request.top_p": + return getTopP(jsonBody); + case "gen_ai.request.temperature": + return getTemperature(jsonBody); + case "gen_ai.request.max_tokens": + return getMaxTokens(jsonBody); + default: + return null; + } + } catch (RuntimeException e) { + return null; + } + } + + private static String getFinishReason(JSONObject body) { + if (body.has("stop_reason")) { + return body.getString("stop_reason"); + } else if (body.has("results")) { + JSONObject result = (JSONObject) body.getJSONArray("results").get(0); + if (result.has("completionReason")) { + return result.getString("completionReason"); + } + } + return null; + } + + private static String getInputTokens(JSONObject body) { + if (body.has("prompt_token_count")) { + return String.valueOf(body.getInt("prompt_token_count")); + } else if (body.has("inputTextTokenCount")) { + return String.valueOf(body.getInt("inputTextTokenCount")); + } else if (body.has("usage")) { + JSONObject usage = (JSONObject) body.get("usage"); + if (usage.has("input_tokens")) { + return String.valueOf(usage.getInt("input_tokens")); + } + } + return null; + } + + private static String getOutputTokens(JSONObject body) { + if (body.has("generation_token_count")) { + return String.valueOf(body.getInt("generation_token_count")); + } else if (body.has("results")) { + JSONObject result = (JSONObject) body.getJSONArray("results").get(0); + if (result.has("tokenCount")) { + return String.valueOf(result.getInt("tokenCount")); + } + } else if (body.has("inputTextTokenCount")) { + return String.valueOf(body.getInt("inputTextTokenCount")); + } else if (body.has("usage")) { + JSONObject usage = (JSONObject) body.get("usage"); + if (usage.has("output_tokens")) { + return String.valueOf(usage.getInt("output_tokens")); + } + } + return null; + } + + private static String getTopP(JSONObject body) { + if (body.has("top_p")) { + return String.valueOf(body.getFloat("top_p")); + } else if (body.has("textGenerationConfig")) { + JSONObject usage = (JSONObject) body.get("textGenerationConfig"); + if (usage.has("topP")) { + return String.valueOf(usage.getFloat("topP")); + } + } + return null; + } + + private static String getTemperature(JSONObject body) { + if (body.has("temperature")) { + return String.valueOf(body.getFloat("temperature")); + } else if (body.has("textGenerationConfig")) { + JSONObject usage = (JSONObject) body.get("textGenerationConfig"); + if (usage.has("temperature")) { + return String.valueOf(usage.getFloat("temperature")); + } + } + return null; + } + + private static String getMaxTokens(JSONObject body) { + if (body.has("max_tokens")) { + return String.valueOf(body.getInt("max_tokens")); + } else if (body.has("max_gen_len")) { + return String.valueOf(body.getInt("max_gen_len")); + } else if (body.has("textGenerationConfig")) { + JSONObject usage = (JSONObject) body.get("textGenerationConfig"); + if (usage.has("maxTokenCount")) { + return String.valueOf(usage.getInt("maxTokenCount")); + } + } + return null; + } }