aws v2 not working

mxiamxia · Jun 25, 2024 · eeeb047 · eeeb047
1 parent 7e5dd96
commit eeeb047
Show file tree

Hide file tree

Showing 3 changed files with 135 additions and 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
@@ -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/...k-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/FieldMapper.java b/...k-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/...dk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/Serializer.java b/...dk-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;
+  }
 }