Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

aws v2 not working draft #7

Open
wants to merge 1 commit into
base: v1_32_1_dev_aws_working
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
Loading