llama-stack 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

llama_stack/providers/remote/inference/tgi/config.py

@@ -5,28 +5,29 @@
 # the root directory of this source tree.
 
 
-from pydantic import BaseModel, Field, SecretStr
+from pydantic import BaseModel, Field, HttpUrl, SecretStr
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 @json_schema_type
 class TGIImplConfig(RemoteInferenceProviderConfig):
     auth_credential: SecretStr | None = Field(default=None, exclude=True)
 
-    url: str = Field(
-        description="The URL for the TGI serving endpoint",
+    base_url: HttpUrl | None = Field(
+        default=None,
+        description="The URL for the TGI serving endpoint (should include /v1 path)",
     )
 
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.TGI_URL:=}",
+        base_url: str = "${env.TGI_URL:=}",
         **kwargs,
     ):
         return {
-            "url": url,
+            "base_url": base_url,
         }
 
 

llama_stack/providers/remote/inference/tgi/tgi.py

@@ -8,14 +8,14 @@
 from collections.abc import Iterable
 
 from huggingface_hub import AsyncInferenceClient, HfApi
-from pydantic import SecretStr
+from pydantic import HttpUrl, SecretStr
 
-from llama_stack.apis.inference import (
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import (
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
-from llama_stack.log import get_logger
-from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import InferenceAPIImplConfig, InferenceEndpointImplConfig, TGIImplConfig
 
@@ -23,7 +23,7 @@ log = get_logger(name=__name__, category="inference::tgi")
 
 
 class _HfAdapter(OpenAIMixin):
-    url: str
+    base_url: HttpUrl
     api_key: SecretStr
 
     hf_client: AsyncInferenceClient
@@ -36,7 +36,7 @@ class _HfAdapter(OpenAIMixin):
         return "NO KEY REQUIRED"
 
     def get_base_url(self):
-        return self.url
+        return self.base_url
 
     async def list_provider_model_ids(self) -> Iterable[str]:
         return [self.model_id]
@@ -50,14 +50,22 @@ class _HfAdapter(OpenAIMixin):
 
 class TGIAdapter(_HfAdapter):
     async def initialize(self, config: TGIImplConfig) -> None:
-        if not config.url:
-            raise ValueError("You must provide a URL in run.yaml (or via the TGI_URL environment variable) to use TGI.")
-        log.info(f"Initializing TGI client with url={config.url}")
-        self.hf_client = AsyncInferenceClient(model=config.url, provider="hf-inference")
+        if not config.base_url:
+            raise ValueError(
+                "You must provide a URL in config.yaml (or via the TGI_URL environment variable) to use TGI."
+            )
+        log.info(f"Initializing TGI client with url={config.base_url}")
+        # Extract base URL without /v1 for HF client initialization
+        base_url_str = str(config.base_url).rstrip("/")
+        if base_url_str.endswith("/v1"):
+            base_url_for_client = base_url_str[:-3]
+        else:
+            base_url_for_client = base_url_str
+        self.hf_client = AsyncInferenceClient(model=base_url_for_client, provider="hf-inference")
         endpoint_info = await self.hf_client.get_endpoint_info()
         self.max_tokens = endpoint_info["max_total_tokens"]
         self.model_id = endpoint_info["model_id"]
-        self.url = f"{config.url.rstrip('/')}/v1"
+        self.base_url = config.base_url
         self.api_key = SecretStr("NO_KEY")
 
 

llama_stack/providers/remote/inference/together/config.py

@@ -6,22 +6,22 @@
 
 from typing import Any
 
-from pydantic import Field
+from pydantic import Field, HttpUrl
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 @json_schema_type
 class TogetherImplConfig(RemoteInferenceProviderConfig):
-    url: str = Field(
-        default="https://api.together.xyz/v1",
+    base_url: HttpUrl | None = Field(
+        default=HttpUrl("https://api.together.xyz/v1"),
         description="The URL for the Together AI server",
     )
 
     @classmethod
     def sample_run_config(cls, **kwargs) -> dict[str, Any]:
         return {
-            "url": "https://api.together.xyz/v1",
+            "base_url": "https://api.together.xyz/v1",
             "api_key": "${env.TOGETHER_API_KEY:=}",
         }

llama_stack/providers/remote/inference/together/together.py

@@ -6,19 +6,19 @@
 
 
 from collections.abc import Iterable
+from typing import Any, cast
 
-from together import AsyncTogether
-from together.constants import BASE_URL
+from together import AsyncTogether  # type: ignore[import-untyped]
 
-from llama_stack.apis.inference import (
-    OpenAIEmbeddingsRequestWithExtraBody,
-    OpenAIEmbeddingsResponse,
-)
-from llama_stack.apis.inference.inference import OpenAIEmbeddingUsage
-from llama_stack.apis.models import Model
 from llama_stack.core.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import (
+    Model,
+    OpenAIEmbeddingsRequestWithExtraBody,
+    OpenAIEmbeddingsResponse,
+    OpenAIEmbeddingUsage,
+)
 
 from .config import TogetherImplConfig
 
@@ -41,7 +41,7 @@ class TogetherInferenceAdapter(OpenAIMixin, NeedsRequestProviderData):
     provider_data_api_key_field: str = "together_api_key"
 
     def get_base_url(self):
-        return BASE_URL
+        return str(self.config.base_url)
 
     def _get_client(self) -> AsyncTogether:
         together_api_key = None
@@ -81,10 +81,11 @@ class TogetherInferenceAdapter(OpenAIMixin, NeedsRequestProviderData):
         if params.dimensions is not None:
             raise ValueError("Together's embeddings endpoint does not support dimensions param.")
 
+        # Cast encoding_format to match OpenAI SDK's expected Literal type
         response = await self.client.embeddings.create(
             model=await self._get_provider_model_id(params.model),
             input=params.input,
-            encoding_format=params.encoding_format,
+            encoding_format=cast(Any, params.encoding_format),
         )
 
         response.model = (
@@ -97,6 +98,8 @@ class TogetherInferenceAdapter(OpenAIMixin, NeedsRequestProviderData):
             logger.warning(
                 f"Together's embedding endpoint for {params.model} did not return usage information, substituting -1s."
             )
-            response.usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
+            # Cast to allow monkey-patching the response object
+            response.usage = cast(Any, OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1))
 
-        return response  # type: ignore[no-any-return]
+        # Together's CreateEmbeddingResponse is compatible with OpenAIEmbeddingsResponse after monkey-patching
+        return cast(OpenAIEmbeddingsResponse, response)

llama_stack/providers/remote/inference/vertexai/config.py

@@ -9,7 +9,7 @@ from typing import Any
 from pydantic import BaseModel, Field, SecretStr
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 class VertexAIProviderDataValidator(BaseModel):

llama_stack/providers/remote/inference/vllm/config.py

@@ -6,15 +6,15 @@
 
 from pathlib import Path
 
-from pydantic import Field, SecretStr, field_validator
+from pydantic import Field, HttpUrl, SecretStr, field_validator
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 @json_schema_type
 class VLLMInferenceAdapterConfig(RemoteInferenceProviderConfig):
-    url: str | None = Field(
+    base_url: HttpUrl | None = Field(
         default=None,
         description="The URL for the vLLM model serving endpoint",
     )
@@ -48,11 +48,11 @@ class VLLMInferenceAdapterConfig(RemoteInferenceProviderConfig):
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.VLLM_URL:=}",
+        base_url: str = "${env.VLLM_URL:=}",
         **kwargs,
     ):
         return {
-            "url": url,
+            "base_url": base_url,
             "max_tokens": "${env.VLLM_MAX_TOKENS:=4096}",
             "api_token": "${env.VLLM_API_TOKEN:=fake}",
             "tls_verify": "${env.VLLM_TLS_VERIFY:=true}",

llama_stack/providers/remote/inference/vllm/vllm.py

@@ -7,22 +7,18 @@ from collections.abc import AsyncIterator
 from urllib.parse import urljoin
 
 import httpx
-from openai.types.chat.chat_completion_chunk import (
-    ChatCompletionChunk as OpenAIChatCompletionChunk,
-)
 from pydantic import ConfigDict
 
-from llama_stack.apis.inference import (
-    OpenAIChatCompletion,
-    OpenAIChatCompletionRequestWithExtraBody,
-    ToolChoice,
-)
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import (
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import (
     HealthResponse,
     HealthStatus,
+    OpenAIChatCompletion,
+    OpenAIChatCompletionChunk,
+    OpenAIChatCompletionRequestWithExtraBody,
+    ToolChoice,
 )
-from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import VLLMInferenceAdapterConfig
 
@@ -34,6 +30,9 @@ class VLLMInferenceAdapter(OpenAIMixin):
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    # vLLM does not support the stream_options parameter
+    supports_stream_options: bool = False
+
     provider_data_api_key_field: str = "vllm_api_token"
 
     def get_api_key(self) -> str | None:
@@ -43,14 +42,14 @@ class VLLMInferenceAdapter(OpenAIMixin):
 
     def get_base_url(self) -> str:
         """Get the base URL from config."""
-        if not self.config.url:
+        if not self.config.base_url:
             raise ValueError("No base URL configured")
-        return self.config.url
+        return str(self.config.base_url)
 
     async def initialize(self) -> None:
-        if not self.config.url:
+        if not self.config.base_url:
             raise ValueError(
-                "You must provide a URL in run.yaml (or via the VLLM_URL environment variable) to use vLLM."
+                "You must provide a URL in config.yaml (or via the VLLM_URL environment variable) to use vLLM."
             )
 
     async def health(self) -> HealthResponse:

llama_stack/providers/remote/inference/watsonx/config.py

@@ -7,10 +7,10 @@
 import os
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, HttpUrl
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 class WatsonXProviderDataValidator(BaseModel):
@@ -23,7 +23,7 @@ class WatsonXProviderDataValidator(BaseModel):
 
 @json_schema_type
 class WatsonXConfig(RemoteInferenceProviderConfig):
-    url: str = Field(
+    base_url: HttpUrl | None = Field(
         default_factory=lambda: os.getenv("WATSONX_BASE_URL", "https://us-south.ml.cloud.ibm.com"),
         description="A base url for accessing the watsonx.ai",
     )
@@ -39,7 +39,7 @@ class WatsonXConfig(RemoteInferenceProviderConfig):
     @classmethod
     def sample_run_config(cls, **kwargs) -> dict[str, Any]:
         return {
-            "url": "${env.WATSONX_BASE_URL:=https://us-south.ml.cloud.ibm.com}",
+            "base_url": "${env.WATSONX_BASE_URL:=https://us-south.ml.cloud.ibm.com}",
             "api_key": "${env.WATSONX_API_KEY:=}",
             "project_id": "${env.WATSONX_PROJECT_ID:=}",
         }

llama_stack/providers/remote/inference/watsonx/watsonx.py

@@ -10,23 +10,20 @@ from typing import Any
 import litellm
 import requests
 
-from llama_stack.apis.inference.inference import (
+from llama_stack.log import get_logger
+from llama_stack.providers.remote.inference.watsonx.config import WatsonXConfig
+from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
+from llama_stack_api import (
+    Model,
+    ModelType,
     OpenAIChatCompletion,
     OpenAIChatCompletionChunk,
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAIChatCompletionUsage,
-    OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
-from llama_stack.apis.models import Model
-from llama_stack.apis.models.models import ModelType
-from llama_stack.log import get_logger
-from llama_stack.providers.remote.inference.watsonx.config import WatsonXConfig
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
-from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
-from llama_stack.providers.utils.telemetry.tracing import get_current_span
 
 logger = get_logger(name=__name__, category="providers::remote::watsonx")
 
@@ -48,57 +45,25 @@ class WatsonXInferenceAdapter(LiteLLMOpenAIMixin):
             openai_compat_api_base=self.get_base_url(),
         )
 
+    def _litellm_extra_request_params(
+        self,
+        params: OpenAIChatCompletionRequestWithExtraBody | OpenAICompletionRequestWithExtraBody,
+    ) -> dict[str, Any]:
+        # These are watsonx-specific parameters used by LiteLLM.
+        return {"timeout": self.config.timeout, "project_id": self.config.project_id}
+
     async def openai_chat_completion(
         self,
         params: OpenAIChatCompletionRequestWithExtraBody,
     ) -> OpenAIChatCompletion | AsyncIterator[OpenAIChatCompletionChunk]:
         """
-        Override parent method to add timeout and inject usage object when missing.
+        Override parent method to inject usage object when missing.
+
         This works around a LiteLLM defect where usage block is sometimes dropped.
+        Note: request parameter construction (including telemetry-driven stream_options injection)
+        is handled by LiteLLMOpenAIMixin via _litellm_extra_request_params().
         """
-
-        # Add usage tracking for streaming when telemetry is active
-        stream_options = params.stream_options
-        if params.stream and get_current_span() is not None:
-            if stream_options is None:
-                stream_options = {"include_usage": True}
-            elif "include_usage" not in stream_options:
-                stream_options = {**stream_options, "include_usage": True}
-
-        model_obj = await self.model_store.get_model(params.model)
-
-        request_params = await prepare_openai_completion_params(
-            model=self.get_litellm_model_name(model_obj.provider_resource_id),
-            messages=params.messages,
-            frequency_penalty=params.frequency_penalty,
-            function_call=params.function_call,
-            functions=params.functions,
-            logit_bias=params.logit_bias,
-            logprobs=params.logprobs,
-            max_completion_tokens=params.max_completion_tokens,
-            max_tokens=params.max_tokens,
-            n=params.n,
-            parallel_tool_calls=params.parallel_tool_calls,
-            presence_penalty=params.presence_penalty,
-            response_format=params.response_format,
-            seed=params.seed,
-            stop=params.stop,
-            stream=params.stream,
-            stream_options=stream_options,
-            temperature=params.temperature,
-            tool_choice=params.tool_choice,
-            tools=params.tools,
-            top_logprobs=params.top_logprobs,
-            top_p=params.top_p,
-            user=params.user,
-            api_key=self.get_api_key(),
-            api_base=self.api_base,
-            # These are watsonx-specific parameters
-            timeout=self.config.timeout,
-            project_id=self.config.project_id,
-        )
-
-        result = await litellm.acompletion(**request_params)
+        result = await super().openai_chat_completion(params)
 
         # If not streaming, check and inject usage if missing
         if not params.stream:
@@ -175,44 +140,6 @@ class WatsonXInferenceAdapter(LiteLLMOpenAIMixin):
             logger.error(f"Error normalizing stream: {e}", exc_info=True)
             raise
 
-    async def openai_completion(
-        self,
-        params: OpenAICompletionRequestWithExtraBody,
-    ) -> OpenAICompletion:
-        """
-        Override parent method to add watsonx-specific parameters.
-        """
-        from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
-
-        model_obj = await self.model_store.get_model(params.model)
-
-        request_params = await prepare_openai_completion_params(
-            model=self.get_litellm_model_name(model_obj.provider_resource_id),
-            prompt=params.prompt,
-            best_of=params.best_of,
-            echo=params.echo,
-            frequency_penalty=params.frequency_penalty,
-            logit_bias=params.logit_bias,
-            logprobs=params.logprobs,
-            max_tokens=params.max_tokens,
-            n=params.n,
-            presence_penalty=params.presence_penalty,
-            seed=params.seed,
-            stop=params.stop,
-            stream=params.stream,
-            stream_options=params.stream_options,
-            temperature=params.temperature,
-            top_p=params.top_p,
-            user=params.user,
-            suffix=params.suffix,
-            api_key=self.get_api_key(),
-            api_base=self.api_base,
-            # These are watsonx-specific parameters
-            timeout=self.config.timeout,
-            project_id=self.config.project_id,
-        )
-        return await litellm.atext_completion(**request_params)
-
     async def openai_embeddings(
         self,
         params: OpenAIEmbeddingsRequestWithExtraBody,
@@ -238,8 +165,8 @@ class WatsonXInferenceAdapter(LiteLLMOpenAIMixin):
         )
 
         # Convert response to OpenAI format
-        from llama_stack.apis.inference import OpenAIEmbeddingUsage
         from llama_stack.providers.utils.inference.litellm_openai_mixin import b64_encode_openai_embeddings_response
+        from llama_stack_api import OpenAIEmbeddingUsage
 
         data = b64_encode_openai_embeddings_response(response.data, params.encoding_format)
 
@@ -255,7 +182,7 @@ class WatsonXInferenceAdapter(LiteLLMOpenAIMixin):
         )
 
     def get_base_url(self) -> str:
-        return self.config.url
+        return str(self.config.base_url)
 
     # Copied from OpenAIMixin
     async def check_model_availability(self, model: str) -> bool:
@@ -316,7 +243,7 @@ class WatsonXInferenceAdapter(LiteLLMOpenAIMixin):
         """
         Retrieves foundation model specifications from the watsonx.ai API.
         """
-        url = f"{self.config.url}/ml/v1/foundation_model_specs?version=2023-10-25"
+        url = f"{str(self.config.base_url)}/ml/v1/foundation_model_specs?version=2023-10-25"
         headers = {
             # Note that there is no authorization header.  Listing models does not require authentication.
             "Content-Type": "application/json",

llama_stack/providers/remote/post_training/nvidia/post_training.py

@@ -10,7 +10,10 @@ from typing import Any, Literal
 import aiohttp
 from pydantic import BaseModel, ConfigDict
 
-from llama_stack.apis.post_training import (
+from llama_stack.providers.remote.post_training.nvidia.config import NvidiaPostTrainingConfig
+from llama_stack.providers.remote.post_training.nvidia.utils import warn_unsupported_params
+from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
+from llama_stack_api import (
     AlgorithmConfig,
     DPOAlignmentConfig,
     JobStatus,
@@ -19,9 +22,6 @@ from llama_stack.apis.post_training import (
     PostTrainingJobStatusResponse,
     TrainingConfig,
 )
-from llama_stack.providers.remote.post_training.nvidia.config import NvidiaPostTrainingConfig
-from llama_stack.providers.remote.post_training.nvidia.utils import warn_unsupported_params
-from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 
 from .models import _MODEL_ENTRIES
 

llama_stack/providers/remote/post_training/nvidia/utils.py

@@ -9,9 +9,9 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from llama_stack.apis.post_training import TrainingConfig
 from llama_stack.log import get_logger
 from llama_stack.providers.remote.post_training.nvidia.config import SFTLoRADefaultConfig
+from llama_stack_api import TrainingConfig
 
 from .config import NvidiaPostTrainingConfig
 

llama_stack/providers/remote/safety/bedrock/bedrock.py

@@ -7,17 +7,17 @@
 import json
 from typing import Any
 
-from llama_stack.apis.inference import OpenAIMessageParam
-from llama_stack.apis.safety import (
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.bedrock.client import create_bedrock_client
+from llama_stack_api import (
+    OpenAIMessageParam,
     RunShieldResponse,
     Safety,
     SafetyViolation,
+    Shield,
+    ShieldsProtocolPrivate,
     ViolationLevel,
 )
-from llama_stack.apis.shields import Shield
-from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import ShieldsProtocolPrivate
-from llama_stack.providers.utils.bedrock.client import create_bedrock_client
 
 from .config import BedrockSafetyConfig
 

llama_stack/providers/remote/safety/bedrock/config.py

@@ -6,7 +6,7 @@
 
 
 from llama_stack.providers.utils.bedrock.config import BedrockBaseConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 @json_schema_type

llama_stack/providers/remote/safety/nvidia/config.py

@@ -8,7 +8,7 @@ from typing import Any
 
 from pydantic import BaseModel, Field
 
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 @json_schema_type

llama_stack/providers/remote/safety/nvidia/nvidia.py

@@ -8,11 +8,17 @@ from typing import Any
 
 import requests
 
-from llama_stack.apis.inference import OpenAIMessageParam
-from llama_stack.apis.safety import ModerationObject, RunShieldResponse, Safety, SafetyViolation, ViolationLevel
-from llama_stack.apis.shields import Shield
 from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import ShieldsProtocolPrivate
+from llama_stack_api import (
+    ModerationObject,
+    OpenAIMessageParam,
+    RunShieldResponse,
+    Safety,
+    SafetyViolation,
+    Shield,
+    ShieldsProtocolPrivate,
+    ViolationLevel,
+)
 
 from .config import NVIDIASafetyConfig
 
@@ -66,7 +72,7 @@ class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate):
         self.shield = NeMoGuardrails(self.config, shield.shield_id)
         return await self.shield.run(messages)
 
-    async def run_moderation(self, input: str | list[str], model: str) -> ModerationObject:
+    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
         raise NotImplementedError("NVIDIA safety provider currently does not implement run_moderation")
 
 

llama_stack/providers/remote/safety/sambanova/config.py

@@ -8,7 +8,7 @@ from typing import Any
 
 from pydantic import BaseModel, Field, SecretStr
 
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 class SambaNovaProviderDataValidator(BaseModel):

llama_stack/providers/remote/safety/sambanova/sambanova.py

@@ -9,17 +9,17 @@ from typing import Any
 import litellm
 import requests
 
-from llama_stack.apis.inference import OpenAIMessageParam
-from llama_stack.apis.safety import (
+from llama_stack.core.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
+from llama_stack_api import (
+    OpenAIMessageParam,
     RunShieldResponse,
     Safety,
     SafetyViolation,
+    Shield,
+    ShieldsProtocolPrivate,
     ViolationLevel,
 )
-from llama_stack.apis.shields import Shield
-from llama_stack.core.request_headers import NeedsRequestProviderData
-from llama_stack.log import get_logger
-from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 
 from .config import SambaNovaSafetyConfig