fix(telemetry): move out of core + fix name for safety

src/llama_stack/telemetry/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

src/llama_stack/telemetry/constants.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+llama_stack_prefix = "llama_stack"
+
+# Safety Attributes
+RUN_SHIELD_OPERATION_NAME = "run_shield"
+
+SAFETY_REQUEST_PREFIX = f"{llama_stack_prefix}.safety.request"
+SAFETY_REQUEST_SHIELD_ID_ATTRIBUTE = f"{SAFETY_REQUEST_PREFIX}.shield_id"
+SAFETY_REQUEST_MESSAGES_ATTRIBUTE = f"{SAFETY_REQUEST_PREFIX}.messages"
+
+SAFETY_RESPONSE_PREFIX = f"{llama_stack_prefix}.safety.response"
+SAFETY_RESPONSE_METADATA_ATTRIBUTE = f"{SAFETY_RESPONSE_PREFIX}.metadata"
+SAFETY_RESPONSE_VIOLATION_LEVEL_ATTRIBUTE = f"{SAFETY_RESPONSE_PREFIX}.violation.level"
+SAFETY_RESPONSE_USER_MESSAGE_ATTRIBUTE = f"{SAFETY_RESPONSE_PREFIX}.violation.user_message"

src/llama_stack/telemetry/helpers.py

@@ -0,0 +1,42 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+
+from opentelemetry import trace
+
+from llama_stack.apis.inference import Message
+from llama_stack.apis.safety import RunShieldResponse
+
+from .constants import (
+    RUN_SHIELD_OPERATION_NAME,
+    SAFETY_REQUEST_MESSAGES_ATTRIBUTE,
+    SAFETY_REQUEST_SHIELD_ID_ATTRIBUTE,
+    SAFETY_RESPONSE_METADATA_ATTRIBUTE,
+    SAFETY_RESPONSE_USER_MESSAGE_ATTRIBUTE,
+    SAFETY_RESPONSE_VIOLATION_LEVEL_ATTRIBUTE,
+)
+
+
+def safety_span_name(shield_id: str) -> str:
+    return f"{RUN_SHIELD_OPERATION_NAME} {shield_id}"
+
+
+# TODO: Consider using Wrapt to automatically instrument code
+# This is the industry standard way to package automatically instrumentation in python.
+def safety_request_span_attributes(shield_id: str, messages: list[Message], response: RunShieldResponse) -> None:
+    span = trace.get_current_span()
+    span.set_attribute(SAFETY_REQUEST_SHIELD_ID_ATTRIBUTE, shield_id)
+    messages_json = json.dumps([msg.model_dump() for msg in messages])
+    span.set_attribute(SAFETY_REQUEST_MESSAGES_ATTRIBUTE, messages_json)
+
+    if response.violation:
+        if response.violation.metadata:
+            metadata_json = json.dumps(response.violation.metadata)
+            span.set_attribute(SAFETY_RESPONSE_METADATA_ATTRIBUTE, metadata_json)
+        if response.violation.user_message:
+            span.set_attribute(SAFETY_RESPONSE_USER_MESSAGE_ATTRIBUTE, response.violation.user_message)
+        span.set_attribute(SAFETY_RESPONSE_VIOLATION_LEVEL_ATTRIBUTE, response.violation.violation_level.value)