Add support for async Bedrock streaming. (#1315)

* Add async bedrock instrumentation.

* Async bedrock tests

* Refactor and add safeguards.

* Remove newline.

* Formatting fixes.

---------

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

@@ -1,3 +1,4 @@
+
 # Copyright 2010 New Relic, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,9 +21,11 @@
 from newrelic.api.external_trace import ExternalTrace
 from newrelic.common.object_wrapper import wrap_function_wrapper
 from newrelic.hooks.external_botocore import (
+    AsyncEventStreamWrapper,
     handle_bedrock_exception,
     run_bedrock_response_extractor,
     run_bedrock_request_extractor,
+    EMBEDDING_STREAMING_UNSUPPORTED_LOG_MESSAGE,
     RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE,
 )
 
@@ -75,21 +78,36 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
     if not hasattr(instance, "_nr_is_bedrock"):
         return await wrapped(*args, **kwargs)
 
-    transaction = instance._nr_txn
+    transaction = getattr(instance, "_nr_txn", None)
     if not transaction:
         return await wrapped(*args, **kwargs)
 
+    settings = getattr(instance, "_nr_settings", None)
+
+    # Early exit if we can't access the shared settings object from invoke_model instrumentation
+    # This settings object helps us determine if AIM was enabled as well as streaming
+    if not (settings and settings.ai_monitoring.enabled):
+        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)
+    stream_extractor = getattr(instance, "_nr_stream_extractor", None)
+    response_streaming = getattr(instance, "_nr_response_streaming", False)
+
     ft = getattr(instance, "_nr_ft", None)
 
-    model = args[1].get("modelId")
-    is_embedding = "embed" in model
-    request_body = args[1].get("body")
+    if len(args) >= 2:
+        model = args[1].get("modelId")
+        request_body = args[1].get("body")
+        is_embedding = "embed" in model
+    else:
+        model = ""
+        request_body = None
+        is_embedding = False
 
     try:
         response = await wrapped(*args, **kwargs)
@@ -98,7 +116,18 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
             exc, is_embedding, model, span_id, trace_id, request_extractor, request_body, ft, transaction
         )
 
-    if not response:
+    if not response or response_streaming and not settings.ai_monitoring.streaming.enabled:
+        if ft:
+            ft.__exit__(None, None, None)
+        return response
+
+    if response_streaming and is_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
+        # if this becomes available.
+        _logger.warning(EMBEDDING_STREAMING_UNSUPPORTED_LOG_MESSAGE)
+        if ft:
+            ft.__exit__(None, None, None)
         return response
 
     response_headers = response.get("ResponseMetadata", {}).get("HTTPHeaders") or {}
@@ -112,13 +141,22 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
     run_bedrock_request_extractor(request_extractor, request_body, bedrock_attrs)
 
     try:
+        if response_streaming:
+            # Wrap EventStream object here to intercept __iter__ method instead of instrumenting class.
+            # This class is used in numerous other services in botocore, and would cause conflicts.
+            response["body"] = body = AsyncEventStreamWrapper(response["body"])
+            body._nr_ft = ft or None
+            body._nr_bedrock_attrs = bedrock_attrs or {}
+            body._nr_model_extractor = stream_extractor or None
+            return response
+
         # 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:

newrelic/hooks/external_botocore.py

@@ -141,9 +141,9 @@ def extract_firehose_agent_attrs(instance, *args, **kwargs):
                 region = instance._client_config.region_name
             if account_id and region:
                 agent_attrs["cloud.platform"] = "aws_kinesis_delivery_streams"
-                agent_attrs["cloud.resource_id"] = (
-                    f"arn:aws:firehose:{region}:{account_id}:deliverystream/{stream_name}"
-                )
+                agent_attrs[
+                    "cloud.resource_id"
+                ] = f"arn:aws:firehose:{region}:{account_id}:deliverystream/{stream_name}"
     except Exception as e:
         _logger.debug("Failed to capture AWS Kinesis Delivery Stream (Firehose) info.", exc_info=True)
     return agent_attrs
@@ -547,7 +547,9 @@ 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):
+def handle_bedrock_exception(
+    exc, is_embedding, model, span_id, trace_id, request_extractor, request_body, ft, transaction
+):
     try:
         bedrock_attrs = {
             "model": model,
@@ -678,6 +680,8 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
         instance._nr_stream_extractor = stream_extractor
         instance._nr_txn = transaction
         instance._nr_ft = ft
+        instance._nr_response_streaming = response_streaming
+        instance._nr_settings = settings
 
         # 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
@@ -686,7 +690,9 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
             # For aioboto3 clients, this will call make_api_call instrumentation in external_aiobotocore
             response = wrapped(*args, **kwargs)
         except Exception as exc:
-            handle_bedrock_exception(exc, is_embedding, model, span_id, trace_id, request_extractor, request_body, ft, transaction)
+            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)
@@ -777,6 +783,42 @@ def close(self):
         return super(GeneratorProxy, self).close()
 
 
+class AsyncEventStreamWrapper(ObjectProxy):
+    def __aiter__(self):
+        g = AsyncGeneratorProxy(self.__wrapped__.__aiter__())
+        g._nr_ft = getattr(self, "_nr_ft", None)
+        g._nr_bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
+        g._nr_model_extractor = getattr(self, "_nr_model_extractor", NULL_EXTRACTOR)
+        return g
+
+
+class AsyncGeneratorProxy(ObjectProxy):
+    def __init__(self, wrapped):
+        super(AsyncGeneratorProxy, self).__init__(wrapped)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        transaction = current_transaction()
+        if not transaction:
+            return await self.__wrapped__.__anext__()
+        return_val = None
+        try:
+            return_val = await self.__wrapped__.__anext__()
+            record_stream_chunk(self, return_val, transaction)
+        except StopAsyncIteration as e:
+            record_events_on_stop_iteration(self, transaction)
+            raise
+        except Exception as exc:
+            record_error(self, transaction, exc)
+            raise
+        return return_val
+
+    async def aclose(self):
+        return await super(AsyncGeneratorProxy, self).aclose()
+
+
 def record_stream_chunk(self, return_val, transaction):
     if return_val:
         try: