Skip to content

Commit d76618e

Browse files
ezhang6811ezhang6811
committed
add GenAI attribute to verify
1 parent 4de0497 commit d76618e

File tree

1 file changed

+34
-27
lines changed

1 file changed

+34
-27
lines changed

‎test/contract-tests/tests/test/amazon/awssdk/awssdk_test.py‎

Lines changed: 34 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
_AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id"
3232
_AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id"
3333
_AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
34+
_GEN_AI_SYSTEM: str = "gen_ai.system"
3435
_GEN_AI_REQUEST_MODEL: str = "gen_ai.request.model"
3536
_GEN_AI_REQUEST_TOP_P: str = "gen_ai.request.top_p"
3637
_GEN_AI_REQUEST_TEMPERATURE: str = "gen_ai.request.temperature"
@@ -109,7 +110,7 @@ def test_s3_create_bucket(self):
109110
remote_operation="PutBucket",
110111
remote_resource_type="AWS::S3::Bucket",
111112
remote_resource_identifier="test-bucket-name",
112-
request_specific_attributes={
113+
request_response_specific_attributes={
113114
SpanAttributes.AWS_S3_BUCKET: "test-bucket-name",
114115
},
115116
span_name="S3.PutBucket",
@@ -126,7 +127,7 @@ def test_s3_create_object(self):
126127
remote_operation="PutObject",
127128
remote_resource_type="AWS::S3::Bucket",
128129
remote_resource_identifier="test-bucket-name",
129-
request_specific_attributes={
130+
request_response_specific_attributes={
130131
SpanAttributes.AWS_S3_BUCKET: "test-bucket-name",
131132
},
132133
span_name="S3.PutObject",
@@ -143,7 +144,7 @@ def test_s3_delete_object(self):
143144
remote_operation="DeleteObject",
144145
remote_resource_type="AWS::S3::Bucket",
145146
remote_resource_identifier="test-bucket-name",
146-
request_specific_attributes={
147+
request_response_specific_attributes={
147148
SpanAttributes.AWS_S3_BUCKET: "test-bucket-name",
148149
},
149150
span_name="S3.DeleteObject",
@@ -160,7 +161,7 @@ def test_dynamodb_create_table(self):
160161
remote_operation="CreateTable",
161162
remote_resource_type="AWS::DynamoDB::Table",
162163
remote_resource_identifier="test_table",
163-
request_specific_attributes={
164+
request_response_specific_attributes={
164165
# SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["test_table"],
165166
"aws.table_name": ["test_table"],
166167
},
@@ -178,7 +179,7 @@ def test_dynamodb_put_item(self):
178179
remote_operation="PutItem",
179180
remote_resource_type="AWS::DynamoDB::Table",
180181
remote_resource_identifier="test_table",
181-
request_specific_attributes={
182+
request_response_specific_attributes={
182183
# SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["test_table"],
183184
"aws.table_name": ["test_table"],
184185
},
@@ -196,7 +197,7 @@ def test_sqs_create_queue(self):
196197
remote_operation="CreateQueue",
197198
remote_resource_type="AWS::SQS::Queue",
198199
remote_resource_identifier="test_queue",
199-
request_specific_attributes={
200+
request_response_specific_attributes={
200201
_AWS_SQS_QUEUE_NAME: "test_queue",
201202
},
202203
span_name="SQS.CreateQueue",
@@ -213,7 +214,7 @@ def test_sqs_send_message(self):
213214
remote_operation="SendMessage",
214215
remote_resource_type="AWS::SQS::Queue",
215216
remote_resource_identifier="test_queue",
216-
request_specific_attributes={
217+
request_response_specific_attributes={
217218
_AWS_SQS_QUEUE_URL: "http://sqs.us-east-1.localstack:4566/000000000000/test_queue",
218219
},
219220
span_name="SQS.SendMessage",
@@ -230,7 +231,7 @@ def test_sqs_receive_message(self):
230231
remote_operation="ReceiveMessage",
231232
remote_resource_type="AWS::SQS::Queue",
232233
remote_resource_identifier="test_queue",
233-
request_specific_attributes={
234+
request_response_specific_attributes={
234235
_AWS_SQS_QUEUE_URL: "http://sqs.us-east-1.localstack:4566/000000000000/test_queue",
235236
},
236237
span_name="SQS.ReceiveMessage",
@@ -247,7 +248,7 @@ def test_kinesis_create_stream(self):
247248
remote_operation="CreateStream",
248249
remote_resource_type="AWS::Kinesis::Stream",
249250
remote_resource_identifier="test_stream",
250-
request_specific_attributes={
251+
request_response_specific_attributes={
251252
_AWS_KINESIS_STREAM_NAME: "test_stream",
252253
},
253254
span_name="Kinesis.CreateStream",
@@ -264,7 +265,7 @@ def test_kinesis_put_record(self):
264265
remote_operation="PutRecord",
265266
remote_resource_type="AWS::Kinesis::Stream",
266267
remote_resource_identifier="test_stream",
267-
request_specific_attributes={
268+
request_response_specific_attributes={
268269
_AWS_KINESIS_STREAM_NAME: "test_stream",
269270
},
270271
span_name="Kinesis.PutRecord",
@@ -281,7 +282,7 @@ def test_kinesis_error(self):
281282
remote_operation="DeleteStream",
282283
remote_resource_type="AWS::Kinesis::Stream",
283284
remote_resource_identifier="test_stream_error",
284-
request_specific_attributes={
285+
request_response_specific_attributes={
285286
_AWS_KINESIS_STREAM_NAME: "test_stream_error",
286287
},
287288
span_name="Kinesis.DeleteStream",
@@ -299,7 +300,7 @@ def test_kinesis_error(self):
299300
# remote_operation="CreateStream",
300301
# remote_resource_type="AWS::Kinesis::Stream",
301302
# remote_resource_identifier="test_stream",
302-
# request_specific_attributes={
303+
# request_response_specific_attributes={
303304
# _AWS_KINESIS_STREAM_NAME: "test_stream",
304305
# },
305306
# span_name="Kinesis.CreateStream",
@@ -317,7 +318,7 @@ def test_bedrock_get_guardrail(self):
317318
remote_operation="GetGuardrail",
318319
remote_resource_type="AWS::Bedrock::Guardrail",
319320
remote_resource_identifier="test-guardrail",
320-
request_specific_attributes={
321+
request_response_specific_attributes={
321322
_AWS_BEDROCK_GUARDRAIL_ID: "test-guardrail",
322323
},
323324
span_name="Bedrock.GetGuardrail",
@@ -335,7 +336,8 @@ def test_bedrock_runtime_invoke_model_titan(self):
335336
remote_operation="InvokeModel",
336337
remote_resource_type="AWS::Bedrock::Model",
337338
remote_resource_identifier="amazon.titan-text-express-v1",
338-
request_specific_attributes={
339+
request_response_specific_attributes={
340+
_GEN_AI_SYSTEM: "aws_bedrock",
339341
_GEN_AI_REQUEST_MODEL: "amazon.titan-text-express-v1",
340342
_GEN_AI_REQUEST_TEMPERATURE: 0.123,
341343
_GEN_AI_REQUEST_TOP_P: 0.456,
@@ -359,7 +361,8 @@ def test_bedrock_runtime_invoke_model_claude(self):
359361
remote_operation="InvokeModel",
360362
remote_resource_type="AWS::Bedrock::Model",
361363
remote_resource_identifier="anthropic.claude-v2:1",
362-
request_specific_attributes={
364+
request_response_specific_attributes={
365+
_GEN_AI_SYSTEM: "aws_bedrock",
363366
_GEN_AI_REQUEST_MODEL: "anthropic.claude-v2:1",
364367
_GEN_AI_REQUEST_TEMPERATURE: 0.123,
365368
_GEN_AI_REQUEST_TOP_P: 0.456,
@@ -383,7 +386,8 @@ def test_bedrock_runtime_invoke_model_llama(self):
383386
remote_operation="InvokeModel",
384387
remote_resource_type="AWS::Bedrock::Model",
385388
remote_resource_identifier="meta.llama3-8b-instruct-v1:0",
386-
request_specific_attributes={
389+
request_response_specific_attributes={
390+
_GEN_AI_SYSTEM: "aws_bedrock",
387391
_GEN_AI_REQUEST_MODEL: "meta.llama3-8b-instruct-v1:0",
388392
_GEN_AI_REQUEST_TEMPERATURE: 0.123,
389393
_GEN_AI_REQUEST_TOP_P: 0.456,
@@ -407,7 +411,8 @@ def test_bedrock_runtime_invoke_model_command(self):
407411
remote_operation="InvokeModel",
408412
remote_resource_type="AWS::Bedrock::Model",
409413
remote_resource_identifier="cohere.command-r-v1:0",
410-
request_specific_attributes={
414+
request_response_specific_attributes={
415+
_GEN_AI_SYSTEM: "aws_bedrock",
411416
_GEN_AI_REQUEST_MODEL: "cohere.command-r-v1:0",
412417
_GEN_AI_REQUEST_TEMPERATURE: 0.123,
413418
_GEN_AI_REQUEST_TOP_P: 0.456,
@@ -431,7 +436,8 @@ def test_bedrock_runtime_invoke_model_jamba(self):
431436
remote_operation="InvokeModel",
432437
remote_resource_type="AWS::Bedrock::Model",
433438
remote_resource_identifier="ai21.jamba-1-5-large-v1:0",
434-
request_specific_attributes={
439+
request_response_specific_attributes={
440+
_GEN_AI_SYSTEM: "aws_bedrock",
435441
_GEN_AI_REQUEST_MODEL: "ai21.jamba-1-5-large-v1:0",
436442
_GEN_AI_REQUEST_TEMPERATURE: 0.123,
437443
_GEN_AI_REQUEST_TOP_P: 0.456,
@@ -455,7 +461,8 @@ def test_bedrock_runtime_invoke_model_mistral(self):
455461
remote_operation="InvokeModel",
456462
remote_resource_type="AWS::Bedrock::Model",
457463
remote_resource_identifier="mistral.mistral-7b-instruct-v0:2",
458-
request_specific_attributes={
464+
request_response_specific_attributes={
465+
_GEN_AI_SYSTEM: "aws_bedrock",
459466
_GEN_AI_REQUEST_MODEL: "mistral.mistral-7b-instruct-v0:2",
460467
_GEN_AI_REQUEST_TEMPERATURE: 0.123,
461468
_GEN_AI_REQUEST_TOP_P: 0.456,
@@ -479,7 +486,7 @@ def test_bedrock_agent_runtime_invoke_agent(self):
479486
remote_operation="InvokeAgent",
480487
remote_resource_type="AWS::Bedrock::Agent",
481488
remote_resource_identifier="test-agent",
482-
request_specific_attributes={
489+
request_response_specific_attributes={
483490
_AWS_BEDROCK_AGENT_ID: "test-agent",
484491
},
485492
span_name="Bedrock Agent Runtime.InvokeAgent",
@@ -497,7 +504,7 @@ def test_bedrock_agent_runtime_retrieve(self):
497504
remote_operation="Retrieve",
498505
remote_resource_type="AWS::Bedrock::KnowledgeBase",
499506
remote_resource_identifier="test-knowledge-base",
500-
request_specific_attributes={
507+
request_response_specific_attributes={
501508
_AWS_BEDROCK_KNOWLEDGE_BASE_ID: "test-knowledge-base",
502509
},
503510
span_name="Bedrock Agent Runtime.Retrieve",
@@ -515,7 +522,7 @@ def test_bedrock_agent_get_agent(self):
515522
remote_operation="GetAgent",
516523
remote_resource_type="AWS::Bedrock::Agent",
517524
remote_resource_identifier="test-agent",
518-
request_specific_attributes={
525+
request_response_specific_attributes={
519526
_AWS_BEDROCK_AGENT_ID: "test-agent",
520527
},
521528
span_name="Bedrock Agent.GetAgent",
@@ -533,7 +540,7 @@ def test_bedrock_agent_get_knowledge_base(self):
533540
remote_operation="GetKnowledgeBase",
534541
remote_resource_type="AWS::Bedrock::KnowledgeBase",
535542
remote_resource_identifier="test-knowledge-base",
536-
request_specific_attributes={
543+
request_response_specific_attributes={
537544
_AWS_BEDROCK_KNOWLEDGE_BASE_ID: "test-knowledge-base",
538545
},
539546
span_name="Bedrock Agent.GetKnowledgeBase",
@@ -551,7 +558,7 @@ def test_bedrock_agent_get_data_source(self):
551558
remote_operation="GetDataSource",
552559
remote_resource_type="AWS::Bedrock::DataSource",
553560
remote_resource_identifier="test-data-source",
554-
request_specific_attributes={
561+
request_response_specific_attributes={
555562
_AWS_BEDROCK_DATA_SOURCE_ID: "test-data-source",
556563
},
557564
span_name="Bedrock Agent.GetDataSource",
@@ -614,7 +621,7 @@ def _assert_semantic_conventions_span_attributes(
614621
kwargs.get("remote_service"),
615622
kwargs.get("remote_operation"),
616623
status_code,
617-
kwargs.get("request_specific_attributes", {}),
624+
kwargs.get("request_response_specific_attributes", {}),
618625
)
619626

620627
# pylint: disable=unidiomatic-typecheck
@@ -625,14 +632,14 @@ def _assert_semantic_conventions_attributes(
625632
service: str,
626633
operation: str,
627634
status_code: int,
628-
request_specific_attributes: dict,
635+
request_response_specific_attributes: dict,
629636
) -> None:
630637
attributes_dict: Dict[str, AnyValue] = self._get_attributes_dict(attributes_list)
631638
self._assert_str_attribute(attributes_dict, SpanAttributes.RPC_METHOD, operation)
632639
self._assert_str_attribute(attributes_dict, SpanAttributes.RPC_SYSTEM, "aws-api")
633640
self._assert_str_attribute(attributes_dict, SpanAttributes.RPC_SERVICE, rpc_service)
634641
self._assert_int_attribute(attributes_dict, SpanAttributes.HTTP_STATUS_CODE, status_code)
635-
for key, value in request_specific_attributes.items():
642+
for key, value in request_response_specific_attributes.items():
636643
if isinstance(value, str):
637644
self._assert_str_attribute(attributes_dict, key, value)
638645
elif isinstance(value, int):

0 commit comments

Comments
 (0)