apply linter

Author: thpierce

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -344,7 +344,9 @@ def patched_process_anthropic_claude_chunk(self, chunk):
             return
 
         if message_type == "message_start":
-            # {'type': 'message_start', 'message': {'id': 'id', 'type': 'message', 'role': 'assistant', 'model': 'claude-2.0', 'content': [], 'stop_reason': None, 'stop_sequence': None, 'usage': {'input_tokens': 18, 'output_tokens': 1}}}
+            # {'type': 'message_start', 'message': {'id': 'id', 'type': 'message', 'role': 'assistant',
+            # 'model': 'claude-2.0', 'content': [], 'stop_reason': None, 'stop_sequence': None,
+            # 'usage': {'input_tokens': 18, 'output_tokens': 1}}}
             if chunk.get("message", {}).get("role") == "assistant":
                 self._record_message = True
                 message = chunk["message"]
@@ -356,7 +358,8 @@ def patched_process_anthropic_claude_chunk(self, chunk):
 
         if message_type == "content_block_start":
             # {'type': 'content_block_start', 'index': 0, 'content_block': {'type': 'text', 'text': ''}}
-            # {'type': 'content_block_start', 'index': 1, 'content_block': {'type': 'tool_use', 'id': 'id', 'name': 'func_name', 'input': {}}}
+            # {'type': 'content_block_start', 'index': 1, 'content_block':
+            # {'type': 'tool_use', 'id': 'id', 'name': 'func_name', 'input': {}}}
             if self._record_message:
                 block = chunk.get("content_block", {})
                 if block.get("type") == "text":
@@ -380,9 +383,7 @@ def patched_process_anthropic_claude_chunk(self, chunk):
             # {'type': 'content_block_stop', 'index': 0}
             if self._tool_json_input_buf:
                 try:
-                    self._content_block["input"] = json.loads(
-                        self._tool_json_input_buf
-                    )
+                    self._content_block["input"] = json.loads(self._tool_json_input_buf)
                 except json.JSONDecodeError:
                     self._content_block["input"] = self._tool_json_input_buf
             self._message["content"].append(self._content_block)
@@ -391,18 +392,16 @@ def patched_process_anthropic_claude_chunk(self, chunk):
             return
 
         if message_type == "message_delta":
-            # {'type': 'message_delta', 'delta': {'stop_reason': 'end_turn', 'stop_sequence': None}, 'usage': {'output_tokens': 123}}
-            if (
-                stop_reason := chunk.get("delta", {}).get("stop_reason")
-            ) is not None:
+            # {'type': 'message_delta', 'delta': {'stop_reason': 'end_turn', 'stop_sequence': None},
+            # 'usage': {'output_tokens': 123}}
+            if (stop_reason := chunk.get("delta", {}).get("stop_reason")) is not None:
                 self._response["stopReason"] = stop_reason
             return
 
         if message_type == "message_stop":
-            # {'type': 'message_stop', 'amazon-bedrock-invocationMetrics': {'inputTokenCount': 18, 'outputTokenCount': 123, 'invocationLatency': 5250, 'firstByteLatency': 290}}
-            if invocation_metrics := chunk.get(
-                "amazon-bedrock-invocationMetrics"
-            ):
+            # {'type': 'message_stop', 'amazon-bedrock-invocationMetrics':
+            # {'inputTokenCount': 18, 'outputTokenCount': 123, 'invocationLatency': 5250, 'firstByteLatency': 290}}
+            if invocation_metrics := chunk.get("amazon-bedrock-invocationMetrics"):
                 self._process_invocation_metrics(invocation_metrics)
 
             if self._record_message:
@@ -415,7 +414,9 @@ def patched_process_anthropic_claude_chunk(self, chunk):
 
     bedrock_utils.ConverseStreamWrapper.__init__ = patched_init
     bedrock_utils.ConverseStreamWrapper._process_event = patched_process_event
-    bedrock_utils.InvokeModelWithResponseStreamWrapper._process_anthropic_claude_chunk = patched_process_anthropic_claude_chunk
+    bedrock_utils.InvokeModelWithResponseStreamWrapper._process_anthropic_claude_chunk = (
+        patched_process_anthropic_claude_chunk
+    )
     bedrock_utils.extract_tool_calls = patched_extract_tool_calls
 
     # END The OpenTelemetry Authors code

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_instrumentation_patch.py

@@ -228,7 +228,6 @@ def _test_unpatched_botocore_instrumentation(self):
         self._test_unpatched_process_anthropic_claude_chunk(None, None)
         self._test_unpatched_process_anthropic_claude_chunk({}, {})
 
-
     def _test_unpatched_gevent_instrumentation(self):
         self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched")
         self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched")
@@ -690,13 +689,19 @@ def _test_patched_extract_tool_calls(self):
         result = bedrock_utils.extract_tool_calls(message_with_string_content, True)
         self.assertIsNone(result)
 
-    def _test_patched_process_anthropic_claude_chunk(self, input_value: Dict[str, str], expected_output: Dict[str, str]):
+    def _test_patched_process_anthropic_claude_chunk(
+        self, input_value: Dict[str, str], expected_output: Dict[str, str]
+    ):
         self._test_process_anthropic_claude_chunk(input_value, expected_output, False)
 
-    def _test_unpatched_process_anthropic_claude_chunk(self, input_value: Dict[str, str], expected_output: Dict[str, str]):
+    def _test_unpatched_process_anthropic_claude_chunk(
+        self, input_value: Dict[str, str], expected_output: Dict[str, str]
+    ):
         self._test_process_anthropic_claude_chunk(input_value, expected_output, True)
 
-    def _test_process_anthropic_claude_chunk(self, input_value: Dict[str, str], expected_output: Dict[str, str], expect_exception: bool):
+    def _test_process_anthropic_claude_chunk(
+        self, input_value: Dict[str, str], expected_output: Dict[str, str], expect_exception: bool
+    ):
         """Test that _process_anthropic_claude_chunk handles various tool_use input formats."""
         wrapper = bedrock_utils.InvokeModelWithResponseStreamWrapper(
             stream=MagicMock(),
@@ -735,9 +740,7 @@ def _test_process_anthropic_claude_chunk(self, input_value: Dict[str, str], expe
 
         # Simulate content_block_stop
         try:
-            wrapper._process_anthropic_claude_chunk(
-                {"type": "content_block_stop", "index": 0}
-            )
+            wrapper._process_anthropic_claude_chunk({"type": "content_block_stop", "index": 0})
         except TypeError:
             if expect_exception:
                 return