diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 91011f2a8b..c99293a63f 100644 --- a/newrelic/hooks/external_botocore.py +++ b/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), ] diff --git a/tests/mlmodel_bedrock/_mock_external_bedrock_server.py b/tests/mlmodel_bedrock/_mock_external_bedrock_server.py index 96598a02a5..c9149e3f1b 100644 --- a/tests/mlmodel_bedrock/_mock_external_bedrock_server.py +++ b/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"}, { diff --git a/tests/mlmodel_bedrock/_test_chat_completion.py b/tests/mlmodel_bedrock/_test_chat_completion.py index 17fa8549cb..3f1f297d4a 100644 --- a/tests/mlmodel_bedrock/_test_chat_completion.py +++ b/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"}, diff --git a/tests/mlmodel_bedrock/test_chat_completion.py b/tests/mlmodel_bedrock/test_chat_completion.py index 7e6c04899d..4da60eaee4 100644 --- a/tests/mlmodel_bedrock/test_chat_completion.py +++ b/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", ], )