Skip to content

Commit d86f003

Browse files
yuliia-fryshkoyuliia-fryshko
committed
added tests for cohere, cohere-r and mistral
1 parent 16a0e5c commit d86f003

4 files changed

+344
-2
lines changed

‎packages/instrumentation-aws-sdk/test/bedrock-runtime.test.ts‎

Lines changed: 146 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -667,7 +667,6 @@ describe('Bedrock', () => {
667667
});
668668

669669
const response = await client.send(command);
670-
console.log('response', response);
671670

672671
let collectedText = '';
673672
if (!response.body) return;
@@ -741,7 +740,6 @@ describe('Bedrock', () => {
741740
}
742741
}
743742
}
744-
745743
expect(collectedText).toBe(
746744
"Hello! I'm doing well, thank you for asking."
747745
);
@@ -813,6 +811,152 @@ describe('Bedrock', () => {
813811
[ATTR_GEN_AI_RESPONSE_FINISH_REASONS]: ['max_tokens'],
814812
});
815813
});
814+
815+
it('adds mistral ai model attributes to span', async () => {
816+
const modelId = 'mistral.mistral-small-2402-v1:0';
817+
818+
const prompt = '\n\nHuman: Hello, How are you today? \n\nAssistant:';
819+
const nativeRequest: any = {
820+
prompt: prompt,
821+
max_tokens: 20,
822+
temperature: 0.8,
823+
top_p: 1,
824+
stop: ['|'],
825+
};
826+
const command = new InvokeModelWithResponseStreamCommand({
827+
modelId,
828+
body: JSON.stringify(nativeRequest),
829+
});
830+
831+
const response = await client.send(command);
832+
833+
let collectedText = '';
834+
if (!response.body) return;
835+
for await (const chunk of response.body) {
836+
if (chunk?.chunk?.bytes instanceof Uint8Array) {
837+
const parsed = JSON.parse(decodeChunk(chunk));
838+
839+
if (parsed.outputs[0].text) {
840+
collectedText += parsed.outputs[0].text;
841+
}
842+
}
843+
}
844+
expect(collectedText).toBe(
845+
" I'm an AI, so I don't have feelings, but I'm functioning well"
846+
);
847+
848+
const invokeModelSpans: ReadableSpan[] =
849+
getInvokeModelWithResponseStreamSpans();
850+
851+
expect(invokeModelSpans.length).toBe(1);
852+
expect(invokeModelSpans[0].attributes).toMatchObject({
853+
[ATTR_GEN_AI_SYSTEM]: GEN_AI_SYSTEM_VALUE_AWS_BEDROCK,
854+
[ATTR_GEN_AI_REQUEST_MODEL]: modelId,
855+
[ATTR_GEN_AI_REQUEST_MAX_TOKENS]: 20,
856+
[ATTR_GEN_AI_REQUEST_TEMPERATURE]: 0.8,
857+
[ATTR_GEN_AI_REQUEST_TOP_P]: 1,
858+
[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES]: ['|'],
859+
[ATTR_GEN_AI_USAGE_INPUT_TOKENS]: 8,
860+
[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: 1,
861+
[ATTR_GEN_AI_RESPONSE_FINISH_REASONS]: ['length'],
862+
});
863+
});
864+
865+
it('adds cohere command r model attributes to span', async () => {
866+
const modelId = 'cohere.command-r-v1:0';
867+
const prompt = 'Say this is a test';
868+
const nativeRequest: any = {
869+
message: prompt,
870+
max_tokens: 10,
871+
temperature: 0.8,
872+
p: 0.99,
873+
stop_sequences: ['|'],
874+
};
875+
876+
const command = new InvokeModelWithResponseStreamCommand({
877+
modelId,
878+
body: JSON.stringify(nativeRequest),
879+
});
880+
881+
const response = await client.send(command);
882+
883+
let collectedText = '';
884+
if (!response.body) return;
885+
for await (const chunk of response.body) {
886+
if (chunk?.chunk?.bytes instanceof Uint8Array) {
887+
const parsed = JSON.parse(decodeChunk(chunk));
888+
if (parsed.text) {
889+
collectedText += parsed.text;
890+
}
891+
}
892+
}
893+
expect(collectedText).toBe("This is indeed a test. Hopefully, it's");
894+
895+
const invokeModelSpans: ReadableSpan[] =
896+
getInvokeModelWithResponseStreamSpans();
897+
898+
expect(invokeModelSpans.length).toBe(1);
899+
expect(invokeModelSpans[0].attributes).toMatchObject({
900+
[ATTR_GEN_AI_SYSTEM]: GEN_AI_SYSTEM_VALUE_AWS_BEDROCK,
901+
[ATTR_GEN_AI_REQUEST_MODEL]: modelId,
902+
[ATTR_GEN_AI_REQUEST_MAX_TOKENS]: 10,
903+
[ATTR_GEN_AI_REQUEST_TEMPERATURE]: 0.8,
904+
[ATTR_GEN_AI_REQUEST_TOP_P]: 0.99,
905+
[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES]: ['|'],
906+
[ATTR_GEN_AI_USAGE_INPUT_TOKENS]: 3,
907+
[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: 1,
908+
[ATTR_GEN_AI_RESPONSE_FINISH_REASONS]: ['MAX_TOKENS'],
909+
});
910+
});
911+
912+
it('adds cohere command model attributes to span', async () => {
913+
const modelId = 'cohere.command-light-text-v14';
914+
const prompt = 'Say this is a test';
915+
const nativeRequest: any = {
916+
prompt: prompt,
917+
max_tokens: 10,
918+
temperature: 0.8,
919+
p: 1,
920+
stop_sequences: ['|'],
921+
};
922+
923+
const command = new InvokeModelWithResponseStreamCommand({
924+
modelId,
925+
body: JSON.stringify(nativeRequest),
926+
});
927+
928+
const response = await client.send(command);
929+
930+
let collectedText = '';
931+
if (!response.body) return;
932+
for await (const chunk of response.body) {
933+
if (chunk?.chunk?.bytes instanceof Uint8Array) {
934+
const parsed = JSON.parse(decodeChunk(chunk));
935+
if (parsed.generations[0].text) {
936+
collectedText += parsed.generations[0].text;
937+
}
938+
}
939+
}
940+
expect(collectedText).toBe(
941+
' Okay, I will follow your instructions and this will'
942+
);
943+
944+
const invokeModelSpans: ReadableSpan[] =
945+
getInvokeModelWithResponseStreamSpans();
946+
947+
expect(invokeModelSpans.length).toBe(1);
948+
expect(invokeModelSpans[0].attributes).toMatchObject({
949+
[ATTR_GEN_AI_SYSTEM]: GEN_AI_SYSTEM_VALUE_AWS_BEDROCK,
950+
[ATTR_GEN_AI_REQUEST_MODEL]: modelId,
951+
[ATTR_GEN_AI_REQUEST_MAX_TOKENS]: 10,
952+
[ATTR_GEN_AI_REQUEST_TEMPERATURE]: 0.8,
953+
[ATTR_GEN_AI_REQUEST_TOP_P]: 1,
954+
[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES]: ['|'],
955+
[ATTR_GEN_AI_USAGE_INPUT_TOKENS]: 3,
956+
[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: 9,
957+
[ATTR_GEN_AI_RESPONSE_FINISH_REASONS]: ['MAX_TOKENS'],
958+
});
959+
});
816960
});
817961

818962
function getInvokeModelWithResponseStreamSpans(): ReadableSpan[] {

‎packages/instrumentation-aws-sdk/test/mock-responses/bedrock-invokemodelwithstreams-adds-cohere-command-model-attributes-to-span.json‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
[
2+
{
3+
"scope": "https://bedrock-runtime.us-east-1.amazonaws.com:443",
4+
"method": "POST",
5+
"path": "/model/cohere.command-light-text-v14/invoke-with-response-stream",
6+
"body": {
7+
"prompt": "Say this is a test",
8+
"max_tokens": 10,
9+
"temperature": 0.8,
10+
"p": 1,
11+
"stop_sequences": [
12+
"|"
13+
]
14+
},
15+
"status": 403,
16+
"response": {
17+
"message": "Signature expired: 20250808T105739Z is now earlier than 20250808T130125Z (20250808T130625Z - 5 min.)"
18+
},
19+
"rawHeaders": [
20+
"Date",
21+
"Fri, 08 Aug 2025 13:06:25 GMT",
22+
"Content-Type",
23+
"application/json",
24+
"Content-Length",
25+
"114",
26+
"Connection",
27+
"keep-alive",
28+
"x-amzn-RequestId",
29+
"c69d0f37-e52d-4a04-97e3-ddc49f98d665",
30+
"x-amzn-ErrorType",
31+
"InvalidSignatureException:http://internal.amazon.com/coral/com.amazon.coral.service/"
32+
],
33+
"responseIsBinary": false
34+
},
35+
{
36+
"scope": "https://bedrock-runtime.us-east-1.amazonaws.com:443",
37+
"method": "POST",
38+
"path": "/model/cohere.command-light-text-v14/invoke-with-response-stream",
39+
"body": {
40+
"prompt": "Say this is a test",
41+
"max_tokens": 10,
42+
"temperature": 0.8,
43+
"p": 1,
44+
"stop_sequences": [
45+
"|"
46+
]
47+
},
48+
"status": 200,
49+
"response": "000002630000004b1776ebe70b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a705a434936496d4e6b5a6a51304e7a686b4c54457a595755744e4467325a4331694d5752694c5467324d4455314d6a4d77596a56684d534973496d646c626d5679595852706232357a496a706265794a705a434936496a526c5a444d305a6d557a4c544d334e4755744e446b7a4d6930354d4755334c5751354d6a646c4d6a63784d7a6c6a4e694973496e526c654851694f6949675432746865537767535342336157787349475a7662477876647942356233567949476c756333527964574e306157397563794268626d5167644768706379423361577873496977695a6d6c7561584e6f58334a6c59584e7662694936496b31425746395554307446546c4d6966563073496e42796232317764434936496c4e686553423061476c7a49476c7a494745676447567a64434973496d467459587076626931695a57527962324e724c576c75646d396a595852706232354e5a58527961574e7a496a7037496d6c7563485630564739725a57354462335675644349364e5377696233563063485630564739725a57354462335675644349364d544173496d6c75646d396a595852706232354d5958526c626d4e35496a6f794e7a5973496d5a70636e4e30516e6c305a5578686447567559336b694f6a49334e6e3139222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a414243227d60d70fc7",
50+
"rawHeaders": [
51+
"Date",
52+
"Fri, 08 Aug 2025 13:06:25 GMT",
53+
"Content-Type",
54+
"application/vnd.amazon.eventstream",
55+
"Transfer-Encoding",
56+
"chunked",
57+
"Connection",
58+
"keep-alive",
59+
"x-amzn-RequestId",
60+
"cdf4478d-13ae-486d-b1db-86055230b5a1",
61+
"X-Amzn-Bedrock-Content-Type",
62+
"application/json"
63+
],
64+
"responseIsBinary": true
65+
}
66+
]

‎packages/instrumentation-aws-sdk/test/mock-responses/bedrock-invokemodelwithstreams-adds-cohere-command-r-model-attributes-to-span.json‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
[
2+
{
3+
"scope": "https://bedrock-runtime.us-east-1.amazonaws.com:443",
4+
"method": "POST",
5+
"path": "/model/cohere.command-r-v1%3A0/invoke-with-response-stream",
6+
"body": {
7+
"message": "Say this is a test",
8+
"max_tokens": 10,
9+
"temperature": 0.8,
10+
"p": 0.99,
11+
"stop_sequences": [
12+
"|"
13+
]
14+
},
15+
"status": 403,
16+
"response": {
17+
"message": "Signature expired: 20250808T084915Z is now earlier than 20250808T105239Z (20250808T105739Z - 5 min.)"
18+
},
19+
"rawHeaders": [
20+
"Date",
21+
"Fri, 08 Aug 2025 10:57:39 GMT",
22+
"Content-Type",
23+
"application/json",
24+
"Content-Length",
25+
"114",
26+
"Connection",
27+
"keep-alive",
28+
"x-amzn-RequestId",
29+
"dcd1bc90-580f-46f5-ac27-a6628c4f70c5",
30+
"x-amzn-ErrorType",
31+
"InvalidSignatureException:http://internal.amazon.com/coral/com.amazon.coral.service/"
32+
],
33+
"responseIsBinary": false
34+
},
35+
{
36+
"scope": "https://bedrock-runtime.us-east-1.amazonaws.com:443",
37+
"method": "POST",
38+
"path": "/model/cohere.command-r-v1%3A0/invoke-with-response-stream",
39+
"body": {
40+
"message": "Say this is a test",
41+
"max_tokens": 10,
42+
"temperature": 0.8,
43+
"p": 0.99,
44+
"stop_sequences": [
45+
"|"
46+
]
47+
},
48+
"status": 200,
49+
"response": "000000fc0000004b09e8fa2d0b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f695647687063794a39222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a4142434445464748494a4b4c4d4e4f505152535455565758595a3031227dfe7356b8000000e70000004b1ed85cbe0b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f6949476c7a496e303d222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a41424344454647227d2aa812da000000cc0000004ba8c942ab0b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f6949476c755a47566c5a434a39222c2270223a226162227de5dc40e9000000d60000004b8299cd880b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f694947456966513d3d222c2270223a226162636465666768696a6b6c6d6e6f70227d3e962bdf000001010000004b7210bd640b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f694948526c6333516966513d3d222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a4142434445464748494a4b4c4d4e4f505152535455565758595a303132227da52a4efc000000d10000004b30b911980b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f694c694a39222c2270223a226162636465666768696a6b6c6d6e6f227dcacaffe7000000d00000004b0dd938280b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f69494568766347566d6457787365534a39222c2270223a226162227dd9e4e3e7000000ee0000004b13c83ecf0b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f694c434a39222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a4142434445464748494a4b4c4d4e4f505152227da9d92024000000d50000004bc539b7580b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f6949476c30496e303d222c2270223a226162636465666768696a6b6c6d6e6f227d291846b2000000fb0000004bbbc8263d0b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a4349365a6d467363325573496d56325a57353058335235634755694f694a305a5868304c57646c626d567959585270623234694c434a305a586830496a6f694a334d6966513d3d222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a4142434445464748494a4b4c4d4e4f505152535455565758595a30227db50235630000041c0000004bbaad7fe40b3a6576656e742d747970650700056368756e6b0d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226279746573223a2265794a706331396d615735706332686c5a43493664484a315a5377695a585a6c626e526664486c775a534936496e4e30636d56686253316c626d51694c434a795a584e776232357a5a53493665794a795a584e776232357a5a5639705a434936496a4d334f57566b4d4445344c7a4a6d4f475978597a64694c5449344d5759744e475a695a5331684e4459334c5452694d4755304e474e6b4d47517a4d534973496e526c654851694f694a5561476c7a49476c7a49476c755a47566c5a4342684948526c63335175494568766347566d64577873655377676158516e63794973496d646c626d56795958527062323566615751694f6949784d4459785a47457959693077597a49314c5451355a474974596d4a694f5330775a475134596d457a597a68685a5451694c434a6a614746305832687063335276636e6b694f6c7437496e4a76624755694f694a5655305653496977696257567a6332466e5a534936496c4e686553423061476c7a49476c7a494745676447567a64434a394c487369636d39735a534936496b4e4951565243543151694c434a745a584e7a5957646c496a6f69564768706379427063794270626d526c5a575167595342305a584e304c6942496233426c5a6e567362486b7349476c304a334d6966563073496d5a70626d6c7a614639795a57467a623234694f694a4e515668665645394c5255355449697769625756305953493665794a6863476c66646d567963326c766269493665794a325a584a7a61573975496a6f694d534a394c434a69615778735a5752666457357064484d694f6e73696157357764585266644739725a57357a496a6f314c434a766458527764585266644739725a57357a496a6f784d483073496e5276613256756379493665794a70626e4231644639306232746c626e4d694f6a63784c434a766458527764585266644739725a57357a496a6f784d483139665377695a6d6c7561584e6f58334a6c59584e7662694936496b31425746395554307446546c4d694c434a686257463662323474596d566b636d396a61793170626e5a76593246306157397554575630636d6c6a6379493665794a70626e4231644652766132567551323931626e51694f6a5573496d393164484231644652766132567551323931626e51694f6a45774c434a70626e5a765932463061573975544746305a57356a655349364d6a41774c434a6d61584a7a64454a356447564d5958526c626d4e35496a6f354d333139222c2270223a226162636465666768696a6b6c6d6e6f707172737475767778797a4142434445464748227dffffe971",
50+
"rawHeaders": [
51+
"Date",
52+
"Fri, 08 Aug 2025 10:57:39 GMT",
53+
"Content-Type",
54+
"application/vnd.amazon.eventstream",
55+
"Transfer-Encoding",
56+
"chunked",
57+
"Connection",
58+
"keep-alive",
59+
"x-amzn-RequestId",
60+
"22fb7d78-9282-4b2d-a946-46d339e711f9",
61+
"X-Amzn-Bedrock-Content-Type",
62+
"application/json"
63+
],
64+
"responseIsBinary": true
65+
}
66+
]

0 commit comments

Comments
 (0)