Instrument chat endpoint for nim. (#1489)

Author: nickolai-dr

custom_model_runner/datarobot_drum/drum/common.py

@@ -182,3 +182,34 @@ def otel_context(tracer, span_name, carrier):
             yield span
     finally:
         context.detach(token)
+
+
+def extract_chat_request_attributes(completion_params):
+    """Extracts otel related attributes from chat request payload
+
+    Used to populate span with relevant monitoring attriubtes.
+    """
+    attributes = {}
+    attributes["gen_ai.request.model"] = completion_params.get("model")
+    for i, m in enumerate(completion_params.get("messages", [])):
+        attributes[f"gen_ai.prompt.{i}.role"] = m.get("role")
+        attributes[f"gen_ai.prompt.{i}.content"] = m.get("content")
+        # last promt wins
+        attributes["gen_ai.prompt"] = m.get("content")
+    return attributes
+
+
+def extract_chat_response_attributes(response):
+    """Extracts otel related attributes from chat response.
+
+    Used to populate span with relevant monitoring attriubtes.
+    """
+    attributes = {}
+    attributes["gen_ai.response.model"] = response.get("model")
+    for i, c in enumerate(response.get("choices", [])):
+        m = c.get("message", {})
+        attributes[f"gen_ai.completion.{i}.role"] = m.get("role")
+        attributes[f"gen_ai.completion.{i}.content"] = m.get("content")
+        # last completion wins
+        attributes["gen_ai.completion"] = m.get("content")
+    return attributes

custom_model_runner/datarobot_drum/drum/root_predictors/prediction_server.py

@@ -39,7 +39,12 @@
     get_flask_app,
 )
 from datarobot_drum.profiler.stats_collector import StatsCollector, StatsOperation
-from datarobot_drum.drum.common import otel_context
+from datarobot_drum.drum.common import (
+    otel_context,
+    extract_chat_request_attributes,
+    extract_chat_response_attributes,
+)
+from opentelemetry.trace.status import StatusCode
 
 logger = logging.getLogger(LOGGER_NAME_PREFIX + "." + __name__)
 
@@ -200,13 +205,17 @@ def predict_unstructured():
         @model_api.route("/v1/chat/completions", methods=["POST"])
         def chat():
             logger.debug("Entering chat endpoint")
-            with otel_context(tracer, "drum.chat.completions", request.headers):
+            with otel_context(tracer, "drum.chat.completions", request.headers) as span:
+                span.set_attributes(extract_chat_request_attributes(request.json))
                 self._pre_predict_and_transform()
                 try:
                     response, response_status = self.do_chat(logger=logger)
                 finally:
                     self._post_predict_and_transform()
 
+                if isinstance(response, dict) and response_status == 200:
+                    span.set_attributes(extract_chat_response_attributes(response))
+
             return response, response_status
 
         # models routes are defined without trailing slash because this is required by the OpenAI python client.