diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py index 018df2d32..4131e2bf7 100644 --- a/newrelic/hooks/external_botocore.py +++ b/newrelic/hooks/external_botocore.py @@ -140,9 +140,29 @@ def extract_bedrock_ai21_j2_model(request_body, response_body): 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) + + message_list = [{"role": "user", "content": request_body.get("prompt", "")}] + message_list.extend( + {"role": "assistant", "content": result["text"]} for result in response_body.get("generations", []) + ) + + chat_completion_summary_dict = { + "request.max_tokens": request_body.get("max_tokens", ""), + "request.temperature": request_body.get("temperature", ""), + "response.choices.finish_reason": response_body["generations"][0]["finish_reason"], + "response.number_of_messages": len(message_list), + "response_id": str(response_body.get("id", "")), + } + return message_list, chat_completion_summary_dict + + MODEL_EXTRACTORS = { "amazon.titan": extract_bedrock_titan_model, "ai21.j2": extract_bedrock_ai21_j2_model, + "cohere": extract_bedrock_cohere_model, } diff --git a/tests/mlmodel_bedrock/_mock_external_bedrock_server.py b/tests/mlmodel_bedrock/_mock_external_bedrock_server.py index de8a91aad..126c663ce 100644 --- a/tests/mlmodel_bedrock/_mock_external_bedrock_server.py +++ b/tests/mlmodel_bedrock/_mock_external_bedrock_server.py @@ -42,6 +42,20 @@ ], }, ], + "cohere.command-text-v14::What is 212 degrees Fahrenheit converted to Celsius?": [ + {"content-type": "application/json", "x-amzn-requestid": "c5188fb5-dc58-4cbe-948d-af173c69ce0d"}, + { + "generations": [ + { + "finish_reason": "MAX_TOKENS", + "id": "0730f5c0-9a49-4f35-af94-cf8f77327740", + "text": " To convert 212 degrees Fahrenheit to Celsius, we can use the conversion factor that Celsius is equal to (Fahrenheit - 32) x 5/9. \\n\\nApplying this formula, we have:\\n212°F = (212°F - 32) x 5/9\\n= (180) x 5/9\\n= 100°C.\\n\\nTherefore, 212 degrees F", + } + ], + "id": "a9cc8ce6-50b6-40b6-bf77-cf24561d8de7", + "prompt": "What is 212 degrees Fahrenheit converted to Celsius?", + }, + ], "ai21.j2-mid-v1::What is 212 degrees Fahrenheit converted to Celsius?": [ {"content-type": "application/json", "x-amzn-requestid": "3bf1bb6b-b6f0-4901-85a1-2fa0e814440e"}, { diff --git a/tests/mlmodel_bedrock/_test_chat_completion.py b/tests/mlmodel_bedrock/_test_chat_completion.py index 9ccbe5b64..17fa8549c 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}', + "cohere.command-text-v14": '{"prompt": "%s", "temperature": %f, "max_tokens": %d}', } chat_completion_expected_events = { @@ -132,4 +133,67 @@ }, ), ], + "cohere.command-text-v14": [ + ( + {"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": "c5188fb5-dc58-4cbe-948d-af173c69ce0d", + "response_id": None, # UUID that varies with each run + "api_key_last_four_digits": "CRET", + "duration": None, # Response time varies each test run + "request.model": "cohere.command-text-v14", + "response.model": "cohere.command-text-v14", + "request.temperature": 0.7, + "request.max_tokens": 100, + "response.choices.finish_reason": "MAX_TOKENS", + "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": "c5188fb5-dc58-4cbe-948d-af173c69ce0d", + "span_id": "span-id", + "trace_id": "trace-id", + "transaction_id": None, + "content": "What is 212 degrees Fahrenheit converted to Celsius?", + "role": "user", + "completion_id": None, + "sequence": 0, + "response.model": "cohere.command-text-v14", + "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": "c5188fb5-dc58-4cbe-948d-af173c69ce0d", + "span_id": "span-id", + "trace_id": "trace-id", + "transaction_id": None, + "content": " To convert 212 degrees Fahrenheit to Celsius, we can use the conversion factor that Celsius is equal to (Fahrenheit - 32) x 5/9. \\n\\nApplying this formula, we have:\\n212°F = (212°F - 32) x 5/9\\n= (180) x 5/9\\n= 100°C.\\n\\nTherefore, 212 degrees F", + "role": "assistant", + "completion_id": None, + "sequence": 1, + "response.model": "cohere.command-text-v14", + "vendor": "bedrock", + "ingest_source": "Python", + }, + ), + ], } diff --git a/tests/mlmodel_bedrock/conftest.py b/tests/mlmodel_bedrock/conftest.py index 67ecfc46b..b0aa70c57 100644 --- a/tests/mlmodel_bedrock/conftest.py +++ b/tests/mlmodel_bedrock/conftest.py @@ -60,7 +60,7 @@ def bedrock_server(): if not _environ_as_bool("NEW_RELIC_TESTING_RECORD_BEDROCK_RESPONSES", False): # Use mocked Bedrock backend and prerecorded responses with MockExternalBedrockServer() as server: - client = boto3.client( + client = boto3.client( # nosec "bedrock-runtime", "us-east-1", endpoint_url="http://localhost:%d" % server.port, diff --git a/tests/mlmodel_bedrock/test_chat_completion.py b/tests/mlmodel_bedrock/test_chat_completion.py index 574776299..50f851f92 100644 --- a/tests/mlmodel_bedrock/test_chat_completion.py +++ b/tests/mlmodel_bedrock/test_chat_completion.py @@ -53,7 +53,7 @@ def is_file_payload(request): "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}}'), - # ("cohere.command-text-v14", '{"prompt": "{prompt}", "max_tokens": {max_tokens:d}, "temperature": {temperature:f}}'), + "cohere.command-text-v14", ], ) def model_id(request):