llama-stack 0.4.3__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

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

@@ -23,6 +23,7 @@ from llama_stack_api import (
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
+    validate_embeddings_input_is_text,
 )
 
 logger = get_logger(name=__name__, category="providers::remote::watsonx")
@@ -147,6 +148,9 @@ class WatsonXInferenceAdapter(LiteLLMOpenAIMixin):
         """
         Override parent method to add watsonx-specific parameters.
         """
+        # Validate that input contains only text, not token arrays
+        validate_embeddings_input_is_text(params)
+
         model_obj = await self.model_store.get_model(params.model)
 
         # Convert input to list if it's a string

llama_stack/providers/remote/post_training/nvidia/README.md

@@ -0,0 +1,151 @@
+# NVIDIA Post-Training Provider for LlamaStack
+
+This provider enables fine-tuning of LLMs using NVIDIA's NeMo Customizer service.
+
+## Features
+
+- Supervised fine-tuning of Llama models
+- LoRA fine-tuning support
+- Job management and status tracking
+
+## Getting Started
+
+### Prerequisites
+
+- LlamaStack with NVIDIA configuration
+- Access to Hosted NVIDIA NeMo Customizer service
+- Dataset registered in the Hosted NVIDIA NeMo Customizer service
+- Base model downloaded and available in the Hosted NVIDIA NeMo Customizer service
+
+### Setup
+
+Build the NVIDIA environment:
+
+```bash
+uv pip install llama-stack-client
+uv run llama stack list-deps nvidia | xargs -L1 uv pip install
+```
+
+### Basic Usage using the LlamaStack Python Client
+
+### Create Customization Job
+
+#### Initialize the client
+
+```python
+import os
+
+os.environ["NVIDIA_API_KEY"] = "your-api-key"
+os.environ["NVIDIA_CUSTOMIZER_URL"] = "http://nemo.test"
+os.environ["NVIDIA_DATASET_NAMESPACE"] = "default"
+os.environ["NVIDIA_PROJECT_ID"] = "test-project"
+os.environ["NVIDIA_OUTPUT_MODEL_DIR"] = "test-example-model@v1"
+
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
+
+client = LlamaStackAsLibraryClient("nvidia")
+client.initialize()
+```
+
+#### Configure fine-tuning parameters
+
+```python
+from llama_stack_client.types.post_training_supervised_fine_tune_params import (
+    TrainingConfig,
+    TrainingConfigDataConfig,
+    TrainingConfigOptimizerConfig,
+)
+from llama_stack_client.types.algorithm_config_param import LoraFinetuningConfig
+```
+
+#### Set up LoRA configuration
+
+```python
+algorithm_config = LoraFinetuningConfig(type="LoRA", adapter_dim=16)
+```
+
+#### Configure training data
+
+```python
+data_config = TrainingConfigDataConfig(
+    dataset_id="your-dataset-id",  # Use client.datasets.list() to see available datasets
+    batch_size=16,
+)
+```
+
+#### Configure optimizer
+
+```python
+optimizer_config = TrainingConfigOptimizerConfig(
+    lr=0.0001,
+)
+```
+
+#### Set up training configuration
+
+```python
+training_config = TrainingConfig(
+    n_epochs=2,
+    data_config=data_config,
+    optimizer_config=optimizer_config,
+)
+```
+
+#### Start fine-tuning job
+
+```python
+training_job = client.post_training.supervised_fine_tune(
+    job_uuid="unique-job-id",
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    checkpoint_dir="",
+    algorithm_config=algorithm_config,
+    training_config=training_config,
+    logger_config={},
+    hyperparam_search_config={},
+)
+```
+
+### List all jobs
+
+```python
+jobs = client.post_training.job.list()
+```
+
+###  Check job status
+
+```python
+job_status = client.post_training.job.status(job_uuid="your-job-id")
+```
+
+### Cancel a job
+
+```python
+client.post_training.job.cancel(job_uuid="your-job-id")
+```
+
+### Inference with the fine-tuned model
+
+#### 1. Register the model
+
+```python
+from llama_stack_api.models import Model, ModelType
+
+client.models.register(
+    model_id="test-example-model@v1",
+    provider_id="nvidia",
+    provider_model_id="test-example-model@v1",
+    model_type=ModelType.llm,
+)
+```
+
+#### 2. Inference with the fine-tuned model
+
+```python
+response = client.completions.create(
+    prompt="Complete the sentence using one word: Roses are red, violets are ",
+    stream=False,
+    model="test-example-model@v1",
+    max_tokens=50,
+)
+print(response.choices[0].text)
+```

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

@@ -14,13 +14,15 @@ from llama_stack.providers.remote.post_training.nvidia.config import NvidiaPostT
 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,
+    CancelTrainingJobRequest,
+    GetTrainingJobArtifactsRequest,
+    GetTrainingJobStatusRequest,
     JobStatus,
     PostTrainingJob,
     PostTrainingJobArtifactsResponse,
     PostTrainingJobStatusResponse,
-    TrainingConfig,
+    PreferenceOptimizeRequest,
+    SupervisedFineTuneRequest,
 )
 
 from .models import _MODEL_ENTRIES
@@ -156,7 +158,9 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
 
         return ListNvidiaPostTrainingJobs(data=jobs)
 
-    async def get_training_job_status(self, job_uuid: str) -> NvidiaPostTrainingJobStatusResponse:
+    async def get_training_job_status(
+        self, request: GetTrainingJobStatusRequest
+    ) -> NvidiaPostTrainingJobStatusResponse:
         """Get the status of a customization job.
         Updated the base class return type from PostTrainingJobResponse to NvidiaPostTrainingJob.
 
@@ -178,8 +182,8 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
         """
         response = await self._make_request(
             "GET",
-            f"/v1/customization/jobs/{job_uuid}/status",
-            params={"job_id": job_uuid},
+            f"/v1/customization/jobs/{request.job_uuid}/status",
+            params={"job_id": request.job_uuid},
         )
 
         api_status = response.pop("status").lower()
@@ -187,18 +191,20 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
 
         return NvidiaPostTrainingJobStatusResponse(
             status=JobStatus(mapped_status),
-            job_uuid=job_uuid,
+            job_uuid=request.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:
+    async def cancel_training_job(self, request: CancelTrainingJobRequest) -> None:
         await self._make_request(
-            method="POST", path=f"/v1/customization/jobs/{job_uuid}/cancel", params={"job_id": job_uuid}
+            method="POST", path=f"/v1/customization/jobs/{request.job_uuid}/cancel", params={"job_id": request.job_uuid}
         )
 
-    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse:
+    async def get_training_job_artifacts(
+        self, request: GetTrainingJobArtifactsRequest
+    ) -> PostTrainingJobArtifactsResponse:
         raise NotImplementedError("Job artifacts are not implemented yet")
 
     async def get_post_training_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse:
@@ -206,13 +212,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
 
     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,
+        request: SupervisedFineTuneRequest,
     ) -> NvidiaPostTrainingJob:
         """
         Fine-tunes a model on a dataset.
@@ -300,13 +300,16 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
             User is informed about unsupported parameters via warnings.
         """
 
+        # Convert training_config to dict for internal processing
+        training_config = request.training_config.model_dump()
+
         # Check for unsupported method parameters
         unsupported_method_params = []
-        if checkpoint_dir:
-            unsupported_method_params.append(f"checkpoint_dir={checkpoint_dir}")
-        if hyperparam_search_config:
+        if request.checkpoint_dir:
+            unsupported_method_params.append(f"checkpoint_dir={request.checkpoint_dir}")
+        if request.hyperparam_search_config:
             unsupported_method_params.append("hyperparam_search_config")
-        if logger_config:
+        if request.logger_config:
             unsupported_method_params.append("logger_config")
 
         if unsupported_method_params:
@@ -344,7 +347,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
 
         # Prepare base job configuration
         job_config = {
-            "config": model,
+            "config": request.model,
             "dataset": {
                 "name": training_config["data_config"]["dataset_id"],
                 "namespace": self.config.dataset_namespace,
@@ -388,14 +391,14 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
             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")
+        if request.algorithm_config:
+            if request.algorithm_config.type == "LoRA":
+                warn_unsupported_params(request.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
+                    k: v for k, v in {"alpha": request.algorithm_config.alpha}.items() if v is not None
                 }
             else:
-                raise NotImplementedError(f"Unsupported algorithm config: {algorithm_config}")
+                raise NotImplementedError(f"Unsupported algorithm config: {request.algorithm_config}")
 
         # Create the customization job
         response = await self._make_request(
@@ -416,12 +419,7 @@ class NvidiaPostTrainingAdapter(ModelRegistryHelper):
 
     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],
+        request: PreferenceOptimizeRequest,
     ) -> PostTrainingJob:
         """Optimize a model based on preference data."""
         raise NotImplementedError("Preference optimization is not implemented yet")

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

@@ -5,12 +5,13 @@
 # the root directory of this source tree.
 
 import json
-from typing import Any
 
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.bedrock.client import create_bedrock_client
+from llama_stack.providers.utils.safety import ShieldToModerationMixin
 from llama_stack_api import (
-    OpenAIMessageParam,
+    GetShieldRequest,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -24,7 +25,7 @@ from .config import BedrockSafetyConfig
 logger = get_logger(name=__name__, category="safety::bedrock")
 
 
-class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
+class BedrockSafetyAdapter(ShieldToModerationMixin, Safety, ShieldsProtocolPrivate):
     def __init__(self, config: BedrockSafetyConfig) -> None:
         self.config = config
         self.registered_shields = []
@@ -55,49 +56,31 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
     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)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.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"]
-        """
+            raise ValueError(f"Shield {request.shield_id} not found")
 
         shield_params = shield.params
-        logger.debug(f"run_shield::{shield_params}::messages={messages}")
+        logger.debug(f"run_shield::{shield_params}::messages={request.messages}")
 
-        # - convert the messages into format Bedrock expects
         content_messages = []
-        for message in messages:
+        for message in request.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
+            source="OUTPUT",
             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(

llama_stack/providers/remote/safety/nvidia/README.md

@@ -0,0 +1,78 @@
+# NVIDIA Safety Provider for LlamaStack
+
+This provider enables safety checks and guardrails for LLM interactions using NVIDIA's NeMo Guardrails service.
+
+## Features
+
+- Run safety checks for messages
+
+## Getting Started
+
+### Prerequisites
+
+- LlamaStack with NVIDIA configuration
+- Access to NVIDIA NeMo Guardrails service
+- NIM for model to use for safety check is deployed
+
+### Setup
+
+Build the NVIDIA environment:
+
+```bash
+uv pip install llama-stack-client
+uv run llama stack list-deps nvidia | xargs -L1 uv pip install
+```
+
+### Basic Usage using the LlamaStack Python Client
+
+#### Initialize the client
+
+```python
+import os
+
+os.environ["NVIDIA_API_KEY"] = "your-api-key"
+os.environ["NVIDIA_GUARDRAILS_URL"] = "http://guardrails.test"
+
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
+
+client = LlamaStackAsLibraryClient("nvidia")
+client.initialize()
+```
+
+#### Create a safety shield
+
+```python
+from llama_stack_api.safety import Shield
+from llama_stack_api.inference import Message
+
+# Create a safety shield
+shield = Shield(
+    shield_id="your-shield-id",
+    provider_resource_id="safety-model-id",  # The model to use for safety checks
+    description="Safety checks for content moderation",
+)
+
+# Register the shield
+await client.safety.register_shield(shield)
+```
+
+#### Run safety checks
+
+```python
+# Messages to check
+messages = [Message(role="user", content="Your message to check")]
+
+# Run safety check
+response = await client.safety.run_shield(
+    shield_id="your-shield-id",
+    messages=messages,
+)
+
+# Check for violations
+if response.violation:
+    print(f"Safety violation detected: {response.violation.user_message}")
+    print(f"Violation level: {response.violation.violation_level}")
+    print(f"Metadata: {response.violation.metadata}")
+else:
+    print("No safety violations detected")
+```

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

@@ -9,9 +9,11 @@ from typing import Any
 import requests
 
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.safety import ShieldToModerationMixin
 from llama_stack_api import (
-    ModerationObject,
+    GetShieldRequest,
     OpenAIMessageParam,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -25,7 +27,7 @@ from .config import NVIDIASafetyConfig
 logger = get_logger(name=__name__, category="safety::nvidia")
 
 
-class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate):
+class NVIDIASafetyAdapter(ShieldToModerationMixin, Safety, ShieldsProtocolPrivate):
     def __init__(self, config: NVIDIASafetyConfig) -> None:
         """
         Initialize the NVIDIASafetyAdapter with a given safety configuration.
@@ -48,32 +50,14 @@ class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate):
     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 = None
-    ) -> RunShieldResponse:
-        """
-        Run a safety shield check against the provided messages.
-
-        Args:
-            shield_id (str): The unique identifier for the shield to be used.
-            messages (List[Message]): A list of Message objects representing the conversation history.
-            params (Optional[dict[str, Any]]): Additional parameters for the shield check.
-
-        Returns:
-            RunShieldResponse: The response containing safety violation details if any.
-
-        Raises:
-            ValueError: If the shield with the provided shield_id is not found.
-        """
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        """Run a safety shield check against the provided messages."""
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Shield {shield_id} not found")
+            raise ValueError(f"Shield {request.shield_id} not found")
 
         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 | None = None) -> ModerationObject:
-        raise NotImplementedError("NVIDIA safety provider currently does not implement run_moderation")
+        return await self.shield.run(request.messages)
 
 
 class NeMoGuardrails:

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

@@ -4,15 +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
-
 import litellm
 import requests
 
 from llama_stack.core.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.safety import ShieldToModerationMixin
 from llama_stack_api import (
-    OpenAIMessageParam,
+    GetShieldRequest,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -28,7 +28,7 @@ logger = get_logger(name=__name__, category="safety::sambanova")
 CANNED_RESPONSE_TEXT = "I can't answer that. Can I help with something else?"
 
 
-class SambaNovaSafetyAdapter(Safety, ShieldsProtocolPrivate, NeedsRequestProviderData):
+class SambaNovaSafetyAdapter(ShieldToModerationMixin, Safety, ShieldsProtocolPrivate, NeedsRequestProviderData):
     def __init__(self, config: SambaNovaSafetyConfig) -> None:
         self.config = config
         self.environment_available_models = []
@@ -69,17 +69,19 @@ class SambaNovaSafetyAdapter(Safety, ShieldsProtocolPrivate, NeedsRequestProvide
     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 = None
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Shield {shield_id} not found")
+            raise ValueError(f"Shield {request.shield_id} not found")
 
         shield_params = shield.params
-        logger.debug(f"run_shield::{shield_params}::messages={messages}")
+        logger.debug(f"run_shield::{shield_params}::messages={request.messages}")
 
-        response = litellm.completion(model=shield.provider_resource_id, messages=messages, api_key=self._get_api_key())
+        response = litellm.completion(
+            model=shield.provider_resource_id,
+            messages=request.messages,
+            api_key=self._get_api_key(),
+        )
         shield_message = response.choices[0].message.content
 
         if "unsafe" in shield_message.lower():

llama_stack/providers/remote/vector_io/elasticsearch/__init__.py

@@ -0,0 +1,17 @@
+# 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_api import Api, ProviderSpec
+
+from .config import ElasticsearchVectorIOConfig
+
+
+async def get_adapter_impl(config: ElasticsearchVectorIOConfig, deps: dict[Api, ProviderSpec]):
+    from .elasticsearch import ElasticsearchVectorIOAdapter
+
+    impl = ElasticsearchVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
+    await impl.initialize()
+    return impl

llama_stack/providers/remote/vector_io/elasticsearch/config.py

@@ -0,0 +1,32 @@
+# 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 typing import Any
+
+from pydantic import BaseModel, Field
+
+from llama_stack.core.storage.datatypes import KVStoreReference
+from llama_stack_api import json_schema_type
+
+
+@json_schema_type
+class ElasticsearchVectorIOConfig(BaseModel):
+    elasticsearch_api_key: str | None = Field(description="The API key for the Elasticsearch instance", default=None)
+    elasticsearch_url: str | None = Field(description="The URL of the Elasticsearch instance", default="localhost:9200")
+    persistence: KVStoreReference | None = Field(
+        description="Config for KV store backend (SQLite only for now)", default=None
+    )
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
+        return {
+            "elasticsearch_url": "${env.ELASTICSEARCH_URL:=localhost:9200}",
+            "elasticsearch_api_key": "${env.ELASTICSEARCH_API_KEY:=}",
+            "persistence": KVStoreReference(
+                backend="kv_default",
+                namespace="vector_io::elasticsearch",
+            ).model_dump(exclude_none=True),
+        }