fix(rebase): resolve errors from rebase

Author: iamemilio

src/llama_stack/core/routers/safety.py

@@ -7,11 +7,11 @@
 from typing import Any
 
 from llama_stack_api import ModerationObject, OpenAIMessageParam, RoutingTable, RunShieldResponse, Safety, Shield
+from opentelemetry import trace
 
 from llama_stack.core.datatypes import SafetyConfig
 from llama_stack.log import get_logger
-
-from opentelemetry import trace
+from llama_stack.telemetry.helpers import safety_request_span_attributes, safety_span_name
 
 logger = get_logger(name=__name__, category="core::routers")
 tracer = trace.get_tracer(__name__)

src/llama_stack/telemetry/helpers.py

@@ -6,11 +6,9 @@
 
 import json
 
+from llama_stack_api import OpenAIMessageParam, RunShieldResponse
 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,
@@ -27,7 +25,9 @@ def safety_span_name(shield_id: str) -> str:
 
 # 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:
+def safety_request_span_attributes(
+    shield_id: str, messages: list[OpenAIMessageParam], 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])