LangChain: Fix message timestamps, add default role assignment, and Bedrock support (#1580)

* Record the request message as the time the request started for LangChain.

* Tracking the original timestamp of the request for input messages that are
recorded as LlmChatCompletionMessage event types.

* First pass at preserving LlmChatCompletionMessage timestamp for the request
with Bedrock methods.

* the `kwargs` was being mapped directly to the OpenAI client
and having timestamp in there caused a problem. As a quick test,
only add the request timestamp after the wrapped function has been
invoked.

* Moved the request timestamp to its own variable instead of part of kwargs.

* OpenAI async request messages were not being assigned the correct timestamp.

* Trying to improve the passing of the request timestamp through for Bedrock.

* Passing too many parameters.

* Set a default role on input/output messages within LangChain.

* [MegaLinter] Apply linters fixes

* Fix request_timestamp for LlmChatCompletionSummary table

* Fix request_timestamp for LlmChatCompletionSummary table

* [MegaLinter] Apply linters fixes

* Bedrock Converse Streaming Support (#1565)

* Add more formatting to custom event validatators

* Add streamed responses to converse mock server

* Add streaming fixtures for testing for converse

* Rename other bedrock test files

* Add tests for converse streaming

* Instrument converse streaming

* Move GeneratorProxy adjacent functions to mixin

* Fix checking of supported models

* Reorganize converse error tests

* Port new converse botocore tests to aiobotocore

* Instrument response streaming in aiobotocore converse

* Fix suggestions from code review

* Port in converse changes from strands PR

* Delete commented code

---------

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

* Bedrock Converse Streaming Support (#1565)

* Add more formatting to custom event validatators

* Add streamed responses to converse mock server

* Add streaming fixtures for testing for converse

* Rename other bedrock test files

* Add tests for converse streaming

* Instrument converse streaming

* Move GeneratorProxy adjacent functions to mixin

* Fix checking of supported models

* Reorganize converse error tests

* Port new converse botocore tests to aiobotocore

* Instrument response streaming in aiobotocore converse

* Fix suggestions from code review

* Port in converse changes from strands PR

* Delete commented code

---------

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

* [MegaLinter] Apply linters fixes

* request_timestamp is now passed across different method

* Fixed gemini model kwargs issue

* [MegaLinter] Apply linters fixes

* Update tests to validate presence of timestamp/ role and fix bugs in instrumentation.

* Update aiobotocore instrumentation to receive request timestamp.

---------

Co-authored-by: Josh Bonczkowski <jbonczkowski@newrelic.com>
Co-authored-by: sgoel-nr <236423107+sgoel-nr@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Timothy Pansino <11214426+TimPansino@users.noreply.github.com>
Co-authored-by: Uma Annamalai <uannamalai@newrelic.com>

Author: 5 people

newrelic/core/custom_event.py

@@ -141,7 +141,7 @@ def create_custom_event(event_type, params, settings=None, is_ml_event=False):
         )
         return None
 
-    intrinsics = {"type": name, "timestamp": int(1000.0 * time.time())}
+    intrinsics = {"type": name, "timestamp": params.get("timestamp") or int(1000.0 * time.time())}
 
     event = [intrinsics, attributes]
     return event

newrelic/hooks/external_aiobotocore.py

@@ -98,6 +98,7 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
     response_extractor = getattr(instance, "_nr_response_extractor", None)
     stream_extractor = getattr(instance, "_nr_stream_extractor", None)
     response_streaming = getattr(instance, "_nr_response_streaming", False)
+    request_timestamp = getattr(instance, "_nr_request_timestamp", None)
     is_converse = getattr(instance, "_nr_is_converse", False)
     ft = getattr(instance, "_nr_ft", None)
 
@@ -125,6 +126,7 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
             transaction,
             bedrock_args,
             is_converse,
+            request_timestamp,
         )
         raise
 
@@ -187,7 +189,9 @@ async def wrap_client__make_api_call(wrapped, instance, args, kwargs):
         if ft:
             ft.__exit__(None, None, None)
             bedrock_attrs["duration"] = ft.duration * 1000
-        run_bedrock_response_extractor(response_extractor, response_body, bedrock_attrs, is_embedding, transaction)
+        run_bedrock_response_extractor(
+            response_extractor, response_body, bedrock_attrs, is_embedding, transaction, request_timestamp
+        )
 
     except Exception:
         _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE, exc_info=True)

newrelic/hooks/external_botocore.py

@@ -17,6 +17,7 @@
 import logging
 import re
 import sys
+import time
 import uuid
 from io import BytesIO
 
@@ -193,6 +194,7 @@ def create_chat_completion_message_event(
     request_id,
     llm_metadata_dict,
     response_id=None,
+    request_timestamp=None,
 ):
     if not transaction:
         return
@@ -227,6 +229,8 @@ def create_chat_completion_message_event(
 
         if settings.ai_monitoring.record_content.enabled:
             chat_completion_message_dict["content"] = content
+        if request_timestamp:
+            chat_completion_message_dict["timestamp"] = request_timestamp
 
         chat_completion_message_dict.update(llm_metadata_dict)
 
@@ -266,6 +270,8 @@ def create_chat_completion_message_event(
 
         if settings.ai_monitoring.record_content.enabled:
             chat_completion_message_dict["content"] = content
+        if request_timestamp:
+            chat_completion_message_dict["timestamp"] = request_timestamp
 
         chat_completion_message_dict.update(llm_metadata_dict)
 
@@ -542,10 +548,22 @@ 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, kwargs, is_converse
+    exc,
+    is_embedding,
+    model,
+    span_id,
+    trace_id,
+    request_extractor,
+    request_body,
+    ft,
+    transaction,
+    kwargs,
+    is_converse,
+    request_timestamp=None,
 ):
     try:
         bedrock_attrs = {"model": model, "span_id": span_id, "trace_id": trace_id}
+
         if is_converse:
             try:
                 input_message_list = [
@@ -589,12 +607,14 @@ def handle_bedrock_exception(
         if is_embedding:
             handle_embedding_event(transaction, error_attributes)
         else:
-            handle_chat_completion_event(transaction, error_attributes)
+            handle_chat_completion_event(transaction, error_attributes, request_timestamp)
     except Exception:
         _logger.warning(EXCEPTION_HANDLING_FAILURE_LOG_MESSAGE, exc_info=True)
 
 
-def run_bedrock_response_extractor(response_extractor, response_body, bedrock_attrs, is_embedding, transaction):
+def run_bedrock_response_extractor(
+    response_extractor, response_body, bedrock_attrs, is_embedding, transaction, request_timestamp=None
+):
     # Run response extractor for non-streaming responses
     try:
         response_extractor(response_body, bedrock_attrs)
@@ -604,7 +624,7 @@ def run_bedrock_response_extractor(response_extractor, response_body, bedrock_at
     if is_embedding:
         handle_embedding_event(transaction, bedrock_attrs)
     else:
-        handle_chat_completion_event(transaction, bedrock_attrs)
+        handle_chat_completion_event(transaction, bedrock_attrs, request_timestamp)
 
 
 def run_bedrock_request_extractor(request_extractor, request_body, bedrock_attrs):
@@ -628,6 +648,8 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
         if not settings.ai_monitoring.enabled:
             return wrapped(*args, **kwargs)
 
+        request_timestamp = int(1000.0 * time.time())
+
         transaction.add_ml_model_info("Bedrock", BOTOCORE_VERSION)
         transaction._add_agent_attribute("llm", True)
 
@@ -683,6 +705,7 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
         instance._nr_ft = ft
         instance._nr_response_streaming = response_streaming
         instance._nr_settings = settings
+        instance._nr_request_timestamp = request_timestamp
 
         # 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
@@ -703,6 +726,7 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
                 transaction,
                 kwargs,
                 is_converse=False,
+                request_timestamp=request_timestamp,
             )
             raise
 
@@ -733,6 +757,8 @@ def _wrap_bedrock_runtime_invoke_model(wrapped, instance, args, kwargs):
         run_bedrock_request_extractor(request_extractor, request_body, bedrock_attrs)
 
         try:
+            bedrock_attrs.pop("timestamp", None)  # The request timestamp is only needed for request extraction
+
             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.
@@ -748,7 +774,14 @@ 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_bedrock_response_extractor(response_extractor, response_body, bedrock_attrs, is_embedding, transaction)
+            run_bedrock_response_extractor(
+                response_extractor,
+                response_body,
+                bedrock_attrs,
+                is_embedding,
+                transaction,
+                request_timestamp=request_timestamp,
+            )
 
         except Exception:
             _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE, exc_info=True)
@@ -770,6 +803,8 @@ def _wrap_bedrock_runtime_converse(wrapped, instance, args, kwargs):
         if not settings.ai_monitoring.enabled:
             return wrapped(*args, **kwargs)
 
+        request_timestamp = int(1000.0 * time.time())
+
         transaction.add_ml_model_info("Bedrock", BOTOCORE_VERSION)
         transaction._add_agent_attribute("llm", True)
 
@@ -800,6 +835,7 @@ def _wrap_bedrock_runtime_converse(wrapped, instance, args, kwargs):
         instance._nr_ft = ft
         instance._nr_response_streaming = response_streaming
         instance._nr_settings = settings
+        instance._nr_request_timestamp = request_timestamp
         instance._nr_is_converse = True
 
         # 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
@@ -810,7 +846,18 @@ def _wrap_bedrock_runtime_converse(wrapped, instance, args, kwargs):
             response = wrapped(*args, **kwargs)
         except Exception as exc:
             handle_bedrock_exception(
-                exc, False, model, span_id, trace_id, request_extractor, {}, ft, transaction, kwargs, is_converse=True
+                exc,
+                False,
+                model,
+                span_id,
+                trace_id,
+                request_extractor,
+                {},
+                ft,
+                transaction,
+                kwargs,
+                is_converse=True,
+                request_timestamp=request_timestamp,
             )
             raise
 
@@ -824,6 +871,7 @@ def _wrap_bedrock_runtime_converse(wrapped, instance, args, kwargs):
 
         response_headers = response.get("ResponseMetadata", {}).get("HTTPHeaders") or {}
         bedrock_attrs = extract_bedrock_converse_attrs(kwargs, response, response_headers, model, span_id, trace_id)
+        bedrock_attrs["timestamp"] = request_timestamp
 
         try:
             if response_streaming:
@@ -838,7 +886,9 @@ def _wrap_bedrock_runtime_converse(wrapped, instance, args, kwargs):
 
             ft.__exit__(None, None, None)
             bedrock_attrs["duration"] = ft.duration * 1000
-            run_bedrock_response_extractor(response_extractor, {}, bedrock_attrs, False, transaction)
+            run_bedrock_response_extractor(
+                response_extractor, {}, bedrock_attrs, False, transaction, request_timestamp=request_timestamp
+            )
 
         except Exception:
             _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE, exc_info=True)
@@ -888,7 +938,7 @@ def extract_bedrock_converse_attrs(kwargs, response, response_headers, model, sp
 
 
 class BedrockRecordEventMixin:
-    def record_events_on_stop_iteration(self, transaction):
+    def record_events_on_stop_iteration(self, transaction, request_timestamp=None):
         if hasattr(self, "_nr_ft"):
             bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
             self._nr_ft.__exit__(None, None, None)
@@ -899,14 +949,14 @@ def record_events_on_stop_iteration(self, transaction):
 
             try:
                 bedrock_attrs["duration"] = self._nr_ft.duration * 1000
-                handle_chat_completion_event(transaction, bedrock_attrs)
+                handle_chat_completion_event(transaction, bedrock_attrs, request_timestamp)
             except Exception:
                 _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE, exc_info=True)
 
             # Clear cached data as this can be very large.
             self._nr_bedrock_attrs.clear()
 
-    def record_error(self, transaction, exc):
+    def record_error(self, transaction, exc, request_timestamp=None):
         if hasattr(self, "_nr_ft"):
             try:
                 ft = self._nr_ft
@@ -929,32 +979,32 @@ def record_error(self, transaction, exc):
                 ft.__exit__(*sys.exc_info())
                 error_attributes["duration"] = ft.duration * 1000
 
-                handle_chat_completion_event(transaction, error_attributes)
+                handle_chat_completion_event(transaction, error_attributes, request_timestamp)
 
                 # Clear cached data as this can be very large.
                 error_attributes.clear()
             except Exception:
                 _logger.warning(EXCEPTION_HANDLING_FAILURE_LOG_MESSAGE, exc_info=True)
 
-    def record_stream_chunk(self, event, transaction):
+    def record_stream_chunk(self, event, transaction, request_timestamp=None):
         if event:
             try:
                 if getattr(self, "_nr_is_converse", False):
                     return self.converse_record_stream_chunk(event, transaction)
                 else:
-                    return self.invoke_record_stream_chunk(event, transaction)
+                    return self.invoke_record_stream_chunk(event, transaction, request_timestamp)
             except Exception:
                 _logger.warning(RESPONSE_EXTRACTOR_FAILURE_LOG_MESSAGE, exc_info=True)
 
-    def invoke_record_stream_chunk(self, event, transaction):
+    def invoke_record_stream_chunk(self, event, transaction, request_timestamp=None):
         bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
         chunk = json.loads(event["chunk"]["bytes"].decode("utf-8"))
         self._nr_model_extractor(chunk, bedrock_attrs)
         # In Langchain, the bedrock iterator exits early if type is "content_block_stop".
         # So we need to call the record events here since stop iteration will not be raised.
         _type = chunk.get("type")
         if _type == "content_block_stop":
-            self.record_events_on_stop_iteration(transaction)
+            self.record_events_on_stop_iteration(transaction, request_timestamp)
 
     def converse_record_stream_chunk(self, event, transaction):
         bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
@@ -984,6 +1034,7 @@ def __iter__(self):
 class GeneratorProxy(BedrockRecordEventMixin, ObjectProxy):
     def __init__(self, wrapped):
         super().__init__(wrapped)
+        self._nr_request_timestamp = int(1000.0 * time.time())
 
     def __iter__(self):
         return self
@@ -996,12 +1047,12 @@ def __next__(self):
         return_val = None
         try:
             return_val = self.__wrapped__.__next__()
-            self.record_stream_chunk(return_val, transaction)
+            self.record_stream_chunk(return_val, transaction, self._nr_request_timestamp)
         except StopIteration:
-            self.record_events_on_stop_iteration(transaction)
+            self.record_events_on_stop_iteration(transaction, self._nr_request_timestamp)
             raise
         except Exception as exc:
-            self.record_error(transaction, exc)
+            self.record_error(transaction, exc, self._nr_request_timestamp)
             raise
         return return_val
 
@@ -1020,6 +1071,10 @@ def __aiter__(self):
 
 
 class AsyncGeneratorProxy(BedrockRecordEventMixin, ObjectProxy):
+    def __init__(self, wrapped):
+        super().__init__(wrapped)
+        self._nr_request_timestamp = int(1000.0 * time.time())
+
     def __aiter__(self):
         return self
 
@@ -1030,12 +1085,12 @@ async def __anext__(self):
         return_val = None
         try:
             return_val = await self.__wrapped__.__anext__()
-            self.record_stream_chunk(return_val, transaction)
+            self.record_stream_chunk(return_val, transaction, self._nr_request_timestamp)
         except StopAsyncIteration:
-            self.record_events_on_stop_iteration(transaction)
+            self.record_events_on_stop_iteration(transaction, self._nr_request_timestamp)
             raise
         except Exception as exc:
-            self.record_error(transaction, exc)
+            self.record_error(transaction, exc, self._nr_request_timestamp)
             raise
         return return_val
 
@@ -1084,7 +1139,7 @@ def handle_embedding_event(transaction, bedrock_attrs):
     transaction.record_custom_event("LlmEmbedding", embedding_dict)
 
 
-def handle_chat_completion_event(transaction, bedrock_attrs):
+def handle_chat_completion_event(transaction, bedrock_attrs, request_timestamp=None):
     chat_completion_id = str(uuid.uuid4())
     # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events
     custom_attrs_dict = transaction._custom_params
@@ -1128,6 +1183,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs):
         "response.number_of_messages": number_of_messages,
         "response.choices.finish_reason": bedrock_attrs.get("response.choices.finish_reason", None),
         "error": bedrock_attrs.get("error", None),
+        "timestamp": request_timestamp or None,
     }
     chat_completion_summary_dict.update(llm_metadata_dict)
     chat_completion_summary_dict = {k: v for k, v in chat_completion_summary_dict.items() if v is not None}
@@ -1144,6 +1200,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs):
         request_id=request_id,
         llm_metadata_dict=llm_metadata_dict,
         response_id=response_id,
+        request_timestamp=request_timestamp,
     )