Fix botocore tests & re-structure

Author: hmstepanek

newrelic/hooks/external_botocore.py

@@ -12,25 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
-import logging
-import uuid
-from io import BytesIO
-
-from botocore.response import StreamingBody
-
 from newrelic.api.datastore_trace import datastore_trace
 from newrelic.api.external_trace import ExternalTrace
-from newrelic.api.function_trace import FunctionTrace
 from newrelic.api.message_trace import message_trace
-from newrelic.api.time_trace import get_trace_linking_metadata
-from newrelic.api.transaction import current_transaction
-from newrelic.common.object_names import callable_name
-from newrelic.common.object_wrapper import function_wrapper, wrap_function_wrapper
-from newrelic.core.config import global_settings
-
-_logger = logging.getLogger(__name__)
-UNSUPPORTED_MODEL_WARNING_SENT = False
+from newrelic.common.object_wrapper import wrap_function_wrapper
 
 
 def extract_sqs(*args, **kwargs):
@@ -55,353 +40,6 @@ def extractor_string(*args, **kwargs):
     return extractor_list
 
 
-def bedrock_error_attributes(exception, request_args, client, extractor):
-    response = getattr(exception, "response", None)
-    if not response:
-        return {}
-    
-    request_body = request_args.get("body", "")
-    error_attributes = extractor(request_body)[1]
-
-    error_attributes.update({
-        "request_id": response.get("ResponseMetadata", {}).get("RequestId", ""),
-        "api_key_last_four_digits": client._request_signer._credentials.access_key[-4:],
-        "request.model": request_args.get("modelId", ""),
-        "vendor": "Bedrock",
-        "ingest_source": "Python",
-        "http.statusCode": response.get("ResponseMetadata", "").get("HTTPStatusCode", ""),
-        "error.message": response.get("Error", "").get("Message", ""),
-        "error.code": response.get("Error", "").get("Code", ""),
-    })
-    return error_attributes
-
-
-def create_chat_completion_message_event(
-    transaction,
-    app_name,
-    message_list,
-    chat_completion_id,
-    span_id,
-    trace_id,
-    request_model,
-    request_id,
-    conversation_id,
-    response_id="",
-):
-    if not transaction:
-        return
-
-    for index, message in enumerate(message_list):
-        if response_id:
-            id_ = "%s-%d" % (response_id, index)  # Response ID was set, append message index to it.
-        else:
-            id_ = str(uuid.uuid4())  # No response IDs, use random UUID
-
-        chat_completion_message_dict = {
-            "id": id_,
-            "appName": app_name,
-            "conversation_id": conversation_id,
-            "request_id": request_id,
-            "span_id": span_id,
-            "trace_id": trace_id,
-            "transaction_id": transaction._transaction_id,
-            "content": message.get("content", ""),
-            "role": message.get("role"),
-            "completion_id": chat_completion_id,
-            "sequence": index,
-            "response.model": request_model,
-            "vendor": "bedrock",
-            "ingest_source": "Python",
-        }
-        transaction.record_ml_event("LlmChatCompletionMessage", chat_completion_message_dict)
-
-
-def extract_bedrock_titan_text_model(request_body, response_body=None):
-    request_body = json.loads(request_body)
-    if response_body:
-        response_body = json.loads(response_body)
-
-    request_config = request_body.get("textGenerationConfig", {})
-
-    chat_completion_summary_dict = {
-        "request.max_tokens": request_config.get("maxTokenCount", ""),
-        "request.temperature": request_config.get("temperature", ""),
-    }
-
-    if response_body:
-        input_tokens = response_body["inputTextTokenCount"]
-        completion_tokens = sum(result["tokenCount"] for result in response_body.get("results", []))
-        total_tokens = input_tokens + completion_tokens
-
-        message_list = [{"role": "user", "content": request_body.get("inputText", "")}]
-        message_list.extend(
-            {"role": "assistant", "content": result["outputText"]} for result in response_body.get("results", [])
-        )
-
-        chat_completion_summary_dict.update({
-            "response.choices.finish_reason": response_body["results"][0]["completionReason"],
-            "response.usage.completion_tokens": completion_tokens,
-            "response.usage.prompt_tokens": input_tokens,
-            "response.usage.total_tokens": total_tokens,
-            "response.number_of_messages": len(message_list),
-        })
-    else:
-        message_list = []
-
-    return message_list, chat_completion_summary_dict
-
-
-def extract_bedrock_titan_embedding_model(request_body, response_body=None):
-    if not response_body:
-        return [], {}  # No extracted information necessary for embedding
-
-    request_body = json.loads(request_body)
-    response_body = json.loads(response_body)
-
-    input_tokens = response_body.get("inputTextTokenCount", None)
-
-    embedding_dict = {
-        "input": request_body.get("inputText", ""),
-        "response.usage.prompt_tokens": input_tokens,
-        "response.usage.total_tokens": input_tokens,
-    }
-    return [], embedding_dict
-
-
-def extract_bedrock_ai21_j2_model(request_body, response_body=None):
-    request_body = json.loads(request_body)
-    if response_body:
-        response_body = json.loads(response_body)
-
-    chat_completion_summary_dict = {
-        "request.max_tokens": request_body.get("maxTokens", ""),
-        "request.temperature": request_body.get("temperature", ""),
-    }
-
-    if response_body:
-        message_list = [{"role": "user", "content": request_body.get("prompt", "")}]
-        message_list.extend(
-            {"role": "assistant", "content": result["data"]["text"]} for result in response_body.get("completions", [])
-        )
-
-        chat_completion_summary_dict.update({
-            "response.choices.finish_reason": response_body["completions"][0]["finishReason"]["reason"],
-            "response.number_of_messages": len(message_list),
-            "response_id": str(response_body.get("id", "")),
-        })
-    else:
-        message_list = []
-
-    return message_list, chat_completion_summary_dict
-
-
-def extract_bedrock_claude_model(request_body, response_body=None):
-    request_body = json.loads(request_body)
-    if response_body:
-        response_body = json.loads(response_body)
-
-    chat_completion_summary_dict = {
-        "request.max_tokens": request_body.get("max_tokens_to_sample", ""),
-        "request.temperature": request_body.get("temperature", ""),
-    }
-
-    if response_body:
-        message_list = [
-            {"role": "user", "content": request_body.get("prompt", "")},
-            {"role": "assistant", "content": response_body.get("completion", "")},
-        ]
-
-        chat_completion_summary_dict.update({
-            "response.choices.finish_reason": response_body.get("stop_reason", ""),
-            "response.number_of_messages": len(message_list),
-        })
-    else:
-        message_list = []
-
-    return message_list, chat_completion_summary_dict
-
-
-def extract_bedrock_cohere_model(request_body, response_body=None):
-    request_body = json.loads(request_body)
-    if response_body:
-        response_body = json.loads(response_body)
-
-    chat_completion_summary_dict = {
-        "request.max_tokens": request_body.get("max_tokens", ""),
-        "request.temperature": request_body.get("temperature", ""),
-    }
-
-    if response_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.update({
-            "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", "")),
-        })
-    else:
-        message_list = []
-
-    return message_list, chat_completion_summary_dict
-
-
-MODEL_EXTRACTORS = [  # Order is important here, avoiding dictionaries
-    ("amazon.titan-embed", extract_bedrock_titan_embedding_model),
-    ("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),
-]
-
-
-@function_wrapper
-def wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
-    # Wrapped function only takes keyword arguments, no need for binding
-
-    transaction = current_transaction()
-
-    if not transaction:
-        return wrapped(*args, **kwargs)
-
-    # Read and replace request file stream bodies
-    request_body = kwargs["body"]
-    if hasattr(request_body, "read"):
-        request_body = request_body.read()
-        kwargs["body"] = request_body
-
-    # Determine model to be used with extractor
-    model = kwargs.get("modelId")
-    if not model:
-        return wrapped(*args, **kwargs)
-
-    # Determine extractor by model type
-    for extractor_name, extractor in MODEL_EXTRACTORS:
-        if model.startswith(extractor_name):
-            break
-    else:
-        # Model was not found in extractor list
-        global UNSUPPORTED_MODEL_WARNING_SENT
-        if not UNSUPPORTED_MODEL_WARNING_SENT:
-            # Only send warning once to avoid spam
-            _logger.warning(
-                "Unsupported Amazon Bedrock model in use (%s). Upgrade to a newer version of the agent, and contact New Relic support if the issue persists.",
-                model,
-            )
-            UNSUPPORTED_MODEL_WARNING_SENT = True
-        
-        extractor = lambda *args: ([], {})  # Empty extractor that returns nothing
-
-    ft_name = callable_name(wrapped)
-    with FunctionTrace(ft_name) as ft:
-        try:
-            response = wrapped(*args, **kwargs)
-        except Exception as exc:
-            try:
-                error_attributes = extractor(request_body)
-                error_attributes = bedrock_error_attributes(exc, kwargs, instance, extractor)
-                ft.notice_error(
-                    attributes=error_attributes,
-                )
-            finally:
-                raise
-
-    if not response:
-        return response
-
-    # Read and replace response streaming bodies
-    response_body = response["body"].read()
-    response["body"] = StreamingBody(BytesIO(response_body), len(response_body))
-    response_headers = response["ResponseMetadata"]["HTTPHeaders"]
-
-    if model.startswith("amazon.titan-embed"):  # Only available embedding models
-        handle_embedding_event(instance, transaction, extractor, model, response_body, response_headers, request_body, ft.duration)
-    else:
-        handle_chat_completion_event(instance, transaction, extractor, model, response_body, response_headers, request_body, ft.duration)
-
-    return response
-
-def handle_embedding_event(client, transaction, extractor, model, response_body, response_headers, request_body, duration):
-    embedding_id = str(uuid.uuid4())
-    available_metadata = get_trace_linking_metadata()
-    span_id = available_metadata.get("span.id", "")
-    trace_id = available_metadata.get("trace.id", "")
-
-    request_id = response_headers.get("x-amzn-requestid", "")
-    settings = transaction.settings if transaction.settings is not None else global_settings()
-
-    _, embedding_dict = extractor(request_body, response_body)
-
-    embedding_dict.update({
-        "vendor": "bedrock",
-        "ingest_source": "Python",
-        "id": embedding_id,
-        "appName": settings.app_name,
-        "span_id": span_id,
-        "trace_id": trace_id,
-        "request_id": request_id,
-        "transaction_id": transaction._transaction_id,
-        "api_key_last_four_digits": client._request_signer._credentials.access_key[-4:],
-        "duration": duration,
-        "request.model": model,
-        "response.model": model,
-    })
-
-    transaction.record_ml_event("LlmEmbedding", embedding_dict)
-
-
-def handle_chat_completion_event(client, transaction, extractor, model, response_body, response_headers, request_body, duration):
-    custom_attrs_dict = transaction._custom_params
-    conversation_id = custom_attrs_dict.get("conversation_id", "")
-
-    chat_completion_id = str(uuid.uuid4())
-    available_metadata = get_trace_linking_metadata()
-    span_id = available_metadata.get("span.id", "")
-    trace_id = available_metadata.get("trace.id", "")
-
-    request_id = response_headers.get("x-amzn-requestid", "")
-    settings = transaction.settings if transaction.settings is not None else global_settings()
-
-    message_list, chat_completion_summary_dict = extractor(request_body, response_body)
-    response_id = chat_completion_summary_dict.get("response_id", "")
-    chat_completion_summary_dict.update(
-        {
-            "vendor": "bedrock",
-            "ingest_source": "Python",
-            "api_key_last_four_digits": client._request_signer._credentials.access_key[-4:],
-            "id": chat_completion_id,
-            "appName": settings.app_name,
-            "conversation_id": conversation_id,
-            "span_id": span_id,
-            "trace_id": trace_id,
-            "transaction_id": transaction._transaction_id,
-            "request_id": request_id,
-            "duration": duration,
-            "request.model": model,
-            "response.model": model,  # Duplicate data required by the UI
-        }
-    )
-
-    transaction.record_ml_event("LlmChatCompletionSummary", chat_completion_summary_dict)
-
-    create_chat_completion_message_event(
-        transaction=transaction,
-        app_name=settings.app_name,
-        message_list=message_list,
-        chat_completion_id=chat_completion_id,
-        span_id=span_id,
-        trace_id=trace_id,
-        request_model=model,
-        request_id=request_id,
-        conversation_id=conversation_id,
-        response_id=response_id,
-    )
-
-
 CUSTOM_TRACE_POINTS = {
     ("sns", "publish"): message_trace("SNS", "Produce", "Topic", extract(("TopicArn", "TargetArn"), "PhoneNumber")),
     ("dynamodb", "put_item"): datastore_trace("DynamoDB", extract("TableName"), "put_item"),
@@ -415,7 +53,6 @@ def handle_chat_completion_event(client, transaction, extractor, model, response
     ("sqs", "send_message"): message_trace("SQS", "Produce", "Queue", extract_sqs),
     ("sqs", "send_message_batch"): message_trace("SQS", "Produce", "Queue", extract_sqs),
     ("sqs", "receive_message"): message_trace("SQS", "Consume", "Queue", extract_sqs),
-    ("bedrock-runtime", "invoke_model"): wrap_bedrock_runtime_invoke_model,
 }
 
 
@@ -447,6 +84,7 @@ def _nr_endpoint_make_request_(wrapped, instance, args, kwargs):
     method = request_dict.get("method", None)
 
     with ExternalTrace(library="botocore", url=url, method=method, source=wrapped) as trace:
+
         try:
             trace._add_agent_attribute("aws.operation", operation_model.name)
         except: