Skip to content

Commit

Permalink
fix: add strict undefined checks
Browse files Browse the repository at this point in the history
  • Loading branch information
@yiyuan-he
yiyuan-he committed Oct 31, 2024
1 parent c730602 commit 3a57302
Showing 1 changed file with 41 additions and 38 deletions.
79 changes: 41 additions & 38 deletions aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -214,63 +214,66 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
if (request.commandInput?.body) {
const requestBody = JSON.parse(request.commandInput.body);
if (modelId.includes('amazon.titan')) {
if (requestBody.textGenerationConfig) {
const config = requestBody.textGenerationConfig;
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = config.temperature;
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = config.topP;
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = config.maxTokenCount;
if (requestBody.textGenerationConfig?.temperature !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.textGenerationConfig.temperature;
}
if (requestBody.textGenerationConfig?.topP !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.textGenerationConfig.topP;
}
if (requestBody.textGenerationConfig?.maxTokenCount != undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.textGenerationConfig.maxTokenCount;
}
} else if (modelId.includes('anthropic.claude')) {
if (requestBody.max_tokens) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature) {
if (requestBody.temperature !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p) {
if (requestBody.top_p !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
} else if (modelId.includes('meta.llama')) {
if (requestBody.max_gen_len) {
if (requestBody.max_gen_len !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_gen_len;
}
if (requestBody.temperature) {
if (requestBody.temperature !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p) {
if (requestBody.top_p !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
} else if (modelId.includes('cohere.command')) {
if (requestBody.max_tokens) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature) {
if (requestBody.temperature !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.p) {
if (requestBody.p !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.p;
}
} else if (modelId.includes('ai21.jamba')) {
if (requestBody.max_tokens) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature) {
if (requestBody.temperature !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p) {
if (requestBody.top_p !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
} else if (modelId.includes('mistral.mistral')) {
if (requestBody.prompt) {
if (requestBody.prompt !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6);
}
if (requestBody.max_tokens) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature) {
if (requestBody.temperature !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p) {
if (requestBody.top_p !== undefined) {
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
}
Expand All @@ -290,72 +293,72 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
const decodedResponseBody = new TextDecoder().decode(response.data.body);
const responseBody = JSON.parse(decodedResponseBody);
if (currentModelId.includes('amazon.titan')) {
if (responseBody.inputTextTokenCount) {
if (responseBody.inputTextTokenCount !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.inputTextTokenCount);
}
if (responseBody.results?.[0]?.tokenCount) {
if (responseBody.results?.[0]?.tokenCount !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.results[0].tokenCount);
}
if (responseBody.results?.[0]?.completionReason) {
if (responseBody.results?.[0]?.completionReason !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.results[0].completionReason,
]);
}
} else if (currentModelId.includes('anthropic.claude')) {
if (responseBody.usage?.input_tokens) {
if (responseBody.usage?.input_tokens !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.input_tokens);
}
if (responseBody.usage?.output_tokens) {
if (responseBody.usage?.output_tokens !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.output_tokens);
}
if (responseBody.stop_reason) {
if (responseBody.stop_reason !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]);
}
} else if (currentModelId.includes('meta.llama')) {
if (responseBody.prompt_token_count) {
if (responseBody.prompt_token_count !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.prompt_token_count);
}
if (responseBody.generation_token_count) {
if (responseBody.generation_token_count !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.generation_token_count);
}
if (responseBody.stop_reason) {
if (responseBody.stop_reason !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]);
}
} else if (currentModelId.includes('cohere.command')) {
if (responseBody.prompt) {
if (responseBody.prompt !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, Math.ceil(responseBody.prompt.length / 6));
}
if (responseBody.generations?.[0]?.text) {
if (responseBody.generations?.[0]?.text !== undefined) {
span.setAttribute(
AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,
Math.ceil(responseBody.generations[0].text.length / 6)
);
}
if (responseBody.generations?.[0]?.finish_reason) {
if (responseBody.generations?.[0]?.finish_reason !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.generations[0].finish_reason,
]);
}
} else if (currentModelId.includes('ai21.jamba')) {
if (responseBody.usage?.prompt_tokens) {
if (responseBody.usage?.prompt_tokens !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.prompt_tokens);
}
if (responseBody.usage?.completion_tokens) {
if (responseBody.usage?.completion_tokens !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.completion_tokens);
}
if (responseBody.choices?.[0]?.finish_reason) {
if (responseBody.choices?.[0]?.finish_reason !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.choices[0].finish_reason,
]);
}
} else if (currentModelId.includes('mistral.mistral')) {
if (responseBody.outputs?.[0]?.text) {
if (responseBody.outputs?.[0]?.text !== undefined) {
span.setAttribute(
AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,
Math.ceil(responseBody.outputs[0].text.length / 6)
);
}
if (responseBody.outputs?.[0]?.stop_reason) {
if (responseBody.outputs?.[0]?.stop_reason !== undefined) {
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.outputs[0].stop_reason,
]);
Expand Down

0 comments on commit 3a57302

Please sign in to comment.