|
15 | 15 | apply_instrumentation_patches,
|
16 | 16 | )
|
17 | 17 | from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
|
18 |
| -from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS |
| 18 | +from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS, bedrock_utils |
19 | 19 | from opentelemetry.propagate import get_global_textmap
|
20 | 20 | from opentelemetry.semconv.trace import SpanAttributes
|
21 | 21 | from opentelemetry.trace.span import Span
|
@@ -84,6 +84,10 @@ def _run_patch_behaviour_tests(self):
|
84 | 84 | self._test_unpatched_botocore_propagator()
|
85 | 85 | self._test_unpatched_gevent_instrumentation()
|
86 | 86 | self._test_unpatched_starlette_instrumentation()
|
| 87 | + # TODO: remove these tests once we bump botocore instrumentation version to 0.56b0 |
| 88 | + # Bedrock Runtime tests |
| 89 | + self._test_unpatched_converse_stream_wrapper() |
| 90 | + self._test_unpatched_extract_tool_calls() |
87 | 91 |
|
88 | 92 | # Apply patches
|
89 | 93 | apply_instrumentation_patches()
|
@@ -219,6 +223,11 @@ def _test_patched_botocore_instrumentation(self):
|
219 | 223 | # Bedrock Agent Operation
|
220 | 224 | self._test_patched_bedrock_agent_instrumentation()
|
221 | 225 |
|
| 226 | + # TODO: remove these tests once we bump botocore instrumentation version to 0.56b0 |
| 227 | + # Bedrock Runtime |
| 228 | + self._test_patched_converse_stream_wrapper() |
| 229 | + self._test_patched_extract_tool_calls() |
| 230 | + |
222 | 231 | # Bedrock Agent Runtime
|
223 | 232 | self.assertTrue("bedrock-agent-runtime" in _KNOWN_EXTENSIONS)
|
224 | 233 | bedrock_agent_runtime_attributes: Dict[str, str] = _do_extract_attributes_bedrock("bedrock-agent-runtime")
|
@@ -470,6 +479,127 @@ def _test_patched_bedrock_instrumentation(self):
|
470 | 479 | self.assertEqual(len(bedrock_sucess_attributes), 1)
|
471 | 480 | self.assertEqual(bedrock_sucess_attributes["aws.bedrock.guardrail.id"], _BEDROCK_GUARDRAIL_ID)
|
472 | 481 |
|
| 482 | + def _test_unpatched_extract_tool_calls(self): |
| 483 | + """Test unpatched extract_tool_calls with string content throws AttributeError""" |
| 484 | + message_with_string_content = {"role": "assistant", "content": "{"} |
| 485 | + with self.assertRaises(AttributeError): |
| 486 | + bedrock_utils.extract_tool_calls(message_with_string_content, True) |
| 487 | + |
| 488 | + def _test_unpatched_converse_stream_wrapper(self): |
| 489 | + """Test unpatched bedrock-runtime where input values remain as numbers""" |
| 490 | + |
| 491 | + mock_stream = MagicMock() |
| 492 | + mock_span = MagicMock() |
| 493 | + mock_stream_error_callback = MagicMock() |
| 494 | + |
| 495 | + wrapper = bedrock_utils.ConverseStreamWrapper(mock_stream, mock_span, mock_stream_error_callback) |
| 496 | + wrapper._record_message = True |
| 497 | + wrapper._message = {"role": "assistant", "content": []} |
| 498 | + |
| 499 | + start_event = { |
| 500 | + "contentBlockStart": { |
| 501 | + "start": { |
| 502 | + "toolUse": { |
| 503 | + "toolUseId": "random_id", |
| 504 | + "name": "example", |
| 505 | + "input": '{"input": 999999999999999999}', |
| 506 | + } |
| 507 | + }, |
| 508 | + "contentBlockIndex": 0, |
| 509 | + } |
| 510 | + } |
| 511 | + wrapper._process_event(start_event) |
| 512 | + |
| 513 | + # Validate that _content_block contains toolUse input that has been JSON decoded |
| 514 | + self.assertIn("toolUse", wrapper._content_block) |
| 515 | + self.assertIn("input", wrapper._content_block["toolUse"]) |
| 516 | + self.assertIn("input", wrapper._content_block["toolUse"]["input"]) |
| 517 | + # Validate that input values are numbers (unpatched behavior) |
| 518 | + self.assertIsInstance(wrapper._content_block["toolUse"]["input"]["input"], int) |
| 519 | + self.assertEqual(wrapper._content_block["toolUse"]["input"]["input"], 999999999999999999) |
| 520 | + |
| 521 | + stop_event = {"contentBlockStop": {"contentBlockIndex": 0}} |
| 522 | + wrapper._process_event(stop_event) |
| 523 | + |
| 524 | + expected_tool_use = { |
| 525 | + "toolUseId": "random_id", |
| 526 | + "name": "example", |
| 527 | + "input": {"input": 999999999999999999}, |
| 528 | + } |
| 529 | + self.assertEqual(len(wrapper._message["content"]), 1) |
| 530 | + self.assertEqual(wrapper._message["content"][0]["toolUse"], expected_tool_use) |
| 531 | + |
| 532 | + def _test_patched_converse_stream_wrapper(self): |
| 533 | + """Test patched bedrock-runtime""" |
| 534 | + |
| 535 | + # Create mock arguments for ConverseStreamWrapper |
| 536 | + mock_stream = MagicMock() |
| 537 | + mock_span = MagicMock() |
| 538 | + mock_stream_error_callback = MagicMock() |
| 539 | + |
| 540 | + # Create real ConverseStreamWrapper with mocked arguments |
| 541 | + wrapper = bedrock_utils.ConverseStreamWrapper(mock_stream, mock_span, mock_stream_error_callback) |
| 542 | + wrapper._record_message = True |
| 543 | + wrapper._message = {"role": "assistant", "content": []} |
| 544 | + |
| 545 | + # Test contentBlockStart |
| 546 | + start_event = { |
| 547 | + "contentBlockStart": { |
| 548 | + "start": { |
| 549 | + "toolUse": { |
| 550 | + "toolUseId": "random_id", |
| 551 | + "name": "example", |
| 552 | + "input": '{"input": 999999999999999999}', |
| 553 | + } |
| 554 | + }, |
| 555 | + "contentBlockIndex": 0, |
| 556 | + } |
| 557 | + } |
| 558 | + |
| 559 | + wrapper._process_event(start_event) |
| 560 | + |
| 561 | + # Validate that _content_block contains toolUse input as literal string (patched behavior) |
| 562 | + self.assertIn("toolUse", wrapper._content_block) |
| 563 | + self.assertIn("input", wrapper._content_block["toolUse"]) |
| 564 | + # Validate that input is a string containing the literal JSON (not decoded) |
| 565 | + self.assertIsInstance(wrapper._content_block["toolUse"]["input"], str) |
| 566 | + self.assertEqual(wrapper._content_block["toolUse"]["input"], '{"input": 999999999999999999}') |
| 567 | + |
| 568 | + # Test contentBlockDelta events |
| 569 | + delta_events = [ |
| 570 | + {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"in'}}, "contentBlockIndex": 0}}, |
| 571 | + {"contentBlockDelta": {"delta": {"toolUse": {"input": 'put": 9'}}, "contentBlockIndex": 0}}, |
| 572 | + {"contentBlockDelta": {"delta": {"toolUse": {"input": "99"}}, "contentBlockIndex": 0}}, |
| 573 | + {"contentBlockDelta": {"delta": {"toolUse": {"input": "99"}}, "contentBlockIndex": 0}}, |
| 574 | + ] |
| 575 | + |
| 576 | + for delta_event in delta_events: |
| 577 | + wrapper._process_event(delta_event) |
| 578 | + |
| 579 | + # Verify accumulated input buffer |
| 580 | + self.assertEqual(wrapper._tool_json_input_buf, '{"input": 99999') |
| 581 | + |
| 582 | + # Test contentBlockStop |
| 583 | + stop_event = {"contentBlockStop": {"contentBlockIndex": 0}} |
| 584 | + wrapper._process_event(stop_event) |
| 585 | + |
| 586 | + # Verify final content_block toolUse value (input becomes the accumulated JSON string) |
| 587 | + expected_tool_use = { |
| 588 | + "toolUseId": "random_id", |
| 589 | + "name": "example", |
| 590 | + "input": '{"input": 99999', |
| 591 | + } |
| 592 | + self.assertEqual(len(wrapper._message["content"]), 1) |
| 593 | + self.assertEqual(wrapper._message["content"][0]["toolUse"], expected_tool_use) |
| 594 | + |
| 595 | + def _test_patched_extract_tool_calls(self): |
| 596 | + """Test patched extract_tool_calls with string content""" |
| 597 | + |
| 598 | + # Test extract_tool_calls with string content (should return None) |
| 599 | + message_with_string_content = {"role": "assistant", "content": "{"} |
| 600 | + result = bedrock_utils.extract_tool_calls(message_with_string_content, True) |
| 601 | + self.assertIsNone(result) |
| 602 | + |
473 | 603 | def _test_patched_bedrock_agent_instrumentation(self):
|
474 | 604 | """For bedrock-agent service, both extract_attributes and on_success provides attributes,
|
475 | 605 | the attributes depend on the API being invoked."""
|
|
0 commit comments