From 148512c77929943c3182aaaf011708d54a1e4762 Mon Sep 17 00:00:00 2001
From: Lu Peng <bolu.peng@oracle.com>
Date: Thu, 7 Nov 2024 14:46:42 -0500
Subject: [PATCH 1/2] Improved aqua telemetry.

---
 ads/aqua/common/decorator.py                    | 6 ++++++
 ads/aqua/extension/aqua_ws_msg_handler.py       | 2 ++
 ads/aqua/extension/base_handler.py              | 2 ++
 tests/unitary/with_extras/aqua/test_handlers.py | 7 +++++++
 4 files changed, 17 insertions(+)

diff --git a/ads/aqua/common/decorator.py b/ads/aqua/common/decorator.py
index b697afff2..7d251df0d 100644
--- a/ads/aqua/common/decorator.py
+++ b/ads/aqua/common/decorator.py
@@ -69,6 +69,12 @@ def inner_function(
                 reason=error.message,
                 service_payload=error.args[0] if error.args else None,
                 exc_info=sys.exc_info(),
+                aqua_api_details=dict(
+                    # __qualname__ gives information of class and name of api
+                    aqua_api_name=func.__qualname__,
+                    oci_api_name=error.operation_name,
+                    service_endpoint=error.request_endpoint
+                )
             )
         except (
             ClientError,
diff --git a/ads/aqua/extension/aqua_ws_msg_handler.py b/ads/aqua/extension/aqua_ws_msg_handler.py
index 1494ce028..04ff651f4 100644
--- a/ads/aqua/extension/aqua_ws_msg_handler.py
+++ b/ads/aqua/extension/aqua_ws_msg_handler.py
@@ -78,10 +78,12 @@ def write_error(self, status_code, **kwargs):
                 logger.warning(reply["message"])
         # telemetry may not be present if there is an error while initializing
         if hasattr(self, "telemetry"):
+            aqua_api_details = kwargs.get("aqua_api_details", {})
             self.telemetry.record_event_async(
                 category="aqua/error",
                 action=str(status_code),
                 value=reason,
+                **aqua_api_details
             )
         response = AquaWsError(
             status=status_code,
diff --git a/ads/aqua/extension/base_handler.py b/ads/aqua/extension/base_handler.py
index d84602bf7..5bd9f7091 100644
--- a/ads/aqua/extension/base_handler.py
+++ b/ads/aqua/extension/base_handler.py
@@ -98,10 +98,12 @@ def write_error(self, status_code, **kwargs):
 
         # telemetry may not be present if there is an error while initializing
         if hasattr(self, "telemetry"):
+            aqua_api_details = kwargs.get("aqua_api_details", {})
             self.telemetry.record_event_async(
                 category="aqua/error",
                 action=str(status_code),
                 value=reason,
+                **aqua_api_details
             )
 
         self.finish(json.dumps(reply))
diff --git a/tests/unitary/with_extras/aqua/test_handlers.py b/tests/unitary/with_extras/aqua/test_handlers.py
index 74b9853b4..a4ae749e9 100644
--- a/tests/unitary/with_extras/aqua/test_handlers.py
+++ b/tests/unitary/with_extras/aqua/test_handlers.py
@@ -129,6 +129,11 @@ def test_finish(self, name, payload, expected_call, mock_super_finish):
                         ),
                         None,
                     ),
+                    aqua_api_details=dict(
+                        aqua_api_name="TestDataset.create",
+                        oci_api_name=TestDataset.mock_service_payload_create["operation_name"],
+                        service_endpoint=TestDataset.mock_service_payload_create["request_endpoint"]
+                    )
                 ),
                 "Authorization Failed: The resource you're looking for isn't accessible. Operation Name: get_job_run.",
             ],
@@ -159,12 +164,14 @@ def test_write_error(self, name, input, expected_msg, mock_uuid, mock_logger):
             "request_id": "1234",
         }
         self.test_instance.finish.assert_called_once_with(json.dumps(expected_reply))
+        aqua_api_details = input.get("aqua_api_details", {})
         self.test_instance.telemetry.record_event_async.assert_called_with(
             category="aqua/error",
             action=str(
                 input.get("status_code"),
             ),
             value=input.get("reason"),
+            **aqua_api_details
         )
 
         mock_logger.warning.assert_called_with(expected_msg)

From 544a223f8c0c9c053d735d183c9d5e5a9762b150 Mon Sep 17 00:00:00 2001
From: Lu Peng <bolu.peng@oracle.com>
Date: Sat, 9 Nov 2024 11:46:10 -0500
Subject: [PATCH 2/2] Updated pr.

---
 ads/aqua/common/decorator.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/ads/aqua/common/decorator.py b/ads/aqua/common/decorator.py
index 7d251df0d..d03805ab2 100644
--- a/ads/aqua/common/decorator.py
+++ b/ads/aqua/common/decorator.py
@@ -72,8 +72,12 @@ def inner_function(
                 aqua_api_details=dict(
                     # __qualname__ gives information of class and name of api
                     aqua_api_name=func.__qualname__,
-                    oci_api_name=error.operation_name,
-                    service_endpoint=error.request_endpoint
+                    oci_api_name=getattr(
+                        error, "operation_name", "Unknown OCI Operation"
+                    ),
+                    service_endpoint=getattr(
+                        error, "request_endpoint", "Unknown Request Endpoint"
+                    )
                 )
             )
         except (