llama-stack 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

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/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/vector_io/pgvector/pgvector.py

@@ -10,6 +10,7 @@ from typing import Any
 import psycopg2
 from numpy.typing import NDArray
 from psycopg2 import sql
+from psycopg2.extensions import cursor
 from psycopg2.extras import Json, execute_values
 from pydantic import BaseModel, TypeAdapter
 
@@ -54,6 +55,17 @@ def check_extension_version(cur):
     return result[0] if result else None
 
 
+def create_vector_extension(cur: cursor) -> None:
+    try:
+        log.info("Vector extension not found, creating...")
+        cur.execute("CREATE EXTENSION vector;")
+        log.info("Vector extension created successfully")
+        log.info(f"Vector extension version: {check_extension_version(cur)}")
+
+    except psycopg2.Error as e:
+        raise RuntimeError(f"Failed to create vector extension for PGVector: {e}") from e
+
+
 def upsert_models(conn, keys_models: list[tuple[str, BaseModel]]):
     with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
         query = sql.SQL(
@@ -364,7 +376,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
                 if version:
                     log.info(f"Vector extension version: {version}")
                 else:
-                    raise RuntimeError("Vector extension is not installed.")
+                    create_vector_extension(cur)
 
                 cur.execute(
                     """

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -25,7 +25,8 @@ from llama_stack_api import (
     OpenAIEmbeddingUsage,
 )
 
-EMBEDDING_MODELS = {}
+EMBEDDING_MODELS: dict[str, "SentenceTransformer"] = {}
+EMBEDDING_MODELS_LOCK = asyncio.Lock()
 
 DARWIN = "Darwin"
 
@@ -76,26 +77,29 @@ class SentenceTransformerEmbeddingMixin:
         )
 
     async def _load_sentence_transformer_model(self, model: str) -> "SentenceTransformer":
-        global EMBEDDING_MODELS
-
         loaded_model = EMBEDDING_MODELS.get(model)
         if loaded_model is not None:
             return loaded_model
 
-        log.info(f"Loading sentence transformer for {model}...")
+        async with EMBEDDING_MODELS_LOCK:
+            loaded_model = EMBEDDING_MODELS.get(model)
+            if loaded_model is not None:
+                return loaded_model
+
+            log.info(f"Loading sentence transformer for {model}...")
 
-        def _load_model():
-            from sentence_transformers import SentenceTransformer
+            def _load_model():
+                from sentence_transformers import SentenceTransformer
 
-            platform_name = platform.system()
-            if platform_name == DARWIN:
-                # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
-                # threads with the default parallel settings, so force a single-threaded CPU run.
-                log.debug(f"Constraining torch threads on {platform_name} to a single worker")
-                torch.set_num_threads(1)
+                platform_name = platform.system()
+                if platform_name == DARWIN:
+                    # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
+                    # threads with the default parallel settings, so force a single-threaded CPU run.
+                    log.debug(f"Constraining torch threads on {platform_name} to a single worker")
+                    torch.set_num_threads(1)
 
-            return SentenceTransformer(model, trust_remote_code=True)
+                return SentenceTransformer(model, trust_remote_code=True)
 
-        loaded_model = await asyncio.to_thread(_load_model)
-        EMBEDDING_MODELS[model] = loaded_model
-        return loaded_model
+            loaded_model = await asyncio.to_thread(_load_model)
+            EMBEDDING_MODELS[model] = loaded_model
+            return loaded_model

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -122,6 +122,39 @@ class OpenAIVectorStoreMixin(ABC):
         # update in-memory cache
         self.openai_vector_stores[store_id] = store_info
 
+    async def _ensure_openai_metadata_exists(self, vector_store: VectorStore, name: str | None = None) -> None:
+        """
+        Ensure OpenAI-compatible metadata exists for a vector store.
+        """
+        if vector_store.identifier not in self.openai_vector_stores:
+            store_info = {
+                "id": vector_store.identifier,
+                "object": "vector_store",
+                "created_at": int(time.time()),
+                "name": name or vector_store.vector_store_name or vector_store.identifier,
+                "usage_bytes": 0,
+                "file_counts": VectorStoreFileCounts(
+                    cancelled=0,
+                    completed=0,
+                    failed=0,
+                    in_progress=0,
+                    total=0,
+                ).model_dump(),
+                "status": "completed",
+                "expires_after": None,
+                "expires_at": None,
+                "last_active_at": int(time.time()),
+                "file_ids": [],
+                "chunking_strategy": None,
+                "metadata": {
+                    "provider_id": vector_store.provider_id,
+                    "provider_vector_store_id": vector_store.provider_resource_id,
+                    "embedding_model": vector_store.embedding_model,
+                    "embedding_dimension": str(vector_store.embedding_dimension),
+                },
+            }
+            await self._save_openai_vector_store(vector_store.identifier, store_info)
+
     async def _load_openai_vector_stores(self) -> dict[str, dict[str, Any]]:
         """Load all vector store metadata from persistent storage."""
         assert self.kvstore

llama_stack/providers/utils/responses/responses_store.py

@@ -89,6 +89,40 @@ class ResponsesStore:
     ) -> None:
         await self._write_response_object(response_object, input, messages)
 
+    async def upsert_response_object(
+        self,
+        response_object: OpenAIResponseObject,
+        input: list[OpenAIResponseInput],
+        messages: list[OpenAIMessageParam],
+    ) -> None:
+        """Upsert response object using INSERT on first call, UPDATE on subsequent calls.
+
+        This method enables incremental persistence during streaming, allowing clients
+        to poll GET /v1/responses/{response_id} and see in-progress turn state.
+
+        :param response_object: The response object to store/update.
+        :param input: The input items for the response.
+        :param messages: The chat completion messages (for conversation continuity).
+        """
+        if self.sql_store is None:
+            raise ValueError("Responses store is not initialized")
+
+        data = response_object.model_dump()
+        data["input"] = [input_item.model_dump() for input_item in input]
+        data["messages"] = [msg.model_dump() for msg in messages]
+
+        await self.sql_store.upsert(
+            table="openai_responses",
+            data={
+                "id": data["id"],
+                "created_at": data["created_at"],
+                "model": data["model"],
+                "response_object": data,
+            },
+            conflict_columns=["id"],
+            update_columns=["response_object"],
+        )
+
     async def _write_response_object(
         self,
         response_object: OpenAIResponseObject,