Add Async Bedrock Instrumentation (#1307)

* Add async bedrock instrumentation.

* Async bedrock tests

* Refactor and add safeguards.

* Remove unused imports.

* Add review feedback.

* Update attr checks.

---------

Co-authored-by: Tim Pansino <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

newrelic/hooks/external_aiobotocore.py

@@ -11,9 +11,35 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
+import traceback
+import sys
+from aiobotocore.response import StreamingBody
+from io import BytesIO
 
 from newrelic.api.external_trace import ExternalTrace
 from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.hooks.external_botocore import (
+    handle_bedrock_exception,
+    run_bedrock_response_extractor,
+    run_bedrock_request_extractor,
+    RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+# Class from https://github.com/aio-libs/aiobotocore/blob/master/tests/test_response.py
+# aiobotocore Apache 2 license: https://github.com/aio-libs/aiobotocore/blob/master/LICENSE
+class AsyncBytesIO(BytesIO):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.content = self
+
+    async def read(self, amt=-1):
+        if amt == -1:  # aiohttp to regular response
+            amt = None
+        return super().read(amt)
 
 
 def _bind_make_request_params(operation_model, request_dict, *args, **kwargs):
@@ -44,5 +70,67 @@ async def wrap_endpoint_make_request(wrapped, instance, args, kwargs):
         return result
 
 
+async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
+    # This instrumentation only applies to bedrock runtimes so exit if this method was hit through a different path
+    if not hasattr(instance, "_nr_is_bedrock"):
+        return await wrapped(*args, **kwargs)
+
+    transaction = instance._nr_txn
+    if not transaction:
+        return await wrapped(*args, **kwargs)
+
+    # Grab all context data from botocore invoke_model instrumentation off the shared instance
+    trace_id = getattr(instance, "_nr_trace_id", "")
+    span_id = getattr(instance, "_nr_span_id", "")
+
+    request_extractor = getattr(instance, "_nr_request_extractor", None)
+    response_extractor = getattr(instance, "_nr_response_extractor", None)
+    ft = getattr(instance, "_nr_ft", None)
+
+    model = args[1].get("modelId")
+    is_embedding = "embed" in model
+    request_body = args[1].get("body")
+
+    try:
+        response = await wrapped(*args, **kwargs)
+    except Exception as exc:
+        handle_bedrock_exception(
+            exc, is_embedding, model, span_id, trace_id, request_extractor, request_body, ft, transaction
+        )
+
+    if not response:
+        return response
+
+    response_headers = response.get("ResponseMetadata", {}).get("HTTPHeaders") or {}
+    bedrock_attrs = {
+        "request_id": response_headers.get("x-amzn-requestid"),
+        "model": model,
+        "span_id": span_id,
+        "trace_id": trace_id,
+    }
+
+    run_bedrock_request_extractor(request_extractor, request_body, bedrock_attrs)
+
+    try:
+        # Read and replace response streaming bodies
+        response_body = await response["body"].read()
+        if ft:
+            ft.__exit__(None, None, None)
+            bedrock_attrs["duration"] = ft.duration * 1000
+        response["body"] = StreamingBody(AsyncBytesIO(response_body), len(response_body))
+
+        run_bedrock_response_extractor(response_extractor, response_body, bedrock_attrs, is_embedding, transaction)
+
+    except Exception:
+        _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
+
+    return response
+
+
 def instrument_aiobotocore_endpoint(module):
     wrap_function_wrapper(module, "AioEndpoint.make_request", wrap_endpoint_make_request)
+
+
+def instrument_aiobotocore_client(module):
+    if hasattr(module, "AioBaseClient"):
+        wrap_function_wrapper(module, "AioBaseClient._make_api_call", wrap_client__make_api_call)

newrelic/hooks/external_botocore.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import inspect
 import logging
 import re
 import sys
@@ -546,12 +547,77 @@ def extract_bedrock_cohere_model_streaming_response(response_body, bedrock_attrs
 ]
 
 
+def handle_bedrock_exception(exc, is_embedding, model, span_id, trace_id, request_extractor, request_body, ft, transaction):
+    try:
+        bedrock_attrs = {
+            "model": model,
+            "span_id": span_id,
+            "trace_id": trace_id,
+        }
+        try:
+            request_extractor(request_body, bedrock_attrs)
+        except json.decoder.JSONDecodeError:
+            pass
+        except Exception:
+            _logger.warning(REQUEST_EXTACTOR_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
+
+        error_attributes = bedrock_error_attributes(exc, bedrock_attrs)
+        notice_error_attributes = {
+            "http.statusCode": error_attributes.get("http.statusCode"),
+            "error.message": error_attributes.get("error.message"),
+            "error.code": error_attributes.get("error.code"),
+        }
+
+        if is_embedding:
+            notice_error_attributes.update({"embedding_id": str(uuid.uuid4())})
+        else:
+            notice_error_attributes.update({"completion_id": str(uuid.uuid4())})
+
+        if ft:
+            ft.notice_error(
+                attributes=notice_error_attributes,
+            )
+
+            ft.__exit__(*sys.exc_info())
+            error_attributes["duration"] = ft.duration * 1000
+
+        if is_embedding:
+            handle_embedding_event(transaction, error_attributes)
+        else:
+            handle_chat_completion_event(transaction, error_attributes)
+    except Exception:
+        _logger.warning(EXCEPTION_HANDLING_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
+
+    raise
+
+
+def run_bedrock_response_extractor(response_extractor, response_body, bedrock_attrs, is_embedding, transaction):
+    # Run response extractor for non-streaming responses
+    try:
+        response_extractor(response_body, bedrock_attrs)
+    except Exception:
+        _logger.warning(RESPONSE_EXTRACTOR_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
+
+    if is_embedding:
+        handle_embedding_event(transaction, bedrock_attrs)
+    else:
+        handle_chat_completion_event(transaction, bedrock_attrs)
+
+
+def run_bedrock_request_extractor(request_extractor, request_body, bedrock_attrs):
+    try:
+        request_extractor(request_body, bedrock_attrs)
+    except json.decoder.JSONDecodeError:
+        pass
+    except Exception:
+        _logger.warning(REQUEST_EXTACTOR_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
+
+
 def wrap_bedrock_runtime_invoke_model(response_streaming=False):
     @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)
 
@@ -604,54 +670,32 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
         span_id = available_metadata.get("span.id")
         trace_id = available_metadata.get("trace.id")
 
+        # Store data on instance to pass context to async instrumentation
+        instance._nr_trace_id = trace_id
+        instance._nr_span_id = span_id
+        instance._nr_request_extractor = request_extractor
+        instance._nr_response_extractor = response_extractor
+        instance._nr_stream_extractor = stream_extractor
+        instance._nr_txn = transaction
+        instance._nr_ft = ft
+
+        # Add a bedrock flag to instance so we can determine when make_api_call instrumentation is hit from non-Bedrock paths and bypass it if so
+        instance._nr_is_bedrock = True
+
         try:
+            # For aioboto3 clients, this will call make_api_call instrumentation in external_aiobotocore
             response = wrapped(*args, **kwargs)
         except Exception as exc:
-            try:
-                bedrock_attrs = {
-                    "model": model,
-                    "span_id": span_id,
-                    "trace_id": trace_id,
-                }
-                try:
-                    request_extractor(request_body, bedrock_attrs)
-                except json.decoder.JSONDecodeError:
-                    pass
-                except Exception:
-                    _logger.warning(REQUEST_EXTACTOR_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
-
-                error_attributes = bedrock_error_attributes(exc, bedrock_attrs)
-                notice_error_attributes = {
-                    "http.statusCode": error_attributes.get("http.statusCode"),
-                    "error.message": error_attributes.get("error.message"),
-                    "error.code": error_attributes.get("error.code"),
-                }
-
-                if is_embedding:
-                    notice_error_attributes.update({"embedding_id": str(uuid.uuid4())})
-                else:
-                    notice_error_attributes.update({"completion_id": str(uuid.uuid4())})
-
-                ft.notice_error(
-                    attributes=notice_error_attributes,
-                )
-
-                ft.__exit__(*sys.exc_info())
-                error_attributes["duration"] = ft.duration * 1000
-
-                if operation == "embedding":
-                    handle_embedding_event(transaction, error_attributes)
-                else:
-                    handle_chat_completion_event(transaction, error_attributes)
-            except Exception:
-                _logger.warning(EXCEPTION_HANDLING_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
-
-            raise
+            handle_bedrock_exception(exc, is_embedding, model, span_id, trace_id, request_extractor, request_body, ft, transaction)
 
         if not response or response_streaming and not settings.ai_monitoring.streaming.enabled:
             ft.__exit__(None, None, None)
             return response
 
+        # Let the instrumentation of make_api_call in the aioboto3 client handle it if we have an async case
+        if inspect.iscoroutine(response):
+            return response
+
         if response_streaming and operation == "embedding":
             # This combination is not supported at time of writing, but may become
             # a supported feature in the future. Instrumentation will need to be written
@@ -668,12 +712,7 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
             "trace_id": trace_id,
         }
 
-        try:
-            request_extractor(request_body, bedrock_attrs)
-        except json.decoder.JSONDecodeError:
-            pass
-        except Exception:
-            _logger.warning(REQUEST_EXTACTOR_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
+        run_bedrock_request_extractor(request_extractor, request_body, bedrock_attrs)
 
         try:
             if response_streaming:
@@ -691,16 +730,7 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
             bedrock_attrs["duration"] = ft.duration * 1000
             response["body"] = StreamingBody(BytesIO(response_body), len(response_body))
 
-            # Run response extractor for non-streaming responses
-            try:
-                response_extractor(response_body, bedrock_attrs)
-            except Exception:
-                _logger.warning(RESPONSE_EXTRACTOR_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
-
-            if operation == "embedding":
-                handle_embedding_event(transaction, bedrock_attrs)
-            else:
-                handle_chat_completion_event(transaction, bedrock_attrs)
+            run_bedrock_response_extractor(response_extractor, response_body, bedrock_attrs, is_embedding, transaction)
 
         except Exception:
             _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE % traceback.format_exception(*sys.exc_info()))
@@ -864,7 +894,6 @@ def handle_chat_completion_event(transaction, bedrock_attrs):
     llm_context_attrs = getattr(transaction, "_llm_context_attrs", None)
     if llm_context_attrs:
         llm_metadata_dict.update(llm_context_attrs)
-
     span_id = bedrock_attrs.get("span_id", None)
     trace_id = bedrock_attrs.get("trace_id", None)
     request_id = bedrock_attrs.get("request_id", None)
@@ -1009,9 +1038,9 @@ def _nr_dynamodb_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
                     partition = "aws-us-gov"
 
             if partition and region and account_id and _target:
-                agent_attrs["cloud.resource_id"] = (
-                    f"arn:{partition}:dynamodb:{region}:{account_id:012d}:table/{_target}"
-                )
+                agent_attrs[
+                    "cloud.resource_id"
+                ] = f"arn:{partition}:dynamodb:{region}:{account_id:012d}:table/{_target}"
                 agent_attrs["db.system"] = "DynamoDB"
 
         except Exception as e:

tests/external_aiobotocore/conftest.py

@@ -17,6 +17,7 @@
 import socket
 import threading
 
+import pytest
 import moto.server
 import werkzeug.serving
 from testing_support.fixture.event_loop import (  # noqa: F401, pylint: disable=W0611
@@ -27,6 +28,14 @@
     collector_available_fixture,
 )
 
+from newrelic.common.package_version_utils import (
+    get_package_version,
+    get_package_version_tuple,
+)
+from external_botocore._mock_external_bedrock_server import MockExternalBedrockServer
+
+BOTOCORE_VERSION = get_package_version("botocore")
+
 PORT = 4443
 AWS_ACCESS_KEY_ID = "AAAAAAAAAAAACCESSKEY"
 AWS_SECRET_ACCESS_KEY = "AAAAAASECRETKEY"  # nosec
@@ -40,6 +49,8 @@
     "transaction_tracer.stack_trace_threshold": 0.0,
     "debug.log_data_collector_payloads": True,
     "debug.record_transaction_failure": True,
+    "custom_insights_events.max_attribute_value": 4096,
+    "ai_monitoring.enabled": True,
 }
 collector_agent_registration = collector_agent_registration_fixture(
     app_name="Python Agent Test (external_aiobotocore)",
@@ -146,3 +157,37 @@ async def _stop(self):
             self._server.shutdown()
 
         self._thread.join()
+
+
+# Bedrock Fixtures
+@pytest.fixture(scope="session")
+def bedrock_server(loop):
+    """
+    This fixture will create a mocked backend for testing purposes.
+    """
+    import aiobotocore
+
+    from newrelic.core.config import _environ_as_bool
+
+    if get_package_version_tuple("botocore") < (1, 31, 57):
+        pytest.skip(reason="Bedrock Runtime not available.")
+
+    if _environ_as_bool("NEW_RELIC_TESTING_RECORD_BEDROCK_RESPONSES", False):
+        raise NotImplementedError("To record test responses, use botocore instead.")
+
+    # Use mocked Bedrock backend and prerecorded responses
+    with MockExternalBedrockServer() as server:
+        session = aiobotocore.session.get_session()
+        client = loop.run_until_complete(
+            session.create_client(
+                "bedrock-runtime",
+                "us-east-1",
+                endpoint_url=f"http://localhost:{server.port}",
+                aws_access_key_id="NOT-A-REAL-SECRET",
+                aws_secret_access_key="NOT-A-REAL-SECRET",
+            ).__aenter__()
+        )
+
+        yield client
+
+        loop.run_until_complete(client.__aexit__(None, None, None))