Skip to content

Commit d9b811e

Browse files
liustveliustve
authored
Merge branch 'main' into aws_resource_support
2 parents d8bffd5 + 27c0f80 commit d9b811e

File tree

6 files changed

+461
-29
lines changed

6 files changed

+461
-29
lines changed

‎aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,22 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
251251
if (requestBody.top_p !== undefined) {
252252
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
253253
}
254+
} else if (modelId.includes('cohere.command-r')) {
255+
if (requestBody.max_tokens !== undefined) {
256+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
257+
}
258+
if (requestBody.temperature !== undefined) {
259+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
260+
}
261+
if (requestBody.p !== undefined) {
262+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.p;
263+
}
264+
if (requestBody.message !== undefined) {
265+
// NOTE: We approximate the token count since this value is not directly available in the body
266+
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
267+
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
268+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.message.length / 6);
269+
}
254270
} else if (modelId.includes('cohere.command')) {
255271
if (requestBody.max_tokens !== undefined) {
256272
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
@@ -261,6 +277,9 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
261277
if (requestBody.p !== undefined) {
262278
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.p;
263279
}
280+
if (requestBody.prompt !== undefined) {
281+
spanAttributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6);
282+
}
264283
} else if (modelId.includes('ai21.jamba')) {
265284
if (requestBody.max_tokens !== undefined) {
266285
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
@@ -271,7 +290,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
271290
if (requestBody.top_p !== undefined) {
272291
spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
273292
}
274-
} else if (modelId.includes('mistral.mistral')) {
293+
} else if (modelId.includes('mistral')) {
275294
if (requestBody.prompt !== undefined) {
276295
// NOTE: We approximate the token count since this value is not directly available in the body
277296
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
@@ -335,13 +354,17 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
335354
if (responseBody.stop_reason !== undefined) {
336355
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]);
337356
}
338-
} else if (currentModelId.includes('cohere.command')) {
339-
if (responseBody.prompt !== undefined) {
357+
} else if (currentModelId.includes('cohere.command-r')) {
358+
if (responseBody.text !== undefined) {
340359
// NOTE: We approximate the token count since this value is not directly available in the body
341360
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
342361
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
343-
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, Math.ceil(responseBody.prompt.length / 6));
362+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.text.length / 6));
344363
}
364+
if (responseBody.finish_reason !== undefined) {
365+
span.setAttribute(AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.finish_reason]);
366+
}
367+
} else if (currentModelId.includes('cohere.command')) {
345368
if (responseBody.generations?.[0]?.text !== undefined) {
346369
span.setAttribute(
347370
AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,
@@ -368,7 +391,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
368391
responseBody.choices[0].finish_reason,
369392
]);
370393
}
371-
} else if (currentModelId.includes('mistral.mistral')) {
394+
} else if (currentModelId.includes('mistral')) {
372395
if (responseBody.outputs?.[0]?.text !== undefined) {
373396
span.setAttribute(
374397
AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,

‎aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -520,6 +520,60 @@ describe('BedrockRuntime', () => {
520520
expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT);
521521
});
522522

523+
it('Add Cohere Command R model attributes to span', async () => {
524+
const modelId: string = 'cohere.command-r-v1:0"';
525+
const prompt: string = "Describe the purpose of a 'hello world' program in one line";
526+
const nativeRequest: any = {
527+
message: prompt,
528+
max_tokens: 512,
529+
temperature: 0.5,
530+
p: 0.65,
531+
};
532+
const mockRequestBody: string = JSON.stringify(nativeRequest);
533+
const mockResponseBody: any = {
534+
finish_reason: 'COMPLETE',
535+
text: 'test-generation-text',
536+
prompt: prompt,
537+
request: {
538+
commandInput: {
539+
modelId: modelId,
540+
},
541+
},
542+
};
543+
544+
nock(`https://bedrock-runtime.${region}.amazonaws.com`)
545+
.post(`/model/${encodeURIComponent(modelId)}/invoke`)
546+
.reply(200, mockResponseBody);
547+
548+
await bedrock
549+
.invokeModel({
550+
modelId: modelId,
551+
body: mockRequestBody,
552+
})
553+
.catch((err: any) => {
554+
console.log('error', err);
555+
});
556+
557+
const testSpans: ReadableSpan[] = getTestSpans();
558+
const invokeModelSpans: ReadableSpan[] = testSpans.filter((s: ReadableSpan) => {
559+
return s.name === 'BedrockRuntime.InvokeModel';
560+
});
561+
expect(invokeModelSpans.length).toBe(1);
562+
const invokeModelSpan = invokeModelSpans[0];
563+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined();
564+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined();
565+
expect(invokeModelSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined();
566+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM]).toBe('aws_bedrock');
567+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]).toBe(modelId);
568+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS]).toBe(512);
569+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE]).toBe(0.5);
570+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P]).toBe(0.65);
571+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS]).toBe(10);
572+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(4);
573+
expect(invokeModelSpan.attributes[AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(['COMPLETE']);
574+
expect(invokeModelSpan.kind).toBe(SpanKind.CLIENT);
575+
});
576+
523577
it('Add Meta Llama model attributes to span', async () => {
524578
const modelId: string = 'meta.llama2-13b-chat-v1';
525579
const prompt: string = 'Describe the purpose of an interpreter program in one line.';

‎contract-tests/images/applications/aws-sdk/server.js‎

Lines changed: 176 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ const { S3Client, CreateBucketCommand, PutObjectCommand, GetObjectCommand } = re
1010
const { DynamoDBClient, CreateTableCommand, PutItemCommand } = require('@aws-sdk/client-dynamodb');
1111
const { SQSClient, CreateQueueCommand, SendMessageCommand, ReceiveMessageCommand } = require('@aws-sdk/client-sqs');
1212
const { KinesisClient, CreateStreamCommand, PutRecordCommand } = require('@aws-sdk/client-kinesis');
13-
const fetch = require('node-fetch');
1413
const { BedrockClient, GetGuardrailCommand } = require('@aws-sdk/client-bedrock');
1514
const { BedrockAgentClient, GetKnowledgeBaseCommand, GetDataSourceCommand, GetAgentCommand } = require('@aws-sdk/client-bedrock-agent');
1615
const { BedrockRuntimeClient, InvokeModelCommand } = require('@aws-sdk/client-bedrock-runtime');
@@ -553,30 +552,190 @@ async function handleBedrockRequest(req, res, path) {
553552
});
554553
res.statusCode = 200;
555554
} else if (path.includes('invokemodel/invoke-model')) {
556-
await withInjected200Success(bedrockRuntimeClient, ['InvokeModelCommand'], {}, async () => {
557-
const modelId = 'amazon.titan-text-premier-v1:0';
558-
const userMessage = "Describe the purpose of a 'hello world' program in one line.";
559-
const prompt = `<s>[INST] ${userMessage} [/INST]`;
560-
561-
const body = JSON.stringify({
562-
inputText: prompt,
563-
textGenerationConfig: {
564-
maxTokenCount: 3072,
565-
stopSequences: [],
566-
temperature: 0.7,
567-
topP: 0.9,
568-
},
569-
});
555+
const get_model_request_response = function () {
556+
const prompt = "Describe the purpose of a 'hello world' program in one line.";
557+
let modelId = ''
558+
let request_body = {}
559+
let response_body = {}
560+
561+
if (path.includes('amazon.titan')) {
562+
563+
modelId = 'amazon.titan-text-premier-v1:0';
564+
565+
request_body = {
566+
inputText: prompt,
567+
textGenerationConfig: {
568+
maxTokenCount: 3072,
569+
stopSequences: [],
570+
temperature: 0.7,
571+
topP: 0.9,
572+
},
573+
};
574+
575+
response_body = {
576+
inputTextTokenCount: 15,
577+
results: [
578+
{
579+
tokenCount: 13,
580+
outputText: 'text-test-response',
581+
completionReason: 'CONTENT_FILTERED',
582+
},
583+
],
584+
}
585+
586+
}
587+
588+
if (path.includes('anthropic.claude')) {
589+
590+
modelId = 'anthropic.claude-v2:1';
591+
592+
request_body = {
593+
anthropic_version: 'bedrock-2023-05-31',
594+
max_tokens: 1000,
595+
temperature: 0.99,
596+
top_p: 1,
597+
messages: [
598+
{
599+
role: 'user',
600+
content: [{ type: 'text', text: prompt }],
601+
},
602+
],
603+
};
604+
605+
response_body = {
606+
stop_reason: 'end_turn',
607+
usage: {
608+
input_tokens: 15,
609+
output_tokens: 13,
610+
},
611+
}
612+
}
613+
614+
if (path.includes('meta.llama')) {
615+
modelId = 'meta.llama2-13b-chat-v1';
616+
617+
request_body = {
618+
prompt,
619+
max_gen_len: 512,
620+
temperature: 0.5,
621+
top_p: 0.9
622+
};
623+
624+
response_body = {
625+
prompt_token_count: 31,
626+
generation_token_count: 49,
627+
stop_reason: 'stop'
628+
}
629+
}
630+
631+
if (path.includes('cohere.command')) {
632+
modelId = 'cohere.command-light-text-v14';
633+
634+
request_body = {
635+
prompt,
636+
max_tokens: 512,
637+
temperature: 0.5,
638+
p: 0.65,
639+
};
640+
641+
response_body = {
642+
generations: [
643+
{
644+
finish_reason: 'COMPLETE',
645+
text: 'test-generation-text',
646+
},
647+
],
648+
prompt: prompt,
649+
};
650+
}
651+
652+
if (path.includes('cohere.command-r')) {
653+
modelId = 'cohere.command-r-v1:0';
654+
655+
request_body = {
656+
message: prompt,
657+
max_tokens: 512,
658+
temperature: 0.5,
659+
p: 0.65,
660+
};
661+
662+
response_body = {
663+
finish_reason: 'COMPLETE',
664+
text: 'test-generation-text',
665+
prompt: prompt,
666+
request: {
667+
commandInput: {
668+
modelId: modelId,
669+
},
670+
},
671+
}
672+
}
673+
674+
if (path.includes('ai21.jamba')) {
675+
modelId = 'ai21.jamba-1-5-large-v1:0';
676+
677+
request_body = {
678+
messages: [
679+
{
680+
role: 'user',
681+
content: prompt,
682+
},
683+
],
684+
top_p: 0.8,
685+
temperature: 0.6,
686+
max_tokens: 512,
687+
};
688+
689+
response_body = {
690+
stop_reason: 'end_turn',
691+
usage: {
692+
prompt_tokens: 21,
693+
completion_tokens: 24,
694+
},
695+
choices: [
696+
{
697+
finish_reason: 'stop',
698+
},
699+
],
700+
}
701+
}
702+
703+
if (path.includes('mistral')) {
704+
modelId = 'mistral.mistral-7b-instruct-v0:2';
705+
706+
request_body = {
707+
prompt,
708+
max_tokens: 4096,
709+
temperature: 0.75,
710+
top_p: 0.99,
711+
};
712+
713+
response_body = {
714+
outputs: [
715+
{
716+
text: 'test-output-text',
717+
stop_reason: 'stop',
718+
},
719+
]
720+
}
721+
}
722+
723+
return [modelId, JSON.stringify(request_body), new TextEncoder().encode(JSON.stringify(response_body))]
724+
}
725+
726+
const [modelId, request_body, response_body] = get_model_request_response();
570727

728+
await withInjected200Success(bedrockRuntimeClient, ['InvokeModelCommand'], { body: response_body }, async () => {
571729
await bedrockRuntimeClient.send(
572730
new InvokeModelCommand({
573-
body: body,
731+
body: request_body,
574732
modelId: modelId,
575733
accept: 'application/json',
576734
contentType: 'application/json',
577735
})
578736
);
579737
});
738+
580739
res.statusCode = 200;
581740
} else {
582741
res.statusCode = 404;
@@ -624,3 +783,4 @@ prepareAwsServer().then(() => {
624783
console.log('Ready');
625784
});
626785
});
786+

0 commit comments

Comments
 (0)