Merge branch 'mark-instrumentation-pts-for-sdk' of github.com:newreli…

…c/newrelic-python-agent into mark-instrumentation-pts-for-sdk
newrelic · Dec 16, 2023 · 97f8e72 · 97f8e72
2 parents 13151ff + 1cf881d
commit 97f8e72
Show file tree

Hide file tree

Showing 4 changed files with 353 additions and 24 deletions.
diff --git a/newrelic/hooks/mlmodel_openai.py b/newrelic/hooks/mlmodel_openai.py
@@ -59,19 +59,37 @@ def wrap_embedding_sync(wrapped, instance, args, kwargs):
         try:
             response = wrapped(*args, **kwargs)
         except Exception as exc:
-            notice_error_attributes = {
-                "http.statusCode": getattr(exc, "http_status", ""),
-                "error.message": getattr(exc, "_message", ""),
-                "error.code": getattr(getattr(exc, "error", ""), "code", ""),
-                "error.param": getattr(exc, "param", ""),
-                "embedding_id": embedding_id,
-            }
-            exc._nr_message = notice_error_attributes.pop("error.message")
+            if OPENAI_V1:
+                response = getattr(exc, "response", "")
+                response_headers = getattr(response, "headers", "")
+                exc_organization = response_headers.get("openai-organization", "") if response_headers else ""
+                # There appears to be a bug here in openai v1 where despite having code,
+                # param, etc in the error response, they are not populated on the exception
+                # object so grab them from the response body object instead.
+                body = getattr(exc, "body", {}) or {}
+                notice_error_attributes = {
+                    "http.statusCode": getattr(exc, "status_code", "") or "",
+                    "error.message": body.get("message", "") or "",
+                    "error.code": body.get("code", "") or "",
+                    "error.param": body.get("param", "") or "",
+                    "embedding_id": embedding_id,
+                }
+            else:
+                exc_organization = getattr(exc, "organization", "")
+                notice_error_attributes = {
+                    "http.statusCode": getattr(exc, "http_status", ""),
+                    "error.message": getattr(exc, "_message", ""),
+                    "error.code": getattr(getattr(exc, "error", ""), "code", ""),
+                    "error.param": getattr(exc, "param", ""),
+                    "embedding_id": embedding_id,
+                }
+            message = notice_error_attributes.pop("error.message")
+            if message:
+                exc._nr_message = message
             ft.notice_error(
                 attributes=notice_error_attributes,
             )
-            # Gather attributes to add to embedding summary event in error context
-            exc_organization = getattr(exc, "organization", "")
+
             error_embedding_dict = {
                 "id": embedding_id,
                 "appName": settings.app_name,
@@ -498,19 +516,37 @@ async def wrap_embedding_async(wrapped, instance, args, kwargs):
         try:
             response = await wrapped(*args, **kwargs)
         except Exception as exc:
-            notice_error_attributes = {
-                "http.statusCode": getattr(exc, "http_status", ""),
-                "error.message": getattr(exc, "_message", ""),
-                "error.code": getattr(getattr(exc, "error", ""), "code", ""),
-                "error.param": getattr(exc, "param", ""),
-                "embedding_id": embedding_id,
-            }
-            exc._nr_message = notice_error_attributes.pop("error.message")
+            if OPENAI_V1:
+                response = getattr(exc, "response", "")
+                response_headers = getattr(response, "headers", "")
+                exc_organization = response_headers.get("openai-organization", "") if response_headers else ""
+                # There appears to be a bug here in openai v1 where despite having code,
+                # param, etc in the error response, they are not populated on the exception
+                # object so grab them from the response body object instead.
+                body = getattr(exc, "body", {}) or {}
+                notice_error_attributes = {
+                    "http.statusCode": getattr(exc, "status_code", "") or "",
+                    "error.message": body.get("message", "") or "",
+                    "error.code": body.get("code", "") or "",
+                    "error.param": body.get("param", "") or "",
+                    "embedding_id": embedding_id,
+                }
+            else:
+                exc_organization = getattr(exc, "organization", "")
+                notice_error_attributes = {
+                    "http.statusCode": getattr(exc, "http_status", ""),
+                    "error.message": getattr(exc, "_message", ""),
+                    "error.code": getattr(getattr(exc, "error", ""), "code", ""),
+                    "error.param": getattr(exc, "param", ""),
+                    "embedding_id": embedding_id,
+                }
+            message = notice_error_attributes.pop("error.message")
+            if message:
+                exc._nr_message = message
             ft.notice_error(
                 attributes=notice_error_attributes,
             )
-            # Gather attributes to add to embedding summary event in error context
-            exc_organization = getattr(exc, "organization", "")
+
             error_embedding_dict = {
                 "id": embedding_id,
                 "appName": settings.app_name,

diff --git a/tests/mlmodel_openai/conftest.py b/tests/mlmodel_openai/conftest.py
@@ -65,7 +65,6 @@
         "test_chat_completion.py",
         "test_get_llm_message_ids.py",
         "test_chat_completion_error.py",
-        "test_embeddings_error_v1.py",
     ]