Fixup: lint errors

newrelic · Dec 20, 2023 · 1c1f408 · 1c1f408
1 parent 368225f
commit 1c1f408
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 11 deletions.
diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py
@@ -158,7 +158,6 @@ def extract_bedrock_titan_text_model(request_body, response_body=None):
 
     input_message_list = [{"role": "user", "content": request_body.get("inputText", "")}]
 
-
     chat_completion_summary_dict = {
         "request.max_tokens": request_config.get("maxTokenCount", ""),
         "request.temperature": request_config.get("temperature", ""),
@@ -170,7 +169,9 @@ def extract_bedrock_titan_text_model(request_body, response_body=None):
         completion_tokens = sum(result["tokenCount"] for result in response_body.get("results", []))
         total_tokens = input_tokens + completion_tokens
 
-        output_message_list = [{"role": "assistant", "content": result["outputText"]} for result in response_body.get("results", [])]
+        output_message_list = [
+            {"role": "assistant", "content": result["outputText"]} for result in response_body.get("results", [])
+        ]
 
         chat_completion_summary_dict.update(
             {
@@ -218,7 +219,9 @@ def extract_bedrock_ai21_j2_model(request_body, response_body=None):
     }
 
     if response_body:
-        output_message_list =[{"role": "assistant", "content": result["data"]["text"]} for result in response_body.get("completions", [])]
+        output_message_list = [
+            {"role": "assistant", "content": result["data"]["text"]} for result in response_body.get("completions", [])
+        ]
 
         chat_completion_summary_dict.update(
             {
@@ -309,7 +312,9 @@ def extract_bedrock_cohere_model(request_body, response_body=None):
     }
 
     if response_body:
-        output_message_list = [{"role": "assistant", "content": result["text"]} for result in response_body.get("generations", [])]
+        output_message_list = [
+            {"role": "assistant", "content": result["text"]} for result in response_body.get("generations", [])
+        ]
         chat_completion_summary_dict.update(
             {
                 "response.choices.finish_reason": response_body["generations"][0]["finish_reason"],
@@ -412,13 +417,31 @@ def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
 
                 if operation == "embedding":  # Only available embedding models
                     handle_embedding_event(
-                        instance, transaction, extractor, model, None, None, request_body,
-                        ft.duration, True, trace_id, span_id
+                        instance,
+                        transaction,
+                        extractor,
+                        model,
+                        None,
+                        None,
+                        request_body,
+                        ft.duration,
+                        True,
+                        trace_id,
+                        span_id
                     )
                 else:
                     handle_chat_completion_event(
-                        instance, transaction, extractor, model, None, None, request_body,
-                        ft.duration, True, trace_id, span_id
+                        instance,
+                        transaction,
+                        extractor,
+                        model,
+                        None,
+                        None,
+                        request_body,
+                        ft.duration,
+                        True,
+                        trace_id,
+                        span_id
                     )
 
             finally:
@@ -465,7 +488,17 @@ def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
 
 
 def handle_embedding_event(
-    client, transaction, extractor, model, response_body, response_headers, request_body, duration, is_error, trace_id, span_id
+    client,
+    transaction,
+    extractor,
+    model,
+    response_body,
+    response_headers,
+    request_body,
+    duration,
+    is_error,
+    trace_id,
+    span_id
 ):
     embedding_id = str(uuid.uuid4())
 
@@ -500,7 +533,17 @@ def handle_embedding_event(
 
 
 def handle_chat_completion_event(
-   client, transaction, extractor, model, response_body, response_headers, request_body, duration, is_error, trace_id, span_id
+    client,
+    transaction,
+    extractor,
+    model,
+    response_body,
+    response_headers,
+    request_body,
+    duration,
+    is_error,
+    trace_id,
+    span_id
 ):
     custom_attrs_dict = transaction._custom_params
     conversation_id = custom_attrs_dict.get("conversation_id", "")

diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py
@@ -23,7 +23,6 @@
     chat_completion_expected_events,
     chat_completion_invalid_access_key_error_events,
     chat_completion_payload_templates,
-    chat_completion_invalid_access_key_error_events,
 )
 from conftest import BOTOCORE_VERSION
 from testing_support.fixtures import (