Add time to first token for Bedrock and OpenAI. (#1689)

* Add time to first token for Bedrock and OpenAI.

* Prevent resetting of time_to_first_token in bedrock

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Tim Pansino <timpansino@gmail.com>
Co-authored-by: Timothy Pansino <11214426+TimPansino@users.noreply.github.com>

Author: 3 people

newrelic/hooks/external_botocore.py

@@ -963,6 +963,8 @@ def record_events_on_stop_iteration(self, transaction, request_timestamp=None):
 
             try:
                 bedrock_attrs["duration"] = self._nr_ft.duration * 1000
+                if hasattr(self, "_nr_time_to_first_token"):
+                    bedrock_attrs["time_to_first_token"] = self._nr_time_to_first_token
                 handle_chat_completion_event(transaction, bedrock_attrs, request_timestamp)
             except Exception:
                 _logger.warning(RESPONSE_PROCESSING_FAILURE_LOG_MESSAGE, exc_info=True)
@@ -1012,8 +1014,23 @@ def record_stream_chunk(self, event, transaction, request_timestamp=None):
 
     def invoke_record_stream_chunk(self, event, transaction, request_timestamp=None):
         bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
+
+        # Store time to first token now, but only attach it if we successfully parse the chunk.
+        # Record the time here since parsing may be a bit slow, and can inflate the metric.
+        time_to_first_token = (
+            int(1000.0 * time.time()) - self._nr_request_timestamp
+            if not hasattr(self, "_nr_time_to_first_token")
+            else None
+        )
+
+        # Load and parse the chunk to extract model specific data
         chunk = json.loads(event["chunk"]["bytes"].decode("utf-8"))
         self._nr_model_extractor(chunk, bedrock_attrs)
+
+        # Attach time to first token after parsing to ensure there are no errors
+        if time_to_first_token is not None:
+            self._nr_time_to_first_token = time_to_first_token
+
         # 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")
@@ -1027,10 +1044,14 @@ def converse_record_stream_chunk(self, event, transaction):
                 return
 
             content = ((event.get("contentBlockDelta") or {}).get("delta") or {}).get("text", "")
+
             if "output_message_list" not in bedrock_attrs:
                 bedrock_attrs["output_message_list"] = [{"role": "assistant", "content": ""}]
             bedrock_attrs["output_message_list"][0]["content"] += content
 
+            if content and not hasattr(self, "_nr_time_to_first_token"):
+                self._nr_time_to_first_token = int(1000.0 * time.time()) - self._nr_request_timestamp
+
         if "messageStop" in event:
             bedrock_attrs["response.choices.finish_reason"] = (event.get("messageStop") or {}).get("stopReason", "")
 
@@ -1198,6 +1219,7 @@ def handle_chat_completion_event(transaction, bedrock_attrs, request_timestamp=N
         "response.choices.finish_reason": bedrock_attrs.get("response.choices.finish_reason", None),
         "error": bedrock_attrs.get("error", None),
         "timestamp": request_timestamp or None,
+        "time_to_first_token": bedrock_attrs.get("time_to_first_token", 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}

newrelic/hooks/mlmodel_openai.py

@@ -571,6 +571,10 @@ def _record_completion_success(
             "timestamp": request_timestamp,
         }
         llm_metadata = _get_llm_attributes(transaction)
+
+        if "time_to_first_token" in kwargs:
+            full_chat_completion_summary_dict["time_to_first_token"] = kwargs.get("time_to_first_token")
+
         full_chat_completion_summary_dict.update(llm_metadata)
         transaction.record_custom_event("LlmChatCompletionSummary", full_chat_completion_summary_dict)
 
@@ -782,6 +786,10 @@ def _record_stream_chunk(self, return_val):
             if choices:
                 delta = choices[0].get("delta") or {}
                 if delta:
+                    if delta.get("content") and "time_to_first_token" not in self._nr_openai_attrs:
+                        self._nr_openai_attrs["time_to_first_token"] = (
+                            int(1000.0 * time.time()) - self._nr_request_timestamp
+                        )
                     self._nr_openai_attrs["content"] = self._nr_openai_attrs.get("content", "") + (
                         delta.get("content") or ""
                     )

tests/external_botocore/_test_bedrock_chat_completion_converse.py

@@ -116,6 +116,7 @@
             "vendor": "bedrock",
             "ingest_source": "Python",
             "response.number_of_messages": 3,
+            "time_to_first_token": None,  # Varies each test run
         },
     ),
     (

tests/external_botocore/_test_bedrock_chat_completion_invoke_model.py

@@ -1073,6 +1073,7 @@
                 "vendor": "bedrock",
                 "ingest_source": "Python",
                 "response.number_of_messages": 2,
+                "time_to_first_token": None,  # Varies each test run
             },
         ),
         (
@@ -1134,6 +1135,7 @@
                 "vendor": "bedrock",
                 "ingest_source": "Python",
                 "response.number_of_messages": 2,
+                "time_to_first_token": None,  # Varies each test run
             },
         ),
         (
@@ -1195,6 +1197,7 @@
                 "vendor": "bedrock",
                 "ingest_source": "Python",
                 "response.number_of_messages": 2,
+                "time_to_first_token": None,  # Varies each test run
             },
         ),
         (
@@ -1255,6 +1258,7 @@
                 "vendor": "bedrock",
                 "ingest_source": "Python",
                 "response.number_of_messages": 2,
+                "time_to_first_token": None,  # Varies each test run
             },
         ),
         (
@@ -1317,6 +1321,7 @@
                 "vendor": "bedrock",
                 "ingest_source": "Python",
                 "response.number_of_messages": 2,
+                "time_to_first_token": None,  # Varies each test run
             },
         ),
         (
@@ -1378,6 +1383,7 @@
                 "vendor": "bedrock",
                 "ingest_source": "Python",
                 "response.number_of_messages": 2,
+                "time_to_first_token": None,  # Varies each test run
             },
         ),
         (

tests/mlmodel_openai/test_chat_completion_stream.py

@@ -68,6 +68,7 @@
             "vendor": "openai",
             "ingest_source": "Python",
             "response.number_of_messages": 3,
+            "time_to_first_token": None,  # Varies each test run
         },
     ),
     (

tests/mlmodel_openai/test_chat_completion_stream_v1.py

@@ -78,6 +78,7 @@
             "vendor": "openai",
             "ingest_source": "Python",
             "response.number_of_messages": 3,
+            "time_to_first_token": None,  # Varies each test run
         },
     ),
     (