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

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

@@ -0,0 +1,430 @@
+# 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 warnings
+from datetime import datetime
+from typing import Any, Literal
+
+import aiohttp
+from pydantic import BaseModel, ConfigDict
+
+from llama_stack.apis.post_training import (
+    AlgorithmConfig,
+    DPOAlignmentConfig,
+    JobStatus,
+    PostTrainingJob,
+    PostTrainingJobArtifactsResponse,
+    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
+
+# Map API status to JobStatus enum
+STATUS_MAPPING = {
+    "running": JobStatus.in_progress.value,
+    "completed": JobStatus.completed.value,
+    "failed": JobStatus.failed.value,
+    "cancelled": JobStatus.cancelled.value,
+    "pending": JobStatus.scheduled.value,
+    "unknown": JobStatus.scheduled.value,
+}
+
+
+class NvidiaPostTrainingJob(PostTrainingJob):
+    """Parse the response from the Customizer API.
+    Inherits job_uuid from PostTrainingJob.
+    Adds status, created_at, updated_at parameters.
+    Passes through all other parameters from data field in the response.
+    """
+
+    model_config = ConfigDict(extra="allow")
+    status: JobStatus
+    created_at: datetime
+    updated_at: datetime
+
+
+class ListNvidiaPostTrainingJobs(BaseModel):
+    data: list[NvidiaPostTrainingJob]
+
+
+class NvidiaPostTrainingJobStatusResponse(PostTrainingJobStatusResponse):
+    model_config = ConfigDict(extra="allow")
+
+
+class NvidiaPostTrainingAdapter(ModelRegistryHelper):
+    def __init__(self, config: NvidiaPostTrainingConfig):
+        self.config = config
+        self.headers = {}
+        if config.api_key:
+            self.headers["Authorization"] = f"Bearer {config.api_key}"
+
+        self.timeout = aiohttp.ClientTimeout(total=config.timeout)
+        # TODO: filter by available models based on /config endpoint
+        ModelRegistryHelper.__init__(self, model_entries=_MODEL_ENTRIES)
+        self.session = None
+
+        self.customizer_url = config.customizer_url
+        if not self.customizer_url:
+            warnings.warn("Customizer URL is not set, using default value: http://nemo.test", stacklevel=2)
+            self.customizer_url = "http://nemo.test"
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        if self.session is None or self.session.closed:
+            self.session = aiohttp.ClientSession(headers=self.headers, timeout=self.timeout)
+        return self.session
+
+    async def _make_request(
+        self,
+        method: str,
+        path: str,
+        headers: dict[str, Any] | None = None,
+        params: dict[str, Any] | None = None,
+        json: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Helper method to make HTTP requests to the Customizer API."""
+        url = f"{self.customizer_url}{path}"
+        request_headers = self.headers.copy()
+
+        if headers:
+            request_headers.update(headers)
+
+        # Add content-type header for JSON requests
+        if json and "Content-Type" not in request_headers:
+            request_headers["Content-Type"] = "application/json"
+
+        session = await self._get_session()
+        for _ in range(self.config.max_retries):
+            async with session.request(method, url, params=params, json=json, **kwargs) as response:
+                if response.status >= 400:
+                    error_data = await response.json()
+                    raise Exception(f"API request failed: {error_data}")
+                return await response.json()
+
+    async def get_training_jobs(
+        self,
+        page: int | None = 1,
+        page_size: int | None = 10,
+        sort: Literal["created_at", "-created_at"] | None = "created_at",
+    ) -> ListNvidiaPostTrainingJobs:
+        """Get all customization jobs.
+        Updated the base class return type from ListPostTrainingJobsResponse to ListNvidiaPostTrainingJobs.
+
+        Returns a ListNvidiaPostTrainingJobs object with the following fields:
+            - data: List[NvidiaPostTrainingJob] - List of NvidiaPostTrainingJob objects
+
+        ToDo: Support for schema input for filtering.
+        """
+        params = {"page": page, "page_size": page_size, "sort": sort}
+
+        response = await self._make_request("GET", "/v1/customization/jobs", params=params)
+
+        jobs = []
+        for job in response.get("data", []):
+            job_id = job.pop("id")
+            job_status = job.pop("status", "scheduled").lower()
+            mapped_status = STATUS_MAPPING.get(job_status, "scheduled")
+
+            # Convert string timestamps to datetime objects
+            created_at = (
+                datetime.fromisoformat(job.pop("created_at"))
+                if "created_at" in job
+                else datetime.now(tz=datetime.timezone.utc)
+            )
+            updated_at = (
+                datetime.fromisoformat(job.pop("updated_at"))
+                if "updated_at" in job
+                else datetime.now(tz=datetime.timezone.utc)
+            )
+
+            # Create NvidiaPostTrainingJob instance
+            jobs.append(
+                NvidiaPostTrainingJob(
+                    job_uuid=job_id,
+                    status=JobStatus(mapped_status),
+                    created_at=created_at,
+                    updated_at=updated_at,
+                    **job,
+                )
+            )
+
+        return ListNvidiaPostTrainingJobs(data=jobs)
+
+    async def get_training_job_status(self, job_uuid: str) -> NvidiaPostTrainingJobStatusResponse:
+        """Get the status of a customization job.
+        Updated the base class return type from PostTrainingJobResponse to NvidiaPostTrainingJob.
+
+        Returns a NvidiaPostTrainingJob object with the following fields:
+            - job_uuid: str - Unique identifier for the job
+            - status: JobStatus - Current status of the job (in_progress, completed, failed, cancelled, scheduled)
+            - created_at: datetime - The time when the job was created
+            - updated_at: datetime - The last time the job status was updated
+
+        Additional fields that may be included:
+            - steps_completed: Optional[int] - Number of training steps completed
+            - epochs_completed: Optional[int] - Number of epochs completed
+            - percentage_done: Optional[float] - Percentage of training completed (0-100)
+            - best_epoch: Optional[int] - The epoch with the best performance
+            - train_loss: Optional[float] - Training loss of the best checkpoint
+            - val_loss: Optional[float] - Validation loss of the best checkpoint
+            - metrics: Optional[Dict] - Additional training metrics
+            - status_logs: Optional[List] - Detailed logs of status changes
+        """
+        response = await self._make_request(
+            "GET",
+            f"/v1/customization/jobs/{job_uuid}/status",
+            params={"job_id": job_uuid},
+        )
+
+        api_status = response.pop("status").lower()
+        mapped_status = STATUS_MAPPING.get(api_status, "scheduled")
+
+        return NvidiaPostTrainingJobStatusResponse(
+            status=JobStatus(mapped_status),
+            job_uuid=job_uuid,
+            started_at=datetime.fromisoformat(response.pop("created_at")),
+            updated_at=datetime.fromisoformat(response.pop("updated_at")),
+            **response,
+        )
+
+    async def cancel_training_job(self, job_uuid: str) -> None:
+        await self._make_request(
+            method="POST", path=f"/v1/customization/jobs/{job_uuid}/cancel", params={"job_id": job_uuid}
+        )
+
+    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse:
+        raise NotImplementedError("Job artifacts are not implemented yet")
+
+    async def get_post_training_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse:
+        raise NotImplementedError("Job artifacts are not implemented yet")
+
+    async def supervised_fine_tune(
+        self,
+        job_uuid: str,
+        training_config: dict[str, Any],
+        hyperparam_search_config: dict[str, Any],
+        logger_config: dict[str, Any],
+        model: str,
+        checkpoint_dir: str | None,
+        algorithm_config: AlgorithmConfig | None = None,
+    ) -> NvidiaPostTrainingJob:
+        """
+        Fine-tunes a model on a dataset.
+        Currently only supports Lora finetuning for standlone docker container.
+        Assumptions:
+            - nemo microservice is running and endpoint is set in config.customizer_url
+            - dataset is registered separately in nemo datastore
+            - model checkpoint is downloaded as per nemo customizer requirements
+
+        Parameters:
+            training_config: TrainingConfig - Configuration for training
+            model: str - NeMo Customizer configuration name
+            algorithm_config: Optional[AlgorithmConfig] - Algorithm-specific configuration
+            checkpoint_dir: Optional[str] - Directory containing model checkpoints, ignored atm
+            job_uuid: str - Unique identifier for the job, ignored atm
+            hyperparam_search_config: Dict[str, Any] - Configuration for hyperparameter search, ignored atm
+            logger_config: Dict[str, Any] - Configuration for logging, ignored atm
+
+        Environment Variables:
+            - NVIDIA_API_KEY: str - API key for the NVIDIA API
+                Default: None
+            - NVIDIA_DATASET_NAMESPACE: str - Namespace of the dataset
+                Default: "default"
+            - NVIDIA_CUSTOMIZER_URL: str - URL of the NeMo Customizer API
+                Default: "http://nemo.test"
+            - NVIDIA_PROJECT_ID: str - ID of the project
+                Default: "test-project"
+            - NVIDIA_OUTPUT_MODEL_DIR: str - Directory to save the output model
+                Default: "test-example-model@v1"
+
+        Supported models:
+            - meta/llama-3.1-8b-instruct
+            - meta/llama-3.2-1b-instruct
+
+        Supported algorithm configs:
+            - LoRA, SFT
+
+        Supported Parameters:
+            - TrainingConfig:
+                - n_epochs: int - Number of epochs to train
+                    Default: 50
+                - data_config: DataConfig - Configuration for the dataset
+                - optimizer_config: OptimizerConfig - Configuration for the optimizer
+                - dtype: str - Data type for training
+                    not supported (users are informed via warnings)
+                - efficiency_config: EfficiencyConfig - Configuration for efficiency
+                    not supported
+                - max_steps_per_epoch: int - Maximum number of steps per epoch
+                    Default: 1000
+                ## NeMo customizer specific parameters
+                - log_every_n_steps: int - Log every n steps
+                    Default: None
+                - val_check_interval: float - Validation check interval
+                    Default: 0.25
+                - sequence_packing_enabled: bool - Sequence packing enabled
+                    Default: False
+                ## NeMo customizer specific SFT parameters
+                - hidden_dropout: float - Hidden dropout
+                    Default: None (0.0-1.0)
+                - attention_dropout: float - Attention dropout
+                    Default: None (0.0-1.0)
+                - ffn_dropout: float - FFN dropout
+                    Default: None (0.0-1.0)
+
+            - DataConfig:
+                - dataset_id: str - Dataset ID
+                - batch_size: int - Batch size
+                    Default: 8
+
+            - OptimizerConfig:
+                - lr: float - Learning rate
+                    Default: 0.0001
+                ## NeMo customizer specific parameter
+                - weight_decay: float - Weight decay
+                    Default: 0.01
+
+            - LoRA config:
+                ## NeMo customizer specific LoRA parameters
+                - alpha: int - Scaling factor for the LoRA update
+                    Default: 16
+            Note:
+                - checkpoint_dir, hyperparam_search_config, logger_config are not supported (users are informed via warnings)
+                - Some parameters from TrainingConfig, DataConfig, OptimizerConfig are not supported (users are informed via warnings)
+
+            User is informed about unsupported parameters via warnings.
+        """
+
+        # Check for unsupported method parameters
+        unsupported_method_params = []
+        if checkpoint_dir:
+            unsupported_method_params.append(f"checkpoint_dir={checkpoint_dir}")
+        if hyperparam_search_config:
+            unsupported_method_params.append("hyperparam_search_config")
+        if logger_config:
+            unsupported_method_params.append("logger_config")
+
+        if unsupported_method_params:
+            warnings.warn(
+                f"Parameters: {', '.join(unsupported_method_params)} are not supported and will be ignored",
+                stacklevel=2,
+            )
+
+        # Define all supported parameters
+        supported_params = {
+            "training_config": {
+                "n_epochs",
+                "data_config",
+                "optimizer_config",
+                "log_every_n_steps",
+                "val_check_interval",
+                "sequence_packing_enabled",
+                "hidden_dropout",
+                "attention_dropout",
+                "ffn_dropout",
+            },
+            "data_config": {"dataset_id", "batch_size"},
+            "optimizer_config": {"lr", "weight_decay"},
+            "lora_config": {"type", "alpha"},
+        }
+
+        # Validate all parameters at once
+        warn_unsupported_params(training_config, supported_params["training_config"], "TrainingConfig")
+        warn_unsupported_params(training_config["data_config"], supported_params["data_config"], "DataConfig")
+        warn_unsupported_params(
+            training_config["optimizer_config"], supported_params["optimizer_config"], "OptimizerConfig"
+        )
+
+        output_model = self.config.output_model_dir
+
+        # Prepare base job configuration
+        job_config = {
+            "config": model,
+            "dataset": {
+                "name": training_config["data_config"]["dataset_id"],
+                "namespace": self.config.dataset_namespace,
+            },
+            "hyperparameters": {
+                "training_type": "sft",
+                "finetuning_type": "lora",
+                **{
+                    k: v
+                    for k, v in {
+                        "epochs": training_config.get("n_epochs"),
+                        "batch_size": training_config["data_config"].get("batch_size"),
+                        "learning_rate": training_config["optimizer_config"].get("lr"),
+                        "weight_decay": training_config["optimizer_config"].get("weight_decay"),
+                        "log_every_n_steps": training_config.get("log_every_n_steps"),
+                        "val_check_interval": training_config.get("val_check_interval"),
+                        "sequence_packing_enabled": training_config.get("sequence_packing_enabled"),
+                    }.items()
+                    if v is not None
+                },
+            },
+            "project": self.config.project_id,
+            # TODO: ignored ownership, add it later
+            # "ownership": {"created_by": self.config.user_id, "access_policies": self.config.access_policies},
+            "output_model": output_model,
+        }
+
+        # Handle SFT-specific optional parameters
+        job_config["hyperparameters"]["sft"] = {
+            k: v
+            for k, v in {
+                "ffn_dropout": training_config.get("ffn_dropout"),
+                "hidden_dropout": training_config.get("hidden_dropout"),
+                "attention_dropout": training_config.get("attention_dropout"),
+            }.items()
+            if v is not None
+        }
+
+        # Remove the sft dictionary if it's empty
+        if not job_config["hyperparameters"]["sft"]:
+            job_config["hyperparameters"].pop("sft")
+
+        # Handle LoRA-specific configuration
+        if algorithm_config:
+            if algorithm_config.type == "LoRA":
+                warn_unsupported_params(algorithm_config, supported_params["lora_config"], "LoRA config")
+                job_config["hyperparameters"]["lora"] = {
+                    k: v for k, v in {"alpha": algorithm_config.alpha}.items() if v is not None
+                }
+            else:
+                raise NotImplementedError(f"Unsupported algorithm config: {algorithm_config}")
+
+        # Create the customization job
+        response = await self._make_request(
+            method="POST",
+            path="/v1/customization/jobs",
+            headers={"Accept": "application/json"},
+            json=job_config,
+        )
+
+        job_uuid = response["id"]
+        response.pop("status")
+        created_at = datetime.fromisoformat(response.pop("created_at"))
+        updated_at = datetime.fromisoformat(response.pop("updated_at"))
+
+        return NvidiaPostTrainingJob(
+            job_uuid=job_uuid, status=JobStatus.in_progress, created_at=created_at, updated_at=updated_at, **response
+        )
+
+    async def preference_optimize(
+        self,
+        job_uuid: str,
+        finetuned_model: str,
+        algorithm_config: DPOAlignmentConfig,
+        training_config: TrainingConfig,
+        hyperparam_search_config: dict[str, Any],
+        logger_config: dict[str, Any],
+    ) -> PostTrainingJob:
+        """Optimize a model based on preference data."""
+        raise NotImplementedError("Preference optimization is not implemented yet")
+
+    async def get_training_job_container_logs(self, job_uuid: str) -> PostTrainingJobStatusResponse:
+        raise NotImplementedError("Job logs are not implemented yet")

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

@@ -0,0 +1,63 @@
+# 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 warnings
+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 .config import NvidiaPostTrainingConfig
+
+logger = get_logger(name=__name__, category="post_training::nvidia")
+
+
+def warn_unsupported_params(config_dict: Any, supported_keys: set[str], config_name: str) -> None:
+    keys = set(config_dict.__annotations__.keys()) if isinstance(config_dict, BaseModel) else config_dict.keys()
+    unsupported_params = [k for k in keys if k not in supported_keys]
+    if unsupported_params:
+        warnings.warn(
+            f"Parameters: {unsupported_params} in `{config_name}` not supported and will be ignored.", stacklevel=2
+        )
+
+
+def validate_training_params(
+    training_config: dict[str, Any], supported_keys: set[str], config_name: str = "TrainingConfig"
+) -> None:
+    """
+    Validates training parameters against supported keys.
+
+    Args:
+        training_config: Dictionary containing training configuration parameters
+        supported_keys: Set of supported parameter keys
+        config_name: Name of the configuration for warning messages
+    """
+    sft_lora_fields = set(SFTLoRADefaultConfig.__annotations__.keys())
+    training_config_fields = set(TrainingConfig.__annotations__.keys())
+
+    # Check for not supported parameters:
+    # - not in either of configs
+    # - in TrainingConfig but not in SFTLoRADefaultConfig
+    unsupported_params = []
+    for key in training_config:
+        if isinstance(key, str) and key not in (supported_keys.union(sft_lora_fields)):
+            if key in (not sft_lora_fields or training_config_fields):
+                unsupported_params.append(key)
+
+    if unsupported_params:
+        warnings.warn(
+            f"Parameters: {unsupported_params} in `{config_name}` are not supported and will be ignored.", stacklevel=2
+        )
+
+
+# ToDo: implement post health checks for customizer are enabled
+async def _get_health(url: str) -> tuple[bool, bool]: ...
+
+
+async def check_health(config: NvidiaPostTrainingConfig) -> None: ...

llama_stack/providers/remote/safety/__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.

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

@@ -0,0 +1,111 @@
+# 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 typing import Any
+
+from llama_stack.apis.inference import OpenAIMessageParam
+from llama_stack.apis.safety import (
+    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.providers.utils.bedrock.client import create_bedrock_client
+
+from .config import BedrockSafetyConfig
+
+logger = get_logger(name=__name__, category="safety::bedrock")
+
+
+class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
+    def __init__(self, config: BedrockSafetyConfig) -> None:
+        self.config = config
+        self.registered_shields = []
+
+    async def initialize(self) -> None:
+        try:
+            self.bedrock_runtime_client = create_bedrock_client(self.config)
+            self.bedrock_client = create_bedrock_client(self.config, "bedrock")
+        except Exception as e:
+            raise RuntimeError("Error initializing BedrockSafetyAdapter") from e
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def register_shield(self, shield: Shield) -> None:
+        response = self.bedrock_client.list_guardrails(
+            guardrailIdentifier=shield.provider_resource_id,
+        )
+        if (
+            not response["guardrails"]
+            or len(response["guardrails"]) == 0
+            or response["guardrails"][0]["version"] != shield.params["guardrailVersion"]
+        ):
+            raise ValueError(
+                f"Shield {shield.provider_resource_id} with version {shield.params['guardrailVersion']} not found in Bedrock"
+            )
+
+    async def unregister_shield(self, identifier: str) -> None:
+        pass
+
+    async def run_shield(
+        self, shield_id: str, messages: list[OpenAIMessageParam], params: dict[str, Any] = None
+    ) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(shield_id)
+        if not shield:
+            raise ValueError(f"Shield {shield_id} not found")
+
+        """
+        This is the implementation for the bedrock guardrails. The input to the guardrails is to be of this format
+        ```content = [
+            {
+                "text": {
+                    "text": "Is the AB503 Product a better investment than the S&P 500?"
+                }
+            }
+        ]```
+        Incoming messages contain content, role . For now we will extract the content and
+        default the "qualifiers": ["query"]
+        """
+
+        shield_params = shield.params
+        logger.debug(f"run_shield::{shield_params}::messages={messages}")
+
+        # - convert the messages into format Bedrock expects
+        content_messages = []
+        for message in messages:
+            content_messages.append({"text": {"text": message.content}})
+        logger.debug(f"run_shield::final:messages::{json.dumps(content_messages, indent=2)}:")
+
+        response = self.bedrock_runtime_client.apply_guardrail(
+            guardrailIdentifier=shield.provider_resource_id,
+            guardrailVersion=shield_params["guardrailVersion"],
+            source="OUTPUT",  # or 'INPUT' depending on your use case
+            content=content_messages,
+        )
+        if response["action"] == "GUARDRAIL_INTERVENED":
+            user_message = ""
+            metadata = {}
+            for output in response["outputs"]:
+                # guardrails returns a list - however for this implementation we will leverage the last values
+                user_message = output["text"]
+            for assessment in response["assessments"]:
+                # guardrails returns a list - however for this implementation we will leverage the last values
+                metadata = dict(assessment)
+
+            return RunShieldResponse(
+                violation=SafetyViolation(
+                    user_message=user_message,
+                    violation_level=ViolationLevel.ERROR,
+                    metadata=metadata,
+                )
+            )
+
+        return RunShieldResponse()

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

@@ -0,0 +1,14 @@
+# 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.
+
+
+from llama_stack.providers.utils.bedrock.config import BedrockBaseConfig
+from llama_stack.schema_utils import json_schema_type
+
+
+@json_schema_type
+class BedrockSafetyConfig(BedrockBaseConfig):
+    pass

llama_stack/providers/{adapters/inference/sample → remote/safety/nvidia}/__init__.py

@@ -4,14 +4,15 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+
 from typing import Any
 
-from .config import SampleConfig
+from .config import NVIDIASafetyConfig
 
 
-async def get_adapter_impl(config: SampleConfig, _deps) -> Any:
-    from .sample import SampleInferenceImpl
+async def get_adapter_impl(config: NVIDIASafetyConfig, _deps) -> Any:
+    from .nvidia import NVIDIASafetyAdapter
 
-    impl = SampleInferenceImpl(config)
+    impl = NVIDIASafetyAdapter(config)
     await impl.initialize()
     return impl

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

@@ -0,0 +1,40 @@
+# 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 os
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from llama_stack.schema_utils import json_schema_type
+
+
+@json_schema_type
+class NVIDIASafetyConfig(BaseModel):
+    """
+    Configuration for the NVIDIA Guardrail microservice endpoint.
+
+    Attributes:
+        guardrails_service_url (str): A base url for accessing the NVIDIA guardrail endpoint, e.g. http://0.0.0.0:7331
+        config_id (str): The ID of the guardrails configuration to use from the configuration store
+         (https://developer.nvidia.com/docs/nemo-microservices/guardrails/source/guides/configuration-store-guide.html)
+
+    """
+
+    guardrails_service_url: str = Field(
+        default_factory=lambda: os.getenv("GUARDRAILS_SERVICE_URL", "http://0.0.0.0:7331"),
+        description="The url for accessing the Guardrails service",
+    )
+    config_id: str | None = Field(
+        default_factory=lambda: os.getenv("NVIDIA_GUARDRAILS_CONFIG_ID", "self-check"),
+        description="Guardrails configuration ID to use from the Guardrails configuration store",
+    )
+
+    @classmethod
+    def sample_run_config(cls, **kwargs) -> dict[str, Any]:
+        return {
+            "guardrails_service_url": "${env.GUARDRAILS_SERVICE_URL:=http://localhost:7331}",
+            "config_id": "${env.NVIDIA_GUARDRAILS_CONFIG_ID:=self-check}",
+        }