OpenAI and Bedrock Preview

newrelic/agent.py

@@ -153,6 +153,10 @@ def __asgi_application(*args, **kwargs):
 from newrelic.api.message_transaction import (
     wrap_message_transaction as __wrap_message_transaction,
 )
+from newrelic.api.ml_model import get_llm_message_ids as __get_llm_message_ids
+from newrelic.api.ml_model import (
+    record_llm_feedback_event as __record_llm_feedback_event,
+)
 from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel
 from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper
 from newrelic.api.profile_trace import profile_trace as __profile_trace
@@ -340,3 +344,5 @@ def __asgi_application(*args, **kwargs):
 insert_html_snippet = __wrap_api_call(__insert_html_snippet, "insert_html_snippet")
 verify_body_exists = __wrap_api_call(__verify_body_exists, "verify_body_exists")
 wrap_mlmodel = __wrap_api_call(__wrap_mlmodel, "wrap_mlmodel")
+get_llm_message_ids = __wrap_api_call(__get_llm_message_ids, "get_llm_message_ids")
+record_llm_feedback_event = __wrap_api_call(__record_llm_feedback_event, "record_llm_feedback_event")

newrelic/api/ml_model.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 
 import sys
+import uuid
+import warnings
 
+from newrelic.api.transaction import current_transaction
 from newrelic.common.object_names import callable_name
 from newrelic.hooks.mlmodel_sklearn import _nr_instrument_model
 
@@ -33,3 +36,49 @@ def wrap_mlmodel(model, name=None, version=None, feature_names=None, label_names
         model._nr_wrapped_label_names = label_names
     if metadata:
         model._nr_wrapped_metadata = metadata
+
+
+def get_llm_message_ids(response_id=None):
+    transaction = current_transaction()
+    if response_id and transaction:
+        nr_message_ids = getattr(transaction, "_nr_message_ids", {})
+        message_id_info = nr_message_ids.pop(response_id, ())
+
+        if not message_id_info:
+            warnings.warn("No message ids found for %s" % response_id)
+            return []
+
+        conversation_id, request_id, ids = message_id_info
+
+        return [{"conversation_id": conversation_id, "request_id": request_id, "message_id": _id} for _id in ids]
+    warnings.warn("No message ids found. get_llm_message_ids must be called within the scope of a transaction.")
+    return []
+
+
+def record_llm_feedback_event(
+    message_id, rating, conversation_id=None, request_id=None, category=None, message=None, metadata=None
+):
+    transaction = current_transaction()
+    if not transaction:
+        warnings.warn(
+            "No message feedback events will be recorded. record_llm_feedback_event must be called within the "
+            "scope of a transaction."
+        )
+        return
+
+    feedback_message_id = str(uuid.uuid4())
+    metadata = metadata or {}
+
+    feedback_message_event = {
+        "id": feedback_message_id,
+        "message_id": message_id,
+        "rating": rating,
+        "conversation_id": conversation_id or "",
+        "request_id": request_id or "",
+        "category": category or "",
+        "message": message or "",
+        "ingest_source": "Python",
+    }
+    feedback_message_event.update(metadata)
+
+    transaction.record_ml_event("LlmFeedbackMessage", feedback_message_event)

newrelic/api/time_trace.py

@@ -29,7 +29,6 @@
 )
 from newrelic.core.config import is_expected_error, should_ignore_error
 from newrelic.core.trace_cache import trace_cache
-
 from newrelic.packages import six
 
 _logger = logging.getLogger(__name__)
@@ -260,6 +259,11 @@ def _observe_exception(self, exc_info=None, ignore=None, expected=None, status_c
         module, name, fullnames, message_raw = parse_exc_info((exc, value, tb))
         fullname = fullnames[0]
 
+        # In case message is in JSON format for OpenAI models
+        # this will result in a "cleaner" message format
+        if getattr(value, "_nr_message", None):
+            message_raw = value._nr_message
+
         # Check to see if we need to strip the message before recording it.
 
         if settings.strip_exception_messages.enabled and fullname not in settings.strip_exception_messages.allowlist:
@@ -422,23 +426,32 @@ def notice_error(self, error=None, attributes=None, expected=None, ignore=None,
                         input_attributes = {}
                         input_attributes.update(transaction._custom_params)
                         input_attributes.update(attributes)
-                        error_group_name_raw = settings.error_collector.error_group_callback(value, {
-                            "traceback": tb,
-                            "error.class": exc,
-                            "error.message": message_raw,
-                            "error.expected": is_expected,
-                            "custom_params": input_attributes,
-                            "transactionName": getattr(transaction, "name", None),
-                            "response.status": getattr(transaction, "_response_code", None),
-                            "request.method": getattr(transaction, "_request_method", None),
-                            "request.uri": getattr(transaction, "_request_uri", None),
-                        })
+                        error_group_name_raw = settings.error_collector.error_group_callback(
+                            value,
+                            {
+                                "traceback": tb,
+                                "error.class": exc,
+                                "error.message": message_raw,
+                                "error.expected": is_expected,
+                                "custom_params": input_attributes,
+                                "transactionName": getattr(transaction, "name", None),
+                                "response.status": getattr(transaction, "_response_code", None),
+                                "request.method": getattr(transaction, "_request_method", None),
+                                "request.uri": getattr(transaction, "_request_uri", None),
+                            },
+                        )
                         if error_group_name_raw:
                             _, error_group_name = process_user_attribute("error.group.name", error_group_name_raw)
                             if error_group_name is None or not isinstance(error_group_name, six.string_types):
-                                raise ValueError("Invalid attribute value for error.group.name. Expected string, got: %s" % repr(error_group_name_raw))
+                                raise ValueError(
+                                    "Invalid attribute value for error.group.name. Expected string, got: %s"
+                                    % repr(error_group_name_raw)
+                                )
                     except Exception:
-                        _logger.error("Encountered error when calling error group callback:\n%s", "".join(traceback.format_exception(*sys.exc_info())))
+                        _logger.error(
+                            "Encountered error when calling error group callback:\n%s",
+                            "".join(traceback.format_exception(*sys.exc_info())),
+                        )
                         error_group_name = None
 
             transaction._create_error_node(
@@ -595,13 +608,11 @@ def update_async_exclusive_time(self, min_child_start_time, exclusive_duration):
     def process_child(self, node, is_async):
         self.children.append(node)
         if is_async:
-
             # record the lowest start time
             self.min_child_start_time = min(self.min_child_start_time, node.start_time)
 
             # if there are no children running, finalize exclusive time
             if self.child_count == len(self.children):
-
                 exclusive_duration = node.end_time - self.min_child_start_time
 
                 self.update_async_exclusive_time(self.min_child_start_time, exclusive_duration)

newrelic/api/transaction.py

@@ -191,6 +191,7 @@ def __init__(self, application, enabled=None, source=None):
         self._frameworks = set()
         self._message_brokers = set()
         self._dispatchers = set()
+        self._ml_models = set()
 
         self._frozen_path = None
 
@@ -559,6 +560,10 @@ def __exit__(self, exc, value, tb):
             for dispatcher, version in self._dispatchers:
                 self.record_custom_metric("Python/Dispatcher/%s/%s" % (dispatcher, version), 1)
 
+        if self._ml_models:
+            for ml_model, version in self._ml_models:
+                self.record_custom_metric("Python/ML/%s/%s" % (ml_model, version), 1)
+
         if self._settings.distributed_tracing.enabled:
             # Sampled and priority need to be computed at the end of the
             # transaction when distributed tracing or span events are enabled.
@@ -1648,7 +1653,7 @@ def record_ml_event(self, event_type, params):
         if not settings.ml_insights_events.enabled:
             return
 
-        event = create_custom_event(event_type, params)
+        event = create_custom_event(event_type, params, is_ml_event=True)
         if event:
             self._ml_events.add(event, priority=self.priority)
 
@@ -1755,6 +1760,10 @@ def add_dispatcher_info(self, name, version=None):
         if name:
             self._dispatchers.add((name, version))
 
+    def add_ml_model_info(self, name, version=None):
+        if name:
+            self._ml_models.add((name, version))
+
     def dump(self, file):
         """Dumps details about the transaction to the file object."""
 

newrelic/config.py

@@ -2037,6 +2037,21 @@ def _process_trace_cache_import_hooks():
 
 
 def _process_module_builtin_defaults():
+    _process_module_definition(
+        "openai.api_resources.embedding",
+        "newrelic.hooks.mlmodel_openai",
+        "instrument_openai_api_resources_embedding",
+    )
+    _process_module_definition(
+        "openai.api_resources.chat_completion",
+        "newrelic.hooks.mlmodel_openai",
+        "instrument_openai_api_resources_chat_completion",
+    )
+    _process_module_definition(
+        "openai.util",
+        "newrelic.hooks.mlmodel_openai",
+        "instrument_openai_util",
+    )
     _process_module_definition(
         "asyncio.base_events",
         "newrelic.hooks.coroutines_asyncio",

newrelic/core/application.py

@@ -932,7 +932,7 @@ def record_ml_event(self, event_type, params):
         if settings is None or not settings.ml_insights_events.enabled:
             return
 
-        event = create_custom_event(event_type, params)
+        event = create_custom_event(event_type, params, is_ml_event=True)
 
         if event:
             with self._stats_custom_lock:

newrelic/core/attribute.py

@@ -89,6 +89,8 @@
 
 MAX_NUM_USER_ATTRIBUTES = 128
 MAX_ATTRIBUTE_LENGTH = 255
+MAX_NUM_ML_USER_ATTRIBUTES = 64
+MAX_ML_ATTRIBUTE_LENGTH = 4095
 MAX_64_BIT_INT = 2**63 - 1
 MAX_LOG_MESSAGE_LENGTH = 32768
 

newrelic/core/custom_event.py

@@ -18,7 +18,7 @@
 
 from newrelic.core.attribute import (check_name_is_string, check_name_length,
         process_user_attribute, NameIsNotStringException, NameTooLongException,
-        MAX_NUM_USER_ATTRIBUTES)
+        MAX_NUM_USER_ATTRIBUTES, MAX_ML_ATTRIBUTE_LENGTH, MAX_NUM_ML_USER_ATTRIBUTES, MAX_ATTRIBUTE_LENGTH)
 
 
 _logger = logging.getLogger(__name__)
@@ -72,7 +72,8 @@ def process_event_type(name):
     else:
         return name
 
-def create_custom_event(event_type, params):
+
+def create_custom_event(event_type, params, is_ml_event=False):
     """Creates a valid custom event.
 
     Ensures that the custom event has a valid name, and also checks
@@ -83,6 +84,8 @@ def create_custom_event(event_type, params):
     Args:
         event_type (str): The type (name) of the custom event.
         params (dict): Attributes to add to the event.
+        is_ml_event (bool): Boolean indicating whether create_custom_event was called from
+        record_ml_event for truncation purposes
 
     Returns:
         Custom event (list of 2 dicts), if successful.
@@ -99,12 +102,18 @@ def create_custom_event(event_type, params):
 
     try:
         for k, v in params.items():
-            key, value = process_user_attribute(k, v)
+            if is_ml_event:
+                max_length = MAX_ML_ATTRIBUTE_LENGTH
+                max_num_attrs = MAX_NUM_ML_USER_ATTRIBUTES
+            else:
+                max_length = MAX_ATTRIBUTE_LENGTH
+                max_num_attrs = MAX_NUM_USER_ATTRIBUTES
+            key, value = process_user_attribute(k, v, max_length=max_length)
             if key:
-                if len(attributes) >= MAX_NUM_USER_ATTRIBUTES:
+                if len(attributes) >= max_num_attrs:
                     _logger.debug('Maximum number of attributes already '
-                            'added to event %r. Dropping attribute: %r=%r',
-                            name, key, value)
+                                  'added to event %r. Dropping attribute: %r=%r',
+                                  name, key, value)
                 else:
                     attributes[key] = value
     except Exception:

newrelic/core/otlp_utils.py

@@ -21,6 +21,7 @@
 
 import logging
 
+from newrelic.api.time_trace import get_service_linking_metadata
 from newrelic.common.encoding_utils import json_encode
 from newrelic.core.config import global_settings
 from newrelic.core.stats_engine import CountStats, TimeStats
@@ -124,8 +125,11 @@ def create_key_values_from_iterable(iterable):
     )
 
 
-def create_resource(attributes=None):
+def create_resource(attributes=None, attach_apm_entity=True):
     attributes = attributes or {"instrumentation.provider": "newrelic-opentelemetry-python-ml"}
+    if attach_apm_entity:
+        metadata = get_service_linking_metadata()
+        attributes.update(metadata)
     return Resource(attributes=create_key_values_from_iterable(attributes))
 
 
@@ -203,7 +207,7 @@ def stats_to_otlp_metrics(metric_data, start_time, end_time):
 
 
 def encode_metric_data(metric_data, start_time, end_time, resource=None, scope=None):
-    resource = resource or create_resource()
+    resource = resource or create_resource(attach_apm_entity=False)
     return MetricsData(
         resource_metrics=[
             ResourceMetrics(
@@ -220,24 +224,45 @@ def encode_metric_data(metric_data, start_time, end_time, resource=None, scope=N
 
 
 def encode_ml_event_data(custom_event_data, agent_run_id):
-    resource = create_resource()
-    ml_events = []
+    # An InferenceEvent is attached to a separate ML Model entity instead
+    # of the APM entity.
+    ml_inference_events = []
+    ml_apm_events = []
     for event in custom_event_data:
         event_info, event_attrs = event
+        event_type = event_info["type"]
         event_attrs.update(
             {
                 "real_agent_id": agent_run_id,
                 "event.domain": "newrelic.ml_events",
-                "event.name": event_info["type"],
+                "event.name": event_type,
             }
         )
         ml_attrs = create_key_values_from_iterable(event_attrs)
         unix_nano_timestamp = event_info["timestamp"] * 1e6
-        ml_events.append(
-            {
-                "time_unix_nano": int(unix_nano_timestamp),
-                "attributes": ml_attrs,
-            }
+        if event_type == "InferenceEvent":
+            ml_inference_events.append(
+                {
+                    "time_unix_nano": int(unix_nano_timestamp),
+                    "attributes": ml_attrs,
+                }
+            )
+        else:
+            ml_apm_events.append(
+                {
+                    "time_unix_nano": int(unix_nano_timestamp),
+                    "attributes": ml_attrs,
+                }
+            )
+
+    resource_logs = []
+    if ml_inference_events:
+        inference_resource = create_resource(attach_apm_entity=False)
+        resource_logs.append(
+            ResourceLogs(resource=inference_resource, scope_logs=[ScopeLogs(log_records=ml_inference_events)])
         )
+    if ml_apm_events:
+        apm_resource = create_resource()
+        resource_logs.append(ResourceLogs(resource=apm_resource, scope_logs=[ScopeLogs(log_records=ml_apm_events)]))
 
-    return LogsData(resource_logs=[ResourceLogs(resource=resource, scope_logs=[ScopeLogs(log_records=ml_events)])])
+    return LogsData(resource_logs=resource_logs)

newrelic/core/stats_engine.py

@@ -724,6 +724,11 @@ def notice_error(self, error=None, attributes=None, expected=None, ignore=None,
         module, name, fullnames, message_raw = parse_exc_info(error)
         fullname = fullnames[0]
 
+        # In the case case of JSON formatting for OpenAI models
+        # this will result in a "cleaner" message format
+        if getattr(value, "_nr_message", None):
+            message_raw = value._nr_message
+
         # Check to see if we need to strip the message before recording it.
 
         if settings.strip_exception_messages.enabled and fullname not in settings.strip_exception_messages.allowlist: