Add bedrock claude model support

newrelic/hooks/external_botocore.py

@@ -154,6 +154,24 @@ def extract_bedrock_ai21_j2_model(request_body, response_body):
     return message_list, chat_completion_summary_dict
 
 
+def extract_bedrock_claude_model(request_body, response_body):
+    response_body = json.loads(response_body)
+    request_body = json.loads(request_body)
+
+    message_list = [
+        {"role": "user", "content": request_body.get("prompt", "")},
+        {"role": "assistant", "content": response_body.get("completion", "")},
+    ]
+
+    chat_completion_summary_dict = {
+        "request.max_tokens": request_body.get("max_tokens_to_sample", ""),
+        "request.temperature": request_body.get("temperature", ""),
+        "response.choices.finish_reason": response_body.get("stop_reason", ""),
+        "response.number_of_messages": len(message_list),
+    }
+    return message_list, chat_completion_summary_dict
+
+
 def extract_bedrock_cohere_model(request_body, response_body):
     response_body = json.loads(response_body)
     request_body = json.loads(request_body)
@@ -178,6 +196,7 @@ def extract_bedrock_cohere_model(request_body, response_body):
     ("amazon.titan", extract_bedrock_titan_text_model),
     ("ai21.j2", extract_bedrock_ai21_j2_model),
     ("cohere", extract_bedrock_cohere_model),
+    ("anthropic.claude", extract_bedrock_claude_model),
 ]
 
 

tests/mlmodel_bedrock/_mock_external_bedrock_server.py

@@ -42,6 +42,14 @@
             ],
         },
     ],
+    "anthropic.claude-instant-v1::Human: What is 212 degrees Fahrenheit converted to Celsius? Assistant:": [
+        {"content-type": "application/json", "x-amzn-requestid": "f354b9a7-9eac-4f50-a8d7-7d5d23566176"},
+        {
+            "completion": " Here are the step-by-step workings:\n1) 212 degrees Fahrenheit \n2) To convert to Celsius, use the formula: C = (F - 32) * 5/9\n3) Plug in the values: C = (212 - 32) * 5/9 = 100 * 5/9 = 100 degrees Celsius\n\nSo, 212 degrees Fahrenheit converted to Celsius is 100 degrees Celsius.",
+            "stop_reason": "stop_sequence",
+            "stop": "\n\nHuman:",
+        },
+    ],
     "cohere.command-text-v14::What is 212 degrees Fahrenheit converted to Celsius?": [
         {"content-type": "application/json", "x-amzn-requestid": "c5188fb5-dc58-4cbe-948d-af173c69ce0d"},
         {

tests/mlmodel_bedrock/_test_chat_completion.py

@@ -1,6 +1,7 @@
 chat_completion_payload_templates = {
     "amazon.titan-text-express-v1": '{ "inputText": "%s", "textGenerationConfig": {"temperature": %f, "maxTokenCount": %d }}',
     "ai21.j2-mid-v1": '{"prompt": "%s", "temperature": %f, "maxTokens": %d}',
+    "anthropic.claude-instant-v1": '{"prompt": "Human: %s Assistant:", "temperature": %f, "max_tokens_to_sample": %d}',
     "cohere.command-text-v14": '{"prompt": "%s", "temperature": %f, "max_tokens": %d}',
 }
 
@@ -133,6 +134,68 @@
             },
         ),
     ],
+    "anthropic.claude-instant-v1": [
+        (
+            {"type": "LlmChatCompletionSummary"},
+            {
+                "id": None,  # UUID that varies with each run
+                "appName": "Python Agent Test (mlmodel_bedrock)",
+                "conversation_id": "my-awesome-id",
+                "transaction_id": None,
+                "span_id": "span-id",
+                "trace_id": "trace-id",
+                "request_id": "f354b9a7-9eac-4f50-a8d7-7d5d23566176",
+                "api_key_last_four_digits": "CRET",
+                "duration": None,  # Response time varies each test run
+                "request.model": "anthropic.claude-instant-v1",
+                "response.model": "anthropic.claude-instant-v1",
+                "request.temperature": 0.7,
+                "request.max_tokens": 100,
+                "response.choices.finish_reason": "stop_sequence",
+                "vendor": "bedrock",
+                "ingest_source": "Python",
+                "response.number_of_messages": 2,
+            },
+        ),
+        (
+            {"type": "LlmChatCompletionMessage"},
+            {
+                "id": None,  # UUID that varies with each run
+                "appName": "Python Agent Test (mlmodel_bedrock)",
+                "conversation_id": "my-awesome-id",
+                "request_id": "f354b9a7-9eac-4f50-a8d7-7d5d23566176",
+                "span_id": "span-id",
+                "trace_id": "trace-id",
+                "transaction_id": None,
+                "content": "Human: What is 212 degrees Fahrenheit converted to Celsius? Assistant:",
+                "role": "user",
+                "completion_id": None,
+                "sequence": 0,
+                "response.model": "anthropic.claude-instant-v1",
+                "vendor": "bedrock",
+                "ingest_source": "Python",
+            },
+        ),
+        (
+            {"type": "LlmChatCompletionMessage"},
+            {
+                "id": None,  # UUID that varies with each run
+                "appName": "Python Agent Test (mlmodel_bedrock)",
+                "conversation_id": "my-awesome-id",
+                "request_id": "f354b9a7-9eac-4f50-a8d7-7d5d23566176",
+                "span_id": "span-id",
+                "trace_id": "trace-id",
+                "transaction_id": None,
+                "content": " Here are the step-by-step workings:\n1) 212 degrees Fahrenheit \n2) To convert to Celsius, use the formula: C = (F - 32) * 5/9\n3) Plug in the values: C = (212 - 32) * 5/9 = 100 * 5/9 = 100 degrees Celsius\n\nSo, 212 degrees Fahrenheit converted to Celsius is",
+                "role": "assistant",
+                "completion_id": None,
+                "sequence": 1,
+                "response.model": "anthropic.claude-instant-v1",
+                "vendor": "bedrock",
+                "ingest_source": "Python",
+            },
+        ),
+    ],
     "cohere.command-text-v14": [
         (
             {"type": "LlmChatCompletionSummary"},

tests/mlmodel_bedrock/test_chat_completion.py

@@ -43,7 +43,7 @@ def is_file_payload(request):
     params=[
         "amazon.titan-text-express-v1",
         "ai21.j2-mid-v1",
-        # ("anthropic.claude-instant-v1", '{"prompt": "Human: {prompt}\n\nAssistant:", "max_tokens_to_sample": {max_tokens:d}}'),
+        "anthropic.claude-instant-v1",
         "cohere.command-text-v14",
     ],
 )